* feat: add reflect mode to LoComo benchmark and improve reflect agent - Replace think mode with reflect mode in LoComo benchmark using reflect_async with Budget.HIGH - Add --question-index CLI flag to run a single question by its index - Track and display original question index in logs and visualizer - Update visualizer to show reflect mode results Reflect agent improvements: - tool_recall: always fetch chunks (max_chunk_tokens=1000 min, non-optional) - tool_search_observations: use include_source_facts=True instead of separate DB query - Use model_dump() throughout to avoid manual error-prone dict conversion - Enforce minimum 1000 tokens for max_tokens and max_chunk_tokens in _execute_tool - Fix NoneType error when LLM passes null for mental_model_ids/observation_ids arrays - Add non-conversational constraint to system prompt to prevent follow-up questions - Fix recall_fn Callable type hint to include max_chunk_tokens parameter - Fix main.py missing reranker_zeroentropy fields in HindsightConfig constructor * fix: update tests for reflect tool API changes - source_memory_ids -> source_fact_ids in test_search_observations (MemoryFact.model_dump() field name) - Remove proof_count check (not in MemoryFact, was ObservationResult-specific) - Remove max_results param from tool_recall call (no longer supported) - Fix recall_result["count"] -> len(recall_result["memories"])
259 lines
11 KiB
Python
259 lines
11 KiB
Python
"""
|
|
Tool schema definitions for the reflect agent.
|
|
|
|
These are OpenAI-format tool definitions used with native tool calling.
|
|
The reflect agent uses a hierarchical retrieval strategy:
|
|
1. search_mental_models - User-curated stored reflect responses (highest quality, if applicable)
|
|
2. search_observations - Consolidated knowledge with freshness awareness
|
|
3. recall - Raw facts (world/experience) as ground truth fallback
|
|
"""
|
|
|
|
# Tool definitions in OpenAI format
|
|
|
|
TOOL_SEARCH_MENTAL_MODELS = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_mental_models",
|
|
"description": (
|
|
"Search user-curated mental models (stored reflect responses). These are high-quality, manually created "
|
|
"summaries about specific topics. Use FIRST when the question might be covered by an "
|
|
"existing mental model. Returns mental models with their content and last refresh time."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"reason": {
|
|
"type": "string",
|
|
"description": "Brief explanation of why you're making this search (for debugging)",
|
|
},
|
|
"query": {
|
|
"type": "string",
|
|
"description": "Search query to find relevant mental models",
|
|
},
|
|
"max_results": {
|
|
"type": "integer",
|
|
"description": "Maximum number of mental models to return (default 5)",
|
|
},
|
|
},
|
|
"required": ["reason", "query"],
|
|
},
|
|
},
|
|
}
|
|
|
|
TOOL_SEARCH_OBSERVATIONS = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_observations",
|
|
"description": (
|
|
"Search consolidated observations (auto-generated knowledge). These are automatically "
|
|
"synthesized from memories. Returns observations with freshness info (updated_at, is_stale). "
|
|
"If an observation is STALE, you should ALSO use recall() to verify with current facts. "
|
|
"IMPORTANT: If search_mental_models is available, you MUST call it FIRST before using this tool."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"reason": {
|
|
"type": "string",
|
|
"description": "Brief explanation of why you're making this search (for debugging)",
|
|
},
|
|
"query": {
|
|
"type": "string",
|
|
"description": "Search query to find relevant observations",
|
|
},
|
|
"max_tokens": {
|
|
"type": "integer",
|
|
"description": "Maximum tokens for results (default 5000). Use higher values for broader searches.",
|
|
},
|
|
},
|
|
"required": ["reason", "query"],
|
|
},
|
|
},
|
|
}
|
|
|
|
TOOL_RECALL = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "recall",
|
|
"description": (
|
|
"Search raw memories (facts and experiences). This is the ground truth data. "
|
|
"Use when: (1) no reflections/mental models exist, (2) mental models are stale, "
|
|
"(3) you need specific details not in synthesized knowledge. "
|
|
"Returns individual memory facts with their timestamps."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"reason": {
|
|
"type": "string",
|
|
"description": "Brief explanation of why you're making this search (for debugging)",
|
|
},
|
|
"query": {
|
|
"type": "string",
|
|
"description": "Search query string",
|
|
},
|
|
"max_tokens": {
|
|
"type": "integer",
|
|
"description": "Optional limit on result size (default 2048). Use higher values for broader searches.",
|
|
},
|
|
"max_chunk_tokens": {
|
|
"type": "integer",
|
|
"description": "Maximum tokens for raw source chunk text included alongside each memory fact (default 1000, min 1000). Chunks provide the surrounding context the fact was extracted from. Increase for broader context.",
|
|
},
|
|
},
|
|
"required": ["reason", "query"],
|
|
},
|
|
},
|
|
}
|
|
|
|
TOOL_EXPAND = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "expand",
|
|
"description": "Get more context for one or more memories. Memory hierarchy: memory -> chunk -> document.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"reason": {
|
|
"type": "string",
|
|
"description": "Brief explanation of why you need more context (for debugging)",
|
|
},
|
|
"memory_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of memory IDs from recall results (batch multiple for efficiency)",
|
|
},
|
|
"depth": {
|
|
"type": "string",
|
|
"enum": ["chunk", "document"],
|
|
"description": "chunk: surrounding text chunk, document: full source document",
|
|
},
|
|
},
|
|
"required": ["reason", "memory_ids", "depth"],
|
|
},
|
|
},
|
|
}
|
|
|
|
TOOL_DONE_ANSWER = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "done",
|
|
"description": "Signal completion with your final answer. Use this when you have gathered enough information to answer the question.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"answer": {
|
|
"type": "string",
|
|
"description": "Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array. LANGUAGE: By default, write in the SAME language as the user's question. However, if a language directive in the system prompt specifies a different language, follow that directive instead.",
|
|
},
|
|
"memory_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of memory IDs that support your answer (put IDs here, NOT in answer text)",
|
|
},
|
|
"mental_model_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of mental model IDs that support your answer",
|
|
},
|
|
"observation_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of observation IDs that support your answer",
|
|
},
|
|
},
|
|
"required": ["answer"],
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def _build_done_tool_with_directives(directive_rules: list[str]) -> dict:
|
|
"""
|
|
Build the done tool schema with directive compliance field.
|
|
|
|
When directives are present, adds a required field that forces the agent
|
|
to confirm compliance with each directive before submitting.
|
|
|
|
Args:
|
|
directive_rules: List of directive rule strings
|
|
"""
|
|
# Build rules list for description
|
|
rules_list = "\n".join(f" {i + 1}. {rule}" for i, rule in enumerate(directive_rules))
|
|
|
|
# Build the tool with directive compliance field
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "done",
|
|
"description": (
|
|
"Signal completion with your final answer. IMPORTANT: You must confirm directive compliance before submitting. "
|
|
"Your answer will be REJECTED if it violates any directive."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"answer": {
|
|
"type": "string",
|
|
"description": (
|
|
"Your response as well-formatted markdown. Use headers, lists, bold/italic, and code blocks for clarity. "
|
|
"NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array. "
|
|
f"MANDATORY: Your answer MUST comply with ALL directives:\n{rules_list}"
|
|
),
|
|
},
|
|
"memory_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of memory IDs that support your answer (put IDs here, NOT in answer text)",
|
|
},
|
|
"mental_model_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of mental model IDs that support your answer",
|
|
},
|
|
"observation_ids": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Array of observation IDs that support your answer",
|
|
},
|
|
"directive_compliance": {
|
|
"type": "string",
|
|
"description": f"REQUIRED: Confirm your answer complies with ALL directives. List each directive and how your answer follows it:\n{rules_list}\n\nFormat: 'Directive 1: [how answer complies]. Directive 2: [how answer complies]...'",
|
|
},
|
|
},
|
|
"required": ["answer", "directive_compliance"],
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def get_reflect_tools(directive_rules: list[str] | None = None) -> list[dict]:
|
|
"""
|
|
Get the list of tools for the reflect agent.
|
|
|
|
The tools support a hierarchical retrieval strategy:
|
|
1. search_mental_models - User-curated stored reflect responses (try first)
|
|
2. search_observations - Consolidated knowledge with freshness
|
|
3. recall - Raw facts as ground truth
|
|
|
|
Args:
|
|
directive_rules: Optional list of directive rule strings. If provided,
|
|
the done() tool will require directive compliance confirmation.
|
|
|
|
Returns:
|
|
List of tool definitions in OpenAI format
|
|
"""
|
|
tools = [
|
|
TOOL_SEARCH_MENTAL_MODELS,
|
|
TOOL_SEARCH_OBSERVATIONS,
|
|
TOOL_RECALL,
|
|
TOOL_EXPAND,
|
|
]
|
|
|
|
# Use directive-aware done tool if directives are present
|
|
if directive_rules:
|
|
tools.append(_build_done_tool_with_directives(directive_rules))
|
|
else:
|
|
tools.append(TOOL_DONE_ANSWER)
|
|
|
|
return tools
|