第一次提交

2026-06-29 23:48:53 +08:00
commit 8714fb07c4
13 changed files with 1471 additions and 0 deletions
@@ -0,0 +1,158 @@
+import asyncio
+from collections.abc import AsyncGenerator
+from datetime import UTC, datetime
+from openai import APIError, AsyncOpenAI
+
+from models import Message, SendMessageResponse, Session, TokenUsage
+from storage import JsonSessionStorage
+
+
+class DeepSeekUpstreamError(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]
+                    stream=False,
+                    extra_body={"thinking": {"type": "disabled"}},
+                )
+                assistant_content = completion.choices[0].message.content
+            except (APIError, IndexError, AttributeError) as exc:
+                raise DeepSeekUpstreamError("DeepSeek request failed") from exc
+
+            if assistant_content is None:
+                raise DeepSeekUpstreamError("DeepSeek 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 ensure_session_exists(self, session_id: str) -> None:
+        await self.storage.read(session_id)
+
+    async def generate_response_stream(
+        self, session_id: str, content: str
+    ) -> AsyncGenerator[dict[str, object], None]:
+        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:
+                stream = await self.client.chat.completions.create(
+                    model=self.model,
+                    messages=api_messages,  # type: ignore[arg-type]
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    extra_body={"thinking": {"type": "disabled"}},
+                )
+                parts: list[str] = []
+                usage: TokenUsage | None = None
+
+                try:
+                    async for chunk in stream:
+                        chunk_usage = getattr(chunk, "usage", None)
+                        if chunk_usage is not None:
+                            usage = self._parse_usage(chunk_usage)
+
+                        for choice in chunk.choices:
+                            delta = choice.delta.content
+                            if delta:
+                                parts.append(delta)
+                                yield {"type": "delta", "content": delta}
+                finally:
+                    await stream.close()
+            except (APIError, AttributeError, TypeError) as exc:
+                raise DeepSeekUpstreamError("DeepSeek stream failed") from exc
+
+            assistant_content = "".join(parts)
+            if not assistant_content:
+                raise DeepSeekUpstreamError("DeepSeek returned an empty response")
+
+            assistant_message = await self._save_exchange(
+                session, content, user_created_at, assistant_content
+            )
+            yield {
+                "type": "done",
+                "message": assistant_message.model_dump(mode="json"),
+                "usage": usage.model_dump() if usage is not None else None,
+            }
+
+    @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]
+        )