4b766cb908
- 4 provider adapters: OpenAI (SDK), Anthropic (SDK), Google (google-genai), Mistral (direct HTTP) - Core router with automatic failover + exponential backoff - Flask blueprint with /api/v1/ai/* endpoints - Auth via token-broker verify endpoint - DB models for ai_providers, ai_model_mapping, ai_router_log - /health endpoint (parallel provider check), /usage stats - 21 unit tests (all passing)
58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
import logging
|
|
from typing import Optional
|
|
|
|
from .base import AIProvider
|
|
|
|
logger = logging.getLogger("ai_router.google")
|
|
|
|
|
|
class GoogleAdapter(AIProvider):
|
|
@property
|
|
def name(self) -> str:
|
|
return "google"
|
|
|
|
def chat(self, messages: list, model: str, api_key: Optional[str] = None, **kwargs) -> dict:
|
|
from google import genai
|
|
|
|
key = api_key or self.get_api_key()
|
|
client = genai.Client(api_key=key)
|
|
|
|
system_instruction = None
|
|
chat_messages = messages
|
|
if messages and messages[0].get("role") == "system":
|
|
system_instruction = messages[0]["content"]
|
|
chat_messages = messages[1:]
|
|
|
|
contents = []
|
|
for m in chat_messages:
|
|
role = "user" if m["role"] in ("user", "system") else "model"
|
|
contents.append({"role": role, "parts": [{"text": m["content"]}]})
|
|
|
|
resp = client.models.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
config={"system_instruction": system_instruction} if system_instruction else None,
|
|
)
|
|
return {
|
|
"content": resp.text or "",
|
|
"model": model,
|
|
"provider": self.name,
|
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
|
}
|
|
|
|
def models(self) -> list:
|
|
from google import genai
|
|
|
|
client = genai.Client(api_key=self.get_api_key())
|
|
return [m.name for m in client.models.list()]
|
|
|
|
def check_health(self) -> dict:
|
|
try:
|
|
from google import genai
|
|
client = genai.Client(api_key=self.get_api_key())
|
|
client.models.list()
|
|
return {"status": "ok", "details": "API reachable"}
|
|
except Exception as e:
|
|
logger.warning(f"Google health check failed: {e}")
|
|
return {"status": "error", "details": str(e)}
|