hindsight-api 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

hindsight_api/engine/reflect/tools.py

@@ -0,0 +1,437 @@
+"""
+Tool implementations for the reflect agent.
+
+Implements hierarchical retrieval:
+1. search_mental_models - User-curated stored reflect responses (highest quality)
+2. search_observations - Consolidated knowledge with freshness
+3. recall - Raw facts as ground truth
+"""
+
+import logging
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from asyncpg import Connection
+
+    from ...api.http import RequestContext
+    from ..memory_engine import MemoryEngine
+
+logger = logging.getLogger(__name__)
+
+# Observation is considered stale if not updated in this many days
+STALE_THRESHOLD_DAYS = 7
+
+
+async def tool_search_mental_models(
+    conn: "Connection",
+    bank_id: str,
+    query: str,
+    query_embedding: list[float],
+    max_results: int = 5,
+    tags: list[str] | None = None,
+    tags_match: str = "any",
+    exclude_ids: list[str] | None = None,
+) -> dict[str, Any]:
+    """
+    Search user-curated mental models by semantic similarity.
+
+    Mental models are high-quality, manually created summaries about specific topics.
+    They should be searched FIRST as they represent the most reliable synthesized knowledge.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        query: Search query (for logging/tracing)
+        query_embedding: Pre-computed embedding for semantic search
+        max_results: Maximum number of mental models to return
+        tags: Optional tags to filter mental models
+        tags_match: How to match tags - "any" (OR), "all" (AND)
+        exclude_ids: Optional list of mental model IDs to exclude (e.g., when refreshing a mental model)
+
+    Returns:
+        Dict with matching mental models including content and freshness info
+    """
+    from ..memory_engine import fq_table
+
+    # Build filters dynamically
+    filters = ""
+    params: list[Any] = [bank_id, str(query_embedding), max_results]
+    next_param = 4
+
+    if tags:
+        if tags_match == "all":
+            filters += f" AND tags @> ${next_param}::varchar[]"
+        else:
+            filters += f" AND (tags && ${next_param}::varchar[] OR tags IS NULL OR tags = '{{}}')"
+        params.append(tags)
+        next_param += 1
+
+    if exclude_ids:
+        filters += f" AND id != ALL(${next_param}::text[])"
+        params.append(exclude_ids)
+        next_param += 1
+
+    # Search mental models by embedding similarity
+    rows = await conn.fetch(
+        f"""
+        SELECT
+            id, name, content,
+            tags, created_at, last_refreshed_at,
+            1 - (embedding <=> $2::vector) as relevance
+        FROM {fq_table("mental_models")}
+        WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
+        ORDER BY embedding <=> $2::vector
+        LIMIT $3
+        """,
+        *params,
+    )
+
+    now = datetime.now(timezone.utc)
+    mental_models = []
+
+    for row in rows:
+        last_refreshed_at = row["last_refreshed_at"]
+        if last_refreshed_at and last_refreshed_at.tzinfo is None:
+            last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
+
+        # Calculate freshness
+        is_stale = False
+        if last_refreshed_at:
+            age = now - last_refreshed_at
+            is_stale = age > timedelta(days=STALE_THRESHOLD_DAYS)
+
+        mental_models.append(
+            {
+                "id": str(row["id"]),
+                "name": row["name"],
+                "content": row["content"],
+                "tags": row["tags"] or [],
+                "relevance": round(row["relevance"], 4),
+                "updated_at": last_refreshed_at.isoformat() if last_refreshed_at else None,
+                "is_stale": is_stale,
+            }
+        )
+
+    return {
+        "query": query,
+        "count": len(mental_models),
+        "mental_models": mental_models,
+    }
+
+
+async def tool_search_observations(
+    memory_engine: "MemoryEngine",
+    bank_id: str,
+    query: str,
+    request_context: "RequestContext",
+    max_tokens: int = 5000,
+    tags: list[str] | None = None,
+    tags_match: str = "any",
+    last_consolidated_at: datetime | None = None,
+    pending_consolidation: int = 0,
+) -> dict[str, Any]:
+    """
+    Search consolidated observations using recall with include_observations.
+
+    Observations are auto-generated from memories. Returns freshness info
+    so the agent knows if it should also verify with recall().
+
+    Args:
+        memory_engine: Memory engine instance
+        bank_id: Bank identifier
+        query: Search query
+        request_context: Request context for authentication
+        max_tokens: Maximum tokens for results (default 5000)
+        tags: Optional tags to filter observations
+        tags_match: How to match tags - "any" (OR), "all" (AND)
+        last_consolidated_at: When consolidation last ran (for staleness check)
+        pending_consolidation: Number of memories waiting to be consolidated
+
+    Returns:
+        Dict with matching observations including freshness info
+    """
+    from ..memory_engine import fq_table
+
+    # Use recall to search observations (they come back in results field when fact_type=["observation"])
+    result = await memory_engine.recall_async(
+        bank_id=bank_id,
+        query=query,
+        fact_type=["observation"],  # Only retrieve observations
+        max_tokens=max_tokens,  # Token budget controls how many observations are returned
+        enable_trace=False,
+        request_context=request_context,
+        tags=tags,
+        tags_match=tags_match,
+        _connection_budget=1,
+        _quiet=True,
+    )
+
+    observations = []
+
+    # When fact_type=["observation"], results come back in `results` field as MemoryFact objects
+    # We need to fetch additional fields (proof_count, source_memory_ids) from the database
+    if result.results:
+        obs_ids = [m.id for m in result.results]
+
+        # Fetch proof_count and source_memory_ids for these observations
+        pool = await memory_engine._get_pool()
+        async with pool.acquire() as conn:
+            obs_rows = await conn.fetch(
+                f"""
+                SELECT id, proof_count, source_memory_ids
+                FROM {fq_table("memory_units")}
+                WHERE id = ANY($1::uuid[])
+                """,
+                obs_ids,
+            )
+            obs_data = {str(row["id"]): row for row in obs_rows}
+
+        for m in result.results:
+            # Get additional data from DB lookup
+            extra = obs_data.get(m.id, {})
+            proof_count = extra.get("proof_count", 1) if extra else 1
+            source_ids = extra.get("source_memory_ids", []) if extra else []
+            # Convert UUIDs to strings
+            source_memory_ids = [str(sid) for sid in (source_ids or [])]
+
+            # Determine staleness
+            is_stale = False
+            staleness_reason = None
+            if pending_consolidation > 0:
+                is_stale = True
+                staleness_reason = f"{pending_consolidation} memories pending consolidation"
+
+            observations.append(
+                {
+                    "id": str(m.id),
+                    "text": m.text,
+                    "proof_count": proof_count,
+                    "source_memory_ids": source_memory_ids,
+                    "tags": m.tags or [],
+                    "is_stale": is_stale,
+                    "staleness_reason": staleness_reason,
+                }
+            )
+
+    # Return freshness info (more understandable than raw pending_consolidation count)
+    if pending_consolidation == 0:
+        freshness = "up_to_date"
+    elif pending_consolidation < 10:
+        freshness = "slightly_stale"
+    else:
+        freshness = "stale"
+
+    return {
+        "query": query,
+        "count": len(observations),
+        "observations": observations,
+        "freshness": freshness,
+    }
+
+
+async def tool_recall(
+    memory_engine: "MemoryEngine",
+    bank_id: str,
+    query: str,
+    request_context: "RequestContext",
+    max_tokens: int = 2048,
+    max_results: int = 50,
+    tags: list[str] | None = None,
+    tags_match: str = "any",
+    connection_budget: int = 1,
+) -> dict[str, Any]:
+    """
+    Search memories using TEMPR retrieval.
+
+    This is the ground truth - raw facts and experiences.
+    Use when mental models/observations don't exist, are stale, or need verification.
+
+    Args:
+        memory_engine: Memory engine instance
+        bank_id: Bank identifier
+        query: Search query
+        request_context: Request context for authentication
+        max_tokens: Maximum tokens for results (default 2048)
+        max_results: Maximum number of results
+        tags: Filter by tags (includes untagged memories)
+        tags_match: How to match tags - "any" (OR), "all" (AND), or "exact"
+        connection_budget: Max DB connections for this recall (default 1 for internal ops)
+
+    Returns:
+        Dict with list of matching memories
+    """
+    result = await memory_engine.recall_async(
+        bank_id=bank_id,
+        query=query,
+        fact_type=["experience", "world"],  # Exclude opinions and observations
+        max_tokens=max_tokens,
+        enable_trace=False,
+        request_context=request_context,
+        tags=tags,
+        tags_match=tags_match,
+        _connection_budget=connection_budget,
+        _quiet=True,  # Suppress logging for internal operations
+    )
+
+    memories = []
+    for m in result.results[:max_results]:
+        memories.append(
+            {
+                "id": str(m.id),
+                "text": m.text,
+                "type": m.fact_type,
+                "entities": m.entities or [],
+                "occurred": m.occurred_start,  # Already ISO format string
+            }
+        )
+
+    return {
+        "query": query,
+        "count": len(memories),
+        "memories": memories,
+    }
+
+
+async def tool_expand(
+    conn: "Connection",
+    bank_id: str,
+    memory_ids: list[str],
+    depth: str,
+) -> dict[str, Any]:
+    """
+    Expand multiple memories to get chunk or document context.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        memory_ids: List of memory unit IDs
+        depth: "chunk" or "document"
+
+    Returns:
+        Dict with results array, each containing memory, chunk, and optionally document data
+    """
+    from ..memory_engine import fq_table
+
+    if not memory_ids:
+        return {"error": "memory_ids is required and must not be empty"}
+
+    # Validate and convert UUIDs
+    valid_uuids: list[uuid.UUID] = []
+    errors: dict[str, str] = {}
+    for mid in memory_ids:
+        try:
+            valid_uuids.append(uuid.UUID(mid))
+        except ValueError:
+            errors[mid] = f"Invalid memory_id format: {mid}"
+
+    if not valid_uuids:
+        return {"error": "No valid memory IDs provided", "details": errors}
+
+    # Batch fetch all memory units
+    memories = await conn.fetch(
+        f"""
+        SELECT id, text, chunk_id, document_id, fact_type, context
+        FROM {fq_table("memory_units")}
+        WHERE id = ANY($1) AND bank_id = $2
+        """,
+        valid_uuids,
+        bank_id,
+    )
+    memory_map = {row["id"]: row for row in memories}
+
+    # Collect chunk_ids and document_ids for batch fetching
+    chunk_ids = [m["chunk_id"] for m in memories if m["chunk_id"]]
+    doc_ids_from_chunks: set[str] = set()
+    doc_ids_direct: set[str] = set()
+
+    # Batch fetch all chunks
+    chunk_map: dict[str, Any] = {}
+    if chunk_ids:
+        chunks = await conn.fetch(
+            f"""
+            SELECT chunk_id, chunk_text, chunk_index, document_id
+            FROM {fq_table("chunks")}
+            WHERE chunk_id = ANY($1)
+            """,
+            chunk_ids,
+        )
+        chunk_map = {row["chunk_id"]: row for row in chunks}
+        if depth == "document":
+            doc_ids_from_chunks = {c["document_id"] for c in chunks if c["document_id"]}
+
+    # Collect direct document IDs (memories without chunks)
+    if depth == "document":
+        for m in memories:
+            if not m["chunk_id"] and m["document_id"]:
+                doc_ids_direct.add(m["document_id"])
+
+    # Batch fetch all documents
+    doc_map: dict[str, Any] = {}
+    all_doc_ids = list(doc_ids_from_chunks | doc_ids_direct)
+    if all_doc_ids:
+        docs = await conn.fetch(
+            f"""
+            SELECT id, original_text, metadata, retain_params
+            FROM {fq_table("documents")}
+            WHERE id = ANY($1) AND bank_id = $2
+            """,
+            all_doc_ids,
+            bank_id,
+        )
+        doc_map = {row["id"]: row for row in docs}
+
+    # Build results
+    results: list[dict[str, Any]] = []
+    for mid, mem_uuid in zip(memory_ids, valid_uuids):
+        if mid in errors:
+            results.append({"memory_id": mid, "error": errors[mid]})
+            continue
+
+        memory = memory_map.get(mem_uuid)
+        if not memory:
+            results.append({"memory_id": mid, "error": f"Memory not found: {mid}"})
+            continue
+
+        item: dict[str, Any] = {
+            "memory_id": mid,
+            "memory": {
+                "id": str(memory["id"]),
+                "text": memory["text"],
+                "type": memory["fact_type"],
+                "context": memory["context"],
+            },
+        }
+
+        # Add chunk if available
+        if memory["chunk_id"] and memory["chunk_id"] in chunk_map:
+            chunk = chunk_map[memory["chunk_id"]]
+            item["chunk"] = {
+                "id": chunk["chunk_id"],
+                "text": chunk["chunk_text"],
+                "index": chunk["chunk_index"],
+                "document_id": chunk["document_id"],
+            }
+            # Add document if depth=document
+            if depth == "document" and chunk["document_id"] in doc_map:
+                doc = doc_map[chunk["document_id"]]
+                item["document"] = {
+                    "id": doc["id"],
+                    "full_text": doc["original_text"],
+                    "metadata": doc["metadata"],
+                    "retain_params": doc["retain_params"],
+                }
+        elif memory["document_id"] and depth == "document" and memory["document_id"] in doc_map:
+            # No chunk, but has document_id
+            doc = doc_map[memory["document_id"]]
+            item["document"] = {
+                "id": doc["id"],
+                "full_text": doc["original_text"],
+                "metadata": doc["metadata"],
+                "retain_params": doc["retain_params"],
+            }
+
+        results.append(item)
+
+    return {"results": results, "count": len(results)}

hindsight_api/engine/reflect/tools_schema.py

@@ -0,0 +1,250 @@
+"""
+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."
+        ),
+        "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.",
+                },
+            },
+            "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 plain text. Do NOT use markdown formatting. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array.",
+                },
+                "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 plain text. Do NOT use markdown formatting. NEVER include memory IDs, UUIDs, or 'Memory references' in this text - put IDs only in memory_ids array.",
+                    },
+                    "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