src/common/llm_client.py

import os
from typing import Optional

import google.generativeai as genai
from google.generativeai.types import GenerationConfig


class LLMClient:
    def __init__(self):
        self.provider = os.getenv("LLM_PROVIDER", "gemini").lower()
        self.model_name = os.getenv("LLM_MODEL", "gemini-3.1-pro-preview")

        if self.provider == "gemini":
            if genai is None:
                raise ImportError("google-generativeai not installed.")
            key = os.getenv("GEMINI_API_KEY")
            if not key:
                raise ValueError("GEMINI_API_KEY not set.")
            genai.configure(api_key=key)
            self.model = genai.GenerativeModel(self.model_name)
        else:
            raise NotImplementedError("Only Gemini provider is wired for now.")

    def call(self, prompt: str, schema: Optional[dict] = None) -> str:
        """
        Call the underlying LLM.

        If `schema` is provided (as a plain JSON schema dict), and provider is Gemini,
        use it as response_schema with JSON mime type.
        """
        if self.provider == "gemini":
            if schema and GenerationConfig is not None:
                config = GenerationConfig(
                    response_schema=schema,
                    response_mime_type="application/json",
                )
                response = self.model.generate_content(
                    prompt,
                    generation_config=config,
                )
            else:
                response = self.model.generate_content(prompt)

            text = getattr(response, "text", "")
            if not text:
                raise RuntimeError("LLM response did not contain text.")
            return text

        raise NotImplementedError("Schema-based calls only wired for Gemini right now.")