""" Memory Engine for Memory Banks. This implements a sophisticated memory architecture that combines: 1. Temporal links: Memories connected by time proximity 2. Semantic links: Memories connected by meaning/similarity 3. Entity links: Memories connected by shared entities (PERSON, ORG, etc.) 4. Spreading activation: Search through the graph with activation decay 5. Dynamic weighting: Recency and frequency-based importance """ import asyncio import contextvars import json import logging import time import uuid from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any from ..config import get_config from ..metrics import get_metrics_collector from ..tracing import create_operation_span from ..utils import mask_network_location from .db_budget import budgeted_operation # Context variable for current schema (async-safe, per-task isolation) # Note: default is None, actual default comes from config via get_current_schema() _current_schema: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_schema", default=None) def get_current_schema() -> str: """Get the current schema from context (falls back to config default).""" schema = _current_schema.get() if schema is None: # Fall back to configured default schema return get_config().database_schema return schema def fq_table(table_name: str) -> str: """ Get fully-qualified table name with current schema. Example: fq_table("memory_units") -> "public.memory_units" fq_table("memory_units") -> "tenant_xyz.memory_units" (if schema is set) """ return f"{get_current_schema()}.{table_name}" # Tables that must be schema-qualified (for runtime validation) _PROTECTED_TABLES = frozenset( [ "memory_units", "memory_links", "unit_entities", "entities", "entity_cooccurrences", "banks", "documents", "chunks", "async_operations", ] ) # Enable runtime SQL validation (can be disabled in production for performance) _VALIDATE_SQL_SCHEMAS = True class UnqualifiedTableError(Exception): """Raised when SQL contains unqualified table references.""" pass def validate_sql_schema(sql: str) -> None: """ Validate that SQL doesn't contain unqualified table references. This is a runtime safety check to prevent cross-tenant data access. Raises UnqualifiedTableError if any protected table is referenced without a schema prefix. Args: sql: The SQL query to validate Raises: UnqualifiedTableError: If unqualified table reference found """ if not _VALIDATE_SQL_SCHEMAS: return import re sql_upper = sql.upper() for table in _PROTECTED_TABLES: table_upper = table.upper() # Pattern: SQL keyword followed by unqualified table name # Matches: FROM memory_units, JOIN memory_units, INTO memory_units, UPDATE memory_units patterns = [ rf"FROM\s+{table_upper}(?:\s|$|,|\)|;)", rf"JOIN\s+{table_upper}(?:\s|$|,|\)|;)", rf"INTO\s+{table_upper}(?:\s|$|\()", rf"UPDATE\s+{table_upper}(?:\s|$)", rf"DELETE\s+FROM\s+{table_upper}(?:\s|$|;)", ] for pattern in patterns: match = re.search(pattern, sql_upper) if match: # Check if it's actually qualified (preceded by schema.) # Look backwards from match to see if there's a dot start = match.start() # Find the table name position in the match table_pos = sql_upper.find(table_upper, start) if table_pos > 0: # Check character before table name (skip whitespace) prefix = sql[:table_pos].rstrip() if not prefix.endswith("."): raise UnqualifiedTableError( f"Unqualified table reference '{table}' in SQL. " f"Use fq_table('{table}') for schema safety. " f"SQL snippet: ...{sql[max(0, start - 10) : start + 50]}..." ) import asyncpg import numpy as np from pydantic import BaseModel, Field from .cross_encoder import CrossEncoderModel from .embeddings import Embeddings, create_embeddings_from_env from .interface import MemoryEngineInterface if TYPE_CHECKING: from hindsight_api.extensions import OperationValidatorExtension, TenantExtension from hindsight_api.models import RequestContext from enum import Enum from ..metrics import get_metrics_collector from ..pg0 import EmbeddedPostgres, parse_pg0_url from .entity_resolver import EntityResolver from .llm_wrapper import LLMConfig from .query_analyzer import QueryAnalyzer from .reflect import run_reflect_agent from .reflect.tools import tool_expand, tool_recall, tool_search_mental_models, tool_search_observations from .response_models import ( VALID_RECALL_FACT_TYPES, EntityObservation, EntityState, LLMCallTrace, MemoryFact, ObservationRef, ReflectResult, TokenUsage, ToolCallTrace, ) from .response_models import RecallResult as RecallResultModel from .retain import bank_utils, embedding_utils from .retain.types import RetainContentDict from .search import think_utils from .search.reranking import CrossEncoderReranker from .search.tags import TagsMatch from .task_backend import BrokerTaskBackend, SyncTaskBackend, TaskBackend class Budget(str, Enum): """Budget levels for recall/reflect operations.""" LOW = "low" MID = "mid" HIGH = "high" def utcnow(): """Get current UTC time with timezone info.""" return datetime.now(UTC) # Logger for memory system logger = logging.getLogger(__name__) import tiktoken from .db_utils import acquire_with_retry # Cache tiktoken encoding for token budget filtering (module-level singleton) _TIKTOKEN_ENCODING = None def _get_tiktoken_encoding(): """Get cached tiktoken encoding (cl100k_base for GPT-4/3.5).""" global _TIKTOKEN_ENCODING if _TIKTOKEN_ENCODING is None: _TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base") return _TIKTOKEN_ENCODING class MemoryEngine(MemoryEngineInterface): """ Advanced memory system using temporal and semantic linking with PostgreSQL. This class provides: - Embedding generation for semantic search - Entity, temporal, and semantic link creation - Think operations for formulating answers with opinions - bank profile and disposition management """ def __init__( self, db_url: str | None = None, memory_llm_provider: str | None = None, memory_llm_api_key: str | None = None, memory_llm_model: str | None = None, memory_llm_base_url: str | None = None, # Per-operation LLM config (optional, falls back to memory_llm_* params) retain_llm_provider: str | None = None, retain_llm_api_key: str | None = None, retain_llm_model: str | None = None, retain_llm_base_url: str | None = None, reflect_llm_provider: str | None = None, reflect_llm_api_key: str | None = None, reflect_llm_model: str | None = None, reflect_llm_base_url: str | None = None, consolidation_llm_provider: str | None = None, consolidation_llm_api_key: str | None = None, consolidation_llm_model: str | None = None, consolidation_llm_base_url: str | None = None, embeddings: Embeddings | None = None, cross_encoder: CrossEncoderModel | None = None, query_analyzer: QueryAnalyzer | None = None, pool_min_size: int | None = None, pool_max_size: int | None = None, db_command_timeout: int | None = None, db_acquire_timeout: int | None = None, task_backend: TaskBackend | None = None, run_migrations: bool = True, operation_validator: "OperationValidatorExtension | None" = None, tenant_extension: "TenantExtension | None" = None, skip_llm_verification: bool | None = None, lazy_reranker: bool | None = None, ): """ Initialize the temporal + semantic memory system. All parameters are optional and will be read from environment variables if not provided. See hindsight_api.config for environment variable names and defaults. Args: db_url: PostgreSQL connection URL. Defaults to HINDSIGHT_API_DATABASE_URL env var or "pg0". Also supports pg0 URLs: "pg0" or "pg0://instance-name" or "pg0://instance-name:port" memory_llm_provider: LLM provider. Defaults to HINDSIGHT_API_LLM_PROVIDER env var or "groq". memory_llm_api_key: API key for the LLM provider. Defaults to HINDSIGHT_API_LLM_API_KEY env var. memory_llm_model: Model name. Defaults to HINDSIGHT_API_LLM_MODEL env var. memory_llm_base_url: Base URL for the LLM API. Defaults based on provider. retain_llm_provider: LLM provider for retain operations. Falls back to memory_llm_provider. retain_llm_api_key: API key for retain LLM. Falls back to memory_llm_api_key. retain_llm_model: Model for retain operations. Falls back to memory_llm_model. retain_llm_base_url: Base URL for retain LLM. Falls back to memory_llm_base_url. reflect_llm_provider: LLM provider for reflect operations. Falls back to memory_llm_provider. reflect_llm_api_key: API key for reflect LLM. Falls back to memory_llm_api_key. reflect_llm_model: Model for reflect operations. Falls back to memory_llm_model. reflect_llm_base_url: Base URL for reflect LLM. Falls back to memory_llm_base_url. consolidation_llm_provider: LLM provider for consolidation operations. Falls back to memory_llm_provider. consolidation_llm_api_key: API key for consolidation LLM. Falls back to memory_llm_api_key. consolidation_llm_model: Model for consolidation operations. Falls back to memory_llm_model. consolidation_llm_base_url: Base URL for consolidation LLM. Falls back to memory_llm_base_url. embeddings: Embeddings implementation. If not provided, created from env vars. cross_encoder: Cross-encoder model. If not provided, created from env vars. query_analyzer: Query analyzer implementation. If not provided, uses DateparserQueryAnalyzer. pool_min_size: Minimum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MIN_SIZE. pool_max_size: Maximum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MAX_SIZE. db_command_timeout: PostgreSQL command timeout in seconds. Defaults to HINDSIGHT_API_DB_COMMAND_TIMEOUT. db_acquire_timeout: Connection acquisition timeout in seconds. Defaults to HINDSIGHT_API_DB_ACQUIRE_TIMEOUT. task_backend: Custom task backend. If not provided, uses BrokerTaskBackend for distributed processing. run_migrations: Whether to run database migrations during initialize(). Default: True operation_validator: Optional extension to validate operations before execution. If provided, retain/recall/reflect operations will be validated. tenant_extension: Optional extension for multi-tenancy and API key authentication. If provided, operations require a RequestContext for authentication. skip_llm_verification: Skip LLM connection verification during initialization. Defaults to HINDSIGHT_API_SKIP_LLM_VERIFICATION env var or False. lazy_reranker: Delay reranker initialization until first use. Useful for retain-only operations that don't need the cross-encoder. Defaults to HINDSIGHT_API_LAZY_RERANKER env var or False. """ # Load config from environment for any missing parameters from ..config import get_config config = get_config() # Apply optimization flags from config if not explicitly provided self._skip_llm_verification = ( skip_llm_verification if skip_llm_verification is not None else config.skip_llm_verification ) self._lazy_reranker = lazy_reranker if lazy_reranker is not None else config.lazy_reranker # Apply defaults from config db_url = db_url or config.database_url memory_llm_provider = memory_llm_provider or config.llm_provider memory_llm_api_key = memory_llm_api_key or config.llm_api_key # Ollama, openai-codex, claude-code, and mock don't require an API key # openai-codex uses OAuth tokens from ~/.codex/auth.json # claude-code uses OAuth tokens from macOS Keychain if not memory_llm_api_key and memory_llm_provider not in ("ollama", "openai-codex", "claude-code", "mock"): raise ValueError("LLM API key is required. Set HINDSIGHT_API_LLM_API_KEY environment variable.") memory_llm_model = memory_llm_model or config.llm_model memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None # Track pg0 instance (if used) self._pg0: EmbeddedPostgres | None = None # Initialize PostgreSQL connection URL # The actual URL will be set during initialize() after starting the server # Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url) if self._use_pg0: self.db_url = None else: self.db_url = db_url # Set default base URL if not provided if memory_llm_base_url is None: if memory_llm_provider.lower() == "groq": memory_llm_base_url = "https://api.groq.com/openai/v1" elif memory_llm_provider.lower() == "ollama": memory_llm_base_url = "http://localhost:11434/v1" else: memory_llm_base_url = "" # Connection pool (will be created in initialize()) self._pool = None self._initialized = False self._pool_min_size = pool_min_size if pool_min_size is not None else config.db_pool_min_size self._pool_max_size = pool_max_size if pool_max_size is not None else config.db_pool_max_size self._db_command_timeout = db_command_timeout if db_command_timeout is not None else config.db_command_timeout self._db_acquire_timeout = db_acquire_timeout if db_acquire_timeout is not None else config.db_acquire_timeout self._run_migrations = run_migrations # Initialize entity resolver (will be created in initialize()) self.entity_resolver = None # Initialize embeddings (from env vars if not provided) if embeddings is not None: self.embeddings = embeddings else: self.embeddings = create_embeddings_from_env() # Initialize query analyzer if query_analyzer is not None: self.query_analyzer = query_analyzer else: from .query_analyzer import DateparserQueryAnalyzer self.query_analyzer = DateparserQueryAnalyzer() # Initialize LLM configuration (default, used as fallback) self._llm_config = LLMConfig( provider=memory_llm_provider, api_key=memory_llm_api_key, base_url=memory_llm_base_url, model=memory_llm_model, ) # Store client and model for convenience (deprecated: use _llm_config.call() instead) self._llm_client = self._llm_config._client self._llm_model = self._llm_config.model # Initialize per-operation LLM configs (fall back to default if not specified) # Retain LLM config - for fact extraction (benefits from strong structured output) retain_provider = retain_llm_provider or config.retain_llm_provider or memory_llm_provider retain_api_key = retain_llm_api_key or config.retain_llm_api_key or memory_llm_api_key retain_model = retain_llm_model or config.retain_llm_model or memory_llm_model retain_base_url = retain_llm_base_url or config.retain_llm_base_url or memory_llm_base_url # Apply provider-specific base URL defaults for retain if retain_base_url is None: if retain_provider.lower() == "groq": retain_base_url = "https://api.groq.com/openai/v1" elif retain_provider.lower() == "ollama": retain_base_url = "http://localhost:11434/v1" else: retain_base_url = "" self._retain_llm_config = LLMConfig( provider=retain_provider, api_key=retain_api_key, base_url=retain_base_url, model=retain_model, ) # Reflect LLM config - for think/observe operations (can use lighter models) reflect_provider = reflect_llm_provider or config.reflect_llm_provider or memory_llm_provider reflect_api_key = reflect_llm_api_key or config.reflect_llm_api_key or memory_llm_api_key reflect_model = reflect_llm_model or config.reflect_llm_model or memory_llm_model reflect_base_url = reflect_llm_base_url or config.reflect_llm_base_url or memory_llm_base_url # Apply provider-specific base URL defaults for reflect if reflect_base_url is None: if reflect_provider.lower() == "groq": reflect_base_url = "https://api.groq.com/openai/v1" elif reflect_provider.lower() == "ollama": reflect_base_url = "http://localhost:11434/v1" else: reflect_base_url = "" self._reflect_llm_config = LLMConfig( provider=reflect_provider, api_key=reflect_api_key, base_url=reflect_base_url, model=reflect_model, ) # Consolidation LLM config - for mental model consolidation (can use efficient models) consolidation_provider = consolidation_llm_provider or config.consolidation_llm_provider or memory_llm_provider consolidation_api_key = consolidation_llm_api_key or config.consolidation_llm_api_key or memory_llm_api_key consolidation_model = consolidation_llm_model or config.consolidation_llm_model or memory_llm_model consolidation_base_url = consolidation_llm_base_url or config.consolidation_llm_base_url or memory_llm_base_url # Apply provider-specific base URL defaults for consolidation if consolidation_base_url is None: if consolidation_provider.lower() == "groq": consolidation_base_url = "https://api.groq.com/openai/v1" elif consolidation_provider.lower() == "ollama": consolidation_base_url = "http://localhost:11434/v1" else: consolidation_base_url = "" self._consolidation_llm_config = LLMConfig( provider=consolidation_provider, api_key=consolidation_api_key, base_url=consolidation_base_url, model=consolidation_model, ) # Initialize cross-encoder reranker (cached for performance) self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder) # Initialize task backend # If no custom backend provided, use BrokerTaskBackend which stores tasks in PostgreSQL # The pool_getter lambda will return the pool once it's initialized self._task_backend = task_backend or BrokerTaskBackend( pool_getter=lambda: self._pool, schema_getter=get_current_schema, ) # Backpressure mechanism: limit concurrent searches to prevent overwhelming the database # Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50) self._search_semaphore = asyncio.Semaphore(get_config().recall_max_concurrent) # Backpressure for put operations: limit concurrent puts to prevent database contention # Each put_batch holds a connection for the entire transaction, so we limit to 5 # concurrent puts to avoid connection pool exhaustion and reduce write contention self._put_semaphore = asyncio.Semaphore(5) # initialize encoding eagerly to avoid delaying the first time _get_tiktoken_encoding() # Store operation validator extension (optional) self._operation_validator = operation_validator # Store tenant extension (always set, use default if none provided) if tenant_extension is None: from ..extensions.builtin.tenant import DefaultTenantExtension tenant_extension = DefaultTenantExtension(config={}) self._tenant_extension = tenant_extension @property def tenant_extension(self) -> "TenantExtension | None": """The configured tenant extension, if any.""" return self._tenant_extension async def _validate_operation(self, validation_coro) -> None: """ Run validation if an operation validator is configured. Args: validation_coro: Coroutine that returns a ValidationResult Raises: OperationValidationError: If validation fails """ if self._operation_validator is None: return from hindsight_api.extensions import OperationValidationError result = await validation_coro if not result.allowed: raise OperationValidationError(result.reason or "Operation not allowed", result.status_code) async def _authenticate_tenant(self, request_context: "RequestContext | None") -> str: """ Authenticate tenant and set schema in context variable. The schema is stored in a contextvar for async-safe, per-task isolation. Use fq_table(table_name) to get fully-qualified table names. Args: request_context: The request context with API key. Required if tenant_extension is configured. Returns: Schema name that was set in the context. Raises: AuthenticationError: If authentication fails or request_context is missing when required. """ from hindsight_api.extensions import AuthenticationError if request_context is None: raise AuthenticationError("RequestContext is required") # For internal/background operations (e.g., worker tasks), skip extension authentication. # The task was already authenticated at submission time, and execute_task sets _current_schema # from the task's _schema field. if request_context.internal: return _current_schema.get() # Authenticate through tenant extension (always set, may be default no-auth extension) tenant_context = await self._tenant_extension.authenticate(request_context) _current_schema.set(tenant_context.schema_name) return tenant_context.schema_name async def _handle_batch_retain(self, task_dict: dict[str, Any]): """ Handler for batch retain tasks. Args: task_dict: Dict with 'bank_id', 'contents' Raises: ValueError: If bank_id is missing Exception: Any exception from retain_batch_async (propagates to execute_task for retry) """ bank_id = task_dict.get("bank_id") if not bank_id: raise ValueError("bank_id is required for batch retain task") contents = task_dict.get("contents", []) logger.info( f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items" ) # Restore tenant_id/api_key_id from task payload so extensions # (e.g., operation validators) can attribute the operation correctly. # internal=True to skip extension auth (worker has no API key), # user_initiated=True so extensions know this originated from a user request. from hindsight_api.models import RequestContext context = RequestContext( internal=True, user_initiated=True, tenant_id=task_dict.get("_tenant_id"), api_key_id=task_dict.get("_api_key_id"), ) await self.retain_batch_async(bank_id=bank_id, contents=contents, request_context=context) logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}") async def _handle_consolidation(self, task_dict: dict[str, Any]): """ Handler for consolidation tasks. Consolidates new memories into mental models for a bank. Args: task_dict: Dict with 'bank_id' Raises: ValueError: If bank_id is missing Exception: Any exception from consolidation (propagates to execute_task for retry) """ bank_id = task_dict.get("bank_id") if not bank_id: raise ValueError("bank_id is required for consolidation task") from hindsight_api.models import RequestContext from .consolidation import run_consolidation_job # Restore tenant_id/api_key_id from task payload so downstream operations # (e.g., mental model refreshes) can attribute usage to the correct org. internal_context = RequestContext( internal=True, tenant_id=task_dict.get("_tenant_id"), api_key_id=task_dict.get("_api_key_id"), ) result = await run_consolidation_job( memory_engine=self, bank_id=bank_id, request_context=internal_context, ) logger.info(f"[CONSOLIDATION] bank={bank_id} completed: {result.get('memories_processed', 0)} processed") async def _handle_refresh_mental_model(self, task_dict: dict[str, Any]): """ Handler for refresh_mental_model tasks. Re-runs the source query through reflect and updates the mental model content. Args: task_dict: Dict with 'bank_id', 'mental_model_id', 'operation_id' Raises: ValueError: If required fields are missing Exception: Any exception from reflect/update (propagates to execute_task for retry) """ bank_id = task_dict.get("bank_id") mental_model_id = task_dict.get("mental_model_id") if not bank_id or not mental_model_id: raise ValueError("bank_id and mental_model_id are required for refresh_mental_model task") logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Starting for bank_id={bank_id}, mental_model_id={mental_model_id}") from hindsight_api.models import RequestContext # Restore tenant_id/api_key_id from task payload so extensions can # attribute the mental_model_refresh operation to the correct org. internal_context = RequestContext( internal=True, tenant_id=task_dict.get("_tenant_id"), api_key_id=task_dict.get("_api_key_id"), ) # Get the current mental model to get source_query mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=internal_context) if not mental_model: raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}") source_query = mental_model["source_query"] # SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching # to ensure it can only access other mental models/memories with the SAME tags. # This prevents cross-tenant/cross-user information leakage by excluding untagged content. tags = mental_model.get("tags") tags_match = "all_strict" if tags else "any" # Run reflect to generate new content, excluding the mental model being refreshed reflect_result = await self.reflect_async( bank_id=bank_id, query=source_query, request_context=internal_context, tags=tags, tags_match=tags_match, exclude_mental_model_ids=[mental_model_id], ) generated_content = reflect_result.text or "No content generated" # Build reflect_response payload to store # based_on contains MemoryFact objects for most types, but plain dicts for directives based_on_serialized: dict[str, list[dict[str, Any]]] = {} for fact_type, facts in reflect_result.based_on.items(): serialized_facts = [] for fact in facts: if isinstance(fact, dict): # Plain dict (e.g., directives with id, name, content) serialized_facts.append( { "id": str(fact["id"]), "text": fact.get("text", fact.get("content", fact.get("name", ""))), "type": fact_type, } ) else: # MemoryFact object with .id and .text attributes serialized_facts.append( { "id": str(fact.id), "text": fact.text, "type": fact_type, } ) based_on_serialized[fact_type] = serialized_facts reflect_response = { "text": reflect_result.text, "based_on": based_on_serialized, } # Update the mental model with the generated content and reflect_response await self.update_mental_model( bank_id=bank_id, mental_model_id=mental_model_id, content=generated_content, reflect_response=reflect_response, request_context=internal_context, ) # Call post-operation hook if validator is configured if self._operation_validator: from hindsight_api.extensions.operation_validator import MentalModelRefreshResult # Count facts and mental models from based_on facts_used = 0 mental_models_used = 0 if reflect_result.based_on: for fact_type, facts in reflect_result.based_on.items(): if facts: if fact_type == "mental_models": mental_models_used += len(facts) else: facts_used += len(facts) # Estimate tokens query_tokens = len(source_query) // 4 if source_query else 0 output_tokens = len(generated_content) // 4 if generated_content else 0 context_tokens = 0 # refresh doesn't use additional context result_ctx = MentalModelRefreshResult( bank_id=bank_id, mental_model_id=mental_model_id, request_context=internal_context, query_tokens=query_tokens, output_tokens=output_tokens, context_tokens=context_tokens, facts_used=facts_used, mental_models_used=mental_models_used, success=True, ) try: await self._operation_validator.on_mental_model_refresh_complete(result_ctx) except Exception as hook_err: logger.warning(f"Post-mental-model-refresh hook error (non-fatal): {hook_err}") logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Completed for bank_id={bank_id}, mental_model_id={mental_model_id}") async def execute_task(self, task_dict: dict[str, Any]): """ Execute a task by routing it to the appropriate handler. This method is called by the task backend to execute tasks. It receives a plain dict that can be serialized and sent over the network. Args: task_dict: Task dictionary with 'type' key and other payload data Example: {'type': 'batch_retain', 'bank_id': '...', 'contents': [...]} """ task_type = task_dict.get("type") operation_id = task_dict.get("operation_id") retry_count = task_dict.get("retry_count", 0) max_retries = 3 # Set schema context for multi-tenant task execution schema = task_dict.pop("_schema", None) if schema: _current_schema.set(schema) # Check if operation was cancelled (only for tasks with operation_id) if operation_id: try: pool = await self._get_pool() async with acquire_with_retry(pool) as conn: result = await conn.fetchrow( f"SELECT operation_id FROM {fq_table('async_operations')} WHERE operation_id = $1", uuid.UUID(operation_id), ) if not result: # Operation was cancelled, skip processing logger.info(f"Skipping cancelled operation: {operation_id}") return except Exception as e: logger.error(f"Failed to check operation status {operation_id}: {e}") # Continue with processing if we can't check status try: if task_type == "batch_retain": await self._handle_batch_retain(task_dict) elif task_type == "consolidation": await self._handle_consolidation(task_dict) elif task_type == "refresh_mental_model": await self._handle_refresh_mental_model(task_dict) else: logger.error(f"Unknown task type: {task_type}") # Don't retry unknown task types if operation_id: await self._delete_operation_record(operation_id) return # Task succeeded - mark operation as completed if operation_id: await self._mark_operation_completed(operation_id) except Exception as e: # Task failed - check if we should retry logger.error( f"Task execution failed (attempt {retry_count + 1}/{max_retries + 1}): {task_type}, error: {e}" ) import traceback error_traceback = traceback.format_exc() traceback.print_exc() if retry_count < max_retries: # Reschedule with incremented retry count task_dict["retry_count"] = retry_count + 1 logger.info(f"Rescheduling task {task_type} (retry {retry_count + 1}/{max_retries})") await self._task_backend.submit_task(task_dict) else: # Max retries exceeded - mark operation as failed logger.error(f"Max retries exceeded for task {task_type}, marking as failed") if operation_id: await self._mark_operation_failed(operation_id, str(e), error_traceback) async def _delete_operation_record(self, operation_id: str): """Helper to delete an operation record from the database.""" try: pool = await self._get_pool() async with acquire_with_retry(pool) as conn: await conn.execute( f"DELETE FROM {fq_table('async_operations')} WHERE operation_id = $1", uuid.UUID(operation_id) ) except Exception as e: logger.error(f"Failed to delete async operation record {operation_id}: {e}") async def _mark_operation_failed(self, operation_id: str, error_message: str, error_traceback: str): """Helper to mark an operation as failed in the database.""" try: pool = await self._get_pool() # Truncate error message to avoid extremely long strings full_error = f"{error_message}\n\nTraceback:\n{error_traceback}" truncated_error = full_error[:5000] if len(full_error) > 5000 else full_error async with acquire_with_retry(pool) as conn: await conn.execute( f""" UPDATE {fq_table("async_operations")} SET status = 'failed', error_message = $2, updated_at = NOW() WHERE operation_id = $1 """, uuid.UUID(operation_id), truncated_error, ) logger.info(f"Marked async operation as failed: {operation_id}") except Exception as e: logger.error(f"Failed to mark operation as failed {operation_id}: {e}") async def _mark_operation_completed(self, operation_id: str): """Helper to mark an operation as completed in the database.""" try: pool = await self._get_pool() async with acquire_with_retry(pool) as conn: await conn.execute( f""" UPDATE {fq_table("async_operations")} SET status = 'completed', updated_at = NOW(), completed_at = NOW() WHERE operation_id = $1 """, uuid.UUID(operation_id), ) logger.info(f"Marked async operation as completed: {operation_id}") except Exception as e: logger.error(f"Failed to mark operation as completed {operation_id}: {e}") async def initialize(self): """Initialize the connection pool, models, and background workers. Loads models (embeddings, cross-encoder) in parallel with pg0 startup for faster overall initialization. """ if self._initialized: return # Run model loading in thread pool (CPU-bound) in parallel with pg0 startup loop = asyncio.get_event_loop() async def start_pg0(): """Start pg0 if configured.""" if self._use_pg0: kwargs = {"name": self._pg0_instance_name} if self._pg0_port is not None: kwargs["port"] = self._pg0_port pg0 = EmbeddedPostgres(**kwargs) # Check if pg0 is already running before we start it was_already_running = await pg0.is_running() self.db_url = await pg0.ensure_running() # Only track pg0 (to stop later) if WE started it if not was_already_running: self._pg0 = pg0 async def init_embeddings(): """Initialize embedding model.""" # For local providers, run in thread pool to avoid blocking event loop if self.embeddings.provider_name == "local": await loop.run_in_executor(None, lambda: asyncio.run(self.embeddings.initialize())) else: await self.embeddings.initialize() async def init_cross_encoder(): """Initialize cross-encoder model.""" cross_encoder = self._cross_encoder_reranker.cross_encoder # For local providers, run in thread pool to avoid blocking event loop if cross_encoder.provider_name == "local": await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize())) else: await cross_encoder.initialize() # Mark reranker as initialized self._cross_encoder_reranker._initialized = True async def init_query_analyzer(): """Initialize query analyzer model.""" # Query analyzer load is sync and CPU-bound await loop.run_in_executor(None, self.query_analyzer.load) async def verify_llm(): """Verify LLM connections are working for all unique configs.""" if not self._skip_llm_verification: # Verify default config await self._llm_config.verify_connection() # Verify retain config if different from default retain_is_different = ( self._retain_llm_config.provider != self._llm_config.provider or self._retain_llm_config.model != self._llm_config.model ) if retain_is_different: await self._retain_llm_config.verify_connection() # Verify reflect config if different from default and retain reflect_is_different = ( self._reflect_llm_config.provider != self._llm_config.provider or self._reflect_llm_config.model != self._llm_config.model ) and ( self._reflect_llm_config.provider != self._retain_llm_config.provider or self._reflect_llm_config.model != self._retain_llm_config.model ) if reflect_is_different: await self._reflect_llm_config.verify_connection() # Verify consolidation config if different from all others consolidation_is_different = ( ( self._consolidation_llm_config.provider != self._llm_config.provider or self._consolidation_llm_config.model != self._llm_config.model ) and ( self._consolidation_llm_config.provider != self._retain_llm_config.provider or self._consolidation_llm_config.model != self._retain_llm_config.model ) and ( self._consolidation_llm_config.provider != self._reflect_llm_config.provider or self._consolidation_llm_config.model != self._reflect_llm_config.model ) ) if consolidation_is_different: await self._consolidation_llm_config.verify_connection() # Build list of initialization tasks init_tasks = [ start_pg0(), init_embeddings(), init_query_analyzer(), ] # Only init cross-encoder eagerly if not using lazy initialization if not self._lazy_reranker: init_tasks.append(init_cross_encoder()) # Only verify LLM if not skipping if not self._skip_llm_verification: init_tasks.append(verify_llm()) # Run pg0 and selected model initializations in parallel await asyncio.gather(*init_tasks) # Run database migrations if enabled if self._run_migrations: from ..migrations import ensure_embedding_dimension, run_migrations if not self.db_url: raise ValueError("Database URL is required for migrations") # Migrate all schemas from the tenant extension # The tenant extension is the single source of truth for which schemas exist logger.info("Running database migrations...") try: tenants = await self._tenant_extension.list_tenants() if tenants: logger.info(f"Running migrations on {len(tenants)} schema(s)...") for tenant in tenants: schema = tenant.schema if schema: try: run_migrations(self.db_url, schema=schema) except Exception as e: logger.warning(f"Failed to migrate schema {schema}: {e}") logger.info("Schema migrations completed") # Ensure embedding column dimension matches the model's dimension # This is done after migrations and after embeddings.initialize() for tenant in tenants: schema = tenant.schema if schema: try: ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=schema) except Exception as e: logger.warning(f"Failed to ensure embedding dimension for schema {schema}: {e}") except Exception as e: logger.warning(f"Failed to run schema migrations: {e}") logger.info(f"Connecting to PostgreSQL at {mask_network_location(self.db_url)}") # Create connection pool # For read-heavy workloads with many parallel think/search operations, # we need a larger pool. Read operations don't need strong isolation. self._pool = await asyncpg.create_pool( self.db_url, min_size=self._pool_min_size, max_size=self._pool_max_size, command_timeout=self._db_command_timeout, statement_cache_size=0, # Disable prepared statement cache timeout=self._db_acquire_timeout, # Connection acquisition timeout (seconds) ) # Initialize entity resolver with pool self.entity_resolver = EntityResolver(self._pool) # Initialize config resolver for hierarchical configuration from ..config_resolver import ConfigResolver self._config_resolver = ConfigResolver(pool=self._pool, tenant_extension=self._tenant_extension) logger.debug("Config resolver initialized for hierarchical configuration") # Set executor for task backend and initialize self._task_backend.set_executor(self.execute_task) await self._task_backend.initialize() self._initialized = True logger.info("Memory system initialized (pool and task backend started)") async def _get_pool(self) -> asyncpg.Pool: """Get the connection pool (must call initialize() first).""" if not self._initialized: await self.initialize() return self._pool async def _acquire_connection(self): """ Acquire a connection from the pool with retry logic. Returns an async context manager that yields a connection. Retries on transient connection errors with exponential backoff. """ pool = await self._get_pool() async def acquire(): return await pool.acquire() return await _retry_with_backoff(acquire) async def health_check(self) -> dict: """ Perform a health check by querying the database. Returns: dict with status and optional error message Note: Returns unhealthy until initialize() has completed successfully. """ # Not healthy until fully initialized if not self._initialized: return {"status": "unhealthy", "reason": "not_initialized"} try: pool = await self._get_pool() async with pool.acquire() as conn: result = await conn.fetchval("SELECT 1") if result == 1: return {"status": "healthy", "database": "connected"} else: return {"status": "unhealthy", "database": "unexpected response"} except Exception as e: return {"status": "unhealthy", "database": "error", "error": str(e)} async def close(self): """Close the connection pool and shutdown background workers.""" logger.info("close() started") # Shutdown task backend await self._task_backend.shutdown() # Close pool if self._pool is not None: self._pool.terminate() self._pool = None self._initialized = False # Stop pg0 if we started it if self._pg0 is not None: logger.info("Stopping pg0...") await self._pg0.stop() self._pg0 = None logger.info("pg0 stopped") async def wait_for_background_tasks(self): """ Wait for all pending background tasks to complete. This is useful in tests to ensure background tasks complete before making assertions. """ if hasattr(self._task_backend, "wait_for_pending_tasks"): await self._task_backend.wait_for_pending_tasks() def _format_readable_date(self, dt: datetime) -> str: """ Format a datetime into a readable string for temporal matching. Examples: - June 2024 - January 15, 2024 - December 2023 This helps queries like "camping in June" match facts that happened in June. Args: dt: datetime object to format Returns: Readable date string """ # Format as "Month Year" for most cases # Could be extended to include day for very specific dates if needed month_name = dt.strftime("%B") # Full month name (e.g., "June") year = dt.strftime("%Y") # Year (e.g., "2024") # For now, use "Month Year" format # Could check if day is significant (not 1st or 15th) and include it return f"{month_name} {year}" async def _find_duplicate_facts_batch( self, conn, bank_id: str, texts: list[str], embeddings: list[list[float]], event_date: datetime, time_window_hours: int = 24, similarity_threshold: float = 0.95, ) -> list[bool]: """ Check which facts are duplicates using semantic similarity + temporal window. For each new fact, checks if a semantically similar fact already exists within the time window. Uses pgvector cosine similarity for efficiency. Args: conn: Database connection bank_id: bank IDentifier texts: List of fact texts to check embeddings: Corresponding embeddings event_date: Event date for temporal filtering time_window_hours: Hours before/after event_date to search (default: 24) similarity_threshold: Minimum cosine similarity to consider duplicate (default: 0.95) Returns: List of booleans - True if fact is a duplicate (should skip), False if new """ if not texts: return [] # Handle edge cases where event_date is at datetime boundaries try: time_lower = event_date - timedelta(hours=time_window_hours) except OverflowError: time_lower = datetime.min try: time_upper = event_date + timedelta(hours=time_window_hours) except OverflowError: time_upper = datetime.max # Fetch ALL existing facts in time window ONCE (much faster than N queries) import time as time_mod fetch_start = time_mod.time() existing_facts = await conn.fetch( f""" SELECT id, text, embedding FROM {fq_table("memory_units")} WHERE bank_id = $1 AND event_date BETWEEN $2 AND $3 """, bank_id, time_lower, time_upper, ) # If no existing facts, nothing is duplicate if not existing_facts: return [False] * len(texts) # Compute similarities in Python (vectorized with numpy) is_duplicate = [] # Convert existing embeddings to numpy for faster computation embedding_arrays = [] for row in existing_facts: raw_emb = row["embedding"] # Handle different pgvector formats if isinstance(raw_emb, str): # Parse string format: "[1.0, 2.0, ...]" import json emb = np.array(json.loads(raw_emb), dtype=np.float32) elif isinstance(raw_emb, (list, tuple)): emb = np.array(raw_emb, dtype=np.float32) else: # Try direct conversion emb = np.array(raw_emb, dtype=np.float32) embedding_arrays.append(emb) if not embedding_arrays: existing_embeddings = np.array([]) elif len(embedding_arrays) == 1: # Single embedding: reshape to (1, dim) existing_embeddings = embedding_arrays[0].reshape(1, -1) else: # Multiple embeddings: vstack existing_embeddings = np.vstack(embedding_arrays) comp_start = time_mod.time() for embedding in embeddings: # Compute cosine similarity with all existing facts emb_array = np.array(embedding) # Cosine similarity = 1 - cosine distance # For normalized vectors: cosine_sim = dot product similarities = np.dot(existing_embeddings, emb_array) # Check if any existing fact is too similar max_similarity = np.max(similarities) if len(similarities) > 0 else 0 is_duplicate.append(max_similarity > similarity_threshold) return is_duplicate def retain( self, bank_id: str, content: str, context: str = "", event_date: datetime | None = None, request_context: "RequestContext | None" = None, ) -> list[str]: """ Store content as memory units (synchronous wrapper). This is a synchronous wrapper around retain_async() for convenience. For best performance, use retain_async() directly. Args: bank_id: Unique identifier for the bank content: Text content to store context: Context about when/why this memory was formed event_date: When the event occurred (defaults to now) request_context: Request context for authentication (optional, uses internal context if not provided) Returns: List of created unit IDs """ # Run async version synchronously from hindsight_api.models import RequestContext as RC ctx = request_context if request_context is not None else RC() return asyncio.run(self.retain_async(bank_id, content, context, event_date, request_context=ctx)) async def retain_async( self, bank_id: str, content: str, context: str = "", event_date: datetime | None = None, document_id: str | None = None, fact_type_override: str | None = None, confidence_score: float | None = None, *, request_context: "RequestContext", ) -> list[str]: """ Store content as memory units with temporal and semantic links (ASYNC version). This is a convenience wrapper around retain_batch_async for a single content item. Args: bank_id: Unique identifier for the bank content: Text content to store context: Context about when/why this memory was formed event_date: When the event occurred (defaults to now) document_id: Optional document ID for tracking (always upserts if document already exists) fact_type_override: Override fact type ('world', 'experience') confidence_score: Confidence score (0.0 to 1.0) request_context: Request context for authentication. Returns: List of created unit IDs """ # Build content dict content_dict: RetainContentDict = {"content": content, "context": context} if event_date: content_dict["event_date"] = event_date if document_id: content_dict["document_id"] = document_id # Use retain_batch_async with a single item (avoids code duplication) result = await self.retain_batch_async( bank_id=bank_id, contents=[content_dict], request_context=request_context, fact_type_override=fact_type_override, confidence_score=confidence_score, ) # Return the first (and only) list of unit IDs return result[0] if result else [] async def retain_batch_async( self, bank_id: str, contents: list[RetainContentDict], *, request_context: "RequestContext", document_id: str | None = None, fact_type_override: str | None = None, confidence_score: float | None = None, document_tags: list[str] | None = None, return_usage: bool = False, ): """ Store multiple content items as memory units in ONE batch operation. This is MUCH more efficient than calling retain_async multiple times: - Extracts facts from all contents in parallel - Generates ALL embeddings in ONE batch - Does ALL database operations in ONE transaction - Automatically chunks large batches to prevent timeouts Args: bank_id: Unique identifier for the bank contents: List of dicts with keys: - "content" (required): Text content to store - "context" (optional): Context about the memory - "event_date" (optional): When the event occurred - "document_id" (optional): Document ID for this specific content item document_id: **DEPRECATED** - Use "document_id" key in each content dict instead. Applies the same document_id to ALL content items that don't specify their own. fact_type_override: Override fact type for all facts ('world', 'experience') confidence_score: Confidence score (0.0 to 1.0) return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility. Returns: If return_usage=False: List of lists of unit IDs (one list per content item) If return_usage=True: Tuple of (unit_ids, TokenUsage) Example (new style - per-content document_id): unit_ids = await memory.retain_batch_async( bank_id="user123", contents=[ {"content": "Alice works at Google", "document_id": "doc1"}, {"content": "Bob loves Python", "document_id": "doc2"}, {"content": "More about Alice", "document_id": "doc1"}, ] ) # Returns: [["unit-id-1"], ["unit-id-2"], ["unit-id-3"]] Example (deprecated style - batch-level document_id): unit_ids = await memory.retain_batch_async( bank_id="user123", contents=[ {"content": "Alice works at Google"}, {"content": "Bob loves Python"}, ], document_id="meeting-2024-01-15" ) # Returns: [["unit-id-1"], ["unit-id-2"]] """ start_time = time.time() if not contents: if return_usage: return [], TokenUsage() return [] # Authenticate tenant and set schema in context (for fq_table()) await self._authenticate_tenant(request_context) # Validate operation if validator is configured contents_copy = [dict(c) for c in contents] # Convert TypedDict to regular dict for extension if self._operation_validator: from hindsight_api.extensions import RetainContext ctx = RetainContext( bank_id=bank_id, contents=contents_copy, request_context=request_context, document_id=document_id, fact_type_override=fact_type_override, confidence_score=confidence_score, ) await self._validate_operation(self._operation_validator.validate_retain(ctx)) # Apply batch-level document_id to contents that don't have their own (backwards compatibility) if document_id: for item in contents: if "document_id" not in item: item["document_id"] = document_id # Auto-chunk large batches by character count to avoid timeouts and memory issues # Calculate total character count total_chars = sum(len(item.get("content", "")) for item in contents) total_usage = TokenUsage() CHARS_PER_BATCH = 600_000 if total_chars > CHARS_PER_BATCH: # Split into smaller batches based on character count logger.info( f"Large batch detected ({total_chars:,} chars from {len(contents)} items). Splitting into sub-batches of ~{CHARS_PER_BATCH:,} chars each..." ) sub_batches = [] current_batch = [] current_batch_chars = 0 for item in contents: item_chars = len(item.get("content", "")) # If adding this item would exceed the limit, start a new batch # (unless current batch is empty - then we must include it even if it's large) if current_batch and current_batch_chars + item_chars > CHARS_PER_BATCH: sub_batches.append(current_batch) current_batch = [item] current_batch_chars = item_chars else: current_batch.append(item) current_batch_chars += item_chars # Add the last batch if current_batch: sub_batches.append(current_batch) logger.info(f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each") # Process each sub-batch all_results = [] for i, sub_batch in enumerate(sub_batches, 1): sub_batch_chars = sum(len(item.get("content", "")) for item in sub_batch) logger.info( f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars" ) sub_results, sub_usage = await self._retain_batch_async_internal( bank_id=bank_id, contents=sub_batch, request_context=request_context, document_id=document_id, is_first_batch=i == 1, # Only upsert on first batch fact_type_override=fact_type_override, confidence_score=confidence_score, document_tags=document_tags, ) all_results.extend(sub_results) total_usage = total_usage + sub_usage total_time = time.time() - start_time logger.info( f"RETAIN_BATCH_ASYNC (chunked) COMPLETE: {len(all_results)} results from {len(contents)} contents in {total_time:.3f}s" ) result = all_results else: # Small batch - use internal method directly result, total_usage = await self._retain_batch_async_internal( bank_id=bank_id, contents=contents, request_context=request_context, document_id=document_id, is_first_batch=True, fact_type_override=fact_type_override, confidence_score=confidence_score, document_tags=document_tags, ) # Call post-operation hook if validator is configured if self._operation_validator: from hindsight_api.extensions import RetainResult result_ctx = RetainResult( bank_id=bank_id, contents=contents_copy, request_context=request_context, document_id=document_id, fact_type_override=fact_type_override, confidence_score=confidence_score, unit_ids=result, success=True, error=None, llm_input_tokens=total_usage.input_tokens, llm_output_tokens=total_usage.output_tokens, llm_total_tokens=total_usage.total_tokens, ) try: await self._operation_validator.on_retain_complete(result_ctx) except Exception as e: logger.warning(f"Post-retain hook error (non-fatal): {e}") # Trigger consolidation as a tracked async operation if enabled # Resolve bank-specific config to check if observations are enabled for this bank config = await self._config_resolver.resolve_full_config(bank_id, request_context) if config.enable_observations: try: await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context) except Exception as e: # Log but don't fail the retain - consolidation is non-critical logger.warning(f"Failed to submit consolidation task for bank {bank_id}: {e}") if return_usage: return result, total_usage return result async def _retain_batch_async_internal( self, bank_id: str, contents: list[RetainContentDict], request_context: "RequestContext", document_id: str | None = None, is_first_batch: bool = True, fact_type_override: str | None = None, confidence_score: float | None = None, document_tags: list[str] | None = None, ) -> tuple[list[list[str]], "TokenUsage"]: """ Internal method for batch processing without chunking logic. Assumes contents are already appropriately sized (< 50k chars). Called by retain_batch_async after chunking large batches. Uses semaphore for backpressure to limit concurrent retains. Args: bank_id: Unique identifier for the bank contents: List of dicts with content, context, event_date request_context: Request context for config resolution document_id: Optional document ID (always upserts if exists) is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch) fact_type_override: Override fact type for all facts confidence_score: Confidence score for opinions document_tags: Tags applied to all items in this batch Returns: Tuple of (unit ID lists, token usage for fact extraction) """ # Backpressure: limit concurrent retains to prevent database contention async with self._put_semaphore: # Use the new modular orchestrator from .retain import orchestrator pool = await self._get_pool() # Resolve bank-specific config for this operation resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context) # Create parent span for retain operation with create_operation_span("retain", bank_id): return await orchestrator.retain_batch( pool=pool, embeddings_model=self.embeddings, llm_config=self._retain_llm_config, entity_resolver=self.entity_resolver, format_date_fn=self._format_readable_date, duplicate_checker_fn=self._find_duplicate_facts_batch, bank_id=bank_id, contents_dicts=contents, document_id=document_id, is_first_batch=is_first_batch, fact_type_override=fact_type_override, confidence_score=confidence_score, document_tags=document_tags, config=resolved_config, ) def recall( self, bank_id: str, query: str, fact_type: str, budget: Budget = Budget.MID, max_tokens: int = 4096, enable_trace: bool = False, ) -> tuple[list[dict[str, Any]], Any | None]: """ Recall memories using 4-way parallel retrieval (synchronous wrapper). This is a synchronous wrapper around recall_async() for convenience. For best performance, use recall_async() directly. Args: bank_id: bank ID to recall for query: Recall query fact_type: Required filter for fact type ('world', 'experience', or 'opinion') budget: Budget level for graph traversal (low=100, mid=300, high=600 units) max_tokens: Maximum tokens to return (counts only 'text' field, default 4096) enable_trace: If True, returns detailed trace object Returns: Tuple of (results, trace) """ # Run async version synchronously - deprecated sync method, passing None for request_context from hindsight_api.models import RequestContext return asyncio.run( self.recall_async( bank_id, query, budget=budget, max_tokens=max_tokens, enable_trace=enable_trace, fact_type=[fact_type], request_context=RequestContext(), ) ) async def recall_async( self, bank_id: str, query: str, *, budget: Budget | None = None, max_tokens: int = 4096, enable_trace: bool = False, fact_type: list[str] | None = None, question_date: datetime | None = None, include_entities: bool = False, max_entity_tokens: int = 500, include_chunks: bool = False, max_chunk_tokens: int = 8192, request_context: "RequestContext", tags: list[str] | None = None, tags_match: TagsMatch = "any", _connection_budget: int | None = None, _quiet: bool = False, ) -> RecallResultModel: """ Recall memories using N*4-way parallel retrieval (N fact types × 4 retrieval methods). This implements the core RECALL operation: 1. Retrieval: For each fact type, run 4 parallel retrievals (semantic vector, BM25 keyword, graph activation, temporal graph) 2. Merge: Combine using Reciprocal Rank Fusion (RRF) 3. Rerank: Score using selected reranker (heuristic or cross-encoder) 4. Diversify: Apply MMR for diversity 5. Token Filter: Return results up to max_tokens budget Args: bank_id: bank ID to recall for query: Recall query fact_type: List of fact types to recall (e.g., ['world', 'experience']) budget: Budget level for graph traversal (low=100, mid=300, high=600 units) max_tokens: Maximum tokens to return (counts only 'text' field, default 4096) Results are returned until token budget is reached, stopping before including a fact that would exceed the limit enable_trace: Whether to return trace for debugging (deprecated) question_date: Optional date when question was asked (for temporal filtering) include_entities: Whether to include entity observations in the response max_entity_tokens: Maximum tokens for entity observations (default 500) include_chunks: Whether to include raw chunks in the response max_chunk_tokens: Maximum tokens for chunks (default 8192) tags: Optional list of tags for visibility filtering (OR matching - returns memories that have at least one matching tag) Returns: RecallResultModel containing: - results: List of MemoryFact objects - trace: Optional trace information for debugging - entities: Optional dict of entity states (if include_entities=True) - chunks: Optional dict of chunks (if include_chunks=True) """ # Authenticate tenant and set schema in context (for fq_table()) await self._authenticate_tenant(request_context) # Default to all fact types if not specified if fact_type is None: fact_type = list(VALID_RECALL_FACT_TYPES) # Filter out 'opinion' early (deprecated, silently ignore) fact_type = [ft for ft in fact_type if ft != "opinion"] # Validate fact types invalid_types = set(fact_type) - VALID_RECALL_FACT_TYPES if invalid_types: raise ValueError( f"Invalid fact type(s): {', '.join(sorted(invalid_types))}. " f"Must be one of: {', '.join(sorted(VALID_RECALL_FACT_TYPES))}" ) if not fact_type: # All requested types were opinions - return empty result return RecallResultModel(results=[], entities={}, chunks={}) # Validate operation if validator is configured if self._operation_validator: from hindsight_api.extensions import RecallContext ctx = RecallContext( bank_id=bank_id, query=query, request_context=request_context, budget=budget, max_tokens=max_tokens, enable_trace=enable_trace, fact_types=list(fact_type), question_date=question_date, include_entities=include_entities, max_entity_tokens=max_entity_tokens, include_chunks=include_chunks, max_chunk_tokens=max_chunk_tokens, ) await self._validate_operation(self._operation_validator.validate_recall(ctx)) # Map budget enum to thinking_budget number (default to MID if None) budget_mapping = {Budget.LOW: 100, Budget.MID: 300, Budget.HIGH: 1000} effective_budget = budget if budget is not None else Budget.MID thinking_budget = budget_mapping[effective_budget] # Log recall start with tags if present (skip if quiet mode for internal operations) if not _quiet: tags_info = f", tags={tags} ({tags_match})" if tags else "" logger.info(f"[RECALL {bank_id[:8]}] Starting recall for query: {query[:50]}...{tags_info}") # Create parent span for recall operation from ..tracing import get_tracer tracer = get_tracer() # Use start_as_current_span to ensure child spans are linked properly recall_span_context = tracer.start_as_current_span("hindsight.recall") recall_span = recall_span_context.__enter__() recall_span.set_attribute("hindsight.bank_id", bank_id) recall_span.set_attribute("hindsight.query", query[:100]) recall_span.set_attribute("hindsight.fact_types", ",".join(fact_type)) recall_span.set_attribute("hindsight.thinking_budget", thinking_budget) recall_span.set_attribute("hindsight.max_tokens", max_tokens) try: # Backpressure: limit concurrent recalls to prevent overwhelming the database result = None error_msg = None semaphore_wait_start = time.time() async with self._search_semaphore: semaphore_wait = time.time() - semaphore_wait_start # Retry loop for connection errors max_retries = 3 for attempt in range(max_retries + 1): try: result = await self._search_with_retries( bank_id, query, fact_type, thinking_budget, max_tokens, enable_trace, question_date, include_entities, max_entity_tokens, include_chunks, max_chunk_tokens, request_context, semaphore_wait=semaphore_wait, tags=tags, tags_match=tags_match, connection_budget=_connection_budget, quiet=_quiet, ) break # Success - exit retry loop except Exception as e: # Check if it's a connection error is_connection_error = ( isinstance(e, asyncpg.TooManyConnectionsError) or isinstance(e, asyncpg.CannotConnectNowError) or (isinstance(e, asyncpg.PostgresError) and "connection" in str(e).lower()) ) if is_connection_error and attempt < max_retries: # Wait with exponential backoff before retry wait_time = 0.5 * (2**attempt) # 0.5s, 1s, 2s logger.warning( f"Connection error on search attempt {attempt + 1}/{max_retries + 1}: {str(e)}. " f"Retrying in {wait_time:.1f}s..." ) await asyncio.sleep(wait_time) else: # Not a connection error or out of retries - call post-hook and raise error_msg = str(e) if self._operation_validator: from hindsight_api.extensions.operation_validator import RecallResult result_ctx = RecallResult( bank_id=bank_id, query=query, request_context=request_context, budget=budget, max_tokens=max_tokens, enable_trace=enable_trace, fact_types=list(fact_type), question_date=question_date, include_entities=include_entities, max_entity_tokens=max_entity_tokens, include_chunks=include_chunks, max_chunk_tokens=max_chunk_tokens, result=None, success=False, error=error_msg, ) try: await self._operation_validator.on_recall_complete(result_ctx) except Exception as hook_err: logger.warning(f"Post-recall hook error (non-fatal): {hook_err}") raise else: # Exceeded max retries error_msg = "Exceeded maximum retries for search due to connection errors." if self._operation_validator: from hindsight_api.extensions.operation_validator import RecallResult result_ctx = RecallResult( bank_id=bank_id, query=query, request_context=request_context, budget=budget, max_tokens=max_tokens, enable_trace=enable_trace, fact_types=list(fact_type), question_date=question_date, include_entities=include_entities, max_entity_tokens=max_entity_tokens, include_chunks=include_chunks, max_chunk_tokens=max_chunk_tokens, result=None, success=False, error=error_msg, ) try: await self._operation_validator.on_recall_complete(result_ctx) except Exception as hook_err: logger.warning(f"Post-recall hook error (non-fatal): {hook_err}") raise Exception(error_msg) # Call post-operation hook for success if self._operation_validator and result is not None: from hindsight_api.extensions.operation_validator import RecallResult result_ctx = RecallResult( bank_id=bank_id, query=query, request_context=request_context, budget=budget, max_tokens=max_tokens, enable_trace=enable_trace, fact_types=list(fact_type), question_date=question_date, include_entities=include_entities, max_entity_tokens=max_entity_tokens, include_chunks=include_chunks, max_chunk_tokens=max_chunk_tokens, result=result, success=True, error=None, ) try: await self._operation_validator.on_recall_complete(result_ctx) except Exception as e: logger.warning(f"Post-recall hook error (non-fatal): {e}") return result finally: recall_span_context.__exit__(None, None, None) async def _search_with_retries( self, bank_id: str, query: str, fact_type: list[str], thinking_budget: int, max_tokens: int, enable_trace: bool, question_date: datetime | None = None, include_entities: bool = False, max_entity_tokens: int = 500, include_chunks: bool = False, max_chunk_tokens: int = 8192, request_context: "RequestContext" = None, semaphore_wait: float = 0.0, tags: list[str] | None = None, tags_match: TagsMatch = "any", connection_budget: int | None = None, quiet: bool = False, ) -> RecallResultModel: """ Search implementation with modular retrieval and reranking. Architecture: 1. Retrieval: 4-way parallel (semantic, keyword, graph, temporal graph) 2. Merge: RRF to combine ranked lists 3. Reranking: Pluggable strategy (heuristic or cross-encoder) 4. Diversity: MMR with λ=0.5 5. Token Filter: Limit results to max_tokens budget Args: bank_id: bank IDentifier query: Search query fact_type: Type of facts to search thinking_budget: Nodes to explore in graph traversal max_tokens: Maximum tokens to return (counts only 'text' field) enable_trace: Whether to return search trace (deprecated) include_entities: Whether to include entity observations max_entity_tokens: Maximum tokens for entity observations include_chunks: Whether to include raw chunks max_chunk_tokens: Maximum tokens for chunks Returns: RecallResultModel with results, trace, optional entities, and optional chunks """ # Initialize tracer if requested from .search.tracer import SearchTracer tracer = ( SearchTracer(query, thinking_budget, max_tokens, tags=tags, tags_match=tags_match) if enable_trace else None ) if tracer: tracer.start() pool = await self._get_pool() recall_start = time.time() # Buffer logs for clean output in concurrent scenarios recall_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}" log_buffer = [] tags_info = f", tags={tags}, tags_match={tags_match}" if tags else "" log_buffer.append( f"[RECALL {recall_id}] Query: '{query[:50]}...' (budget={thinking_budget}, max_tokens={max_tokens}{tags_info})" ) # Import tracing utilities from ..tracing import get_tracer tracer_otel = get_tracer() try: # Step 1: Generate query embedding (for semantic search) step_start = time.time() embedding_span = tracer_otel.start_span("hindsight.recall_embedding") embedding_span.set_attribute("hindsight.bank_id", bank_id) embedding_span.set_attribute("hindsight.query", query[:100]) try: query_embedding = embedding_utils.generate_embedding(self.embeddings, query) step_duration = time.time() - step_start log_buffer.append(f" [1] Generate query embedding: {step_duration:.3f}s") finally: embedding_span.end() if tracer: tracer.record_query_embedding(query_embedding) tracer.add_phase_metric("generate_query_embedding", step_duration) # Step 2: Optimized parallel retrieval using batched queries # - Semantic + BM25 combined in 1 CTE query for ALL fact types # - Graph runs per fact type (complex traversal) # - Temporal runs per fact type (if constraint detected) step_start = time.time() query_embedding_str = str(query_embedding) from .search.retrieval import ( get_default_graph_retriever, retrieve_all_fact_types_parallel, ) # Track each retrieval start time retrieval_start = time.time() retrieval_span = tracer_otel.start_span("hindsight.recall_retrieval") retrieval_span.set_attribute("hindsight.bank_id", bank_id) retrieval_span.set_attribute("hindsight.fact_types", ",".join(fact_type)) retrieval_span.set_attribute("hindsight.thinking_budget", thinking_budget) try: # Run optimized retrieval with connection budget config = get_config() effective_connection_budget = ( connection_budget if connection_budget is not None else config.recall_connection_budget ) async with budgeted_operation( max_connections=effective_connection_budget, operation_id=f"recall-{recall_id}", ) as op: budgeted_pool = op.wrap_pool(pool) parallel_start = time.time() multi_result = await retrieve_all_fact_types_parallel( budgeted_pool, query, query_embedding_str, bank_id, fact_type, # Pass all fact types at once thinking_budget, question_date, self.query_analyzer, tags=tags, tags_match=tags_match, ) parallel_duration = time.time() - parallel_start finally: retrieval_span.end() # Combine all results from all fact types and aggregate timings semantic_results = [] bm25_results = [] graph_results = [] temporal_results = [] aggregated_timings = { "semantic": 0.0, "bm25": 0.0, "graph": 0.0, "temporal": 0.0, "temporal_extraction": 0.0, } all_mpfp_timings = [] detected_temporal_constraint = None max_conn_wait = multi_result.max_conn_wait for ft in fact_type: retrieval_result = multi_result.results_by_fact_type.get(ft) if not retrieval_result: continue # Log fact types in this retrieval batch logger.debug( f"[RECALL {recall_id}] Fact type '{ft}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}" ) semantic_results.extend(retrieval_result.semantic) bm25_results.extend(retrieval_result.bm25) graph_results.extend(retrieval_result.graph) if retrieval_result.temporal: temporal_results.extend(retrieval_result.temporal) # Track max timing for each method (since they run in parallel across fact types) for method, duration in retrieval_result.timings.items(): aggregated_timings[method] = max(aggregated_timings.get(method, 0.0), duration) # Capture temporal constraint (same across all fact types) if retrieval_result.temporal_constraint: detected_temporal_constraint = retrieval_result.temporal_constraint # If no temporal results from any fact type, set to None if not temporal_results: temporal_results = None # Sort combined results by score (descending) so higher-scored results # get better ranks in the trace, regardless of fact type semantic_results.sort(key=lambda r: r.similarity if hasattr(r, "similarity") else 0, reverse=True) bm25_results.sort(key=lambda r: r.bm25_score if hasattr(r, "bm25_score") else 0, reverse=True) graph_results.sort(key=lambda r: r.activation if hasattr(r, "activation") else 0, reverse=True) if temporal_results: temporal_results.sort( key=lambda r: r.combined_score if hasattr(r, "combined_score") else 0, reverse=True ) retrieval_duration = time.time() - retrieval_start step_duration = time.time() - step_start total_retrievals = len(fact_type) * (4 if temporal_results else 3) # Format per-method timings timing_parts = [ f"semantic={len(semantic_results)}({aggregated_timings['semantic']:.3f}s)", f"bm25={len(bm25_results)}({aggregated_timings['bm25']:.3f}s)", f"graph={len(graph_results)}({aggregated_timings['graph']:.3f}s)", f"temporal_extraction={aggregated_timings['temporal_extraction']:.3f}s", ] temporal_info = "" if detected_temporal_constraint: start_dt, end_dt = detected_temporal_constraint temporal_count = len(temporal_results) if temporal_results else 0 timing_parts.append(f"temporal={temporal_count}({aggregated_timings['temporal']:.3f}s)") temporal_info = f" | temporal_range={start_dt.strftime('%Y-%m-%d')} to {end_dt.strftime('%Y-%m-%d')}" log_buffer.append( f" [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{temporal_info}" ) # Log graph retriever timing breakdown if available if all_mpfp_timings: retriever_name = get_default_graph_retriever().name.upper() mpfp_total = all_mpfp_timings[0] # Take first fact type's timing as representative mpfp_parts = [ f"db_queries={mpfp_total.db_queries}", f"edge_load={mpfp_total.edge_load_time:.3f}s", f"edges={mpfp_total.edge_count}", f"patterns={mpfp_total.pattern_count}", ] if mpfp_total.seeds_time > 0.01: mpfp_parts.append(f"seeds={mpfp_total.seeds_time:.3f}s") if mpfp_total.fusion > 0.001: mpfp_parts.append(f"fusion={mpfp_total.fusion:.3f}s") if mpfp_total.fetch > 0.001: mpfp_parts.append(f"fetch={mpfp_total.fetch:.3f}s") log_buffer.append(f" [{retriever_name}] {', '.join(mpfp_parts)}") # Log detailed hop timing for debugging slow queries if mpfp_total.hop_details: for hd in mpfp_total.hop_details: log_buffer.append( f" hop{hd['hop']}: exec={hd.get('exec_time', 0) * 1000:.0f}ms, " f"uncached={hd.get('uncached_after_filter', 0)}, " f"load={hd.get('load_time', 0) * 1000:.0f}ms, " f"edges={hd.get('edges_loaded', 0)}" ) # Record temporal constraint in tracer if detected if tracer and detected_temporal_constraint: start_dt, end_dt = detected_temporal_constraint tracer.record_temporal_constraint(start_dt, end_dt) # Record retrieval results for tracer - per fact type if tracer: # Convert RetrievalResult to old tuple format for tracer def to_tuple_format(results): return [(r.id, r.__dict__) for r in results] # Add retrieval results per fact type (to show parallel execution in UI) for ft_name in fact_type: rr = multi_result.results_by_fact_type.get(ft_name) if not rr: continue # Add semantic retrieval results for this fact type tracer.add_retrieval_results( method_name="semantic", results=to_tuple_format(rr.semantic), duration_seconds=rr.timings.get("semantic", 0.0), score_field="similarity", metadata={"limit": thinking_budget}, fact_type=ft_name, ) # Add BM25 retrieval results for this fact type tracer.add_retrieval_results( method_name="bm25", results=to_tuple_format(rr.bm25), duration_seconds=rr.timings.get("bm25", 0.0), score_field="bm25_score", metadata={"limit": thinking_budget}, fact_type=ft_name, ) # Add graph retrieval results for this fact type tracer.add_retrieval_results( method_name="graph", results=to_tuple_format(rr.graph), duration_seconds=rr.timings.get("graph", 0.0), score_field="activation", metadata={"budget": thinking_budget}, fact_type=ft_name, ) # Add temporal retrieval results for this fact type # Show temporal even with 0 results if constraint was detected if rr.temporal is not None or rr.temporal_constraint is not None: temporal_metadata = {"budget": thinking_budget} if rr.temporal_constraint: start_dt, end_dt = rr.temporal_constraint temporal_metadata["constraint"] = { "start": start_dt.isoformat() if start_dt else None, "end": end_dt.isoformat() if end_dt else None, } tracer.add_retrieval_results( method_name="temporal", results=to_tuple_format(rr.temporal or []), duration_seconds=rr.timings.get("temporal", 0.0), score_field="temporal_score", metadata=temporal_metadata, fact_type=ft_name, ) # Record entry points (from semantic results) for legacy graph view for rank, retrieval in enumerate(semantic_results[:10], start=1): # Top 10 as entry points tracer.add_entry_point(retrieval.id, retrieval.text, retrieval.similarity or 0.0, rank) tracer.add_phase_metric( "parallel_retrieval", step_duration, { "semantic_count": len(semantic_results), "bm25_count": len(bm25_results), "graph_count": len(graph_results), "temporal_count": len(temporal_results) if temporal_results else 0, }, ) # Step 3: Merge with RRF step_start = time.time() from .search.fusion import reciprocal_rank_fusion fusion_span = tracer_otel.start_span("hindsight.recall_fusion") fusion_span.set_attribute("hindsight.bank_id", bank_id) fusion_span.set_attribute("hindsight.semantic_count", len(semantic_results)) fusion_span.set_attribute("hindsight.bm25_count", len(bm25_results)) fusion_span.set_attribute("hindsight.graph_count", len(graph_results)) fusion_span.set_attribute("hindsight.temporal_count", len(temporal_results) if temporal_results else 0) try: # Merge 3 or 4 result lists depending on temporal constraint if temporal_results: merged_candidates = reciprocal_rank_fusion( [semantic_results, bm25_results, graph_results, temporal_results] ) else: merged_candidates = reciprocal_rank_fusion([semantic_results, bm25_results, graph_results]) step_duration = time.time() - step_start log_buffer.append( f" [3] RRF merge: {len(merged_candidates)} unique candidates in {step_duration:.3f}s" ) finally: fusion_span.set_attribute("hindsight.merged_count", len(merged_candidates)) fusion_span.end() if tracer: # Convert MergedCandidate to old tuple format for tracer tracer_merged = [ (mc.id, mc.retrieval.__dict__, {"rrf_score": mc.rrf_score, **mc.source_ranks}) for mc in merged_candidates ] tracer.add_rrf_merged(tracer_merged) tracer.add_phase_metric("rrf_merge", step_duration, {"candidates_merged": len(merged_candidates)}) # Step 4: Rerank using cross-encoder (MergedCandidate -> ScoredResult) step_start = time.time() reranker_instance = self._cross_encoder_reranker rerank_span = tracer_otel.start_span("hindsight.recall_rerank") rerank_span.set_attribute("hindsight.bank_id", bank_id) rerank_span.set_attribute("hindsight.candidates_count", len(merged_candidates)) try: # Ensure reranker is initialized (for lazy initialization mode) await reranker_instance.ensure_initialized() # Pre-filter candidates to reduce reranking cost (RRF already provides good ranking) # This is especially important for remote rerankers with network latency reranker_max_candidates = get_config().reranker_max_candidates pre_filtered_count = 0 if len(merged_candidates) > reranker_max_candidates: # Sort by RRF score and take top candidates merged_candidates.sort(key=lambda mc: mc.rrf_score, reverse=True) pre_filtered_count = len(merged_candidates) - reranker_max_candidates merged_candidates = merged_candidates[:reranker_max_candidates] # Rerank using cross-encoder scored_results = await reranker_instance.rerank(query, merged_candidates) step_duration = time.time() - step_start pre_filter_note = f" (pre-filtered {pre_filtered_count})" if pre_filtered_count > 0 else "" log_buffer.append( f" [4] Reranking: {len(scored_results)} candidates scored in {step_duration:.3f}s{pre_filter_note}" ) finally: rerank_span.set_attribute("hindsight.scored_count", len(scored_results)) if pre_filtered_count > 0: rerank_span.set_attribute("hindsight.pre_filtered_count", pre_filtered_count) rerank_span.end() # Step 4.5: Combine cross-encoder score with retrieval signals # This preserves retrieval work (RRF, temporal, recency) instead of pure cross-encoder ranking if scored_results: # Normalize RRF scores to [0, 1] range using min-max normalization rrf_scores = [sr.candidate.rrf_score for sr in scored_results] max_rrf = max(rrf_scores) if rrf_scores else 0.0 min_rrf = min(rrf_scores) if rrf_scores else 0.0 rrf_range = max_rrf - min_rrf # Don't force to 1.0, let fallback handle it # Calculate recency based on occurred_start (more recent = higher score) now = utcnow() for sr in scored_results: # Normalize RRF score (0-1 range, 0.5 if all same) if rrf_range > 0: sr.rrf_normalized = (sr.candidate.rrf_score - min_rrf) / rrf_range else: # All RRF scores are the same, use neutral value sr.rrf_normalized = 0.5 # Calculate recency (decay over 365 days, minimum 0.1) sr.recency = 0.5 # default for missing dates if sr.retrieval.occurred_start: occurred = sr.retrieval.occurred_start if hasattr(occurred, "tzinfo") and occurred.tzinfo is None: occurred = occurred.replace(tzinfo=UTC) days_ago = (now - occurred).total_seconds() / 86400 sr.recency = max(0.1, 1.0 - (days_ago / 365)) # Linear decay over 1 year # Get temporal proximity if available (already 0-1) sr.temporal = ( sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5 ) # Weighted combination # Cross-encoder: 60% (semantic relevance) # RRF: 20% (retrieval consensus) # Temporal proximity: 10% (time relevance for temporal queries) # Recency: 10% (prefer recent facts) sr.combined_score = ( 0.6 * sr.cross_encoder_score_normalized + 0.2 * sr.rrf_normalized + 0.1 * sr.temporal + 0.1 * sr.recency ) sr.weight = sr.combined_score # Update weight for final ranking # Re-sort by combined score scored_results.sort(key=lambda x: x.weight, reverse=True) log_buffer.append( " [4.6] Combined scoring: cross_encoder(0.6) + rrf(0.2) + temporal(0.1) + recency(0.1)" ) # Add reranked results to tracer AFTER combined scoring (so normalized values are included) if tracer: results_dict = [sr.to_dict() for sr in scored_results] tracer_merged = [ (mc.id, mc.retrieval.__dict__, {"rrf_score": mc.rrf_score, **mc.source_ranks}) for mc in merged_candidates ] tracer.add_reranked(results_dict, tracer_merged) tracer.add_phase_metric( "reranking", step_duration, {"reranker_type": "cross-encoder", "candidates_reranked": len(scored_results)}, ) # Step 5: Truncate to thinking_budget * 2 for token filtering rerank_limit = thinking_budget * 2 top_scored = scored_results[:rerank_limit] log_buffer.append(f" [5] Truncated to top {len(top_scored)} results") # Step 6: Token budget filtering step_start = time.time() # Convert to dict for token filtering (backward compatibility) top_dicts = [sr.to_dict() for sr in top_scored] filtered_dicts, total_tokens = self._filter_by_token_budget(top_dicts, max_tokens) # Convert back to list of IDs and filter scored_results filtered_ids = {d["id"] for d in filtered_dicts} top_scored = [sr for sr in top_scored if sr.id in filtered_ids] step_duration = time.time() - step_start log_buffer.append( f" [6] Token filtering: {len(top_scored)} results, {total_tokens}/{max_tokens} tokens in {step_duration:.3f}s" ) if tracer: tracer.add_phase_metric( "token_filtering", step_duration, {"results_selected": len(top_scored), "tokens_used": total_tokens, "max_tokens": max_tokens}, ) # Record visits for all retrieved nodes if tracer: for sr in scored_results: tracer.visit_node( node_id=sr.id, text=sr.retrieval.text, context=sr.retrieval.context or "", event_date=sr.retrieval.occurred_start, is_entry_point=(sr.id in [ep.node_id for ep in tracer.entry_points]), parent_node_id=None, # In parallel retrieval, there's no clear parent link_type=None, link_weight=None, activation=sr.candidate.rrf_score, # Use RRF score as activation semantic_similarity=sr.retrieval.similarity or 0.0, recency=sr.recency, frequency=0.0, final_weight=sr.weight, ) # Log fact_type distribution in results fact_type_counts = {} for sr in top_scored: ft = sr.retrieval.fact_type fact_type_counts[ft] = fact_type_counts.get(ft, 0) + 1 fact_type_summary = ", ".join([f"{ft}={count}" for ft, count in sorted(fact_type_counts.items())]) # Convert ScoredResult to dicts with ISO datetime strings top_results_dicts = [] for sr in top_scored: result_dict = sr.to_dict() # Convert datetime objects to ISO strings for JSON serialization if result_dict.get("occurred_start"): occurred_start = result_dict["occurred_start"] result_dict["occurred_start"] = ( occurred_start.isoformat() if hasattr(occurred_start, "isoformat") else occurred_start ) if result_dict.get("occurred_end"): occurred_end = result_dict["occurred_end"] result_dict["occurred_end"] = ( occurred_end.isoformat() if hasattr(occurred_end, "isoformat") else occurred_end ) if result_dict.get("mentioned_at"): mentioned_at = result_dict["mentioned_at"] result_dict["mentioned_at"] = ( mentioned_at.isoformat() if hasattr(mentioned_at, "isoformat") else mentioned_at ) top_results_dicts.append(result_dict) # Convert results to MemoryFact objects memory_facts = [] for result_dict in top_results_dicts: memory_facts.append( MemoryFact( id=str(result_dict.get("id")), text=result_dict.get("text"), fact_type=result_dict.get("fact_type", "world"), entities=None, # Entity observations removed context=result_dict.get("context"), occurred_start=result_dict.get("occurred_start"), occurred_end=result_dict.get("occurred_end"), mentioned_at=result_dict.get("mentioned_at"), document_id=result_dict.get("document_id"), chunk_id=result_dict.get("chunk_id"), tags=result_dict.get("tags"), ) ) # Entity observations removed - always set to None entities_dict = None # Fetch chunks if requested chunks_dict = None total_chunk_tokens = 0 if include_chunks and top_scored: from .response_models import ChunkInfo # Collect chunk_ids in order of fact relevance (preserving order from top_scored) # Use a list to maintain order, but track seen chunks to avoid duplicates chunk_ids_ordered = [] seen_chunk_ids = set() for sr in top_scored: chunk_id = sr.retrieval.chunk_id if chunk_id and chunk_id not in seen_chunk_ids: chunk_ids_ordered.append(chunk_id) seen_chunk_ids.add(chunk_id) if chunk_ids_ordered: # Fetch chunk data from database using chunk_ids (no ORDER BY to preserve input order) async with acquire_with_retry(pool) as conn: chunks_rows = await conn.fetch( f""" SELECT chunk_id, chunk_text, chunk_index FROM {fq_table("chunks")} WHERE chunk_id = ANY($1::text[]) """, chunk_ids_ordered, ) # Create a lookup dict for fast access chunks_lookup = {row["chunk_id"]: row for row in chunks_rows} # Apply token limit and build chunks_dict in the order of chunk_ids_ordered chunks_dict = {} encoding = _get_tiktoken_encoding() for chunk_id in chunk_ids_ordered: if chunk_id not in chunks_lookup: continue row = chunks_lookup[chunk_id] chunk_text = row["chunk_text"] chunk_tokens = len(encoding.encode(chunk_text)) # Check if adding this chunk would exceed the limit if total_chunk_tokens + chunk_tokens > max_chunk_tokens: # Truncate the chunk to fit within the remaining budget remaining_tokens = max_chunk_tokens - total_chunk_tokens if remaining_tokens > 0: # Truncate to remaining tokens truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens]) chunks_dict[chunk_id] = ChunkInfo( chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True ) total_chunk_tokens = max_chunk_tokens # Stop adding more chunks once we hit the limit break else: chunks_dict[chunk_id] = ChunkInfo( chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False ) total_chunk_tokens += chunk_tokens # Finalize trace if enabled trace_dict = None if tracer: trace = tracer.finalize(top_results_dicts) trace_dict = trace.to_dict() if trace else None # Log final recall stats total_time = time.time() - recall_start num_chunks = len(chunks_dict) if chunks_dict else 0 # Include wait times in log if significant wait_parts = [] if semaphore_wait > 0.01: wait_parts.append(f"sem={semaphore_wait:.3f}s") if max_conn_wait > 0.01: wait_parts.append(f"conn={max_conn_wait:.3f}s") wait_info = f" | waits: {', '.join(wait_parts)}" if wait_parts else "" log_buffer.append( f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok) | {fact_type_summary} | {total_time:.3f}s{wait_info}" ) if not quiet: logger.info("\n" + "\n".join(log_buffer)) return RecallResultModel(results=memory_facts, trace=trace_dict, entities=entities_dict, chunks=chunks_dict) except Exception as e: log_buffer.append(f"[RECALL {recall_id}] ERROR after {time.time() - recall_start:.3f}s: {str(e)}") if not quiet: logger.error("\n" + "\n".join(log_buffer)) raise Exception(f"Failed to search memories: {str(e)}") def _filter_by_token_budget( self, results: list[dict[str, Any]], max_tokens: int ) -> tuple[list[dict[str, Any]], int]: """ Filter results to fit within token budget. Counts tokens only for the 'text' field using tiktoken (cl100k_base encoding). Stops before including a fact that would exceed the budget. Args: results: List of search results max_tokens: Maximum tokens allowed Returns: Tuple of (filtered_results, total_tokens_used) """ encoding = _get_tiktoken_encoding() filtered_results = [] total_tokens = 0 for result in results: text = result.get("text", "") text_tokens = len(encoding.encode(text)) # Check if adding this result would exceed budget if total_tokens + text_tokens <= max_tokens: filtered_results.append(result) total_tokens += text_tokens else: # Stop before including a fact that would exceed limit break return filtered_results, total_tokens async def get_document( self, document_id: str, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, Any] | None: """ Retrieve document metadata and statistics. Args: document_id: Document ID to retrieve bank_id: bank ID that owns the document request_context: Request context for authentication. Returns: Dictionary with document info or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: doc = await conn.fetchrow( f""" SELECT d.id, d.bank_id, d.original_text, d.content_hash, d.created_at, d.updated_at, d.tags, COUNT(mu.id) as unit_count FROM {fq_table("documents")} d LEFT JOIN {fq_table("memory_units")} mu ON mu.document_id = d.id WHERE d.id = $1 AND d.bank_id = $2 GROUP BY d.id, d.bank_id, d.original_text, d.content_hash, d.created_at, d.updated_at, d.tags """, document_id, bank_id, ) if not doc: return None return { "id": doc["id"], "bank_id": doc["bank_id"], "original_text": doc["original_text"], "content_hash": doc["content_hash"], "memory_unit_count": doc["unit_count"], "created_at": doc["created_at"].isoformat() if doc["created_at"] else None, "updated_at": doc["updated_at"].isoformat() if doc["updated_at"] else None, "tags": list(doc["tags"]) if doc["tags"] else [], } async def delete_document( self, document_id: str, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, int]: """ Delete a document and all its associated memory units and links. Args: document_id: Document ID to delete bank_id: bank ID that owns the document request_context: Request context for authentication. Returns: Dictionary with counts of deleted items """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Get memory unit IDs before deletion (for mental model invalidation) unit_rows = await conn.fetch( f"SELECT id FROM {fq_table('memory_units')} WHERE document_id = $1", document_id ) unit_ids = [str(row["id"]) for row in unit_rows] units_count = len(unit_ids) # Delete document (cascades to memory_units and all their links) deleted = await conn.fetchval( f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id, ) # Invalidate deleted fact IDs from mental models if deleted and unit_ids: await self._invalidate_facts_from_mental_models(conn, bank_id, unit_ids) return {"document_deleted": 1 if deleted else 0, "memory_units_deleted": units_count if deleted else 0} async def delete_memory_unit( self, unit_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """ Delete a single memory unit and all its associated links. Due to CASCADE DELETE constraints, this will automatically delete: - All links from this unit (memory_links where from_unit_id = unit_id) - All links to this unit (memory_links where to_unit_id = unit_id) - All entity associations (unit_entities where unit_id = unit_id) Args: unit_id: UUID of the memory unit to delete request_context: Request context for authentication. Returns: Dictionary with deletion result """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Get bank_id before deletion (for mental model invalidation) bank_id = await conn.fetchval(f"SELECT bank_id FROM {fq_table('memory_units')} WHERE id = $1", unit_id) # Delete the memory unit (cascades to links and associations) deleted = await conn.fetchval( f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 RETURNING id", unit_id ) # Invalidate deleted fact ID from mental models if deleted and bank_id: await self._invalidate_facts_from_mental_models(conn, bank_id, [str(deleted)]) return { "success": deleted is not None, "unit_id": str(deleted) if deleted else None, "message": "Memory unit and all its links deleted successfully" if deleted else "Memory unit not found", } async def delete_bank( self, bank_id: str, fact_type: str | None = None, *, request_context: "RequestContext", ) -> dict[str, int]: """ Delete all data for a specific agent (multi-tenant cleanup). This is much more efficient than dropping all tables and allows multiple agents to coexist in the same database. Deletes (with CASCADE): - All memory units for this bank (optionally filtered by fact_type) - All entities for this bank (if deleting all memory units) - All associated links, unit-entity associations, and co-occurrences Args: bank_id: bank ID to delete fact_type: Optional fact type filter (world, experience, opinion). If provided, only deletes memories of that type. request_context: Request context for authentication. Returns: Dictionary with counts of deleted items """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Ensure connection is not in read-only mode (can happen with connection poolers) await conn.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE") async with conn.transaction(): try: if fact_type: # Delete only memories of a specific fact type units_count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2", bank_id, fact_type, ) await conn.execute( f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2", bank_id, fact_type, ) # Note: We don't delete entities when fact_type is specified, # as they may be referenced by other memory units return {"memory_units_deleted": units_count, "entities_deleted": 0} else: # Delete all data for the bank units_count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1", bank_id ) entities_count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('entities')} WHERE bank_id = $1", bank_id ) documents_count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('documents')} WHERE bank_id = $1", bank_id ) # Delete documents (cascades to chunks) await conn.execute(f"DELETE FROM {fq_table('documents')} WHERE bank_id = $1", bank_id) # Delete memory units (cascades to unit_entities, memory_links) await conn.execute(f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1", bank_id) # Delete entities (cascades to unit_entities, entity_cooccurrences, memory_links with entity_id) await conn.execute(f"DELETE FROM {fq_table('entities')} WHERE bank_id = $1", bank_id) # Delete the bank profile itself await conn.execute(f"DELETE FROM {fq_table('banks')} WHERE bank_id = $1", bank_id) return { "memory_units_deleted": units_count, "entities_deleted": entities_count, "documents_deleted": documents_count, "bank_deleted": True, } except Exception as e: raise Exception(f"Failed to delete agent data: {str(e)}") async def clear_observations( self, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, int]: """ Clear all observations for a bank (consolidated knowledge). Args: bank_id: Bank ID to clear observations for request_context: Request context for authentication. Returns: Dictionary with count of deleted observations """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Count observations before deletion count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", bank_id, ) # Delete all observations await conn.execute( f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'", bank_id, ) # Reset consolidation timestamp await conn.execute( f"UPDATE {fq_table('banks')} SET last_consolidated_at = NULL WHERE bank_id = $1", bank_id, ) return {"deleted_count": count or 0} async def run_consolidation( self, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, int]: """ Run memory consolidation to create/update mental models. Args: bank_id: Bank ID to run consolidation for request_context: Request context for authentication. Returns: Dictionary with consolidation stats """ await self._authenticate_tenant(request_context) from .consolidation import run_consolidation_job # Create parent span for consolidation operation with create_operation_span("consolidation", bank_id): result = await run_consolidation_job( memory_engine=self, bank_id=bank_id, request_context=request_context, ) return { "processed": result.get("processed", 0), "created": result.get("created", 0), "updated": result.get("updated", 0), "skipped": result.get("skipped", 0), } async def get_graph_data( self, bank_id: str | None = None, fact_type: str | None = None, *, limit: int = 1000, request_context: "RequestContext", ): """ Get graph data for visualization. Args: bank_id: Filter by bank ID fact_type: Filter by fact type (world, experience, opinion) limit: Maximum number of items to return (default: 1000) request_context: Request context for authentication. Returns: Dict with nodes, edges, table_rows, total_units, and limit """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Get memory units, optionally filtered by bank_id and fact_type query_conditions = [] query_params = [] param_count = 0 if bank_id: param_count += 1 query_conditions.append(f"bank_id = ${param_count}") query_params.append(bank_id) if fact_type: param_count += 1 query_conditions.append(f"fact_type = ${param_count}") query_params.append(fact_type) where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" # Get total count first total_count_result = await conn.fetchrow( f""" SELECT COUNT(*) as total FROM {fq_table("memory_units")} {where_clause} """, *query_params, ) total_count = total_count_result["total"] if total_count_result else 0 # Get units with limit param_count += 1 units = await conn.fetch( f""" SELECT id, text, event_date, context, occurred_start, occurred_end, mentioned_at, document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids FROM {fq_table("memory_units")} {where_clause} ORDER BY mentioned_at DESC NULLS LAST, event_date DESC LIMIT ${param_count} """, *query_params, limit, ) # Get links, filtering to only include links between units of the selected agent # Use DISTINCT ON with LEAST/GREATEST to deduplicate bidirectional links unit_ids = [row["id"] for row in units] unit_id_set = set(unit_ids) # Collect source memory IDs from observations source_memory_ids = [] for unit in units: if unit["source_memory_ids"]: source_memory_ids.extend(unit["source_memory_ids"]) source_memory_ids = list(set(source_memory_ids)) # Deduplicate # Fetch links involving both visible units AND source memories all_relevant_ids = unit_ids + source_memory_ids if all_relevant_ids: links = await conn.fetch( f""" SELECT DISTINCT ON (LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight, e.canonical_name as entity_name FROM {fq_table("memory_links")} ml LEFT JOIN {fq_table("entities")} e ON ml.entity_id = e.id WHERE ml.from_unit_id = ANY($1::uuid[]) OR ml.to_unit_id = ANY($1::uuid[]) ORDER BY LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid), ml.weight DESC """, all_relevant_ids, ) else: links = [] # Copy links from source memories to observations # Observations inherit links from their source memories via source_memory_ids # Build a map from source_id to observation_ids source_to_observations = {} for unit in units: if unit["source_memory_ids"]: for source_id in unit["source_memory_ids"]: if source_id not in source_to_observations: source_to_observations[source_id] = [] source_to_observations[source_id].append(unit["id"]) copied_links = [] for link in links: from_id = link["from_unit_id"] to_id = link["to_unit_id"] # Get observations that should inherit this link from_observations = source_to_observations.get(from_id, []) to_observations = source_to_observations.get(to_id, []) # If from_id is a source memory, copy links to its observations if from_observations: for obs_id in from_observations: # Only include if the target is visible if to_id in unit_id_set or to_observations: target = to_observations[0] if to_observations and to_id not in unit_id_set else to_id if target in unit_id_set: copied_links.append( { "from_unit_id": obs_id, "to_unit_id": target, "link_type": link["link_type"], "weight": link["weight"], "entity_name": link["entity_name"], } ) # If to_id is a source memory, copy links to its observations if to_observations and from_id in unit_id_set: for obs_id in to_observations: copied_links.append( { "from_unit_id": from_id, "to_unit_id": obs_id, "link_type": link["link_type"], "weight": link["weight"], "entity_name": link["entity_name"], } ) # Keep only direct links between visible nodes direct_links = [ link for link in links if link["from_unit_id"] in unit_id_set and link["to_unit_id"] in unit_id_set ] # Get entity information unit_entities = await conn.fetch(f""" SELECT ue.unit_id, e.canonical_name FROM {fq_table("unit_entities")} ue JOIN {fq_table("entities")} e ON ue.entity_id = e.id ORDER BY ue.unit_id """) # Build entity mapping entity_map = {} for row in unit_entities: unit_id = row["unit_id"] entity_name = row["canonical_name"] if unit_id not in entity_map: entity_map[unit_id] = [] entity_map[unit_id].append(entity_name) # For observations, inherit entities from source memories for unit in units: if unit["source_memory_ids"] and unit["id"] not in entity_map: # Collect entities from all source memories source_entities = [] for source_id in unit["source_memory_ids"]: if source_id in entity_map: source_entities.extend(entity_map[source_id]) if source_entities: # Deduplicate while preserving order entity_map[unit["id"]] = list(dict.fromkeys(source_entities)) # Build nodes nodes = [] for row in units: unit_id = row["id"] text = row["text"] event_date = row["event_date"] context = row["context"] entities = entity_map.get(unit_id, []) entity_count = len(entities) # Color by entity count if entity_count == 0: color = "#e0e0e0" elif entity_count == 1: color = "#90caf9" else: color = "#42a5f5" nodes.append( { "data": { "id": str(unit_id), "label": f"{text[:30]}..." if len(text) > 30 else text, "text": text, "date": event_date.isoformat() if event_date else "", "context": context if context else "", "entities": ", ".join(entities) if entities else "None", "color": color, } } ) # Build edges (combine direct links and copied links from sources) edges = [] all_links = direct_links + copied_links for row in all_links: from_id = str(row["from_unit_id"]) to_id = str(row["to_unit_id"]) link_type = row["link_type"] weight = row["weight"] entity_name = row.get("entity_name") # Color by link type if link_type == "temporal": color = "#00bcd4" line_style = "dashed" elif link_type == "semantic": color = "#ff69b4" line_style = "solid" elif link_type == "entity": color = "#ffd700" line_style = "solid" else: color = "#999999" line_style = "solid" edges.append( { "data": { "id": f"{from_id}-{to_id}-{link_type}", "source": from_id, "target": to_id, "linkType": link_type, "weight": weight, "entityName": entity_name if entity_name else "", "color": color, "lineStyle": line_style, } } ) # Build table rows table_rows = [] for row in units: unit_id = row["id"] entities = entity_map.get(unit_id, []) table_rows.append( { "id": str(unit_id), "text": row["text"], "context": row["context"] if row["context"] else "N/A", "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, "date": row["event_date"].strftime("%Y-%m-%d %H:%M") if row["event_date"] else "N/A", # Deprecated, kept for backwards compatibility "entities": ", ".join(entities) if entities else "None", "document_id": row["document_id"], "chunk_id": row["chunk_id"] if row["chunk_id"] else None, "fact_type": row["fact_type"], "tags": list(row["tags"]) if row["tags"] else [], "created_at": row["created_at"].isoformat() if row["created_at"] else None, "proof_count": row["proof_count"] if row["proof_count"] else None, } ) return {"nodes": nodes, "edges": edges, "table_rows": table_rows, "total_units": total_count, "limit": limit} async def list_memory_units( self, bank_id: str, *, fact_type: str | None = None, search_query: str | None = None, limit: int = 100, offset: int = 0, request_context: "RequestContext", ): """ List memory units for table view with optional full-text search. Args: bank_id: Filter by bank ID fact_type: Filter by fact type (world, experience, opinion) search_query: Full-text search query (searches text and context fields) limit: Maximum number of results to return offset: Offset for pagination request_context: Request context for authentication. Returns: Dict with items (list of memory units) and total count """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build query conditions query_conditions = [] query_params = [] param_count = 0 if bank_id: param_count += 1 query_conditions.append(f"bank_id = ${param_count}") query_params.append(bank_id) if fact_type: param_count += 1 query_conditions.append(f"fact_type = ${param_count}") query_params.append(fact_type) if search_query: # Full-text search on text and context fields using ILIKE param_count += 1 query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})") query_params.append(f"%{search_query}%") where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" # Get total count count_query = f""" SELECT COUNT(*) as total FROM {fq_table("memory_units")} {where_clause} """ count_result = await conn.fetchrow(count_query, *query_params) total = count_result["total"] # Get units with limit and offset param_count += 1 limit_param = f"${param_count}" query_params.append(limit) param_count += 1 offset_param = f"${param_count}" query_params.append(offset) units = await conn.fetch( f""" SELECT id, text, event_date, context, fact_type, mentioned_at, occurred_start, occurred_end, chunk_id FROM {fq_table("memory_units")} {where_clause} ORDER BY mentioned_at DESC NULLS LAST, created_at DESC LIMIT {limit_param} OFFSET {offset_param} """, *query_params, ) # Get entity information for these units if units: unit_ids = [row["id"] for row in units] unit_entities = await conn.fetch( f""" SELECT ue.unit_id, e.canonical_name FROM {fq_table("unit_entities")} ue JOIN {fq_table("entities")} e ON ue.entity_id = e.id WHERE ue.unit_id = ANY($1::uuid[]) ORDER BY ue.unit_id """, unit_ids, ) else: unit_entities = [] # Build entity mapping entity_map = {} for row in unit_entities: unit_id = row["unit_id"] entity_name = row["canonical_name"] if unit_id not in entity_map: entity_map[unit_id] = [] entity_map[unit_id].append(entity_name) # Build result items items = [] for row in units: unit_id = row["id"] entities = entity_map.get(unit_id, []) items.append( { "id": str(unit_id), "text": row["text"], "context": row["context"] if row["context"] else "", "date": row["event_date"].isoformat() if row["event_date"] else "", "fact_type": row["fact_type"], "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, "entities": ", ".join(entities) if entities else "", "chunk_id": row["chunk_id"] if row["chunk_id"] else None, } ) return {"items": items, "total": total, "limit": limit, "offset": offset} async def get_memory_unit( self, bank_id: str, memory_id: str, request_context: "RequestContext", ): """ Get a single memory unit by ID. Args: bank_id: Bank ID memory_id: Memory unit ID request_context: Request context for authentication. Returns: Dict with memory unit data or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Get the memory unit (include source_memory_ids for mental models) row = await conn.fetchrow( f""" SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, source_memory_ids FROM {fq_table("memory_units")} WHERE id = $1 AND bank_id = $2 """, memory_id, bank_id, ) if not row: return None # Get entity information entities_rows = await conn.fetch( f""" SELECT e.canonical_name FROM {fq_table("unit_entities")} ue JOIN {fq_table("entities")} e ON ue.entity_id = e.id WHERE ue.unit_id = $1 """, row["id"], ) entities = [r["canonical_name"] for r in entities_rows] result = { "id": str(row["id"]), "text": row["text"], "context": row["context"] if row["context"] else "", "date": row["event_date"].isoformat() if row["event_date"] else "", "type": row["fact_type"], "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None, "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None, "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None, "entities": entities, "document_id": row["document_id"] if row["document_id"] else None, "chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None, "tags": row["tags"] if row["tags"] else [], } # For observations, include source_memory_ids and fetch source_memories if row["fact_type"] == "observation" and row["source_memory_ids"]: source_ids = row["source_memory_ids"] result["source_memory_ids"] = [str(sid) for sid in source_ids] # Fetch source memories source_rows = await conn.fetch( f""" SELECT id, text, fact_type, context, occurred_start, mentioned_at FROM {fq_table("memory_units")} WHERE id = ANY($1::uuid[]) ORDER BY mentioned_at DESC NULLS LAST """, source_ids, ) result["source_memories"] = [ { "id": str(r["id"]), "text": r["text"], "type": r["fact_type"], "context": r["context"], "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None, "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None, } for r in source_rows ] return result async def list_documents( self, bank_id: str, *, search_query: str | None = None, limit: int = 100, offset: int = 0, request_context: "RequestContext", ): """ List documents with optional search and pagination. Args: bank_id: bank ID (required) search_query: Search in document ID limit: Maximum number of results offset: Offset for pagination request_context: Request context for authentication. Returns: Dict with items (list of documents without original_text) and total count """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build query conditions query_conditions = [] query_params = [] param_count = 0 param_count += 1 query_conditions.append(f"bank_id = ${param_count}") query_params.append(bank_id) if search_query: # Search in document ID param_count += 1 query_conditions.append(f"id ILIKE ${param_count}") query_params.append(f"%{search_query}%") where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else "" # Get total count count_query = f""" SELECT COUNT(*) as total FROM {fq_table("documents")} {where_clause} """ count_result = await conn.fetchrow(count_query, *query_params) total = count_result["total"] # Get documents with limit and offset (without original_text for performance) param_count += 1 limit_param = f"${param_count}" query_params.append(limit) param_count += 1 offset_param = f"${param_count}" query_params.append(offset) documents = await conn.fetch( f""" SELECT id, bank_id, content_hash, created_at, updated_at, LENGTH(original_text) as text_length, retain_params, tags FROM {fq_table("documents")} {where_clause} ORDER BY created_at DESC LIMIT {limit_param} OFFSET {offset_param} """, *query_params, ) # Get memory unit count for each document if documents: doc_ids = [(row["id"], row["bank_id"]) for row in documents] # Create placeholders for the query placeholders = [] params_for_count = [] for i, (doc_id, bank_id_val) in enumerate(doc_ids): idx_doc = i * 2 + 1 idx_agent = i * 2 + 2 placeholders.append(f"(document_id = ${idx_doc} AND bank_id = ${idx_agent})") params_for_count.extend([doc_id, bank_id_val]) where_clause_count = " OR ".join(placeholders) unit_counts = await conn.fetch( f""" SELECT document_id, bank_id, COUNT(*) as unit_count FROM {fq_table("memory_units")} WHERE {where_clause_count} GROUP BY document_id, bank_id """, *params_for_count, ) else: unit_counts = [] # Build count mapping count_map = {(row["document_id"], row["bank_id"]): row["unit_count"] for row in unit_counts} # Build result items items = [] for row in documents: doc_id = row["id"] bank_id_val = row["bank_id"] unit_count = count_map.get((doc_id, bank_id_val), 0) items.append( { "id": doc_id, "bank_id": bank_id_val, "content_hash": row["content_hash"], "created_at": row["created_at"].isoformat() if row["created_at"] else "", "updated_at": row["updated_at"].isoformat() if row["updated_at"] else "", "text_length": row["text_length"] or 0, "memory_unit_count": unit_count, "retain_params": row["retain_params"] if row["retain_params"] else None, "tags": row["tags"] if row["tags"] else [], } ) return {"items": items, "total": total, "limit": limit, "offset": offset} async def get_chunk( self, chunk_id: str, *, request_context: "RequestContext", ): """ Get a specific chunk by its ID. Args: chunk_id: Chunk ID (format: bank_id_document_id_chunk_index) request_context: Request context for authentication. Returns: Dict with chunk details including chunk_text, or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: chunk = await conn.fetchrow( f""" SELECT chunk_id, document_id, bank_id, chunk_index, chunk_text, created_at FROM {fq_table("chunks")} WHERE chunk_id = $1 """, chunk_id, ) if not chunk: return None return { "chunk_id": chunk["chunk_id"], "document_id": chunk["document_id"], "bank_id": chunk["bank_id"], "chunk_index": chunk["chunk_index"], "chunk_text": chunk["chunk_text"], "created_at": chunk["created_at"].isoformat() if chunk["created_at"] else "", } # ==================== bank profile Methods ==================== async def get_bank_profile( self, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """ Get bank profile (name, disposition + mission). Auto-creates agent with default values if not exists. Args: bank_id: bank IDentifier request_context: Request context for authentication. Returns: Dict with name, disposition traits, and mission """ await self._authenticate_tenant(request_context) pool = await self._get_pool() profile = await bank_utils.get_bank_profile(pool, bank_id) disposition = profile["disposition"] return { "bank_id": bank_id, "name": profile["name"], "disposition": disposition, "mission": profile["mission"], } async def update_bank_disposition( self, bank_id: str, disposition: dict[str, int], *, request_context: "RequestContext", ) -> None: """ Update bank disposition traits. Args: bank_id: bank IDentifier disposition: Dict with skepticism, literalism, empathy (all 1-5) request_context: Request context for authentication. """ await self._authenticate_tenant(request_context) pool = await self._get_pool() await bank_utils.update_bank_disposition(pool, bank_id, disposition) async def set_bank_mission( self, bank_id: str, mission: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """ Set the mission for a bank. Args: bank_id: bank IDentifier mission: The mission text request_context: Request context for authentication. Returns: Dict with bank_id and mission. """ await self._authenticate_tenant(request_context) pool = await self._get_pool() await bank_utils.set_bank_mission(pool, bank_id, mission) return {"bank_id": bank_id, "mission": mission} async def merge_bank_mission( self, bank_id: str, new_info: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """ Merge new mission information with existing mission using LLM. Normalizes to first person ("I") and resolves conflicts. Args: bank_id: bank IDentifier new_info: New mission information to add/merge request_context: Request context for authentication. Returns: Dict with 'mission' (str) key """ await self._authenticate_tenant(request_context) pool = await self._get_pool() return await bank_utils.merge_bank_mission(pool, self._reflect_llm_config, bank_id, new_info) async def list_banks( self, *, request_context: "RequestContext", ) -> list[dict[str, Any]]: """ List all agents in the system. Args: request_context: Request context for authentication. Returns: List of dicts with bank_id, name, disposition, mission, created_at, updated_at """ await self._authenticate_tenant(request_context) pool = await self._get_pool() return await bank_utils.list_banks(pool) # ==================== Reflect Methods ==================== async def reflect_async( self, bank_id: str, query: str, *, budget: Budget | None = None, context: str | None = None, max_tokens: int = 4096, response_schema: dict | None = None, request_context: "RequestContext", tags: list[str] | None = None, tags_match: TagsMatch = "any", exclude_mental_model_ids: list[str] | None = None, _skip_span: bool = False, ) -> ReflectResult: """ Reflect and formulate an answer using an agentic loop with tools. The reflect agent iteratively uses tools to: 1. lookup: Get mental models (synthesized knowledge) 2. recall: Search facts (semantic + temporal retrieval) 3. learn: Create/update mental models with new insights 4. expand: Get chunk/document context for memories The agent starts with empty context and must call tools to gather information. On the last iteration, tools are removed to force a final text response. Args: bank_id: bank identifier query: Question to answer budget: Budget level (currently unused, reserved for future) context: Additional context string to include in agent prompt max_tokens: Max tokens (currently unused, reserved for future) response_schema: Optional JSON Schema for structured output (not yet supported) tags: Optional tags to filter memories tags_match: How to match tags - "any" (OR), "all" (AND) exclude_mental_model_ids: Optional list of mental model IDs to exclude from search (used when refreshing a mental model to avoid circular reference) Returns: ReflectResult containing: - text: Plain text answer - based_on: Empty dict (agent retrieves facts dynamically) - structured_output: None (not yet supported for agentic reflect) """ # Use cached LLM config if self._reflect_llm_config is None: raise ValueError("Memory LLM API key not set. Set HINDSIGHT_API_LLM_API_KEY environment variable.") # Authenticate tenant and set schema in context (for fq_table()) await self._authenticate_tenant(request_context) # Validate operation if validator is configured if self._operation_validator: from hindsight_api.extensions import ReflectContext ctx = ReflectContext( bank_id=bank_id, query=query, request_context=request_context, budget=budget, context=context, ) await self._validate_operation(self._operation_validator.validate_reflect(ctx)) reflect_start = time.time() reflect_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}" tags_info = f", tags={tags} ({tags_match})" if tags else "" logger.info(f"[REFLECT {reflect_id}] Starting agentic reflect for query: {query[:50]}...{tags_info}") # Get bank profile for agent identity profile = await self.get_bank_profile(bank_id, request_context=request_context) # NOTE: Mental models are NOT pre-loaded to keep the initial prompt small. # The agent can call lookup() to list available models if needed. # This is critical for banks with many mental models to avoid huge prompts. # Compute max iterations based on budget config = get_config() base_max_iterations = config.reflect_max_iterations # Budget multipliers: low=0.5x, mid=1x, high=2x budget_multipliers = {Budget.LOW: 0.5, Budget.MID: 1.0, Budget.HIGH: 2.0} effective_budget = budget or Budget.LOW max_iterations = max(1, int(base_max_iterations * budget_multipliers.get(effective_budget, 1.0))) # Run agentic loop - acquire connections only when needed for DB operations # (not held during LLM calls which can be slow) pool = await self._get_pool() # Get bank stats for freshness info bank_stats = await self.get_bank_stats(bank_id, request_context=request_context) last_consolidated_at = bank_stats.last_consolidated_at if hasattr(bank_stats, "last_consolidated_at") else None pending_consolidation = bank_stats.pending_consolidation if hasattr(bank_stats, "pending_consolidation") else 0 # Create tool callbacks that acquire connections only when needed from .retain import embedding_utils async def search_mental_models_fn(q: str, max_results: int = 5) -> dict[str, Any]: # Generate embedding for the query embeddings = await embedding_utils.generate_embeddings_batch(self.embeddings, [q]) query_embedding = embeddings[0] async with pool.acquire() as conn: return await tool_search_mental_models( conn, bank_id, q, query_embedding, max_results=max_results, tags=tags, tags_match=tags_match, exclude_ids=exclude_mental_model_ids, ) async def search_observations_fn(q: str, max_tokens: int = 5000) -> dict[str, Any]: return await tool_search_observations( self, bank_id, q, request_context, max_tokens=max_tokens, tags=tags, tags_match=tags_match, last_consolidated_at=last_consolidated_at, pending_consolidation=pending_consolidation, ) async def recall_fn(q: str, max_tokens: int = 4096) -> dict[str, Any]: return await tool_recall( self, bank_id, q, request_context, max_tokens=max_tokens, tags=tags, tags_match=tags_match ) async def expand_fn(memory_ids: list[str], depth: str) -> dict[str, Any]: async with pool.acquire() as conn: return await tool_expand(conn, bank_id, memory_ids, depth) # Load directives from the dedicated directives table # Directives are hard rules that must be followed in all responses # Use isolation_mode=True to prevent tag-scoped directives from leaking into untagged operations directives_raw = await self.list_directives( bank_id=bank_id, tags=tags, tags_match=tags_match, active_only=True, request_context=request_context, isolation_mode=True, ) # Convert directive format to the expected format for reflect agent # The agent expects: name, description (optional), observations (list of {title, content}) directives = [ { "name": d["name"], "description": d["content"], # Use content as description "observations": [], # Directives use content directly, not observations } for d in directives_raw ] if directives: logger.info(f"[REFLECT {reflect_id}] Loaded {len(directives)} directives") # Check if the bank has any mental models async with pool.acquire() as conn: mental_model_count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('mental_models')} WHERE bank_id = $1", bank_id, ) has_mental_models = mental_model_count > 0 if has_mental_models: logger.info(f"[REFLECT {reflect_id}] Bank has {mental_model_count} mental models") # Run the agent with parent span for reflect operation (skip if called from another operation) if not _skip_span: span_context = create_operation_span("reflect", bank_id) span_context.__enter__() else: span_context = None try: agent_result = await run_reflect_agent( llm_config=self._reflect_llm_config, bank_id=bank_id, query=query, bank_profile=profile, search_mental_models_fn=search_mental_models_fn, search_observations_fn=search_observations_fn, recall_fn=recall_fn, expand_fn=expand_fn, context=context, max_iterations=max_iterations, max_tokens=max_tokens, response_schema=response_schema, directives=directives, has_mental_models=has_mental_models, budget=effective_budget, ) total_time = time.time() - reflect_start logger.info( f"[REFLECT {reflect_id}] Complete: {len(agent_result.text)} chars, " f"{agent_result.iterations} iterations, {agent_result.tools_called} tool calls | {total_time:.3f}s" ) # Convert agent tool trace to ToolCallTrace objects tool_trace_result = [ ToolCallTrace( tool=tc.tool, reason=tc.reason, input=tc.input, output=tc.output, duration_ms=tc.duration_ms, iteration=tc.iteration, ) for tc in agent_result.tool_trace ] # Convert agent LLM trace to LLMCallTrace objects llm_trace_result = [ LLMCallTrace(scope=lc.scope, duration_ms=lc.duration_ms) for lc in agent_result.llm_trace ] # Extract memories from recall tool outputs - only include memories the agent actually used # agent_result.used_memory_ids contains validated IDs from the done action used_memory_ids_set = set(agent_result.used_memory_ids) if agent_result.used_memory_ids else set() # based_on stores facts, mental models, and directives # Note: directives list stores raw directive dicts (not MemoryFact), which will be converted to Directive objects based_on: dict[str, list[MemoryFact] | list[dict[str, Any]]] = { "world": [], "experience": [], "opinion": [], "observation": [], "mental-models": [], "directives": [], } seen_memory_ids: set[str] = set() for tc in agent_result.tool_trace: if tc.tool == "recall" and "memories" in tc.output: for memory_data in tc.output["memories"]: memory_id = memory_data.get("id") # Only include memories that the agent declared as used (or all if none specified) if memory_id and memory_id not in seen_memory_ids: if used_memory_ids_set and memory_id not in used_memory_ids_set: continue # Skip memories not actually used by the agent seen_memory_ids.add(memory_id) fact_type = memory_data.get("type", "world") if fact_type in based_on: based_on[fact_type].append( MemoryFact( id=memory_id, text=memory_data.get("text", ""), fact_type=fact_type, context=None, occurred_start=memory_data.get("occurred"), occurred_end=memory_data.get("occurred"), ) ) # Extract mental models from tool outputs - only include models the agent actually used # agent_result.used_mental_model_ids contains validated IDs from the done action used_model_ids_set = ( set(agent_result.used_mental_model_ids) if agent_result.used_mental_model_ids else set() ) based_on["mental-models"] = [] seen_model_ids: set[str] = set() for tc in agent_result.tool_trace: if tc.tool == "get_mental_model": # Single model lookup (with full details) if tc.output.get("found") and "model" in tc.output: model = tc.output["model"] model_id = model.get("id") if model_id and model_id not in seen_model_ids: # Only include models that the agent declared as used (or all if none specified) if used_model_ids_set and model_id not in used_model_ids_set: continue # Skip models not actually used by the agent seen_model_ids.add(model_id) # Add to based_on as MemoryFact with type "mental-models" model_name = model.get("name", "") model_summary = model.get("summary") or model.get("description", "") based_on["mental-models"].append( MemoryFact( id=model_id, text=f"{model_name}: {model_summary}", fact_type="mental-models", context=f"{model.get('type', 'concept')} ({model.get('subtype', 'structural')})", occurred_start=None, occurred_end=None, ) ) elif tc.tool == "search_mental_models": # Search mental models - include all returned models (filtered by used_model_ids_set if specified) for model in tc.output.get("mental_models", []): model_id = model.get("id") if model_id and model_id not in seen_model_ids: # Only include models that the agent declared as used (or all if none specified) if used_model_ids_set and model_id not in used_model_ids_set: continue # Skip models not actually used by the agent seen_model_ids.add(model_id) # Add to based_on as MemoryFact with type "mental-models" model_name = model.get("name", "") model_summary = model.get("summary") or model.get("description", "") based_on["mental-models"].append( MemoryFact( id=model_id, text=f"{model_name}: {model_summary}", fact_type="mental-models", context=f"{model.get('type', 'concept')} ({model.get('subtype', 'structural')})", occurred_start=None, occurred_end=None, ) ) elif tc.tool == "search_mental_models": # Search mental models - include all returned mental models (filtered by used_mental_model_ids_set if specified) used_mental_model_ids_set = ( set(agent_result.used_mental_model_ids) if agent_result.used_mental_model_ids else set() ) for mental_model in tc.output.get("mental_models", []): mental_model_id = mental_model.get("id") if mental_model_id and mental_model_id not in seen_model_ids: # Only include mental models that the agent declared as used (or all if none specified) if used_mental_model_ids_set and mental_model_id not in used_mental_model_ids_set: continue # Skip mental models not actually used by the agent seen_model_ids.add(mental_model_id) # Add to based_on as MemoryFact with type "mental-models" (mental models are synthesized knowledge) mental_model_name = mental_model.get("name", "") mental_model_content = mental_model.get("content", "") based_on["mental-models"].append( MemoryFact( id=mental_model_id, text=f"{mental_model_name}: {mental_model_content}", fact_type="mental-models", context="mental model (user-curated)", occurred_start=None, occurred_end=None, ) ) # List all models lookup - don't add to based_on (too verbose, just a listing) # Add directives to based_on["directives"] # Store raw directive dicts (with id, name, content) for http.py to convert to ReflectDirective for directive_raw in directives_raw: based_on["directives"].append( { "id": directive_raw["id"], "name": directive_raw["name"], "content": directive_raw["content"], } ) # Build directives_applied from agent result from hindsight_api.engine.response_models import DirectiveRef directives_applied_result = [ DirectiveRef(id=d.id, name=d.name, content=d.content) for d in agent_result.directives_applied ] # Convert agent usage to TokenUsage format from hindsight_api.engine.response_models import TokenUsage usage = TokenUsage( input_tokens=agent_result.usage.input_tokens, output_tokens=agent_result.usage.output_tokens, total_tokens=agent_result.usage.total_tokens, ) # Return response (compatible with existing API) result = ReflectResult( text=agent_result.text, based_on=based_on, structured_output=agent_result.structured_output, usage=usage, tool_trace=tool_trace_result, llm_trace=llm_trace_result, directives_applied=directives_applied_result, ) # Call post-operation hook if validator is configured if self._operation_validator: from hindsight_api.extensions.operation_validator import ReflectResultContext result_ctx = ReflectResultContext( bank_id=bank_id, query=query, request_context=request_context, budget=budget, context=context, result=result, success=True, error=None, ) try: await self._operation_validator.on_reflect_complete(result_ctx) except Exception as e: logger.warning(f"Post-reflect hook error (non-fatal): {e}") return result finally: if span_context: span_context.__exit__(None, None, None) async def list_entities( self, bank_id: str, *, limit: int = 100, offset: int = 0, request_context: "RequestContext", ) -> dict[str, Any]: """ List all entities for a bank with pagination. Args: bank_id: bank IDentifier limit: Maximum number of entities to return offset: Offset for pagination request_context: Request context for authentication. Returns: Dict with items, total, limit, offset """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Get total count total_row = await conn.fetchrow( f""" SELECT COUNT(*) as total FROM {fq_table("entities")} WHERE bank_id = $1 """, bank_id, ) total = total_row["total"] if total_row else 0 # Get paginated entities rows = await conn.fetch( f""" SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata FROM {fq_table("entities")} WHERE bank_id = $1 ORDER BY mention_count DESC, last_seen DESC, id ASC LIMIT $2 OFFSET $3 """, bank_id, limit, offset, ) entities = [] for row in rows: # Handle metadata - may be dict, JSON string, or None metadata = row["metadata"] if metadata is None: metadata = {} elif isinstance(metadata, str): import json try: metadata = json.loads(metadata) except json.JSONDecodeError: metadata = {} entities.append( { "id": str(row["id"]), "canonical_name": row["canonical_name"], "mention_count": row["mention_count"], "first_seen": row["first_seen"].isoformat() if row["first_seen"] else None, "last_seen": row["last_seen"].isoformat() if row["last_seen"] else None, "metadata": metadata, } ) return { "items": entities, "total": total, "limit": limit, "offset": offset, } async def list_tags( self, bank_id: str, *, pattern: str | None = None, limit: int = 100, offset: int = 0, request_context: "RequestContext", ) -> dict[str, Any]: """ List all unique tags for a bank with usage counts. Use this to discover available tags or expand wildcard patterns. Supports '*' as wildcard for flexible matching (case-insensitive): - 'user:*' matches user:alice, user:bob - '*-admin' matches role-admin, super-admin - 'env*-prod' matches env-prod, environment-prod Args: bank_id: Bank identifier pattern: Wildcard pattern to filter tags (use '*' as wildcard, case-insensitive) limit: Maximum number of tags to return offset: Offset for pagination request_context: Request context for authentication. Returns: Dict with items (list of {tag, count}), total, limit, offset """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build pattern filter if provided (convert * to % for ILIKE) pattern_clause = "" params: list[Any] = [bank_id] if pattern: # Convert wildcard pattern: * -> % for SQL ILIKE sql_pattern = pattern.replace("*", "%") pattern_clause = "AND tag ILIKE $2" params.append(sql_pattern) # Get total count of distinct tags matching pattern total_row = await conn.fetchrow( f""" SELECT COUNT(DISTINCT tag) as total FROM {fq_table("memory_units")}, unnest(tags) AS tag WHERE bank_id = $1 AND tags IS NOT NULL AND tags != '{{}}' {pattern_clause} """, *params, ) total = total_row["total"] if total_row else 0 # Get paginated tags with counts, ordered by frequency limit_param = len(params) + 1 offset_param = len(params) + 2 params.extend([limit, offset]) rows = await conn.fetch( f""" SELECT tag, COUNT(*) as count FROM {fq_table("memory_units")}, unnest(tags) AS tag WHERE bank_id = $1 AND tags IS NOT NULL AND tags != '{{}}' {pattern_clause} GROUP BY tag ORDER BY count DESC, tag ASC LIMIT ${limit_param} OFFSET ${offset_param} """, *params, ) items = [{"tag": row["tag"], "count": row["count"]} for row in rows] return { "items": items, "total": total, "limit": limit, "offset": offset, } async def get_entity_state( self, bank_id: str, entity_id: str, entity_name: str, *, limit: int = 10, request_context: "RequestContext", ) -> EntityState: """ Get the current state of an entity. NOTE: Entity observations/summaries have been moved to mental models. This method returns an entity with empty observations. Args: bank_id: bank IDentifier entity_id: Entity UUID entity_name: Canonical name of the entity limit: Maximum number of observations to include (kept for backwards compat) request_context: Request context for authentication. Returns: EntityState with empty observations (summaries now in mental models) """ await self._authenticate_tenant(request_context) return EntityState(entity_id=entity_id, canonical_name=entity_name, observations=[]) # ========================================================================= # Statistics & Operations (for HTTP API layer) # ========================================================================= async def get_bank_stats( self, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """Get statistics about memory nodes and links for a bank.""" await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Get node counts by fact_type node_stats = await conn.fetch( f""" SELECT fact_type, COUNT(*) as count FROM {fq_table("memory_units")} WHERE bank_id = $1 GROUP BY fact_type """, bank_id, ) # Get link counts by link_type link_stats = await conn.fetch( f""" SELECT ml.link_type, COUNT(*) as count FROM {fq_table("memory_links")} ml JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id WHERE mu.bank_id = $1 GROUP BY ml.link_type """, bank_id, ) # Get link counts by fact_type (from nodes) link_fact_type_stats = await conn.fetch( f""" SELECT mu.fact_type, COUNT(*) as count FROM {fq_table("memory_links")} ml JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id WHERE mu.bank_id = $1 GROUP BY mu.fact_type """, bank_id, ) # Get link counts by fact_type AND link_type link_breakdown_stats = await conn.fetch( f""" SELECT mu.fact_type, ml.link_type, COUNT(*) as count FROM {fq_table("memory_links")} ml JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id WHERE mu.bank_id = $1 GROUP BY mu.fact_type, ml.link_type """, bank_id, ) # Get pending and failed operations counts ops_stats = await conn.fetch( f""" SELECT status, COUNT(*) as count FROM {fq_table("async_operations")} WHERE bank_id = $1 GROUP BY status """, bank_id, ) return { "bank_id": bank_id, "node_counts": {row["fact_type"]: row["count"] for row in node_stats}, "link_counts": {row["link_type"]: row["count"] for row in link_stats}, "link_counts_by_fact_type": {row["fact_type"]: row["count"] for row in link_fact_type_stats}, "link_breakdown": [ {"fact_type": row["fact_type"], "link_type": row["link_type"], "count": row["count"]} for row in link_breakdown_stats ], "operations": {row["status"]: row["count"] for row in ops_stats}, } async def get_entity( self, bank_id: str, entity_id: str, *, request_context: "RequestContext", ) -> dict[str, Any] | None: """Get entity details including metadata and observations.""" await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: entity_row = await conn.fetchrow( f""" SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata FROM {fq_table("entities")} WHERE bank_id = $1 AND id = $2 """, bank_id, uuid.UUID(entity_id), ) if not entity_row: return None return { "id": str(entity_row["id"]), "canonical_name": entity_row["canonical_name"], "mention_count": entity_row["mention_count"], "first_seen": entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None, "last_seen": entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None, "metadata": entity_row["metadata"] or {}, "observations": [], } def _parse_observations(self, observations_raw: list): """Parse raw observation dicts into typed Observation models. Returns list of Observation models with computed trend/evidence_span/evidence_count. """ from .reflect.observations import Observation, ObservationEvidence observations: list[Observation] = [] for obs in observations_raw: if not isinstance(obs, dict): continue try: parsed = Observation( title=obs.get("title", ""), content=obs.get("content", ""), evidence=[ ObservationEvidence( memory_id=ev.get("memory_id", ""), quote=ev.get("quote", ""), relevance=ev.get("relevance", ""), timestamp=ev.get("timestamp"), ) for ev in obs.get("evidence", []) if isinstance(ev, dict) ], created_at=obs.get("created_at"), ) observations.append(parsed) except Exception as e: logger.warning(f"Failed to parse observation: {e}") continue return observations async def _count_memories_since( self, bank_id: str, since_timestamp: str | None, pool=None, ) -> int: """ Count memories created after a given timestamp. Args: bank_id: Bank identifier since_timestamp: ISO timestamp string. If None, returns total count. pool: Optional database pool (uses default if not provided) Returns: Number of memories created since the timestamp """ if pool is None: pool = await self._get_pool() async with acquire_with_retry(pool) as conn: if since_timestamp: # Parse the timestamp from datetime import datetime try: ts = datetime.fromisoformat(since_timestamp.replace("Z", "+00:00")) except ValueError: # Invalid timestamp, return total count ts = None if ts: count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND created_at > $2", bank_id, ts, ) return count or 0 # No timestamp or invalid, return total count count = await conn.fetchval( f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1", bank_id, ) return count or 0 async def _invalidate_facts_from_mental_models( self, conn, bank_id: str, fact_ids: list[str], ) -> int: """ Remove fact IDs from observation source_memory_ids when memories are deleted. Observations are stored in memory_units with fact_type='observation' and have a source_memory_ids column (UUID[]) tracking their source memories. Args: conn: Database connection bank_id: Bank identifier fact_ids: List of fact IDs to remove from observations Returns: Number of observations updated """ if not fact_ids: return 0 # Convert string IDs to UUIDs for the array comparison import uuid as uuid_module fact_uuids = [uuid_module.UUID(fid) for fid in fact_ids] # Update observations (memory_units with fact_type='observation') # by removing the deleted fact IDs from source_memory_ids # Use array subtraction: source_memory_ids - deleted_ids result = await conn.execute( f""" UPDATE {fq_table("memory_units")} SET source_memory_ids = ( SELECT COALESCE(array_agg(elem), ARRAY[]::uuid[]) FROM unnest(source_memory_ids) AS elem WHERE elem != ALL($2::uuid[]) ), updated_at = NOW() WHERE bank_id = $1 AND fact_type = 'observation' AND source_memory_ids && $2::uuid[] """, bank_id, fact_uuids, ) # Parse the result to get number of updated rows updated_count = int(result.split()[-1]) if result and "UPDATE" in result else 0 if updated_count > 0: logger.info( f"[OBSERVATIONS] Invalidated {len(fact_ids)} fact IDs from {updated_count} observations in bank {bank_id}" ) return updated_count # ========================================================================= # MENTAL MODELS (CONSOLIDATED) - Read-only access to auto-consolidated mental models # ========================================================================= async def list_mental_models_consolidated( self, bank_id: str, *, tags: list[str] | None = None, tags_match: str = "any", limit: int = 100, offset: int = 0, request_context: "RequestContext", ) -> list[dict[str, Any]]: """List auto-consolidated observations for a bank. Observations are stored in memory_units with fact_type='observation'. They are automatically created and updated by the consolidation engine. Args: bank_id: Bank identifier tags: Optional tags to filter by tags_match: How to match tags - 'any', 'all', or 'exact' limit: Maximum number of results offset: Offset for pagination request_context: Request context for authentication Returns: List of observation dicts """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build tag filter tag_filter = "" params: list[Any] = [bank_id, limit, offset] if tags: if tags_match == "all": tag_filter = " AND tags @> $4::varchar[]" elif tags_match == "exact": tag_filter = " AND tags = $4::varchar[]" else: # any tag_filter = " AND tags && $4::varchar[]" params.append(tags) rows = await conn.fetch( f""" SELECT id, bank_id, text, proof_count, history, tags, source_memory_ids, created_at, updated_at FROM {fq_table("memory_units")} WHERE bank_id = $1 AND fact_type = 'observation' {tag_filter} ORDER BY updated_at DESC NULLS LAST LIMIT $2 OFFSET $3 """, *params, ) return [self._row_to_observation_consolidated(row) for row in rows] async def get_observation_consolidated( self, bank_id: str, observation_id: str, *, include_source_memories: bool = True, request_context: "RequestContext", ) -> dict[str, Any] | None: """Get a single observation by ID. Args: bank_id: Bank identifier observation_id: Observation ID include_source_memories: Whether to include full source memory details request_context: Request context for authentication Returns: Observation dict or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" SELECT id, bank_id, text, proof_count, history, tags, source_memory_ids, created_at, updated_at FROM {fq_table("memory_units")} WHERE bank_id = $1 AND id = $2 AND fact_type = 'observation' """, bank_id, observation_id, ) if not row: return None result = self._row_to_observation_consolidated(row) # Fetch source memories if requested and source_memory_ids exist if include_source_memories and result.get("source_memory_ids"): source_ids = [uuid.UUID(sid) if isinstance(sid, str) else sid for sid in result["source_memory_ids"]] source_rows = await conn.fetch( f""" SELECT id, text, fact_type, context, occurred_start, mentioned_at FROM {fq_table("memory_units")} WHERE id = ANY($1::uuid[]) ORDER BY mentioned_at DESC NULLS LAST """, source_ids, ) result["source_memories"] = [ { "id": str(r["id"]), "text": r["text"], "type": r["fact_type"], "context": r["context"], "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None, "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None, } for r in source_rows ] return result def _row_to_observation_consolidated(self, row: Any) -> dict[str, Any]: """Convert a database row to an observation dict.""" import json history = row["history"] if isinstance(history, str): history = json.loads(history) elif history is None: history = [] # Convert source_memory_ids to strings source_memory_ids = row.get("source_memory_ids") or [] source_memory_ids = [str(sid) for sid in source_memory_ids] return { "id": str(row["id"]), "bank_id": row["bank_id"], "text": row["text"], "proof_count": row["proof_count"] or 1, "history": history, "tags": row["tags"] or [], "source_memory_ids": source_memory_ids, "source_memories": [], # Populated separately when fetching full details "created_at": row["created_at"].isoformat() if row["created_at"] else None, "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None, } # ========================================================================= # MENTAL MODELS CRUD # ========================================================================= async def list_mental_models( self, bank_id: str, *, tags: list[str] | None = None, tags_match: str = "any", limit: int = 100, offset: int = 0, request_context: "RequestContext", ) -> list[dict[str, Any]]: """List pinned mental models for a bank. Args: bank_id: Bank identifier tags: Optional tags to filter by tags_match: How to match tags - 'any', 'all', or 'exact' limit: Maximum number of results offset: Offset for pagination request_context: Request context for authentication Returns: List of pinned mental model dicts """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build tag filter tag_filter = "" params: list[Any] = [bank_id, limit, offset] if tags: if tags_match == "all": tag_filter = " AND tags @> $4::varchar[]" elif tags_match == "exact": tag_filter = " AND tags = $4::varchar[]" else: # any tag_filter = " AND tags && $4::varchar[]" params.append(tags) rows = await conn.fetch( f""" SELECT id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, max_tokens, trigger FROM {fq_table("mental_models")} WHERE bank_id = $1 {tag_filter} ORDER BY last_refreshed_at DESC LIMIT $2 OFFSET $3 """, *params, ) return [self._row_to_mental_model(row) for row in rows] async def get_mental_model( self, bank_id: str, mental_model_id: str, *, request_context: "RequestContext", ) -> dict[str, Any] | None: """Get a single pinned mental model by ID. Args: bank_id: Bank identifier mental_model_id: Pinned mental model UUID request_context: Request context for authentication Returns: Pinned mental model dict or None if not found """ await self._authenticate_tenant(request_context) # Pre-operation validation (credit check / usage metering) if self._operation_validator: from hindsight_api.extensions.operation_validator import MentalModelGetContext ctx = MentalModelGetContext( bank_id=bank_id, mental_model_id=mental_model_id, request_context=request_context, ) await self._validate_operation(self._operation_validator.validate_mental_model_get(ctx)) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" SELECT id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, max_tokens, trigger FROM {fq_table("mental_models")} WHERE bank_id = $1 AND id = $2 """, bank_id, mental_model_id, ) result = self._row_to_mental_model(row) if row else None # Post-operation hook (usage recording) if result and self._operation_validator: from hindsight_api.extensions.operation_validator import MentalModelGetResult content = result.get("content", "") output_tokens = len(content) // 4 if content else 0 result_ctx = MentalModelGetResult( bank_id=bank_id, mental_model_id=mental_model_id, request_context=request_context, output_tokens=output_tokens, success=True, ) try: await self._operation_validator.on_mental_model_get_complete(result_ctx) except Exception as hook_err: logger.warning(f"Post-mental-model-get hook error (non-fatal): {hook_err}") return result async def create_mental_model( self, bank_id: str, name: str, source_query: str, content: str, *, mental_model_id: str | None = None, tags: list[str] | None = None, max_tokens: int | None = None, trigger: dict[str, Any] | None = None, request_context: "RequestContext", ) -> dict[str, Any]: """Create a new pinned mental model. Args: bank_id: Bank identifier name: Human-readable name for the mental model source_query: The query that generated this mental model content: The synthesized content mental_model_id: Optional UUID for the mental model (auto-generated if not provided) tags: Optional tags for scoped visibility max_tokens: Token limit for content generation during refresh trigger: Trigger settings (e.g., refresh_after_consolidation) request_context: Request context for authentication Returns: The created pinned mental model dict """ await self._authenticate_tenant(request_context) pool = await self._get_pool() # Generate embedding for the content embedding_text = f"{name} {content}" embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text]) # Convert embedding to string for asyncpg vector type embedding_str = str(embedding[0]) if embedding else None async with acquire_with_retry(pool) as conn: if mental_model_id: row = await conn.fetchrow( f""" INSERT INTO {fq_table("mental_models")} (id, bank_id, name, source_query, content, embedding, tags, max_tokens, trigger) VALUES ($1, $2, $3, $4, $5, $6, $7, COALESCE($8, 2048), COALESCE($9, '{{"refresh_after_consolidation": false}}'::jsonb)) RETURNING id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, max_tokens, trigger """, mental_model_id, bank_id, name, source_query, content, embedding_str, tags or [], max_tokens, json.dumps(trigger) if trigger else None, ) else: row = await conn.fetchrow( f""" INSERT INTO {fq_table("mental_models")} (bank_id, name, source_query, content, embedding, tags, max_tokens, trigger) VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, 2048), COALESCE($8, '{{"refresh_after_consolidation": false}}'::jsonb)) RETURNING id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, max_tokens, trigger """, bank_id, name, source_query, content, embedding_str, tags or [], max_tokens, json.dumps(trigger) if trigger else None, ) logger.info(f"[MENTAL_MODELS] Created pinned mental model '{name}' for bank {bank_id}") return self._row_to_mental_model(row) async def refresh_mental_model( self, bank_id: str, mental_model_id: str, *, request_context: "RequestContext", ) -> dict[str, Any] | None: """Refresh a pinned mental model by re-running its source query. This method: 1. Gets the pinned mental model 2. Runs the source_query through reflect 3. Updates the content with the new synthesis 4. Updates last_refreshed_at Args: bank_id: Bank identifier mental_model_id: Pinned mental model UUID request_context: Request context for authentication Returns: Updated pinned mental model dict or None if not found """ await self._authenticate_tenant(request_context) # Get the current mental model mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=request_context) if not mental_model: return None # Create parent span for mental model refresh operation with create_operation_span("mental_model_refresh", bank_id): # SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching # to ensure it can only access other mental models/memories with the SAME tags. # This prevents cross-tenant/cross-user information leakage by excluding untagged content. tags = mental_model.get("tags") tags_match = "all_strict" if tags else "any" # Run reflect with the source query, excluding the mental model being refreshed # Skip creating a nested "hindsight.reflect" span since we already have "hindsight.mental_model_refresh" reflect_result = await self.reflect_async( bank_id=bank_id, query=mental_model["source_query"], request_context=request_context, tags=tags, tags_match=tags_match, exclude_mental_model_ids=[mental_model_id], _skip_span=True, ) # Build reflect_response payload to store # based_on contains MemoryFact objects for most types, but plain dicts for directives based_on_serialized_payload: dict[str, list[dict[str, Any]]] = {} for fact_type, facts in reflect_result.based_on.items(): serialized_facts = [] for fact in facts: if isinstance(fact, dict): # Plain dict (e.g., directives with id, name, content) serialized_facts.append( { "id": str(fact["id"]), "text": fact.get("text", fact.get("content", fact.get("name", ""))), "type": fact_type, "context": fact.get("context", None), } ) else: # MemoryFact object with .id, .text, .context attributes serialized_facts.append( { "id": str(fact.id), "text": fact.text, "type": fact_type, "context": fact.context, } ) based_on_serialized_payload[fact_type] = serialized_facts reflect_response_payload = { "text": reflect_result.text, "based_on": based_on_serialized_payload, "mental_models": [], # Mental models are included in based_on["mental-models"] } # Update the mental model with new content and reflect_response return await self.update_mental_model( bank_id, mental_model_id, content=reflect_result.text, reflect_response=reflect_response_payload, request_context=request_context, ) async def update_mental_model( self, bank_id: str, mental_model_id: str, *, name: str | None = None, content: str | None = None, source_query: str | None = None, max_tokens: int | None = None, tags: list[str] | None = None, trigger: dict[str, Any] | None = None, reflect_response: dict[str, Any] | None = None, request_context: "RequestContext", ) -> dict[str, Any] | None: """Update a pinned mental model. Args: bank_id: Bank identifier mental_model_id: Pinned mental model UUID name: New name (if changing) content: New content (if changing) source_query: New source query (if changing) max_tokens: New max tokens (if changing) tags: New tags (if changing) trigger: New trigger settings (if changing) reflect_response: Full reflect API response payload (if changing) request_context: Request context for authentication Returns: Updated pinned mental model dict or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build dynamic update updates = [] params: list[Any] = [bank_id, mental_model_id] param_idx = 3 if name is not None: updates.append(f"name = ${param_idx}") params.append(name) param_idx += 1 if content is not None: updates.append(f"content = ${param_idx}") params.append(content) param_idx += 1 updates.append("last_refreshed_at = NOW()") # Also update embedding (convert to string for asyncpg vector type) embedding_text = f"{name or ''} {content}" embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text]) if embedding: updates.append(f"embedding = ${param_idx}") params.append(str(embedding[0])) param_idx += 1 if reflect_response is not None: updates.append(f"reflect_response = ${param_idx}") params.append(json.dumps(reflect_response)) param_idx += 1 if source_query is not None: updates.append(f"source_query = ${param_idx}") params.append(source_query) param_idx += 1 if max_tokens is not None: updates.append(f"max_tokens = ${param_idx}") params.append(max_tokens) param_idx += 1 if tags is not None: updates.append(f"tags = ${param_idx}") params.append(tags) param_idx += 1 if trigger is not None: updates.append(f"trigger = ${param_idx}") params.append(json.dumps(trigger)) param_idx += 1 if not updates: return None query = f""" UPDATE {fq_table("mental_models")} SET {", ".join(updates)} WHERE bank_id = $1 AND id = $2 RETURNING id, bank_id, name, source_query, content, tags, last_refreshed_at, created_at, reflect_response, max_tokens, trigger """ row = await conn.fetchrow(query, *params) return self._row_to_mental_model(row) if row else None async def delete_mental_model( self, bank_id: str, mental_model_id: str, *, request_context: "RequestContext", ) -> bool: """Delete a pinned mental model. Args: bank_id: Bank identifier mental_model_id: Pinned mental model UUID request_context: Request context for authentication Returns: True if deleted, False if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: result = await conn.execute( f"DELETE FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2", bank_id, mental_model_id, ) return result == "DELETE 1" def _row_to_mental_model(self, row) -> dict[str, Any]: """Convert a database row to a mental model dict.""" reflect_response = row.get("reflect_response") # Parse JSON string to dict if needed (asyncpg may return JSONB as string) if isinstance(reflect_response, str): try: reflect_response = json.loads(reflect_response) except json.JSONDecodeError: reflect_response = None trigger = row.get("trigger") if isinstance(trigger, str): try: trigger = json.loads(trigger) except json.JSONDecodeError: trigger = None return { "id": str(row["id"]), "bank_id": row["bank_id"], "name": row["name"], "source_query": row["source_query"], "content": row["content"], "tags": row["tags"] or [], "max_tokens": row.get("max_tokens"), "trigger": trigger, "last_refreshed_at": row["last_refreshed_at"].isoformat() if row["last_refreshed_at"] else None, "created_at": row["created_at"].isoformat() if row["created_at"] else None, "reflect_response": reflect_response, } # ========================================================================= # Directives - Hard rules injected into prompts # ========================================================================= async def list_directives( self, bank_id: str, *, tags: list[str] | None = None, tags_match: str = "any", active_only: bool = True, limit: int = 100, offset: int = 0, request_context: "RequestContext", isolation_mode: bool = False, ) -> list[dict[str, Any]]: """List directives for a bank. Args: bank_id: Bank identifier tags: Optional tags to filter by tags_match: How to match tags - 'any', 'all', or 'exact' active_only: Only return active directives (default True) limit: Maximum number of results offset: Offset for pagination request_context: Request context for authentication isolation_mode: When True and tags=None, only return directives with no tags. This prevents tag-scoped directives from leaking into untagged operations. Default False (normal API behavior - returns all directives when tags=None) Returns: List of directive dicts """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build filters from .search.tags import build_tags_where_clause filters = ["bank_id = $1"] params: list[Any] = [bank_id] param_idx = 2 if active_only: filters.append("is_active = TRUE") # Apply tags filter: # - If tags provided: use standard filtering (with strict modes support) # - If tags=None and isolation_mode=True: only include directives with NO tags # (prevents tag-scoped directives from leaking into untagged reflect/refresh) # - If tags=None and isolation_mode=False: no filtering (normal API behavior) if tags: tags_clause, tags_params, param_idx = build_tags_where_clause( tags=tags, param_offset=param_idx, table_alias="", match=tags_match ) if tags_clause: # Remove leading "AND " from clause since we're building filters list filters.append(tags_clause.replace("AND ", "", 1)) params.extend(tags_params) elif isolation_mode: # Isolation mode: only include directives with empty/null tags # This ensures tag-scoped directives don't apply to untagged operations filters.append("(tags IS NULL OR tags = '{}')") params.extend([limit, offset]) rows = await conn.fetch( f""" SELECT id, bank_id, name, content, priority, is_active, tags, created_at, updated_at FROM {fq_table("directives")} WHERE {" AND ".join(filters)} ORDER BY priority DESC, created_at DESC LIMIT ${param_idx} OFFSET ${param_idx + 1} """, *params, ) return [self._row_to_directive(row) for row in rows] async def get_directive( self, bank_id: str, directive_id: str, *, request_context: "RequestContext", ) -> dict[str, Any] | None: """Get a single directive by ID. Args: bank_id: Bank identifier directive_id: Directive UUID request_context: Request context for authentication Returns: Directive dict or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" SELECT id, bank_id, name, content, priority, is_active, tags, created_at, updated_at FROM {fq_table("directives")} WHERE bank_id = $1 AND id = $2 """, bank_id, directive_id, ) return self._row_to_directive(row) if row else None async def create_directive( self, bank_id: str, name: str, content: str, *, priority: int = 0, is_active: bool = True, tags: list[str] | None = None, request_context: "RequestContext", ) -> dict[str, Any]: """Create a new directive. Args: bank_id: Bank identifier name: Human-readable name for the directive content: The directive text to inject into prompts priority: Higher priority directives are injected first (default 0) is_active: Whether this directive is active (default True) tags: Optional tags for filtering request_context: Request context for authentication Returns: The created directive dict """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" INSERT INTO {fq_table("directives")} (bank_id, name, content, priority, is_active, tags) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id, bank_id, name, content, priority, is_active, tags, created_at, updated_at """, bank_id, name, content, priority, is_active, tags or [], ) logger.info(f"[DIRECTIVES] Created directive '{name}' for bank {bank_id}") return self._row_to_directive(row) async def update_directive( self, bank_id: str, directive_id: str, *, name: str | None = None, content: str | None = None, priority: int | None = None, is_active: bool | None = None, tags: list[str] | None = None, request_context: "RequestContext", ) -> dict[str, Any] | None: """Update a directive. Args: bank_id: Bank identifier directive_id: Directive UUID name: New name (optional) content: New content (optional) priority: New priority (optional) is_active: New active status (optional) tags: New tags (optional) request_context: Request context for authentication Returns: Updated directive dict or None if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() # Build update query dynamically updates = ["updated_at = now()"] params: list[Any] = [] param_idx = 1 if name is not None: updates.append(f"name = ${param_idx}") params.append(name) param_idx += 1 if content is not None: updates.append(f"content = ${param_idx}") params.append(content) param_idx += 1 if priority is not None: updates.append(f"priority = ${param_idx}") params.append(priority) param_idx += 1 if is_active is not None: updates.append(f"is_active = ${param_idx}") params.append(is_active) param_idx += 1 if tags is not None: updates.append(f"tags = ${param_idx}") params.append(tags) param_idx += 1 params.extend([bank_id, directive_id]) async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" UPDATE {fq_table("directives")} SET {", ".join(updates)} WHERE bank_id = ${param_idx} AND id = ${param_idx + 1} RETURNING id, bank_id, name, content, priority, is_active, tags, created_at, updated_at """, *params, ) return self._row_to_directive(row) if row else None async def delete_directive( self, bank_id: str, directive_id: str, *, request_context: "RequestContext", ) -> bool: """Delete a directive. Args: bank_id: Bank identifier directive_id: Directive UUID request_context: Request context for authentication Returns: True if deleted, False if not found """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: result = await conn.execute( f"DELETE FROM {fq_table('directives')} WHERE bank_id = $1 AND id = $2", bank_id, directive_id, ) return result == "DELETE 1" def _row_to_directive(self, row) -> dict[str, Any]: """Convert a database row to a directive dict.""" return { "id": str(row["id"]), "bank_id": row["bank_id"], "name": row["name"], "content": row["content"], "priority": row["priority"], "is_active": row["is_active"], "tags": row["tags"] or [], "created_at": row["created_at"].isoformat() if row["created_at"] else None, "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None, } async def list_operations( self, bank_id: str, *, status: str | None = None, limit: int = 20, offset: int = 0, request_context: "RequestContext", ) -> dict[str, Any]: """List async operations for a bank with optional filtering and pagination. Args: bank_id: Bank identifier status: Optional status filter (pending, completed, failed) limit: Maximum number of operations to return (default 20) offset: Number of operations to skip (default 0) request_context: Request context for authentication Returns: Dict with total count and list of operations, sorted by most recent first """ await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: # Build WHERE clause where_conditions = ["bank_id = $1"] params: list[Any] = [bank_id] if status: # Map API status to DB statuses (pending includes processing) if status == "pending": where_conditions.append("status IN ('pending', 'processing')") else: where_conditions.append(f"status = ${len(params) + 1}") params.append(status) where_clause = " AND ".join(where_conditions) # Get total count (with filter) total_row = await conn.fetchrow( f"SELECT COUNT(*) as total FROM {fq_table('async_operations')} WHERE {where_clause}", *params, ) total = total_row["total"] if total_row else 0 # Get operations with pagination operations = await conn.fetch( f""" SELECT operation_id, operation_type, created_at, status, error_message FROM {fq_table("async_operations")} WHERE {where_clause} ORDER BY created_at DESC LIMIT ${len(params) + 1} OFFSET ${len(params) + 2} """, *params, limit, offset, ) return { "total": total, "operations": [ { "id": str(row["operation_id"]), "task_type": row["operation_type"], "items_count": 0, "document_id": None, "created_at": row["created_at"].isoformat(), # Map DB status to API status (processing -> pending for simplicity) "status": "pending" if row["status"] in ("pending", "processing") else row["status"], "error_message": row["error_message"], } for row in operations ], } async def get_operation_status( self, bank_id: str, operation_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """Get the status of a specific async operation. Returns: - status: "pending", "completed", or "failed" - updated_at: last update timestamp - completed_at: completion timestamp (if completed) """ await self._authenticate_tenant(request_context) pool = await self._get_pool() op_uuid = uuid.UUID(operation_id) async with acquire_with_retry(pool) as conn: row = await conn.fetchrow( f""" SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message FROM {fq_table("async_operations")} WHERE operation_id = $1 AND bank_id = $2 """, op_uuid, bank_id, ) if row: # Map DB status to API status (processing -> pending for simplicity) db_status = row["status"] api_status = "pending" if db_status in ("pending", "processing") else db_status return { "operation_id": operation_id, "status": api_status, "operation_type": row["operation_type"], "created_at": row["created_at"].isoformat() if row["created_at"] else None, "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None, "completed_at": row["completed_at"].isoformat() if row["completed_at"] else None, "error_message": row["error_message"], } else: # Operation not found return { "operation_id": operation_id, "status": "not_found", "operation_type": None, "created_at": None, "updated_at": None, "completed_at": None, "error_message": None, } async def cancel_operation( self, bank_id: str, operation_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """Cancel a pending async operation.""" await self._authenticate_tenant(request_context) pool = await self._get_pool() op_uuid = uuid.UUID(operation_id) async with acquire_with_retry(pool) as conn: # Check if operation exists and belongs to this memory bank result = await conn.fetchrow( f"SELECT bank_id FROM {fq_table('async_operations')} WHERE operation_id = $1 AND bank_id = $2", op_uuid, bank_id, ) if not result: raise ValueError(f"Operation {operation_id} not found for bank {bank_id}") # Delete the operation await conn.execute(f"DELETE FROM {fq_table('async_operations')} WHERE operation_id = $1", op_uuid) return { "success": True, "message": f"Operation {operation_id} cancelled", "operation_id": operation_id, "bank_id": bank_id, } async def update_bank( self, bank_id: str, *, name: str | None = None, mission: str | None = None, request_context: "RequestContext", ) -> dict[str, Any]: """Update bank name and/or mission.""" await self._authenticate_tenant(request_context) pool = await self._get_pool() async with acquire_with_retry(pool) as conn: if name is not None: await conn.execute( f""" UPDATE {fq_table("banks")} SET name = $2, updated_at = NOW() WHERE bank_id = $1 """, bank_id, name, ) if mission is not None: await conn.execute( f""" UPDATE {fq_table("banks")} SET mission = $2, updated_at = NOW() WHERE bank_id = $1 """, bank_id, mission, ) # Return updated profile return await self.get_bank_profile(bank_id, request_context=request_context) async def _submit_async_operation( self, bank_id: str, operation_type: str, task_type: str, task_payload: dict[str, Any], *, result_metadata: dict[str, Any] | None = None, dedupe_by_bank: bool = False, ) -> dict[str, Any]: """Generic helper to submit an async operation. Args: bank_id: Bank identifier operation_type: Operation type for the async_operations record (e.g., 'consolidation', 'retain') task_type: Task type for the task payload (e.g., 'consolidation', 'batch_retain') task_payload: Additional task payload fields (operation_id and bank_id are added automatically) result_metadata: Optional metadata to store with the operation record dedupe_by_bank: If True, skip creating a new task if one is already pending for this bank+operation_type Returns: Dict with operation_id and optionally deduplicated=True if an existing task was found """ import json pool = await self._get_pool() # Check for existing pending task if deduplication is enabled # Note: We only check 'pending', not 'processing', because a processing task # uses a watermark from when it started - new memories added after that point # would need another consolidation run to be processed. if dedupe_by_bank: async with acquire_with_retry(pool) as conn: existing = await conn.fetchrow( f""" SELECT operation_id FROM {fq_table("async_operations")} WHERE bank_id = $1 AND operation_type = $2 AND status = 'pending' LIMIT 1 """, bank_id, operation_type, ) if existing: logger.debug( f"{operation_type} task already pending for bank_id={bank_id}, " f"skipping duplicate (existing operation_id={existing['operation_id']})" ) return { "operation_id": str(existing["operation_id"]), "deduplicated": True, } operation_id = uuid.uuid4() # Insert operation record into database async with acquire_with_retry(pool) as conn: await conn.execute( f""" INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata) VALUES ($1, $2, $3, $4) """, operation_id, bank_id, operation_type, json.dumps(result_metadata or {}), ) # Build and submit task payload full_payload = { "type": task_type, "operation_id": str(operation_id), "bank_id": bank_id, **task_payload, } await self._task_backend.submit_task(full_payload) logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}") return { "operation_id": str(operation_id), } async def submit_async_retain( self, bank_id: str, contents: list[dict[str, Any]], *, request_context: "RequestContext", document_tags: list[str] | None = None, ) -> dict[str, Any]: """Submit a batch retain operation to run asynchronously.""" await self._authenticate_tenant(request_context) task_payload: dict[str, Any] = {"contents": contents} if document_tags: task_payload["document_tags"] = document_tags # Pass tenant_id and api_key_id through task payload so the worker # can propagate request context to downstream operations (e.g., # consolidation and mental model refreshes triggered after retain). if request_context.tenant_id: task_payload["_tenant_id"] = request_context.tenant_id if request_context.api_key_id: task_payload["_api_key_id"] = request_context.api_key_id result = await self._submit_async_operation( bank_id=bank_id, operation_type="retain", task_type="batch_retain", task_payload=task_payload, result_metadata={"items_count": len(contents)}, dedupe_by_bank=False, ) result["items_count"] = len(contents) return result async def submit_async_consolidation( self, bank_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """Submit a consolidation operation to run asynchronously. Deduplicates by bank_id - if there's already a pending consolidation for this bank, returns the existing operation_id instead of creating a new one. Args: bank_id: Bank identifier request_context: Request context for authentication Returns: Dict with operation_id """ await self._authenticate_tenant(request_context) # Pass tenant_id and api_key_id through task payload so the worker # can provide request context to extension hooks (e.g., usage metering # for mental model refreshes triggered by consolidation). task_payload: dict[str, Any] = {} if request_context.tenant_id: task_payload["_tenant_id"] = request_context.tenant_id if request_context.api_key_id: task_payload["_api_key_id"] = request_context.api_key_id return await self._submit_async_operation( bank_id=bank_id, operation_type="consolidation", task_type="consolidation", task_payload=task_payload, dedupe_by_bank=True, ) async def submit_async_refresh_mental_model( self, bank_id: str, mental_model_id: str, *, request_context: "RequestContext", ) -> dict[str, Any]: """Submit an async mental model refresh operation. This schedules a background task to re-run the source query and update the content. Args: bank_id: Bank identifier mental_model_id: Mental model UUID to refresh request_context: Request context for authentication Returns: Dict with operation_id """ await self._authenticate_tenant(request_context) # Pre-operation validation (credit check) if self._operation_validator: from hindsight_api.extensions.operation_validator import MentalModelRefreshContext ctx = MentalModelRefreshContext( bank_id=bank_id, mental_model_id=mental_model_id, request_context=request_context, ) await self._validate_operation(self._operation_validator.validate_mental_model_refresh(ctx)) # Verify mental model exists mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=request_context) if not mental_model: raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}") # Pass tenant_id and api_key_id through task payload so the worker # can provide request context to extension hooks. task_payload: dict[str, Any] = { "mental_model_id": mental_model_id, } if request_context.tenant_id: task_payload["_tenant_id"] = request_context.tenant_id if request_context.api_key_id: task_payload["_api_key_id"] = request_context.api_key_id return await self._submit_async_operation( bank_id=bank_id, operation_type="refresh_mental_model", task_type="refresh_mental_model", task_payload=task_payload, result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]}, dedupe_by_bank=False, )