161 lines
5.2 KiB
Python
161 lines
5.2 KiB
Python
"""GigaChat provider for VoIdea."""
|
|
|
|
import base64
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Optional
|
|
|
|
import httpx
|
|
|
|
from app.core.config import get_settings
|
|
from app.integrations.ai.base import AIProvider, AIResult
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
class GigaChatProvider(AIProvider):
|
|
"""GigaChat AI provider (Sber).
|
|
|
|
Uses OAuth client credentials for authentication.
|
|
Token is refreshed automatically on 401.
|
|
"""
|
|
|
|
name = "gigachat"
|
|
model = "GigaChat"
|
|
timeout = settings.ai_timeout
|
|
max_retries = settings.ai_max_retries
|
|
|
|
def __init__(self, api_key: str = ""):
|
|
super().__init__(api_key)
|
|
self.base_url = settings.ai_gigachat_url.rstrip("/")
|
|
self.client_id = settings.ai_gigachat_client_id
|
|
self.client_secret = settings.ai_gigachat_secret
|
|
self._client = httpx.AsyncClient(timeout=self.timeout)
|
|
self._access_token: Optional[str] = None
|
|
self._token_expires_at: float = 0
|
|
|
|
async def _get_access_token(self) -> Optional[str]:
|
|
"""Get or refresh GigaChat access token."""
|
|
if self._access_token and time.time() < self._token_expires_at:
|
|
return self._access_token
|
|
|
|
if not self.client_id or not self.client_secret:
|
|
return None
|
|
|
|
auth_str = base64.b64encode(
|
|
f"{self.client_id}:{self.client_secret}".encode()
|
|
).decode()
|
|
|
|
try:
|
|
response = await self._client.post(
|
|
"https://ngw.devices.sber.ru/token",
|
|
headers={
|
|
"Authorization": f"Basic {auth_str}",
|
|
"Content-Type": "application/x-www-form-urlencoded",
|
|
"RqUID": str(time.time()),
|
|
},
|
|
data={"scope": "GIGACHAT_API_PERS"},
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
self._access_token = data.get("access_token")
|
|
expires_in = data.get("expires_in", 1800)
|
|
self._token_expires_at = time.time() + expires_in - 60
|
|
return self._access_token
|
|
|
|
except Exception:
|
|
return None
|
|
|
|
return None
|
|
|
|
async def analyze(self, prompt: str, **kwargs: Any) -> AIResult:
|
|
"""Send prompt to GigaChat and return result."""
|
|
start = time.time()
|
|
temperature = kwargs.get("temperature", 0.7)
|
|
max_tokens = kwargs.get("max_tokens", 2000)
|
|
|
|
token = await self._get_access_token()
|
|
if not token:
|
|
return AIResult(
|
|
success=False,
|
|
error="GigaChat authentication failed",
|
|
provider=self.name,
|
|
)
|
|
|
|
parts = prompt.split("\n\n", 1)
|
|
system_text = parts[0]
|
|
user_text = parts[1] if len(parts) > 1 else ""
|
|
|
|
body = {
|
|
"model": self.model,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"messages": [
|
|
{"role": "system", "content": system_text},
|
|
{"role": "user", "content": user_text},
|
|
],
|
|
}
|
|
|
|
try:
|
|
response = await self._client.post(
|
|
f"{self.base_url}/api/v1/chat/completion",
|
|
json=body,
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
|
|
duration = int((time.time() - start) * 1000)
|
|
|
|
if response.status_code == 401:
|
|
self._access_token = None
|
|
token = await self._get_access_token()
|
|
if token:
|
|
response = await self._client.post(
|
|
f"{self.base_url}/api/v1/chat/completion",
|
|
json=body,
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
duration = int((time.time() - start) * 1000)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
result_text = (
|
|
data.get("choices", [{}])[0]
|
|
.get("message", {})
|
|
.get("content", "")
|
|
)
|
|
|
|
return AIResult(
|
|
success=True,
|
|
content=result_text,
|
|
model=self.model,
|
|
provider=self.name,
|
|
duration_ms=duration,
|
|
)
|
|
else:
|
|
error_body = response.text[:500]
|
|
return AIResult(
|
|
success=False,
|
|
error=f"GigaChat {response.status_code}: {error_body}",
|
|
provider=self.name,
|
|
duration_ms=duration,
|
|
)
|
|
|
|
except httpx.TimeoutException:
|
|
return AIResult(
|
|
success=False,
|
|
error="GigaChat timeout",
|
|
provider=self.name,
|
|
duration_ms=int((time.time() - start) * 1000),
|
|
)
|
|
except Exception as e:
|
|
return AIResult(
|
|
success=False,
|
|
error=f"GigaChat error: {str(e)}",
|
|
provider=self.name,
|
|
duration_ms=int((time.time() - start) * 1000),
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
await self._client.aclose()
|