90e370ef35
* fix: misc fixes for observations and mental models * feat: improve graph retrieval for observations - Update LinkExpansionRetriever to traverse through source_memory_ids for observation entity connections (avoiding data duplication) - Remove entity link copy from world facts to observations in consolidator - Add tests for link expansion graph retrieval - Add directives_applied field to ReflectResult - Include user's other changes (CLI, docs, client updates) * fix: CI test failures - Add mental_model_id parameter to create_mental_model function - Fix ToolCallTrace not including reason field from ToolCall - Improve test_link_expansion_observation_graph_retrieval to wait for consolidation with retry * chore: reduce link expansion log verbosity * Revert "chore: reduce link expansion log verbosity" This reverts commit 3ce759391cead1012157785fa78fef16ef9bfe3b. * feat: add semantic/temporal/entity links as fallback in graph retrieval - Add fallback query for semantic, temporal, and entity links from memory_links - Check both directions (outgoing and incoming links) - Weight fallback results at 0.5x to prioritize entity links via unit_entities - Fixes graph retrieval returning 0 when data has cross-cluster temporal connections * fix: enable observations fixture for link expansion test - Add enable_observations fixture to ensure observations are created - Increase wait time from 10 to 30 seconds for CI reliability
109 lines
5.1 KiB
Python
109 lines
5.1 KiB
Python
"""
|
|
Pydantic models for the reflect agent.
|
|
"""
|
|
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class ObservationSection(BaseModel):
|
|
"""A section within an observation with its supporting memories."""
|
|
|
|
title: str = Field(description="Section header (can be empty for intro)")
|
|
text: str = Field(description="Section content - no headers, use lists/tables/bold")
|
|
memory_ids: list[str] = Field(default_factory=list, description="Memory IDs supporting this section")
|
|
|
|
|
|
class ReflectAction(BaseModel):
|
|
"""Single action the reflect agent can take."""
|
|
|
|
tool: Literal["list_observations", "get_observation", "recall", "expand", "done"] = Field(
|
|
description="Tool to invoke: list_observations, get_observation, recall, expand, or done"
|
|
)
|
|
# Tool-specific parameters
|
|
observation_id: str | None = Field(default=None, description="Observation ID for get_observation")
|
|
query: str | None = Field(default=None, description="Search query for recall")
|
|
max_tokens: int | None = Field(default=None, description="Max tokens for recall results (default 2048)")
|
|
memory_ids: list[str] | None = Field(default=None, description="Memory unit IDs for expand (batched)")
|
|
depth: Literal["chunk", "document"] | None = Field(default=None, description="Expansion depth for expand")
|
|
observation_sections: list[ObservationSection] | None = Field(
|
|
default=None, description="Observation sections for done action (when output_mode=observations)"
|
|
)
|
|
# Plain text answer fields (for output_mode=answer)
|
|
answer: str | None = Field(default=None, description="Plain text answer for done action (no markdown)")
|
|
answer_memory_ids: list[str] | None = Field(
|
|
default=None, description="Memory IDs supporting the answer", alias="memory_ids"
|
|
)
|
|
answer_model_ids: list[str] | None = Field(
|
|
default=None, description="Mental model IDs supporting the answer", alias="model_ids"
|
|
)
|
|
reasoning: str | None = Field(default=None, description="Brief reasoning for this action")
|
|
|
|
|
|
class ReflectActionBatch(BaseModel):
|
|
"""Batch of actions for parallel execution."""
|
|
|
|
actions: list[ReflectAction] = Field(description="List of actions to execute in parallel")
|
|
|
|
|
|
class ToolCall(BaseModel):
|
|
"""A single tool call made during reflect."""
|
|
|
|
tool: str = Field(description="Tool name: lookup, recall, expand")
|
|
reason: str | None = Field(default=None, description="Agent's reasoning for making this tool call")
|
|
input: dict = Field(description="Tool input parameters")
|
|
output: dict = Field(description="Tool output/result")
|
|
duration_ms: int = Field(description="Execution time in milliseconds")
|
|
iteration: int = Field(default=0, description="Iteration number (1-based) when this tool was called")
|
|
|
|
|
|
class LLMCall(BaseModel):
|
|
"""A single LLM call made during reflect."""
|
|
|
|
scope: str = Field(description="Call scope: agent_1, agent_2, final, etc.")
|
|
duration_ms: int = Field(description="Execution time in milliseconds")
|
|
input_tokens: int = Field(default=0, description="Input tokens used")
|
|
output_tokens: int = Field(default=0, description="Output tokens used")
|
|
|
|
|
|
class DirectiveInfo(BaseModel):
|
|
"""Information about a directive that was applied during reflect."""
|
|
|
|
id: str = Field(description="Directive mental model ID")
|
|
name: str = Field(description="Directive name")
|
|
content: str = Field(description="Directive content")
|
|
|
|
|
|
class TokenUsageSummary(BaseModel):
|
|
"""Total token usage across all LLM calls."""
|
|
|
|
input_tokens: int = Field(default=0, description="Total input tokens used")
|
|
output_tokens: int = Field(default=0, description="Total output tokens used")
|
|
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
|
|
|
|
|
|
class ReflectAgentResult(BaseModel):
|
|
"""Result from the reflect agent."""
|
|
|
|
text: str = Field(description="Final answer text")
|
|
structured_output: dict[str, Any] | None = Field(
|
|
default=None, description="Structured output parsed according to provided response_schema"
|
|
)
|
|
iterations: int = Field(default=0, description="Number of iterations taken")
|
|
tools_called: int = Field(default=0, description="Total number of tool calls made")
|
|
tool_trace: list[ToolCall] = Field(default_factory=list, description="Trace of all tool calls made")
|
|
llm_trace: list[LLMCall] = Field(default_factory=list, description="Trace of all LLM calls made")
|
|
usage: TokenUsageSummary = Field(
|
|
default_factory=TokenUsageSummary, description="Total token usage across all LLM calls"
|
|
)
|
|
used_memory_ids: list[str] = Field(default_factory=list, description="Validated memory IDs actually used in answer")
|
|
used_mental_model_ids: list[str] = Field(
|
|
default_factory=list, description="Validated mental model IDs actually used in answer"
|
|
)
|
|
used_observation_ids: list[str] = Field(
|
|
default_factory=list, description="Validated observation IDs actually used in answer"
|
|
)
|
|
directives_applied: list[DirectiveInfo] = Field(
|
|
default_factory=list, description="Directive mental models that affected this reflection"
|
|
)
|