File size: 996 Bytes
befaa94
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
# chat_model_wrapper.py

import os
import google.generativeai as genai
from dotenv import load_dotenv

load_dotenv(override=True)

class ChatRefiner:
    def __init__(self, model_name: str = "gemini-1.5-flash"):
        """
        Wraps the Gemini model for prompt→text generation.
        """
        api_key = os.getenv("GEMINI_API_KEY")
        api_key="AIzaSyDAC_1r8FuL9JckNno5hCI5vMD2MWFYEfg"
        if not api_key:
            raise ValueError("GEMINI_API_KEY not set in environment or .env file.")

        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel(model_name)

    def answer(self, prompt: str) -> str:
        """
        Sends `prompt` to Gemini and returns the generated text.
        """
        try:
            response = self.model.generate_content(prompt)
            return response.text.strip()
        except Exception as e:
            # Log or handle as you prefer
            print(f"[ChatRefiner.answer] Error: {e}")
            raise