113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import httpx
|
|
|
|
from app.exceptions import ProviderAPICallError
|
|
from app.models import (
|
|
LLMChoice,
|
|
LLMMessage,
|
|
LLMProvider,
|
|
LLMRequest,
|
|
LLMResponse,
|
|
LLMRole,
|
|
)
|
|
from app.providers.base import LLMProviderClient
|
|
|
|
|
|
class GeminiProvider(LLMProviderClient):
|
|
name = LLMProvider.GEMINI.value
|
|
api_key_env = "GEMINI_API_KEY"
|
|
base_url = "https://generativelanguage.googleapis.com/v1beta"
|
|
|
|
async def chat(
|
|
self, request: LLMRequest, client: httpx.AsyncClient
|
|
) -> LLMResponse:
|
|
self.ensure_stream_supported(request.stream)
|
|
|
|
system_instruction, contents = self._convert_messages(request.messages)
|
|
config = {
|
|
"temperature": request.temperature,
|
|
"topP": request.top_p,
|
|
"maxOutputTokens": request.max_tokens,
|
|
}
|
|
|
|
payload: Dict[str, Any] = self.merge_payload(
|
|
{"contents": contents}, request.extra_params
|
|
)
|
|
|
|
generation_config = {k: v for k, v in config.items() if v is not None}
|
|
if generation_config:
|
|
payload["generationConfig"] = generation_config
|
|
if system_instruction:
|
|
payload["systemInstruction"] = {
|
|
"role": "system",
|
|
"parts": [{"text": system_instruction}],
|
|
}
|
|
|
|
endpoint = f"{self.base_url}/models/{request.model}:generateContent?key={self.api_key}"
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
response = await client.post(endpoint, json=payload, headers=headers)
|
|
response.raise_for_status()
|
|
except httpx.HTTPError as exc:
|
|
raise ProviderAPICallError(f"Gemini request failed: {exc}") from exc
|
|
|
|
data: Dict[str, Any] = response.json()
|
|
choices = self._build_choices(data.get("candidates", []))
|
|
|
|
return LLMResponse(
|
|
provider=LLMProvider.GEMINI,
|
|
model=request.model,
|
|
choices=choices,
|
|
raw=data,
|
|
)
|
|
|
|
@staticmethod
|
|
def _convert_messages(
|
|
messages: List[LLMMessage],
|
|
) -> Tuple[str | None, List[dict[str, Any]]]:
|
|
system_parts: List[str] = []
|
|
contents: List[dict[str, Any]] = []
|
|
|
|
for msg in messages:
|
|
if msg.role == LLMRole.SYSTEM:
|
|
system_parts.append(msg.content)
|
|
continue
|
|
|
|
role = "user" if msg.role == LLMRole.USER else "model"
|
|
contents.append({"role": role, "parts": [{"text": msg.content}]})
|
|
|
|
system_instruction = "\n\n".join(system_parts) if system_parts else None
|
|
return system_instruction, contents
|
|
|
|
@staticmethod
|
|
def _build_choices(candidates: List[dict[str, Any]]) -> List[LLMChoice]:
|
|
choices: List[LLMChoice] = []
|
|
for idx, candidate in enumerate(candidates):
|
|
content = candidate.get("content", {})
|
|
parts = content.get("parts", [])
|
|
text_parts = [
|
|
part.get("text", "")
|
|
for part in parts
|
|
if isinstance(part, dict) and part.get("text")
|
|
]
|
|
text = "\n\n".join(text_parts)
|
|
choices.append(
|
|
LLMChoice(
|
|
index=candidate.get("index", idx),
|
|
message=LLMMessage(role="assistant", content=text),
|
|
)
|
|
)
|
|
if not choices:
|
|
choices.append(
|
|
LLMChoice(
|
|
index=0,
|
|
message=LLMMessage(role="assistant", content=""),
|
|
)
|
|
)
|
|
return choices
|