saksham
1a4c80958f
FastAPI + Celery + Next.js + Postgres/Redis app with company monitoring, source collection, LLM-based change analysis, enrichment, and account security (Turnstile, escalating lockout, email verification).
74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
"""Gemini provider: structured output via the SDK's native
|
|
`response_schema` support (the model is constrained to the schema and the
|
|
SDK parses the result into an instance of it directly), with the same
|
|
bounded repair loop on a malformed/unparsed response as
|
|
`anthropic_provider.py`. Chosen as the production LLM_PROVIDER option
|
|
alongside Anthropic because Gemini has an actual free rate-limited API
|
|
tier (gemini-2.0-flash), unlike OpenAI's expiring trial credits.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from google import genai
|
|
from google.genai import types
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from app.analysis.llm.base import LLMResponseError
|
|
from app.core.config import Settings
|
|
|
|
|
|
class GeminiLLMProvider:
|
|
provider_name = "gemini"
|
|
|
|
def __init__(self, settings: Settings) -> None:
|
|
self._settings = settings
|
|
self._client = genai.Client(api_key=settings.gemini_api_key)
|
|
|
|
async def generate_structured[T: BaseModel](
|
|
self, system_prompt: str, user_prompt: str, response_model: type[T]
|
|
) -> T:
|
|
last_error: Exception | None = None
|
|
prompt = user_prompt
|
|
|
|
for attempt in range(self._settings.llm_max_retries + 1):
|
|
if attempt > 0 and last_error is not None:
|
|
prompt = (
|
|
f"{user_prompt}\n\nYour previous response did not match the required "
|
|
f"schema: {last_error}. Please try again."
|
|
)
|
|
|
|
response = await self._client.aio.models.generate_content(
|
|
model=self._settings.gemini_model,
|
|
contents=prompt,
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_prompt,
|
|
response_mime_type="application/json",
|
|
response_schema=response_model,
|
|
max_output_tokens=self._settings.llm_max_tokens_per_request,
|
|
),
|
|
)
|
|
if response.parsed is None:
|
|
last_error = ValueError("Gemini did not return a parsed structured response")
|
|
continue
|
|
try:
|
|
return response_model.model_validate(response.parsed)
|
|
except ValidationError as exc:
|
|
last_error = exc
|
|
continue
|
|
|
|
raise LLMResponseError(
|
|
f"Gemini provider failed to produce a valid {response_model.__name__} after "
|
|
f"{self._settings.llm_max_retries + 1} attempt(s): {last_error}"
|
|
)
|
|
|
|
async def generate_text(self, system_prompt: str, user_prompt: str) -> str:
|
|
response = await self._client.aio.models.generate_content(
|
|
model=self._settings.gemini_model,
|
|
contents=user_prompt,
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_prompt,
|
|
max_output_tokens=self._settings.llm_max_tokens_per_request,
|
|
),
|
|
)
|
|
return response.text or ""
|