增加工具使用功能

models.py

@@ -1,15 +1,46 @@
 from datetime import datetime
-from typing import Literal
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 
-class Message(BaseModel):
-    role: Literal["user", "assistant"]
+class ToolFunctionCall(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: Literal["function"] = "function"
+    function: ToolFunctionCall
+
+
+class UserMessage(BaseModel):
+    role: Literal["user"] = "user"
     content: str
     created_at: datetime
 
 
+class AssistantMessage(BaseModel):
+    role: Literal["assistant"] = "assistant"
+    content: str | None = None
+    tool_calls: list[ToolCall] = Field(default_factory=list)
+    created_at: datetime
+
+
+class ToolMessage(BaseModel):
+    role: Literal["tool"] = "tool"
+    content: str
+    tool_call_id: str
+    name: str
+    is_error: bool = False
+    created_at: datetime
+
+
+Message = Annotated[
+    UserMessage | AssistantMessage | ToolMessage,
+    Field(discriminator="role"),
+]
 class Session(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -72,9 +103,16 @@ class TokenUsage(BaseModel):
     total_tokens: int
 
 
+class FinalAssistantMessage(BaseModel):
+    role: Literal["assistant"] = "assistant"
+    content: str
+    created_at: datetime
+
+
 class SendMessageResponse(BaseModel):
     session_id: str
-    message: Message
+    message: FinalAssistantMessage
+    tools_use: list[str] = Field(default_factory=list)
     usage: TokenUsage | None