74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
import requests
|
|
|
|
from app.modules.shared.gigachat.errors import GigaChatError
|
|
from app.modules.shared.gigachat.settings import GigaChatSettings
|
|
from app.modules.shared.gigachat.token_provider import GigaChatTokenProvider
|
|
|
|
|
|
class GigaChatClient:
|
|
def __init__(self, settings: GigaChatSettings, token_provider: GigaChatTokenProvider) -> None:
|
|
self._settings = settings
|
|
self._tokens = token_provider
|
|
|
|
def complete(self, system_prompt: str, user_prompt: str) -> str:
|
|
token = self._tokens.get_access_token()
|
|
payload = {
|
|
"model": self._settings.model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
}
|
|
try:
|
|
response = requests.post(
|
|
f"{self._settings.api_url.rstrip('/')}/chat/completions",
|
|
json=payload,
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=90,
|
|
verify=self._settings.ssl_verify,
|
|
)
|
|
except requests.RequestException as exc:
|
|
raise GigaChatError(f"GigaChat completion request failed: {exc}") from exc
|
|
|
|
if response.status_code >= 400:
|
|
raise GigaChatError(f"GigaChat completion error {response.status_code}: {response.text}")
|
|
|
|
data = response.json()
|
|
choices = data.get("choices") or []
|
|
if not choices:
|
|
return ""
|
|
message = choices[0].get("message") or {}
|
|
return str(message.get("content") or "")
|
|
|
|
def embed(self, texts: list[str]) -> list[list[float]]:
|
|
token = self._tokens.get_access_token()
|
|
payload = {
|
|
"model": self._settings.embedding_model,
|
|
"input": texts,
|
|
}
|
|
try:
|
|
response = requests.post(
|
|
f"{self._settings.api_url.rstrip('/')}/embeddings",
|
|
json=payload,
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=90,
|
|
verify=self._settings.ssl_verify,
|
|
)
|
|
except requests.RequestException as exc:
|
|
raise GigaChatError(f"GigaChat embeddings request failed: {exc}") from exc
|
|
|
|
if response.status_code >= 400:
|
|
raise GigaChatError(f"GigaChat embeddings error {response.status_code}: {response.text}")
|
|
|
|
data = response.json()
|
|
items = data.get("data")
|
|
if not isinstance(items, list):
|
|
raise GigaChatError("Unexpected GigaChat embeddings response")
|
|
return [list(map(float, x.get("embedding") or [])) for x in items]
|