Spaces:
Running
Running
File size: 1,453 Bytes
2aa4095 |
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 34 35 36 37 38 39 40 41 |
import os
from google import genai
from dotenv import load_dotenv
load_dotenv()
class GeminiLLM:
def __init__(self, api_key=None):
self.api_key = api_key or os.environ.get("GEMINI_API_Key")
if not self.api_key:
raise ValueError("API key must be provided or set in the environment variable 'GEMINI_API_Key'")
try:
self.client = genai.Client(api_key=self.api_key)
except Exception as e:
raise RuntimeError(f"Failed to initialize Gemini client: {str(e)}")
self.model_id = "gemini-2.0-flash"
def generate_response(self, prompt: str) -> str:
try:
response = self.client.models.generate_content(
model=self.model_id,
contents=prompt,
)
if hasattr(response, "text"):
return response.text
raise RuntimeError("No response text received from model")
except Exception as e:
raise RuntimeError(f"Error generating response: {str(e)}")
def stream_response(self, prompt: str):
try:
stream = self.client.models.generate_content_stream(
model=self.model_id,
contents=prompt,
)
for chunk in stream:
if hasattr(chunk, "text") and chunk.text:
yield chunk.text
except Exception as e:
yield f"Error during streaming: {str(e)}"
|