123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
import asyncio
|
|
from datetime import UTC, datetime
|
|
from openai import APIError, AsyncOpenAI
|
|
|
|
from models import (
|
|
Message,
|
|
SendMessageResponse,
|
|
Session,
|
|
SessionHistoryResponse,
|
|
TokenUsage,
|
|
)
|
|
from storage import JsonSessionStorage
|
|
|
|
|
|
class ChatProviderError(Exception):
|
|
pass
|
|
|
|
|
|
class ChatService:
|
|
def __init__(
|
|
self,
|
|
storage: JsonSessionStorage,
|
|
client: AsyncOpenAI,
|
|
model: str,
|
|
) -> None:
|
|
self.storage = storage
|
|
self.client = client
|
|
self.model = model
|
|
self._locks: dict[str, asyncio.Lock] = {}
|
|
|
|
async def generate_response(
|
|
self, session_id: str, content: str
|
|
) -> SendMessageResponse:
|
|
lock = self._locks.setdefault(session_id, asyncio.Lock())
|
|
async with lock:
|
|
user_created_at = datetime.now(UTC)
|
|
session = await self.storage.read(session_id)
|
|
api_messages = self._build_api_messages(session, content)
|
|
|
|
try:
|
|
completion = await self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=api_messages, # type: ignore[arg-type]
|
|
extra_body={"thinking": {"type": "disabled"}},
|
|
)
|
|
assistant_content = completion.choices[0].message.content
|
|
except (APIError, IndexError, AttributeError) as exc:
|
|
raise ChatProviderError("Upstream chat request failed") from exc
|
|
|
|
if assistant_content is None:
|
|
raise ChatProviderError("Upstream chat provider returned an empty response")
|
|
|
|
assistant_message = await self._save_exchange(
|
|
session, content, user_created_at, assistant_content
|
|
)
|
|
usage = self._parse_usage(completion.usage)
|
|
|
|
return SendMessageResponse(
|
|
session_id=session_id,
|
|
message=assistant_message,
|
|
usage=usage,
|
|
)
|
|
|
|
async def get_session_history(self, session_id: str) -> SessionHistoryResponse:
|
|
lock = self._locks.setdefault(session_id, asyncio.Lock())
|
|
async with lock:
|
|
session = await self.storage.read(session_id)
|
|
return SessionHistoryResponse(
|
|
session_id=session.session_id,
|
|
system_prompt=session.system_prompt,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
messages=session.messages,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_api_messages(
|
|
session: Session, content: str
|
|
) -> list[dict[str, str]]:
|
|
messages = [{"role": "system", "content": session.system_prompt}]
|
|
messages.extend(
|
|
{"role": message.role, "content": message.content}
|
|
for message in session.messages
|
|
)
|
|
messages.append({"role": "user", "content": content})
|
|
return messages
|
|
|
|
async def _save_exchange(
|
|
self,
|
|
session: Session,
|
|
user_content: str,
|
|
user_created_at: datetime,
|
|
assistant_content: str,
|
|
) -> Message:
|
|
assistant_message = Message(
|
|
role="assistant",
|
|
content=assistant_content,
|
|
created_at=datetime.now(UTC),
|
|
)
|
|
session.messages.extend(
|
|
[
|
|
Message(
|
|
role="user",
|
|
content=user_content,
|
|
created_at=user_created_at,
|
|
),
|
|
assistant_message,
|
|
]
|
|
)
|
|
session.updated_at = assistant_message.created_at
|
|
await self.storage.write(session)
|
|
return assistant_message
|
|
|
|
@staticmethod
|
|
def _parse_usage(usage: object | None) -> TokenUsage | None:
|
|
if usage is None:
|
|
return None
|
|
return TokenUsage(
|
|
prompt_tokens=usage.prompt_tokens, # type: ignore[attr-defined]
|
|
completion_tokens=usage.completion_tokens, # type: ignore[attr-defined]
|
|
total_tokens=usage.total_tokens, # type: ignore[attr-defined]
|
|
)
|