Spaces:
Sleeping
Sleeping
| """Agent graph state, its reducers, and the response shape. | |
| AgentAnswer.to_dict() is a strict SUPERSET of rag.answer.Answer.to_dict(), so the | |
| frontend renders both the fast path and the agent with the same code. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| from typing import Annotated, Any, TypedDict | |
| # --------------------------------------------------------------------------- # | |
| # reducers | |
| # --------------------------------------------------------------------------- # | |
| def merge_evidence(left: list[dict], right: list[dict] | None) -> list[dict]: | |
| """Accumulate retrieved passages across steps, deduped by (page_number, text). | |
| The same page legitimately arrives twice — once from a search hit, once from a | |
| read_page — and we keep the higher score so the compose ordering stays sane. | |
| `right is None` RESETS. The checkpointer persists state across turns of a | |
| conversation, so without a reset the second question would inherit (and cite) | |
| the first question's evidence. `prepare` sends None at the start of every turn. | |
| """ | |
| if right is None: | |
| return [] | |
| if not right: | |
| return left | |
| merged: dict[tuple[int, str], dict] = {} | |
| for item in [*(left or []), *right]: | |
| key = (item.get("page_number", -1), item.get("text", "")) | |
| prev = merged.get(key) | |
| if prev is None or item.get("score", 0.0) > prev.get("score", 0.0): | |
| merged[key] = item | |
| return list(merged.values()) | |
| def append_turns(left: list[dict], right: list[dict]) -> list[dict]: | |
| """Conversation history. Plain append — the checkpointer persists it per thread. | |
| Deliberately has no reset: this is the memory the agent exists to provide. | |
| """ | |
| return [*(left or []), *(right or [])] | |
| def extend_trace(left: list[dict], right: list[dict] | None) -> list[dict]: | |
| """Per-step audit trail. Resets per turn on None, like merge_evidence.""" | |
| if right is None: | |
| return [] | |
| return [*(left or []), *right] | |
| # --------------------------------------------------------------------------- # | |
| # state | |
| # --------------------------------------------------------------------------- # | |
| class AgentState(TypedDict, total=False): | |
| # -- inputs | |
| doc_id: int | |
| question: str # the raw user turn | |
| persona: str | None # "cook" -> in-character Forster (rag/llm.py:_SYSTEM_COOK) | |
| max_steps: int | |
| # -- derived | |
| standalone_question: str # coreference-resolved; what actually gets embedded | |
| history: Annotated[list[dict], append_turns] # [{role, content}], survives via checkpointer | |
| evidence: Annotated[list[dict], merge_evidence] # [{page_number, text, score, via}] | |
| trace: Annotated[list[dict], extend_trace] # per-step audit, shown in the UI | |
| books: dict[str, Any] # book map: {"1": [1, 132], ...} | |
| # -- control | |
| steps: int | |
| rewrites: int | |
| verify_attempts: int | |
| deadline: float # time.monotonic() ceiling, checked in every node | |
| next_action: dict # the action `plan` chose, consumed by `act` | |
| stop_reason: str # why the loop ended: answer | max_steps | deadline | llm_unavailable | |
| # -- outputs | |
| answer: str | |
| source_pages: list[int] | |
| grounded: bool | |
| class AgentAnswer: | |
| """Superset of rag.answer.Answer so the frontend can share one renderer.""" | |
| question: str | |
| answer: str | |
| source_pages: list[int] = field(default_factory=list) | |
| persona: str | None = None | |
| contexts: list[dict] = field(default_factory=list) | |
| # agent-only | |
| thread_id: str | None = None | |
| standalone_question: str | None = None | |
| steps: int = 0 | |
| grounded: bool = True | |
| stop_reason: str = "answer" | |
| trace: list[dict] = field(default_factory=list) | |
| books: dict[str, Any] = field(default_factory=dict) | |
| def to_dict(self) -> dict: | |
| return { | |
| # --- shared with Answer.to_dict() | |
| "question": self.question, | |
| "answer": self.answer, | |
| "source_pages": self.source_pages, | |
| "persona": self.persona, | |
| "in_character": bool(self.persona), | |
| "contexts": self.contexts, | |
| # --- agent extras | |
| "mode": "agent", | |
| "thread_id": self.thread_id, | |
| "standalone_question": self.standalone_question, | |
| "steps": self.steps, | |
| "grounded": self.grounded, | |
| "stop_reason": self.stop_reason, | |
| "trace": self.trace, | |
| "books": self.books, | |
| } | |