92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
import os
|
|
import json
|
|
import time
|
|
import httpx
|
|
import logging
|
|
from typing import Dict, Any, Optional
|
|
from core.metrics import AI_REQUESTS_TOTAL, AI_LATENCY_SECONDS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class LLMClient:
|
|
def __init__(
|
|
self,
|
|
provider: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
):
|
|
self.provider = provider or os.getenv("AI_PROVIDER", "gemini").lower()
|
|
self.api_key = api_key or os.getenv("AI_API_KEY", "")
|
|
self.model = model or os.getenv("AI_MODEL", "gemini-1.5-flash" if self.provider == "gemini" else "gpt-4o-mini")
|
|
self.base_url = base_url or os.getenv("AI_BASE_URL")
|
|
|
|
async def generate_json(self, prompt: str, system_prompt: Optional[str] = None, action_name: str = "general") -> Dict[str, Any]:
|
|
"""Send prompt to LLM and parse JSON response."""
|
|
start_time = time.time()
|
|
status = "error"
|
|
try:
|
|
if self.provider == "gemini":
|
|
result = await self._call_gemini(prompt, system_prompt)
|
|
else:
|
|
result = await self._call_openai(prompt, system_prompt)
|
|
status = "success"
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"LLM generation failed ({self.provider}/{self.model}): {e}")
|
|
raise
|
|
finally:
|
|
duration = time.time() - start_time
|
|
AI_LATENCY_SECONDS.labels(action=action_name).observe(duration)
|
|
AI_REQUESTS_TOTAL.labels(action=action_name, status=status).inc()
|
|
|
|
async def _call_gemini(self, prompt: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent?key={self.api_key}"
|
|
payload: Dict[str, Any] = {
|
|
"contents": [
|
|
{
|
|
"parts": [{"text": prompt}]
|
|
}
|
|
],
|
|
"generationConfig": {
|
|
"responseMimeType": "application/json",
|
|
"temperature": 0.2,
|
|
}
|
|
}
|
|
if system_prompt:
|
|
payload["systemInstruction"] = {
|
|
"parts": [{"text": system_prompt}]
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
resp = await client.post(url, json=payload)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
raw_text = data["candidates"][0]["content"]["parts"][0]["text"]
|
|
return json.loads(raw_text)
|
|
|
|
async def _call_openai(self, prompt: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
|
|
url = self.base_url or "https://api.openai.com/v1/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"response_format": {"type": "json_object"},
|
|
"temperature": 0.2,
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
resp = await client.post(url, headers=headers, json=payload)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
content = data["choices"][0]["message"]["content"]
|
|
return json.loads(content)
|