4130 lines
167 KiB
Python
4130 lines
167 KiB
Python
"""
|
||
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 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
|
||
|
||
# Context variable for current schema (async-safe, per-task isolation)
|
||
_current_schema: contextvars.ContextVar[str] = contextvars.ContextVar("current_schema", default="public")
|
||
|
||
|
||
def get_current_schema() -> str:
|
||
"""Get the current schema from context (default: 'public')."""
|
||
return _current_schema.get()
|
||
|
||
|
||
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 ..pg0 import EmbeddedPostgres, parse_pg0_url
|
||
from .entity_resolver import EntityResolver
|
||
from .llm_wrapper import LLMConfig
|
||
from .query_analyzer import QueryAnalyzer
|
||
from .response_models import (
|
||
VALID_RECALL_FACT_TYPES,
|
||
EntityObservation,
|
||
EntityState,
|
||
MemoryFact,
|
||
ReflectResult,
|
||
TokenUsage,
|
||
)
|
||
from .response_models import RecallResult as RecallResultModel
|
||
from .retain import bank_utils, embedding_utils
|
||
from .retain.types import RetainContentDict
|
||
from .search import observation_utils, think_utils
|
||
from .search.reranking import CrossEncoderReranker
|
||
from .task_backend import AsyncIOQueueBackend, 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,
|
||
embeddings: Embeddings | None = None,
|
||
cross_encoder: CrossEncoderModel | None = None,
|
||
query_analyzer: QueryAnalyzer | None = None,
|
||
pool_min_size: int = 5,
|
||
pool_max_size: int = 100,
|
||
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.
|
||
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 (default: 5)
|
||
pool_max_size: Maximum number of connections in the pool (default: 100)
|
||
task_backend: Custom task backend. If not provided, uses AsyncIOQueueBackend.
|
||
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 doesn't require an API key
|
||
if not memory_llm_api_key and memory_llm_provider != "ollama":
|
||
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
|
||
self._pool_max_size = pool_max_size
|
||
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
|
||
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 cross-encoder reranker (cached for performance)
|
||
self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder)
|
||
|
||
# Initialize task backend
|
||
self._task_backend = task_backend or AsyncIOQueueBackend(batch_size=100, batch_interval=1.0)
|
||
|
||
# Backpressure mechanism: limit concurrent searches to prevent overwhelming the database
|
||
# Limit concurrent searches to prevent connection pool exhaustion
|
||
# Each search can use 2-4 connections, so with 10 concurrent searches
|
||
# we use ~20-40 connections max, staying well within pool limits
|
||
self._search_semaphore = asyncio.Semaphore(10)
|
||
|
||
# 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 (optional)
|
||
self._tenant_extension = 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.
|
||
"""
|
||
if self._tenant_extension is None:
|
||
_current_schema.set("public")
|
||
return "public"
|
||
|
||
from hindsight_api.extensions import AuthenticationError
|
||
|
||
if request_context is None:
|
||
raise AuthenticationError("RequestContext is required when tenant extension is configured")
|
||
|
||
# Let AuthenticationError propagate - HTTP layer will convert to 401
|
||
tenant_context = await self._tenant_extension.authenticate(request_context)
|
||
|
||
_current_schema.set(tenant_context.schema_name)
|
||
return tenant_context.schema_name
|
||
|
||
async def _handle_access_count_update(self, task_dict: dict[str, Any]):
|
||
"""
|
||
Handler for access count update tasks.
|
||
|
||
Args:
|
||
task_dict: Dict with 'node_ids' key containing list of node IDs to update
|
||
|
||
Raises:
|
||
Exception: Any exception from database operations (propagates to execute_task for retry)
|
||
"""
|
||
node_ids = task_dict.get("node_ids", [])
|
||
if not node_ids:
|
||
return
|
||
|
||
pool = await self._get_pool()
|
||
# Convert string UUIDs to UUID type for faster matching
|
||
uuid_list = [uuid.UUID(nid) for nid in node_ids]
|
||
async with acquire_with_retry(pool) as conn:
|
||
await conn.execute(
|
||
f"UPDATE {fq_table('memory_units')} SET access_count = access_count + 1 WHERE id = ANY($1::uuid[])",
|
||
uuid_list,
|
||
)
|
||
|
||
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"
|
||
)
|
||
|
||
# Use internal request context for background tasks
|
||
from hindsight_api.models import RequestContext
|
||
|
||
internal_context = RequestContext()
|
||
await self.retain_batch_async(bank_id=bank_id, contents=contents, request_context=internal_context)
|
||
|
||
logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_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': 'access_count_update', 'node_ids': [...]}
|
||
"""
|
||
task_type = task_dict.get("type")
|
||
operation_id = task_dict.get("operation_id")
|
||
retry_count = task_dict.get("retry_count", 0)
|
||
max_retries = 3
|
||
|
||
# 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 == "access_count_update":
|
||
await self._handle_access_count_update(task_dict)
|
||
elif task_type == "reinforce_opinion":
|
||
await self._handle_reinforce_opinion(task_dict)
|
||
elif task_type == "form_opinion":
|
||
await self._handle_form_opinion(task_dict)
|
||
elif task_type == "batch_retain":
|
||
await self._handle_batch_retain(task_dict)
|
||
elif task_type == "regenerate_observations":
|
||
await self._handle_regenerate_observations(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 - delete operation record
|
||
if operation_id:
|
||
await self._delete_operation_record(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
|
||
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 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) # type: ignore[invalid-argument-type] - dict 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 connection is working."""
|
||
if not self._skip_llm_verification:
|
||
await self._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")
|
||
logger.info("Running database migrations...")
|
||
run_migrations(self.db_url)
|
||
|
||
# Ensure embedding column dimension matches the model's dimension
|
||
# This is done after migrations and after embeddings.initialize()
|
||
ensure_embedding_dimension(self.db_url, self.embeddings.dimension)
|
||
|
||
logger.info(f"Connecting to PostgreSQL at {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=60,
|
||
statement_cache_size=0, # Disable prepared statement cache
|
||
timeout=30, # Connection acquisition timeout (seconds)
|
||
)
|
||
|
||
# Initialize entity resolver with pool
|
||
self.entity_resolver = EntityResolver(self._pool)
|
||
|
||
# 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 (like opinion reinforcement)
|
||
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', 'opinion')
|
||
confidence_score: Confidence score for opinions (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} # type: ignore[typeddict-item] - building incrementally
|
||
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,
|
||
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', 'opinion')
|
||
confidence_score: Confidence score for opinions (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 using internal method (skip chunking check)
|
||
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,
|
||
document_id=document_id,
|
||
is_first_batch=i == 1, # Only upsert on first batch
|
||
fact_type_override=fact_type_override,
|
||
confidence_score=confidence_score,
|
||
)
|
||
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,
|
||
document_id=document_id,
|
||
is_first_batch=True,
|
||
fact_type_override=fact_type_override,
|
||
confidence_score=confidence_score,
|
||
)
|
||
|
||
# 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,
|
||
)
|
||
try:
|
||
await self._operation_validator.on_retain_complete(result_ctx)
|
||
except Exception as e:
|
||
logger.warning(f"Post-retain hook error (non-fatal): {e}")
|
||
|
||
if return_usage:
|
||
return result, total_usage
|
||
return result
|
||
|
||
async def _retain_batch_async_internal(
|
||
self,
|
||
bank_id: str,
|
||
contents: list[RetainContentDict],
|
||
document_id: str | None = None,
|
||
is_first_batch: bool = True,
|
||
fact_type_override: str | None = None,
|
||
confidence_score: float | 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
|
||
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
|
||
|
||
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()
|
||
return await orchestrator.retain_batch(
|
||
pool=pool,
|
||
embeddings_model=self.embeddings,
|
||
llm_config=self._llm_config,
|
||
entity_resolver=self.entity_resolver,
|
||
task_backend=self._task_backend,
|
||
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,
|
||
)
|
||
|
||
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",
|
||
) -> 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)
|
||
|
||
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)
|
||
|
||
# Validate fact types early
|
||
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))}"
|
||
)
|
||
|
||
# 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]
|
||
|
||
# Backpressure: limit concurrent recalls to prevent overwhelming the database
|
||
result = None
|
||
error_msg = None
|
||
async with self._search_semaphore:
|
||
# 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,
|
||
)
|
||
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
|
||
|
||
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,
|
||
) -> 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) 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 = []
|
||
log_buffer.append(
|
||
f"[RECALL {recall_id}] Query: '{query[:50]}...' (budget={thinking_budget}, max_tokens={max_tokens})"
|
||
)
|
||
|
||
try:
|
||
# Step 1: Generate query embedding (for semantic search)
|
||
step_start = time.time()
|
||
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")
|
||
|
||
if tracer:
|
||
tracer.record_query_embedding(query_embedding)
|
||
tracer.add_phase_metric("generate_query_embedding", step_duration)
|
||
|
||
# Step 2: N*4-Way Parallel Retrieval (N fact types × 4 retrieval methods)
|
||
step_start = time.time()
|
||
query_embedding_str = str(query_embedding)
|
||
|
||
from .search.retrieval import retrieve_parallel
|
||
|
||
# Track each retrieval start time
|
||
retrieval_start = time.time()
|
||
|
||
# Run retrieval for each fact type in parallel
|
||
retrieval_tasks = [
|
||
retrieve_parallel(
|
||
pool, query, query_embedding_str, bank_id, ft, thinking_budget, question_date, self.query_analyzer
|
||
)
|
||
for ft in fact_type
|
||
]
|
||
all_retrievals = await asyncio.gather(*retrieval_tasks)
|
||
|
||
# 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}
|
||
|
||
detected_temporal_constraint = None
|
||
for idx, retrieval_result in enumerate(all_retrievals):
|
||
# Log fact types in this retrieval batch
|
||
ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
|
||
logger.debug(
|
||
f"[RECALL {recall_id}] Fact type '{ft_name}': 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)",
|
||
]
|
||
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] {total_retrievals}-way retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {step_duration:.3f}s{temporal_info}"
|
||
)
|
||
|
||
# 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 idx, rr in enumerate(all_retrievals):
|
||
ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
|
||
|
||
# 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 (even if empty, to show it ran)
|
||
if rr.temporal is not None:
|
||
tracer.add_retrieval_results(
|
||
method_name="temporal",
|
||
results=to_tuple_format(rr.temporal),
|
||
duration_seconds=rr.timings.get("temporal", 0.0),
|
||
score_field="temporal_score",
|
||
metadata={"budget": thinking_budget},
|
||
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
|
||
|
||
# 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")
|
||
|
||
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
|
||
|
||
# Ensure reranker is initialized (for lazy initialization mode)
|
||
await reranker_instance.ensure_initialized()
|
||
|
||
# Rerank using cross-encoder
|
||
scored_results = reranker_instance.rerank(query, merged_candidates)
|
||
|
||
step_duration = time.time() - step_start
|
||
log_buffer.append(f" [4] Reranking: {len(scored_results)} candidates scored in {step_duration:.3f}s")
|
||
|
||
# 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,
|
||
access_count=sr.retrieval.access_count,
|
||
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,
|
||
)
|
||
|
||
# Step 8: Queue access count updates for visited nodes
|
||
visited_ids = list(set([sr.id for sr in scored_results[:50]])) # Top 50
|
||
if visited_ids:
|
||
await self._task_backend.submit_task({"type": "access_count_update", "node_ids": visited_ids})
|
||
log_buffer.append(f" [7] Queued access count updates for {len(visited_ids)} nodes")
|
||
|
||
# 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)
|
||
|
||
# Get entities for each fact if include_entities is requested
|
||
fact_entity_map = {} # unit_id -> list of (entity_id, entity_name)
|
||
if include_entities and top_scored:
|
||
unit_ids = [uuid.UUID(sr.id) for sr in top_scored]
|
||
if unit_ids:
|
||
async with acquire_with_retry(pool) as entity_conn:
|
||
entity_rows = await entity_conn.fetch(
|
||
f"""
|
||
SELECT ue.unit_id, e.id as entity_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[])
|
||
""",
|
||
unit_ids,
|
||
)
|
||
for row in entity_rows:
|
||
unit_id = str(row["unit_id"])
|
||
if unit_id not in fact_entity_map:
|
||
fact_entity_map[unit_id] = []
|
||
fact_entity_map[unit_id].append(
|
||
{"entity_id": str(row["entity_id"]), "canonical_name": row["canonical_name"]}
|
||
)
|
||
|
||
# Convert results to MemoryFact objects
|
||
memory_facts = []
|
||
for result_dict in top_results_dicts:
|
||
result_id = str(result_dict.get("id"))
|
||
# Get entity names for this fact
|
||
entity_names = None
|
||
if include_entities and result_id in fact_entity_map:
|
||
entity_names = [e["canonical_name"] for e in fact_entity_map[result_id]]
|
||
|
||
memory_facts.append(
|
||
MemoryFact(
|
||
id=result_id,
|
||
text=result_dict.get("text"),
|
||
fact_type=result_dict.get("fact_type", "world"),
|
||
entities=entity_names,
|
||
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"),
|
||
)
|
||
)
|
||
|
||
# Fetch entity observations if requested
|
||
entities_dict = None
|
||
total_entity_tokens = 0
|
||
total_chunk_tokens = 0
|
||
if include_entities and fact_entity_map:
|
||
# Collect unique entities in order of fact relevance (preserving order from top_scored)
|
||
# Use a list to maintain order, but track seen entities to avoid duplicates
|
||
entities_ordered = [] # list of (entity_id, entity_name) tuples
|
||
seen_entity_ids = set()
|
||
|
||
# Iterate through facts in relevance order
|
||
for sr in top_scored:
|
||
unit_id = sr.id
|
||
if unit_id in fact_entity_map:
|
||
for entity in fact_entity_map[unit_id]:
|
||
entity_id = entity["entity_id"]
|
||
entity_name = entity["canonical_name"]
|
||
if entity_id not in seen_entity_ids:
|
||
entities_ordered.append((entity_id, entity_name))
|
||
seen_entity_ids.add(entity_id)
|
||
|
||
# Fetch observations for each entity (respect token budget, in order)
|
||
entities_dict = {}
|
||
encoding = _get_tiktoken_encoding()
|
||
|
||
for entity_id, entity_name in entities_ordered:
|
||
if total_entity_tokens >= max_entity_tokens:
|
||
break
|
||
|
||
observations = await self.get_entity_observations(
|
||
bank_id, entity_id, limit=5, request_context=request_context
|
||
)
|
||
|
||
# Calculate tokens for this entity's observations
|
||
entity_tokens = 0
|
||
included_observations = []
|
||
for obs in observations:
|
||
obs_tokens = len(encoding.encode(obs.text))
|
||
if total_entity_tokens + entity_tokens + obs_tokens <= max_entity_tokens:
|
||
included_observations.append(obs)
|
||
entity_tokens += obs_tokens
|
||
else:
|
||
break
|
||
|
||
if included_observations:
|
||
entities_dict[entity_name] = EntityState(
|
||
entity_id=entity_id, canonical_name=entity_name, observations=included_observations
|
||
)
|
||
total_entity_tokens += entity_tokens
|
||
|
||
# Fetch chunks if requested
|
||
chunks_dict = None
|
||
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
|
||
num_entities = len(entities_dict) if entities_dict else 0
|
||
log_buffer.append(
|
||
f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok), {num_entities} entities ({total_entity_tokens} tok) | {fact_type_summary} | {total_time:.3f}s"
|
||
)
|
||
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)}")
|
||
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, 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
|
||
""",
|
||
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,
|
||
}
|
||
|
||
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():
|
||
# Count units before deletion
|
||
units_count = await conn.fetchval(
|
||
f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE document_id = $1", document_id
|
||
)
|
||
|
||
# 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,
|
||
)
|
||
|
||
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():
|
||
# 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
|
||
)
|
||
|
||
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 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
|
||
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]
|
||
if unit_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[]) AND 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
|
||
""",
|
||
unit_ids,
|
||
)
|
||
else:
|
||
links = []
|
||
|
||
# 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)
|
||
|
||
# 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
|
||
edges = []
|
||
for row in 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["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"],
|
||
}
|
||
)
|
||
|
||
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 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
|
||
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,
|
||
}
|
||
)
|
||
|
||
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 "",
|
||
}
|
||
|
||
async def _evaluate_opinion_update_async(
|
||
self,
|
||
opinion_text: str,
|
||
opinion_confidence: float,
|
||
new_event_text: str,
|
||
entity_name: str,
|
||
) -> dict[str, Any] | None:
|
||
"""
|
||
Evaluate if an opinion should be updated based on a new event.
|
||
|
||
Args:
|
||
opinion_text: Current opinion text (includes reasons)
|
||
opinion_confidence: Current confidence score (0.0-1.0)
|
||
new_event_text: Text of the new event
|
||
entity_name: Name of the entity this opinion is about
|
||
|
||
Returns:
|
||
Dict with 'action' ('keep'|'update'), 'new_confidence', 'new_text' (if action=='update')
|
||
or None if no changes needed
|
||
"""
|
||
|
||
class OpinionEvaluation(BaseModel):
|
||
"""Evaluation of whether an opinion should be updated."""
|
||
|
||
action: str = Field(description="Action to take: 'keep' (no change) or 'update' (modify opinion)")
|
||
reasoning: str = Field(description="Brief explanation of why this action was chosen")
|
||
new_confidence: float = Field(
|
||
description="New confidence score (0.0-1.0). Can be higher, lower, or same as before."
|
||
)
|
||
new_opinion_text: str | None = Field(
|
||
default=None,
|
||
description="If action is 'update', the revised opinion text that acknowledges the previous view. Otherwise None.",
|
||
)
|
||
|
||
evaluation_prompt = f"""You are evaluating whether an existing opinion should be updated based on new information.
|
||
|
||
ENTITY: {entity_name}
|
||
|
||
EXISTING OPINION:
|
||
{opinion_text}
|
||
Current confidence: {opinion_confidence:.2f}
|
||
|
||
NEW EVENT:
|
||
{new_event_text}
|
||
|
||
Evaluate whether this new event:
|
||
1. REINFORCES the opinion (increase confidence, keep text)
|
||
2. WEAKENS the opinion (decrease confidence, keep text)
|
||
3. CHANGES the opinion (update both text and confidence, noting "Previously I thought X, but now Y...")
|
||
4. IRRELEVANT (keep everything as is)
|
||
|
||
Guidelines:
|
||
- Only suggest 'update' action if the new event genuinely contradicts or significantly modifies the opinion
|
||
- If updating the text, acknowledge the previous opinion and explain the change
|
||
- Confidence should reflect accumulated evidence (0.0 = no confidence, 1.0 = very confident)
|
||
- Small changes in confidence are normal; large jumps should be rare"""
|
||
|
||
try:
|
||
result = await self._llm_config.call(
|
||
messages=[
|
||
{"role": "system", "content": "You evaluate and update opinions based on new information."},
|
||
{"role": "user", "content": evaluation_prompt},
|
||
],
|
||
response_format=OpinionEvaluation,
|
||
scope="memory_evaluate_opinion",
|
||
temperature=0.3, # Lower temperature for more consistent evaluation
|
||
)
|
||
|
||
# Only return updates if something actually changed
|
||
if result.action == "keep" and abs(result.new_confidence - opinion_confidence) < 0.01:
|
||
return None
|
||
|
||
return {
|
||
"action": result.action,
|
||
"reasoning": result.reasoning,
|
||
"new_confidence": result.new_confidence,
|
||
"new_text": result.new_opinion_text if result.action == "update" else None,
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Failed to evaluate opinion update: {str(e)}")
|
||
return None
|
||
|
||
async def _handle_form_opinion(self, task_dict: dict[str, Any]):
|
||
"""
|
||
Handler for form opinion tasks.
|
||
|
||
Args:
|
||
task_dict: Dict with keys: 'bank_id', 'answer_text', 'query', 'tenant_id'
|
||
"""
|
||
bank_id = task_dict["bank_id"]
|
||
answer_text = task_dict["answer_text"]
|
||
query = task_dict["query"]
|
||
tenant_id = task_dict.get("tenant_id")
|
||
|
||
await self._extract_and_store_opinions_async(
|
||
bank_id=bank_id, answer_text=answer_text, query=query, tenant_id=tenant_id
|
||
)
|
||
|
||
async def _handle_reinforce_opinion(self, task_dict: dict[str, Any]):
|
||
"""
|
||
Handler for reinforce opinion tasks.
|
||
|
||
Args:
|
||
task_dict: Dict with keys: 'bank_id', 'created_unit_ids', 'unit_texts', 'unit_entities'
|
||
"""
|
||
bank_id = task_dict["bank_id"]
|
||
created_unit_ids = task_dict["created_unit_ids"]
|
||
unit_texts = task_dict["unit_texts"]
|
||
unit_entities = task_dict["unit_entities"]
|
||
|
||
await self._reinforce_opinions_async(
|
||
bank_id=bank_id, created_unit_ids=created_unit_ids, unit_texts=unit_texts, unit_entities=unit_entities
|
||
)
|
||
|
||
async def _reinforce_opinions_async(
|
||
self,
|
||
bank_id: str,
|
||
created_unit_ids: list[str],
|
||
unit_texts: list[str],
|
||
unit_entities: list[list[dict[str, str]]],
|
||
):
|
||
"""
|
||
Background task to reinforce opinions based on newly ingested events.
|
||
|
||
This runs asynchronously and does not block the put operation.
|
||
|
||
Args:
|
||
bank_id: bank ID
|
||
created_unit_ids: List of newly created memory unit IDs
|
||
unit_texts: Texts of the newly created units
|
||
unit_entities: Entities extracted from each unit
|
||
"""
|
||
try:
|
||
# Extract all unique entity names from the new units
|
||
entity_names = set()
|
||
for entities_list in unit_entities:
|
||
for entity in entities_list:
|
||
# Handle both Entity objects and dicts
|
||
if hasattr(entity, "text"):
|
||
entity_names.add(entity.text)
|
||
elif isinstance(entity, dict):
|
||
entity_names.add(entity["text"])
|
||
|
||
if not entity_names:
|
||
return
|
||
|
||
pool = await self._get_pool()
|
||
async with acquire_with_retry(pool) as conn:
|
||
# Find all opinions related to these entities
|
||
opinions = await conn.fetch(
|
||
f"""
|
||
SELECT DISTINCT mu.id, mu.text, mu.confidence_score, e.canonical_name
|
||
FROM {fq_table("memory_units")} mu
|
||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||
WHERE mu.bank_id = $1
|
||
AND mu.fact_type = 'opinion'
|
||
AND e.canonical_name = ANY($2::text[])
|
||
""",
|
||
bank_id,
|
||
list(entity_names),
|
||
)
|
||
|
||
if not opinions:
|
||
return
|
||
|
||
# Use cached LLM config
|
||
if self._llm_config is None:
|
||
logger.error("[REINFORCE] LLM config not available, skipping opinion reinforcement")
|
||
return
|
||
|
||
# Evaluate each opinion against the new events
|
||
updates_to_apply = []
|
||
for opinion in opinions:
|
||
opinion_id = str(opinion["id"])
|
||
opinion_text = opinion["text"]
|
||
opinion_confidence = opinion["confidence_score"]
|
||
entity_name = opinion["canonical_name"]
|
||
|
||
# Find all new events mentioning this entity
|
||
relevant_events = []
|
||
for unit_text, entities_list in zip(unit_texts, unit_entities):
|
||
if any(e["text"] == entity_name for e in entities_list):
|
||
relevant_events.append(unit_text)
|
||
|
||
if not relevant_events:
|
||
continue
|
||
|
||
# Combine all relevant events
|
||
combined_events = "\n".join(relevant_events)
|
||
|
||
# Evaluate if opinion should be updated
|
||
evaluation = await self._evaluate_opinion_update_async(
|
||
opinion_text, opinion_confidence, combined_events, entity_name
|
||
)
|
||
|
||
if evaluation:
|
||
updates_to_apply.append({"opinion_id": opinion_id, "evaluation": evaluation})
|
||
|
||
# Apply all updates in a single transaction
|
||
if updates_to_apply:
|
||
async with conn.transaction():
|
||
for update in updates_to_apply:
|
||
opinion_id = update["opinion_id"]
|
||
evaluation = update["evaluation"]
|
||
|
||
if evaluation["action"] == "update" and evaluation["new_text"]:
|
||
# Update both text and confidence
|
||
await conn.execute(
|
||
f"""
|
||
UPDATE {fq_table("memory_units")}
|
||
SET text = $1, confidence_score = $2, updated_at = NOW()
|
||
WHERE id = $3
|
||
""",
|
||
evaluation["new_text"],
|
||
evaluation["new_confidence"],
|
||
uuid.UUID(opinion_id),
|
||
)
|
||
else:
|
||
# Only update confidence
|
||
await conn.execute(
|
||
f"""
|
||
UPDATE {fq_table("memory_units")}
|
||
SET confidence_score = $1, updated_at = NOW()
|
||
WHERE id = $2
|
||
""",
|
||
evaluation["new_confidence"],
|
||
uuid.UUID(opinion_id),
|
||
)
|
||
|
||
else:
|
||
pass # No opinions to update
|
||
|
||
except Exception as e:
|
||
logger.error(f"[REINFORCE] Error during opinion reinforcement: {str(e)}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
|
||
# ==================== bank profile Methods ====================
|
||
|
||
async def get_bank_profile(
|
||
self,
|
||
bank_id: str,
|
||
*,
|
||
request_context: "RequestContext",
|
||
) -> dict[str, Any]:
|
||
"""
|
||
Get bank profile (name, disposition + background).
|
||
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 background
|
||
"""
|
||
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,
|
||
"background": profile["background"],
|
||
}
|
||
|
||
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 merge_bank_background(
|
||
self,
|
||
bank_id: str,
|
||
new_info: str,
|
||
*,
|
||
update_disposition: bool = True,
|
||
request_context: "RequestContext",
|
||
) -> dict[str, Any]:
|
||
"""
|
||
Merge new background information with existing background using LLM.
|
||
Normalizes to first person ("I") and resolves conflicts.
|
||
Optionally infers disposition traits from the merged background.
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
new_info: New background information to add/merge
|
||
update_disposition: If True, infer Big Five traits from background (default: True)
|
||
request_context: Request context for authentication.
|
||
|
||
Returns:
|
||
Dict with 'background' (str) and optionally 'disposition' (dict) keys
|
||
"""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
return await bank_utils.merge_bank_background(pool, self._llm_config, bank_id, new_info, update_disposition)
|
||
|
||
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, background, 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",
|
||
) -> ReflectResult:
|
||
"""
|
||
Reflect and formulate an answer using bank identity, world facts, and opinions.
|
||
|
||
This method:
|
||
1. Retrieves experience (conversations and events)
|
||
2. Retrieves world facts (general knowledge)
|
||
3. Retrieves existing opinions (bank's formed perspectives)
|
||
4. Uses LLM to formulate an answer
|
||
5. Extracts and stores any new opinions formed during reflection
|
||
6. Optionally generates structured output based on response_schema
|
||
7. Returns plain text answer and the facts used
|
||
|
||
Args:
|
||
bank_id: bank identifier
|
||
query: Question to answer
|
||
budget: Budget level for memory exploration (low=100, mid=300, high=600 units)
|
||
context: Additional context string to include in LLM prompt (not used in recall)
|
||
response_schema: Optional JSON Schema for structured output
|
||
|
||
Returns:
|
||
ReflectResult containing:
|
||
- text: Plain text answer (no markdown)
|
||
- based_on: Dict with 'world', 'experience', and 'opinion' fact lists (MemoryFact objects)
|
||
- new_opinions: List of newly formed opinions
|
||
- structured_output: Optional dict if response_schema was provided
|
||
"""
|
||
# Use cached LLM config
|
||
if self._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}"
|
||
log_buffer = []
|
||
log_buffer.append(f"[REFLECT {reflect_id}] Query: '{query[:50]}...'")
|
||
|
||
# Steps 1-3: Run multi-fact-type search (12-way retrieval: 4 methods × 3 fact types)
|
||
recall_start = time.time()
|
||
metrics = get_metrics_collector()
|
||
with metrics.record_operation(
|
||
"recall", bank_id=bank_id, source="reflect", budget=budget.value if budget else None
|
||
):
|
||
search_result = await self.recall_async(
|
||
bank_id=bank_id,
|
||
query=query,
|
||
budget=budget,
|
||
max_tokens=4096,
|
||
enable_trace=False,
|
||
fact_type=["experience", "world", "opinion"],
|
||
include_entities=True,
|
||
request_context=request_context,
|
||
)
|
||
recall_time = time.time() - recall_start
|
||
|
||
all_results = search_result.results
|
||
|
||
# Split results by fact type for structured response
|
||
agent_results = [r for r in all_results if r.fact_type == "experience"]
|
||
world_results = [r for r in all_results if r.fact_type == "world"]
|
||
opinion_results = [r for r in all_results if r.fact_type == "opinion"]
|
||
|
||
log_buffer.append(
|
||
f"[REFLECT {reflect_id}] Recall: {len(all_results)} facts (experience={len(agent_results)}, world={len(world_results)}, opinion={len(opinion_results)}) in {recall_time:.3f}s"
|
||
)
|
||
|
||
# Format facts for LLM
|
||
agent_facts_text = think_utils.format_facts_for_prompt(agent_results)
|
||
world_facts_text = think_utils.format_facts_for_prompt(world_results)
|
||
opinion_facts_text = think_utils.format_facts_for_prompt(opinion_results)
|
||
|
||
# Get bank profile (name, disposition + background)
|
||
profile = await self.get_bank_profile(bank_id, request_context=request_context)
|
||
name = profile["name"]
|
||
disposition = profile["disposition"] # Typed as DispositionTraits
|
||
background = profile["background"]
|
||
|
||
# Build the prompt
|
||
prompt = think_utils.build_think_prompt(
|
||
agent_facts_text=agent_facts_text,
|
||
world_facts_text=world_facts_text,
|
||
opinion_facts_text=opinion_facts_text,
|
||
query=query,
|
||
name=name,
|
||
disposition=disposition,
|
||
background=background,
|
||
context=context,
|
||
)
|
||
|
||
log_buffer.append(f"[REFLECT {reflect_id}] Prompt: {len(prompt)} chars")
|
||
|
||
system_message = think_utils.get_system_message(disposition)
|
||
messages = [{"role": "system", "content": system_message}, {"role": "user", "content": prompt}]
|
||
|
||
# Prepare response_format if schema provided
|
||
response_format = None
|
||
if response_schema is not None:
|
||
# Wrapper class to provide Pydantic-like interface for raw JSON schemas
|
||
class JsonSchemaWrapper:
|
||
def __init__(self, schema: dict):
|
||
self._schema = schema
|
||
|
||
def model_json_schema(self):
|
||
return self._schema
|
||
|
||
response_format = JsonSchemaWrapper(response_schema)
|
||
|
||
llm_start = time.time()
|
||
llm_result, usage = await self._llm_config.call(
|
||
messages=messages,
|
||
scope="memory_reflect",
|
||
max_completion_tokens=max_tokens,
|
||
response_format=response_format,
|
||
skip_validation=True if response_format else False,
|
||
# Don't enforce strict_schema - not all providers support it and may retry forever
|
||
# Soft enforcement (schema in prompt + json_object mode) is sufficient
|
||
strict_schema=False,
|
||
return_usage=True,
|
||
)
|
||
llm_time = time.time() - llm_start
|
||
|
||
# Handle response based on whether structured output was requested
|
||
if response_schema is not None:
|
||
structured_output = llm_result
|
||
answer_text = "" # Empty for backward compatibility
|
||
log_buffer.append(f"[REFLECT {reflect_id}] Structured output generated")
|
||
else:
|
||
structured_output = None
|
||
answer_text = llm_result.strip()
|
||
|
||
# Submit form_opinion task for background processing
|
||
# Pass tenant_id from request context for internal authentication in background task
|
||
await self._task_backend.submit_task(
|
||
{
|
||
"type": "form_opinion",
|
||
"bank_id": bank_id,
|
||
"answer_text": answer_text,
|
||
"query": query,
|
||
"tenant_id": getattr(request_context, "tenant_id", None) if request_context else None,
|
||
}
|
||
)
|
||
|
||
total_time = time.time() - reflect_start
|
||
log_buffer.append(
|
||
f"[REFLECT {reflect_id}] Complete: {len(answer_text)} chars response, LLM {llm_time:.3f}s, total {total_time:.3f}s"
|
||
)
|
||
logger.info("\n" + "\n".join(log_buffer))
|
||
|
||
# Return response with facts split by type
|
||
result = ReflectResult(
|
||
text=answer_text,
|
||
based_on={"world": world_results, "experience": agent_results, "opinion": opinion_results},
|
||
new_opinions=[], # Opinions are being extracted asynchronously
|
||
structured_output=structured_output,
|
||
usage=usage,
|
||
)
|
||
|
||
# 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
|
||
|
||
async def _extract_and_store_opinions_async(
|
||
self, bank_id: str, answer_text: str, query: str, tenant_id: str | None = None
|
||
):
|
||
"""
|
||
Background task to extract and store opinions from think response.
|
||
|
||
This runs asynchronously and does not block the think response.
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
answer_text: The generated answer text
|
||
query: The original query
|
||
tenant_id: Tenant identifier for internal authentication
|
||
"""
|
||
try:
|
||
# Extract opinions from the answer
|
||
new_opinions = await think_utils.extract_opinions_from_text(self._llm_config, text=answer_text, query=query)
|
||
|
||
# Store new opinions
|
||
if new_opinions:
|
||
from datetime import datetime
|
||
|
||
current_time = datetime.now(UTC)
|
||
# Use internal context with tenant_id for background authentication
|
||
# Extension can check internal=True to bypass normal auth
|
||
from hindsight_api.models import RequestContext
|
||
|
||
internal_context = RequestContext(tenant_id=tenant_id, internal=True)
|
||
for opinion in new_opinions:
|
||
await self.retain_async(
|
||
bank_id=bank_id,
|
||
content=opinion.opinion,
|
||
context=f"formed during thinking about: {query}",
|
||
event_date=current_time,
|
||
fact_type_override="opinion",
|
||
confidence_score=opinion.confidence,
|
||
request_context=internal_context,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.warning(f"[REFLECT] Failed to extract/store opinions: {str(e)}")
|
||
|
||
async def get_entity_observations(
|
||
self,
|
||
bank_id: str,
|
||
entity_id: str,
|
||
*,
|
||
limit: int = 10,
|
||
request_context: "RequestContext",
|
||
) -> list[Any]:
|
||
"""
|
||
Get observations linked to an entity.
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
entity_id: Entity UUID to get observations for
|
||
limit: Maximum number of observations to return
|
||
request_context: Request context for authentication.
|
||
|
||
Returns:
|
||
List of EntityObservation objects
|
||
"""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
async with acquire_with_retry(pool) as conn:
|
||
rows = await conn.fetch(
|
||
f"""
|
||
SELECT mu.text, mu.mentioned_at
|
||
FROM {fq_table("memory_units")} mu
|
||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||
WHERE mu.bank_id = $1
|
||
AND mu.fact_type = 'observation'
|
||
AND ue.entity_id = $2
|
||
ORDER BY mu.mentioned_at DESC
|
||
LIMIT $3
|
||
""",
|
||
bank_id,
|
||
uuid.UUID(entity_id),
|
||
limit,
|
||
)
|
||
|
||
observations = []
|
||
for row in rows:
|
||
mentioned_at = row["mentioned_at"].isoformat() if row["mentioned_at"] else None
|
||
observations.append(EntityObservation(text=row["text"], mentioned_at=mentioned_at))
|
||
return observations
|
||
|
||
async def list_entities(
|
||
self,
|
||
bank_id: str,
|
||
*,
|
||
limit: int = 100,
|
||
request_context: "RequestContext",
|
||
) -> list[dict[str, Any]]:
|
||
"""
|
||
List all entities for a bank.
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
limit: Maximum number of entities to return
|
||
request_context: Request context for authentication.
|
||
|
||
Returns:
|
||
List of entity dicts with id, canonical_name, mention_count, first_seen, last_seen
|
||
"""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
async with acquire_with_retry(pool) as conn:
|
||
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
|
||
LIMIT $2
|
||
""",
|
||
bank_id,
|
||
limit,
|
||
)
|
||
|
||
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 entities
|
||
|
||
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 (mental model) of an entity.
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
entity_id: Entity UUID
|
||
entity_name: Canonical name of the entity
|
||
limit: Maximum number of observations to include
|
||
request_context: Request context for authentication.
|
||
|
||
Returns:
|
||
EntityState with observations
|
||
"""
|
||
observations = await self.get_entity_observations(
|
||
bank_id, entity_id, limit=limit, request_context=request_context
|
||
)
|
||
return EntityState(entity_id=entity_id, canonical_name=entity_name, observations=observations)
|
||
|
||
async def regenerate_entity_observations(
|
||
self,
|
||
bank_id: str,
|
||
entity_id: str,
|
||
entity_name: str,
|
||
*,
|
||
version: str | None = None,
|
||
conn=None,
|
||
request_context: "RequestContext",
|
||
) -> None:
|
||
"""
|
||
Regenerate observations for an entity by:
|
||
1. Checking version for deduplication (if provided)
|
||
2. Searching all facts mentioning the entity
|
||
3. Using LLM to synthesize observations (no personality)
|
||
4. Deleting old observations for this entity
|
||
5. Storing new observations linked to the entity
|
||
|
||
Args:
|
||
bank_id: bank IDentifier
|
||
entity_id: Entity UUID
|
||
entity_name: Canonical name of the entity
|
||
version: Entity's last_seen timestamp when task was created (for deduplication)
|
||
conn: Optional database connection (for transactional atomicity with caller)
|
||
request_context: Request context for authentication.
|
||
"""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
entity_uuid = uuid.UUID(entity_id)
|
||
|
||
# Helper to run a query with provided conn or acquire one
|
||
async def fetch_with_conn(query, *args):
|
||
if conn is not None:
|
||
return await conn.fetch(query, *args)
|
||
else:
|
||
async with acquire_with_retry(pool) as acquired_conn:
|
||
return await acquired_conn.fetch(query, *args)
|
||
|
||
async def fetchval_with_conn(query, *args):
|
||
if conn is not None:
|
||
return await conn.fetchval(query, *args)
|
||
else:
|
||
async with acquire_with_retry(pool) as acquired_conn:
|
||
return await acquired_conn.fetchval(query, *args)
|
||
|
||
# Step 1: Check version for deduplication
|
||
if version:
|
||
current_last_seen = await fetchval_with_conn(
|
||
f"""
|
||
SELECT last_seen
|
||
FROM {fq_table("entities")}
|
||
WHERE id = $1 AND bank_id = $2
|
||
""",
|
||
entity_uuid,
|
||
bank_id,
|
||
)
|
||
|
||
if current_last_seen and current_last_seen.isoformat() != version:
|
||
return []
|
||
|
||
# Step 2: Get all facts mentioning this entity (exclude observations themselves)
|
||
rows = await fetch_with_conn(
|
||
f"""
|
||
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
|
||
FROM {fq_table("memory_units")} mu
|
||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||
WHERE mu.bank_id = $1
|
||
AND ue.entity_id = $2
|
||
AND mu.fact_type IN ('world', 'experience')
|
||
ORDER BY mu.occurred_start DESC
|
||
LIMIT 50
|
||
""",
|
||
bank_id,
|
||
entity_uuid,
|
||
)
|
||
|
||
if not rows:
|
||
return []
|
||
|
||
# Convert to MemoryFact objects for the observation extraction
|
||
facts = []
|
||
for row in rows:
|
||
occurred_start = row["occurred_start"].isoformat() if row["occurred_start"] else None
|
||
facts.append(
|
||
MemoryFact(
|
||
id=str(row["id"]),
|
||
text=row["text"],
|
||
fact_type=row["fact_type"],
|
||
context=row["context"],
|
||
occurred_start=occurred_start,
|
||
)
|
||
)
|
||
|
||
# Step 3: Extract observations using LLM (no personality)
|
||
observations = await observation_utils.extract_observations_from_facts(self._llm_config, entity_name, facts)
|
||
|
||
if not observations:
|
||
return []
|
||
|
||
# Step 4: Delete old observations and insert new ones
|
||
# If conn provided, we're already in a transaction - don't start another
|
||
# If conn is None, acquire one and start a transaction
|
||
async def do_db_operations(db_conn):
|
||
# Delete old observations for this entity
|
||
await db_conn.execute(
|
||
f"""
|
||
DELETE FROM {fq_table("memory_units")}
|
||
WHERE id IN (
|
||
SELECT mu.id
|
||
FROM {fq_table("memory_units")} mu
|
||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||
WHERE mu.bank_id = $1
|
||
AND mu.fact_type = 'observation'
|
||
AND ue.entity_id = $2
|
||
)
|
||
""",
|
||
bank_id,
|
||
entity_uuid,
|
||
)
|
||
|
||
# Generate embeddings for new observations
|
||
embeddings = await embedding_utils.generate_embeddings_batch(self.embeddings, observations)
|
||
|
||
# Insert new observations
|
||
current_time = utcnow()
|
||
created_ids = []
|
||
|
||
for obs_text, embedding in zip(observations, embeddings):
|
||
result = await db_conn.fetchrow(
|
||
f"""
|
||
INSERT INTO {fq_table("memory_units")} (
|
||
bank_id, text, embedding, context, event_date,
|
||
occurred_start, occurred_end, mentioned_at,
|
||
fact_type, access_count
|
||
)
|
||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'observation', 0)
|
||
RETURNING id
|
||
""",
|
||
bank_id,
|
||
obs_text,
|
||
str(embedding),
|
||
f"observation about {entity_name}",
|
||
current_time,
|
||
current_time,
|
||
current_time,
|
||
current_time,
|
||
)
|
||
obs_id = str(result["id"])
|
||
created_ids.append(obs_id)
|
||
|
||
# Link observation to entity
|
||
await db_conn.execute(
|
||
f"""
|
||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||
VALUES ($1, $2)
|
||
""",
|
||
uuid.UUID(obs_id),
|
||
entity_uuid,
|
||
)
|
||
|
||
return created_ids
|
||
|
||
if conn is not None:
|
||
# Use provided connection (already in a transaction)
|
||
return await do_db_operations(conn)
|
||
else:
|
||
# Acquire connection and start our own transaction
|
||
async with acquire_with_retry(pool) as acquired_conn:
|
||
async with acquired_conn.transaction():
|
||
return await do_db_operations(acquired_conn)
|
||
|
||
async def _regenerate_observations_sync(
|
||
self,
|
||
bank_id: str,
|
||
entity_ids: list[str],
|
||
min_facts: int | None = None,
|
||
conn=None,
|
||
request_context: "RequestContext | None" = None,
|
||
) -> None:
|
||
"""
|
||
Regenerate observations for entities synchronously (called during retain).
|
||
|
||
Processes entities in PARALLEL for faster execution.
|
||
|
||
Args:
|
||
bank_id: Bank identifier
|
||
entity_ids: List of entity IDs to process
|
||
min_facts: Minimum facts required to regenerate observations (uses config default if None)
|
||
conn: Optional database connection (for transactional atomicity)
|
||
"""
|
||
if not bank_id or not entity_ids:
|
||
return
|
||
|
||
# Use config default if min_facts not specified
|
||
if min_facts is None:
|
||
min_facts = get_config().observation_min_facts
|
||
|
||
# Convert to UUIDs
|
||
entity_uuids = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in entity_ids]
|
||
|
||
# Use provided connection or acquire a new one
|
||
if conn is not None:
|
||
# Use the provided connection (transactional with caller)
|
||
entity_rows = await conn.fetch(
|
||
f"""
|
||
SELECT id, canonical_name FROM {fq_table("entities")}
|
||
WHERE id = ANY($1) AND bank_id = $2
|
||
""",
|
||
entity_uuids,
|
||
bank_id,
|
||
)
|
||
entity_names = {row["id"]: row["canonical_name"] for row in entity_rows}
|
||
|
||
fact_counts = await conn.fetch(
|
||
f"""
|
||
SELECT ue.entity_id, COUNT(*) as cnt
|
||
FROM {fq_table("unit_entities")} ue
|
||
JOIN {fq_table("memory_units")} mu ON ue.unit_id = mu.id
|
||
WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
|
||
GROUP BY ue.entity_id
|
||
""",
|
||
entity_uuids,
|
||
bank_id,
|
||
)
|
||
entity_fact_counts = {row["entity_id"]: row["cnt"] for row in fact_counts}
|
||
else:
|
||
# Acquire a new connection (standalone call)
|
||
pool = await self._get_pool()
|
||
async with pool.acquire() as acquired_conn:
|
||
entity_rows = await acquired_conn.fetch(
|
||
f"""
|
||
SELECT id, canonical_name FROM {fq_table("entities")}
|
||
WHERE id = ANY($1) AND bank_id = $2
|
||
""",
|
||
entity_uuids,
|
||
bank_id,
|
||
)
|
||
entity_names = {row["id"]: row["canonical_name"] for row in entity_rows}
|
||
|
||
fact_counts = await acquired_conn.fetch(
|
||
f"""
|
||
SELECT ue.entity_id, COUNT(*) as cnt
|
||
FROM {fq_table("unit_entities")} ue
|
||
JOIN {fq_table("memory_units")} mu ON ue.unit_id = mu.id
|
||
WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
|
||
GROUP BY ue.entity_id
|
||
""",
|
||
entity_uuids,
|
||
bank_id,
|
||
)
|
||
entity_fact_counts = {row["entity_id"]: row["cnt"] for row in fact_counts}
|
||
|
||
# Filter entities that meet the threshold
|
||
entities_to_process = []
|
||
for entity_id in entity_ids:
|
||
entity_uuid = uuid.UUID(entity_id) if isinstance(entity_id, str) else entity_id
|
||
if entity_uuid not in entity_names:
|
||
continue
|
||
fact_count = entity_fact_counts.get(entity_uuid, 0)
|
||
if fact_count >= min_facts:
|
||
entities_to_process.append((entity_id, entity_names[entity_uuid]))
|
||
|
||
if not entities_to_process:
|
||
return
|
||
|
||
# Use internal context if not provided (for internal/background calls)
|
||
from hindsight_api.models import RequestContext as RC
|
||
|
||
ctx = request_context if request_context is not None else RC()
|
||
|
||
# Process all entities in PARALLEL (LLM calls are the bottleneck)
|
||
async def process_entity(entity_id: str, entity_name: str):
|
||
try:
|
||
await self.regenerate_entity_observations(
|
||
bank_id, entity_id, entity_name, version=None, conn=conn, request_context=ctx
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"[OBSERVATIONS] Error processing entity {entity_id}: {e}")
|
||
|
||
await asyncio.gather(*[process_entity(eid, name) for eid, name in entities_to_process])
|
||
|
||
async def _handle_regenerate_observations(self, task_dict: dict[str, Any]):
|
||
"""
|
||
Handler for regenerate_observations tasks.
|
||
|
||
Args:
|
||
task_dict: Dict with 'bank_id' and either:
|
||
- 'entity_ids' (list): Process multiple entities
|
||
- 'entity_id', 'entity_name': Process single entity (legacy)
|
||
|
||
Raises:
|
||
ValueError: If required fields are missing
|
||
Exception: Any exception from regenerate_entity_observations (propagates to execute_task for retry)
|
||
"""
|
||
bank_id = task_dict.get("bank_id")
|
||
# Use internal request context for background tasks
|
||
from hindsight_api.models import RequestContext
|
||
|
||
internal_context = RequestContext()
|
||
|
||
# New format: multiple entity_ids
|
||
if "entity_ids" in task_dict:
|
||
entity_ids = task_dict.get("entity_ids", [])
|
||
min_facts = task_dict.get("min_facts", 5)
|
||
|
||
if not bank_id or not entity_ids:
|
||
raise ValueError(f"[OBSERVATIONS] Missing required fields in task: {task_dict}")
|
||
|
||
# Process each entity
|
||
pool = await self._get_pool()
|
||
async with pool.acquire() as conn:
|
||
for entity_id in entity_ids:
|
||
try:
|
||
# Fetch entity name and check fact count
|
||
import uuid as uuid_module
|
||
|
||
entity_uuid = uuid_module.UUID(entity_id) if isinstance(entity_id, str) else entity_id
|
||
|
||
# First check if entity exists
|
||
entity_exists = await conn.fetchrow(
|
||
f"SELECT canonical_name FROM {fq_table('entities')} WHERE id = $1 AND bank_id = $2",
|
||
entity_uuid,
|
||
bank_id,
|
||
)
|
||
|
||
if not entity_exists:
|
||
logger.debug(f"[OBSERVATIONS] Entity {entity_id} not yet in bank {bank_id}, skipping")
|
||
continue
|
||
|
||
entity_name = entity_exists["canonical_name"]
|
||
|
||
# Count facts linked to this entity
|
||
fact_count = (
|
||
await conn.fetchval(
|
||
f"SELECT COUNT(*) FROM {fq_table('unit_entities')} WHERE entity_id = $1",
|
||
entity_uuid,
|
||
)
|
||
or 0
|
||
)
|
||
|
||
# Only regenerate if entity has enough facts
|
||
if fact_count >= min_facts:
|
||
await self.regenerate_entity_observations(
|
||
bank_id, entity_id, entity_name, version=None, request_context=internal_context
|
||
)
|
||
else:
|
||
logger.debug(
|
||
f"[OBSERVATIONS] Skipping {entity_name} ({fact_count} facts < {min_facts} threshold)"
|
||
)
|
||
|
||
except Exception as e:
|
||
# Log but continue processing other entities - individual entity failures
|
||
# shouldn't fail the whole batch
|
||
logger.error(f"[OBSERVATIONS] Error processing entity {entity_id}: {e}")
|
||
continue
|
||
|
||
# Legacy format: single entity
|
||
else:
|
||
entity_id = task_dict.get("entity_id")
|
||
entity_name = task_dict.get("entity_name")
|
||
version = task_dict.get("version")
|
||
|
||
if not all([bank_id, entity_id, entity_name]):
|
||
raise ValueError(f"[OBSERVATIONS] Missing required fields in task: {task_dict}")
|
||
|
||
# Type assertions after validation
|
||
assert isinstance(bank_id, str) and isinstance(entity_id, str) and isinstance(entity_name, str)
|
||
await self.regenerate_entity_observations(
|
||
bank_id, entity_id, entity_name, version=version, request_context=internal_context
|
||
)
|
||
|
||
# =========================================================================
|
||
# 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
|
||
|
||
# Get observations for the entity
|
||
observations = await self.get_entity_observations(bank_id, entity_id, limit=20, request_context=request_context)
|
||
|
||
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": observations,
|
||
}
|
||
|
||
async def list_operations(
|
||
self,
|
||
bank_id: str,
|
||
*,
|
||
request_context: "RequestContext",
|
||
) -> list[dict[str, Any]]:
|
||
"""List async operations for a bank."""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
|
||
async with acquire_with_retry(pool) as conn:
|
||
operations = await conn.fetch(
|
||
f"""
|
||
SELECT operation_id, bank_id, operation_type, created_at, status, error_message, result_metadata
|
||
FROM {fq_table("async_operations")}
|
||
WHERE bank_id = $1
|
||
ORDER BY created_at DESC
|
||
""",
|
||
bank_id,
|
||
)
|
||
|
||
def parse_metadata(metadata):
|
||
if metadata is None:
|
||
return {}
|
||
if isinstance(metadata, str):
|
||
import json
|
||
|
||
return json.loads(metadata)
|
||
return metadata
|
||
|
||
return [
|
||
{
|
||
"id": str(row["operation_id"]),
|
||
"task_type": row["operation_type"],
|
||
"items_count": parse_metadata(row["result_metadata"]).get("items_count", 0),
|
||
"document_id": parse_metadata(row["result_metadata"]).get("document_id"),
|
||
"created_at": row["created_at"].isoformat(),
|
||
"status": row["status"],
|
||
"error_message": row["error_message"],
|
||
}
|
||
for row in operations
|
||
]
|
||
|
||
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,
|
||
background: str | None = None,
|
||
request_context: "RequestContext",
|
||
) -> dict[str, Any]:
|
||
"""Update bank name and/or background."""
|
||
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 background is not None:
|
||
await conn.execute(
|
||
f"""
|
||
UPDATE {fq_table("banks")}
|
||
SET background = $2, updated_at = NOW()
|
||
WHERE bank_id = $1
|
||
""",
|
||
bank_id,
|
||
background,
|
||
)
|
||
|
||
# Return updated profile
|
||
return await self.get_bank_profile(bank_id, request_context=request_context)
|
||
|
||
async def submit_async_retain(
|
||
self,
|
||
bank_id: str,
|
||
contents: list[dict[str, Any]],
|
||
*,
|
||
request_context: "RequestContext",
|
||
) -> dict[str, Any]:
|
||
"""Submit a batch retain operation to run asynchronously."""
|
||
await self._authenticate_tenant(request_context)
|
||
pool = await self._get_pool()
|
||
|
||
import json
|
||
|
||
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,
|
||
"retain",
|
||
json.dumps({"items_count": len(contents)}),
|
||
)
|
||
|
||
# Submit task to background queue
|
||
await self._task_backend.submit_task(
|
||
{
|
||
"type": "batch_retain",
|
||
"operation_id": str(operation_id),
|
||
"bank_id": bank_id,
|
||
"contents": contents,
|
||
}
|
||
)
|
||
|
||
logger.info(f"Retain task queued for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}")
|
||
|
||
return {
|
||
"operation_id": str(operation_id),
|
||
"items_count": len(contents),
|
||
}
|