73 lines
1.7 KiB
Python
73 lines
1.7 KiB
Python
from datetime import datetime
|
|
from typing import Annotated, Literal
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
|
|
|
|
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")
|
|
|
|
user_id: str | None = None
|
|
session_id: str
|
|
system_prompt: str
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
api_calls: int = 0
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
total_tokens: int = 0
|
|
messages: list[Message] = Field(default_factory=list)
|
|
|
|
@model_validator(mode="before")
|
|
@classmethod
|
|
def add_timestamp_to_legacy_messages(cls, data: object) -> object:
|
|
if not isinstance(data, dict):
|
|
return data
|
|
|
|
fallback = data.get("created_at")
|
|
messages = data.get("messages")
|
|
if fallback is not None and isinstance(messages, list):
|
|
for message in messages:
|
|
if isinstance(message, dict):
|
|
message.setdefault("created_at", fallback)
|
|
return data
|