Spaces:
Runtime error
Runtime error
File size: 3,184 Bytes
0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 a076686 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 02422e3 0d599d9 | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
gemini_adapter.py — Concrete LLMAdapter for Google Gemini.
V3 (Sidecar): adds stream_chat() for sentence-level streaming.
"""
import os
from typing import AsyncGenerator, Optional
import google.generativeai as genai
from llm_adapter import LLMAdapter
class GeminiAdapter(LLMAdapter):
"""
Wraps Google Gemini API as an LLMAdapter.
Supports both blocking `chat()` and streaming `stream_chat()`.
"""
def __init__(
self,
api_key: Optional[str] = None,
model_name: str = "gemini-3.1-flash-lite",
system_prompt: Optional[str] = None,
):
self.api_key = api_key or os.getenv("GEMINI_API_KEY")
if not self.api_key:
raise ValueError(
"Gemini API key not found. Set GEMINI_API_KEY environment variable "
"or pass api_key directly."
)
self.model_name = model_name
self.system_prompt = system_prompt or "You are a helpful AI assistant."
genai.configure(api_key=self.api_key)
self._model = genai.GenerativeModel(
model_name=self.model_name,
system_instruction=self.system_prompt,
)
# ------------------------------------------------------------------
# Blocking interface (existing — unchanged)
# ------------------------------------------------------------------
def chat(self, prompt: str, system_prompt: Optional[str] = None) -> str:
"""Send prompt, return full response text (blocking)."""
model = self._get_model(system_prompt)
response = model.generate_content(prompt)
return response.text
# ------------------------------------------------------------------
# Streaming interface (NEW)
# ------------------------------------------------------------------
async def stream_chat(
self,
prompt: str,
system_prompt: Optional[str] = None,
) -> AsyncGenerator[str, None]:
"""
Yield token chunks from Gemini as they arrive.
This is a synchronous SDK call wrapped in an async generator —
Gemini's Python SDK streams synchronously, so we iterate the
response object directly and yield each text chunk.
"""
model = self._get_model(system_prompt)
# generate_content with stream=True returns a synchronous iterator
response = model.generate_content(prompt, stream=True)
for chunk in response:
text = getattr(chunk, "text", None)
if text:
yield text
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _get_model(self, system_prompt: Optional[str]):
"""Return model instance, recreating if system prompt differs."""
if system_prompt and system_prompt != self.system_prompt:
return genai.GenerativeModel(
model_name=self.model_name,
system_instruction=system_prompt,
)
return self._model
def get_model_name(self) -> str:
return self.model_name
|