wandering
9043886b6f
- src/doc/ 拆分为 src/core/extraction/, matching/, validation/（核心业务逻辑） - src/bot/ 重命名为 src/infra/browser/（浏览器自动化基础设施） - fill_consumable_doc.py → src/infra/documents/consumable.py - 新增 Agent 调度模块：coordinator.py, events.py, session.py，重构 orchestrator.py - 更新 AGENTS.md、README.md 及所有子目录 README
102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
"""Agent 会话管理
|
|
|
|
负责会话状态的定义、序列化和持久化。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass, field
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .. import get_logger
|
|
|
|
log = get_logger("agent.session")
|
|
|
|
# ------------------------------------------------------------------
|
|
# 状态枚举
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
class AgentState(StrEnum):
|
|
IDLE = "idle"
|
|
EXTRACTING = "extracting"
|
|
AWAITING_SUPPLEMENT = "awaiting_supplement"
|
|
READY = "ready"
|
|
SUBMITTING = "submitting"
|
|
DONE = "done"
|
|
ERROR = "error"
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Agent 会话数据模型
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class AgentSession:
|
|
"""Agent 会话状态"""
|
|
|
|
session_id: str
|
|
state: AgentState = AgentState.IDLE
|
|
rounds: int = 0
|
|
max_rounds: int = 5
|
|
invoice_type: str = "travel" # "travel" 或 "normal"
|
|
extracted_info: dict[str, Any] = field(default_factory=dict)
|
|
validation_reports: list[dict[str, Any]] = field(default_factory=list)
|
|
user_supplements: list[str] = field(default_factory=list)
|
|
error_message: str = ""
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> AgentSession:
|
|
# 兼容旧版本：state 可能是字符串
|
|
if "state" in data and isinstance(data["state"], str):
|
|
data["state"] = AgentState(data["state"])
|
|
return cls(**data)
|
|
|
|
def is_terminal(self) -> bool:
|
|
"""判断会话是否处于终态（已提交或已完成）"""
|
|
return self.state in (AgentState.DONE, AgentState.SUBMITTING, AgentState.READY)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# 持久化
|
|
# ------------------------------------------------------------------
|
|
|
|
AGENT_STATE_FILE = "agent_state.json"
|
|
|
|
|
|
def save_agent_state(session_dir: Path, session: AgentSession) -> None:
|
|
"""将 Agent 会话状态持久化到 session 目录。"""
|
|
state_path = session_dir / AGENT_STATE_FILE
|
|
tmp_path = session_dir / (AGENT_STATE_FILE + ".tmp")
|
|
with open(tmp_path, "w", encoding="utf-8") as f:
|
|
json.dump(session.to_dict(), f, ensure_ascii=False, indent=2)
|
|
tmp_path.replace(state_path)
|
|
|
|
|
|
def load_agent_state(session_dir: Path) -> AgentSession | None:
|
|
"""从 session 目录加载 Agent 会话状态。"""
|
|
state_path = session_dir / AGENT_STATE_FILE
|
|
if not state_path.exists():
|
|
return None
|
|
try:
|
|
with open(state_path, encoding="utf-8") as f:
|
|
return AgentSession.from_dict(json.load(f))
|
|
except Exception as e:
|
|
log.warning("加载 Agent 状态失败: %s", e)
|
|
return None
|
|
|
|
|
|
def create_agent_session(session_id: str, invoice_type: str = "travel") -> AgentSession:
|
|
"""创建新的 Agent 会话。"""
|
|
return AgentSession(
|
|
session_id=session_id,
|
|
invoice_type=invoice_type,
|
|
)
|