"""Beaver session 子系统的数据模型。 这层只定义数据结构，不放数据库读写逻辑。目的是把： 1. SQLite 行结构 2. 运行时会话对象 3. 对外暴露的 conversation message 三件事分开，避免后续所有地方都直接和裸字典耦合。 """ from __future__ import annotations import json from dataclasses import dataclass, field from typing import Any @dataclass(slots=True) class SessionUsage: """会话维度的 usage/cost 统计。""" input_tokens: int = 0 output_tokens: int = 0 cache_read_tokens: int = 0 cache_write_tokens: int = 0 reasoning_tokens: int = 0 estimated_cost_usd: float = 0.0 actual_cost_usd: float | None = None def to_dict(self) -> dict[str, Any]: return { "input_tokens": self.input_tokens, "output_tokens": self.output_tokens, "cache_read_tokens": self.cache_read_tokens, "cache_write_tokens": self.cache_write_tokens, "reasoning_tokens": self.reasoning_tokens, "estimated_cost_usd": self.estimated_cost_usd, "actual_cost_usd": self.actual_cost_usd, } @dataclass(slots=True) class MessageRecord: """单条会话事件的结构化表示。 当前仍然沿用 `messages` 这张表名，但语义已经开始向 event stream 收拢： 1. 普通 user/assistant/tool 消息本身就是事件 2. 运行时的 system snapshot / run lifecycle 也可写成隐藏事件 3. 是否进入模型上下文由 `context_visible` 决定，而不是简单看 role """ role: str content: str | None = None timestamp: float | None = None message_id: int | None = None run_id: str | None = None event_type: str | None = None event_payload: dict[str, Any] | None = None context_visible: bool = True tool_name: str | None = None tool_calls: list[dict[str, Any]] | None = None tool_call_id: str | None = None finish_reason: str | None = None reasoning: str | None = None reasoning_details: Any | None = None codex_reasoning_items: Any | None = None def to_conversation_message(self) -> dict[str, Any]: """转成 provider / context builder 可直接消费的消息格式。""" if not self.context_visible: raise ValueError("Hidden session events cannot be converted into conversation messages") payload: dict[str, Any] = { "role": self.role, "content": self.content, } if self.timestamp is not None: payload["timestamp"] = self.timestamp if self.run_id: payload["run_id"] = self.run_id if self.event_payload: if self.event_payload.get("task_id"): payload["task_id"] = self.event_payload.get("task_id") if self.event_payload.get("task_status"): payload["task_status"] = self.event_payload.get("task_status") if self.event_payload.get("validation_status"): payload["validation_status"] = self.event_payload.get("validation_status") if self.event_payload.get("feedback_state"): payload["feedback_state"] = self.event_payload.get("feedback_state") if self.event_payload.get("feedback_error"): payload["feedback_error"] = self.event_payload.get("feedback_error") for key in ( "message_type", "scheduled_job_id", "scheduled_run_id", "cron_job_name", "mode", ): if self.event_payload.get(key): payload[key] = self.event_payload.get(key) if self.tool_name: payload["tool_name"] = self.tool_name if self.tool_calls: payload["tool_calls"] = self.tool_calls if self.tool_call_id: payload["tool_call_id"] = self.tool_call_id if self.finish_reason: payload["finish_reason"] = self.finish_reason if self.reasoning: payload["reasoning"] = self.reasoning if self.reasoning_details is not None: payload["reasoning_details"] = self.reasoning_details if self.codex_reasoning_items is not None: payload["codex_reasoning_items"] = self.codex_reasoning_items return payload @classmethod def from_row(cls, row: dict[str, Any]) -> "MessageRecord": """从 SQLite row/dict 恢复消息模型。""" tool_calls = row.get("tool_calls") if isinstance(tool_calls, str): try: tool_calls = json.loads(tool_calls) except json.JSONDecodeError: tool_calls = [] reasoning_details = row.get("reasoning_details") if isinstance(reasoning_details, str): try: reasoning_details = json.loads(reasoning_details) except json.JSONDecodeError: reasoning_details = None codex_reasoning_items = row.get("codex_reasoning_items") if isinstance(codex_reasoning_items, str): try: codex_reasoning_items = json.loads(codex_reasoning_items) except json.JSONDecodeError: codex_reasoning_items = None event_payload = row.get("event_payload") if isinstance(event_payload, str): try: event_payload = json.loads(event_payload) except json.JSONDecodeError: event_payload = None return cls( message_id=row.get("id"), run_id=row.get("run_id"), role=row["role"], content=row.get("content"), event_type=row.get("event_type") or row.get("role"), event_payload=event_payload, context_visible=bool(row.get("context_visible", 1)), tool_name=row.get("tool_name"), tool_calls=tool_calls, tool_call_id=row.get("tool_call_id"), timestamp=row.get("timestamp"), finish_reason=row.get("finish_reason"), reasoning=row.get("reasoning"), reasoning_details=reasoning_details, codex_reasoning_items=codex_reasoning_items, ) @dataclass(slots=True) class SessionRecord: """单个 session 的结构化表示。""" session_id: str source: str started_at: float last_active: float user_id: str | None = None title: str | None = None model: str | None = None system_prompt: str | None = None parent_session_id: str | None = None ended_at: float | None = None end_reason: str | None = None message_count: int = 0 tool_call_count: int = 0 preview: str | None = None usage: SessionUsage = field(default_factory=SessionUsage) def to_dict(self) -> dict[str, Any]: payload = { "id": self.session_id, "source": self.source, "user_id": self.user_id, "title": self.title, "model": self.model, "system_prompt": self.system_prompt, "parent_session_id": self.parent_session_id, "started_at": self.started_at, "last_active": self.last_active, "ended_at": self.ended_at, "end_reason": self.end_reason, "message_count": self.message_count, "tool_call_count": self.tool_call_count, "preview": self.preview, } payload.update(self.usage.to_dict()) return payload @classmethod def from_row(cls, row: dict[str, Any]) -> "SessionRecord": return cls( session_id=row["id"], source=row["source"], user_id=row.get("user_id"), title=row.get("title"), model=row.get("model"), system_prompt=row.get("system_prompt"), parent_session_id=row.get("parent_session_id"), started_at=row["started_at"], last_active=row["last_active"], ended_at=row.get("ended_at"), end_reason=row.get("end_reason"), message_count=row.get("message_count", 0), tool_call_count=row.get("tool_call_count", 0), preview=row.get("preview"), usage=SessionUsage( input_tokens=row.get("input_tokens", 0), output_tokens=row.get("output_tokens", 0), cache_read_tokens=row.get("cache_read_tokens", 0), cache_write_tokens=row.get("cache_write_tokens", 0), reasoning_tokens=row.get("reasoning_tokens", 0), estimated_cost_usd=row.get("estimated_cost_usd", 0.0) or 0.0, actual_cost_usd=row.get("actual_cost_usd"), ), )