gemini_adapter.py

"""
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