From 7935b0accd0a395ac4a09d8f994e1bf8c3f152f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 12 Jan 2026 18:58:05 +0100 Subject: [PATCH] fix: improve mpfp retrieval (#146) * fix: improve mpfp retrieval * fix: improve mpfp retrieval * fix: improve embeddings service performances * fix: improve embeddings service performances * fix: improve embeddings service performances * fix: improve embeddings service performances --- CLAUDE.md | 46 +++ ...c4d5e6_add_memory_links_composite_index.py | 44 +++ hindsight-api/hindsight_api/api/http.py | 24 +- hindsight-api/hindsight_api/config.py | 20 +- .../hindsight_api/engine/cross_encoder.py | 199 ++++++++++- .../hindsight_api/engine/memory_engine.py | 78 ++++- .../hindsight_api/engine/query_analyzer.py | 7 +- .../engine/search/link_expansion_retrieval.py | 230 +++++++++++++ .../engine/search/mpfp_retrieval.py | 291 +++++++++++----- .../hindsight_api/engine/search/retrieval.py | 176 +++++----- .../hindsight_api/engine/search/types.py | 2 + hindsight-api/hindsight_api/main.py | 27 +- hindsight-api/hindsight_api/server.py | 11 +- hindsight-api/pyproject.toml | 2 + hindsight-api/tests/test_mpfp_retrieval.py | 310 ++++++++++++++++-- hindsight-api/tests/test_tei_cross_encoder.py | 240 ++++++++++++++ .../docs/developer/configuration.md | 29 +- hindsight-docs/docs/developer/retrieval.md | 39 +-- scripts/dev/start-api.sh | 2 +- uv.lock | 93 ++++++ 20 files changed, 1615 insertions(+), 255 deletions(-) create mode 100644 hindsight-api/hindsight_api/alembic/versions/f1a2b3c4d5e6_add_memory_links_composite_index.py create mode 100644 hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py diff --git a/CLAUDE.md b/CLAUDE.md index 9293fabc..eb0ebab5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -108,6 +108,52 @@ PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-ap Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links` +### Adding Database Migrations + +1. **Create a new migration file** in `hindsight-api/hindsight_api/alembic/versions/`: + - File name format: `_.py` (e.g., `f1a2b3c4d5e6_add_new_index.py`) + - Use a unique hex revision ID (12 chars) + - Set `down_revision` to the previous migration's revision ID + +2. **Migration template**: + ```python + """Description of the migration + + Revision ID: f1a2b3c4d5e6 + Revises: + Create Date: YYYY-MM-DD + """ + from collections.abc import Sequence + from alembic import context, op + + revision: str = "f1a2b3c4d5e6" + down_revision: str | Sequence[str] | None = "" + branch_labels: str | Sequence[str] | None = None + depends_on: str | Sequence[str] | None = None + + def _get_schema_prefix() -> str: + """Get schema prefix for table names (required for multi-tenant support).""" + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + def upgrade() -> None: + schema = _get_schema_prefix() + op.execute(f"CREATE INDEX ... ON {schema}table_name(...)") + + def downgrade() -> None: + schema = _get_schema_prefix() + op.execute(f"DROP INDEX IF EXISTS {schema}index_name") + ``` + +3. **Run migrations locally**: + ```bash + # Set database URL and run migrations + uv run hindsight-admin run-db-migration + + # Run on a specific tenant schema + uv run hindsight-admin run-db-migration --schema tenant_xyz + ``` + ## Key Conventions ### Code Quality diff --git a/hindsight-api/hindsight_api/alembic/versions/f1a2b3c4d5e6_add_memory_links_composite_index.py b/hindsight-api/hindsight_api/alembic/versions/f1a2b3c4d5e6_add_memory_links_composite_index.py new file mode 100644 index 00000000..09100aeb --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/f1a2b3c4d5e6_add_memory_links_composite_index.py @@ -0,0 +1,44 @@ +"""add_memory_links_from_type_weight_index + +Revision ID: f1a2b3c4d5e6 +Revises: e0a1b2c3d4e5 +Create Date: 2025-01-12 + +Add composite index on memory_links (from_unit_id, link_type, weight DESC) +to optimize MPFP graph traversal queries that need top-k edges per type. +""" + +from collections.abc import Sequence + +from alembic import context, op + +# revision identifiers, used by Alembic. +revision: str = "f1a2b3c4d5e6" +down_revision: str | Sequence[str] | None = "e0a1b2c3d4e5" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _get_schema_prefix() -> str: + """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public).""" + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def upgrade() -> None: + """Add composite index for efficient MPFP edge loading.""" + schema = _get_schema_prefix() + # Create composite index for efficient top-k per (from_node, link_type) queries + # This enables LATERAL joins to use index-only scans with early termination + # Note: Not using CONCURRENTLY here as it requires running outside a transaction + # For production with large tables, consider running this manually with CONCURRENTLY + op.execute( + f"CREATE INDEX IF NOT EXISTS idx_memory_links_from_type_weight " + f"ON {schema}memory_links(from_unit_id, link_type, weight DESC)" + ) + + +def downgrade() -> None: + """Remove the composite index.""" + schema = _get_schema_prefix() + op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_links_from_type_weight") diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index afdb980e..35058543 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -1196,6 +1196,9 @@ def _register_routes(app: FastAPI): bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context) ): """Run a recall and return results with trace.""" + import time + + handler_start = time.time() metrics = get_metrics_collector() try: @@ -1221,10 +1224,12 @@ def _register_routes(app: FastAPI): include_chunks = request.include.chunks is not None max_chunk_tokens = request.include.chunks.max_tokens if include_chunks else 8192 + pre_recall = time.time() - handler_start # Run recall with tracing (record metrics) with metrics.record_operation( "recall", bank_id=bank_id, source="api", budget=request.budget.value, max_tokens=request.max_tokens ): + recall_start = time.time() core_result = await app.state.memory.recall_async( bank_id=bank_id, query=request.query, @@ -1283,9 +1288,21 @@ def _register_routes(app: FastAPI): ], ) - return RecallResponse( + response = RecallResponse( results=recall_results, trace=core_result.trace, entities=entities_response, chunks=chunks_response ) + + handler_duration = time.time() - handler_start + recall_duration = time.time() - recall_start + post_recall = handler_duration - pre_recall - recall_duration + if handler_duration > 1.0: + logging.info( + f"[RECALL HTTP] bank={bank_id} handler_total={handler_duration:.3f}s " + f"pre={pre_recall:.3f}s recall={recall_duration:.3f}s post={post_recall:.3f}s " + f"results={len(recall_results)} entities={len(entities_response) if entities_response else 0}" + ) + + return response except HTTPException: raise except OperationValidationError as e: @@ -1295,8 +1312,11 @@ def _register_routes(app: FastAPI): except Exception as e: import traceback + handler_duration = time.time() - handler_start error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}" - logger.error(f"Error in /v1/default/banks/{bank_id}/memories/recall: {error_detail}") + logger.error( + f"[RECALL ERROR] bank={bank_id} handler_duration={handler_duration:.3f}s error={str(e)}\n{error_detail}" + ) raise HTTPException(status_code=500, detail=str(e)) @app.post( diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 63bda959..2788e1f4 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -52,12 +52,18 @@ ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL" ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE" ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT" +ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES" +ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL" +ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR" ENV_HOST = "HINDSIGHT_API_HOST" ENV_PORT = "HINDSIGHT_API_PORT" ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL" +ENV_WORKERS = "HINDSIGHT_API_WORKERS" ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED" ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER" +ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS" +ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT" ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID" ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS" @@ -107,6 +113,9 @@ DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2" DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound reranking to prevent thrashing DEFAULT_RERANKER_TEI_BATCH_SIZE = 128 DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8 +DEFAULT_RERANKER_MAX_CANDIDATES = 300 +DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality +DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0" DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0" @@ -114,8 +123,11 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0" DEFAULT_HOST = "0.0.0.0" DEFAULT_PORT = 8888 DEFAULT_LOG_LEVEL = "info" +DEFAULT_WORKERS = 1 DEFAULT_MCP_ENABLED = True -DEFAULT_GRAPH_RETRIEVER = "mpfp" # Options: "mpfp", "bfs" +DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp", "bfs" +DEFAULT_MPFP_TOP_K_NEIGHBORS = 20 # Fan-out limit per node in MPFP graph traversal +DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker DEFAULT_MCP_LOCAL_BANK_ID = "mcp" # Observation thresholds @@ -217,6 +229,7 @@ class HindsightConfig: reranker_tei_url: str | None reranker_tei_batch_size: int reranker_tei_max_concurrent: int + reranker_max_candidates: int # Server host: str @@ -226,6 +239,8 @@ class HindsightConfig: # Recall graph_retriever: str + mpfp_top_k_neighbors: int + recall_max_concurrent: int # Observation thresholds observation_min_facts: int @@ -290,6 +305,7 @@ class HindsightConfig: reranker_tei_max_concurrent=int( os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT)) ), + reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))), # Server host=os.getenv(ENV_HOST, DEFAULT_HOST), port=int(os.getenv(ENV_PORT, DEFAULT_PORT)), @@ -297,6 +313,8 @@ class HindsightConfig: mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true", # Recall graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER), + mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))), + recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))), # Optimization flags skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true", lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true", diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 0370ad70..1189d131 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -16,6 +16,8 @@ import httpx from ..config import ( DEFAULT_RERANKER_COHERE_MODEL, + DEFAULT_RERANKER_FLASHRANK_CACHE_DIR, + DEFAULT_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT, DEFAULT_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_PROVIDER, @@ -23,6 +25,8 @@ from ..config import ( DEFAULT_RERANKER_TEI_MAX_CONCURRENT, ENV_COHERE_API_KEY, ENV_RERANKER_COHERE_MODEL, + ENV_RERANKER_FLASHRANK_CACHE_DIR, + ENV_RERANKER_FLASHRANK_MODEL, ENV_RERANKER_LOCAL_MAX_CONCURRENT, ENV_RERANKER_LOCAL_MODEL, ENV_RERANKER_PROVIDER, @@ -172,8 +176,13 @@ class RemoteTEICrossEncoder(CrossEncoderModel): Note: The TEI server must be running a cross-encoder/reranker model. Requests are made in parallel with configurable batch size and max concurrency (backpressure). + Uses a GLOBAL semaphore to limit concurrent requests across ALL recall operations. """ + # Global semaphore shared across all instances and calls to prevent thundering herd + _global_semaphore: asyncio.Semaphore | None = None + _global_max_concurrent: int = DEFAULT_RERANKER_TEI_MAX_CONCURRENT + def __init__( self, base_url: str, @@ -190,7 +199,8 @@ class RemoteTEICrossEncoder(CrossEncoderModel): base_url: Base URL of the TEI server (e.g., "http://localhost:8080") timeout: Request timeout in seconds (default: 30.0) batch_size: Maximum batch size for rerank requests (default: 128) - max_concurrent: Maximum concurrent requests for backpressure (default: 8) + max_concurrent: Maximum concurrent requests for backpressure (default: 8). + This is a GLOBAL limit across all parallel recall operations. max_retries: Maximum number of retries for failed requests (default: 3) retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5) """ @@ -203,6 +213,14 @@ class RemoteTEICrossEncoder(CrossEncoderModel): self._async_client: httpx.AsyncClient | None = None self._model_id: str | None = None + # Update global semaphore if max_concurrent changed + if ( + RemoteTEICrossEncoder._global_semaphore is None + or RemoteTEICrossEncoder._global_max_concurrent != max_concurrent + ): + RemoteTEICrossEncoder._global_max_concurrent = max_concurrent + RemoteTEICrossEncoder._global_semaphore = asyncio.Semaphore(max_concurrent) + @property def provider_name(self) -> str: return "tei" @@ -327,9 +345,10 @@ class RemoteTEICrossEncoder(CrossEncoderModel): batch_texts = texts[i : i + self.batch_size] tasks_info.append((query, batch_indices, batch_texts)) - # Run all requests in parallel with semaphore for backpressure + # Run all requests in parallel with GLOBAL semaphore for backpressure + # This ensures max_concurrent is respected across ALL parallel recall operations all_scores = [0.0] * len(pairs) - semaphore = asyncio.Semaphore(self.max_concurrent) + semaphore = RemoteTEICrossEncoder._global_semaphore tasks = [ self._rerank_query_group(self._async_client, semaphore, query, texts) for query, _, texts in tasks_info @@ -458,6 +477,170 @@ class CohereCrossEncoder(CrossEncoderModel): return all_scores +class RRFPassthroughCrossEncoder(CrossEncoderModel): + """ + Passthrough cross-encoder that preserves RRF scores without neural reranking. + + This is useful for: + - Testing retrieval quality without reranking overhead + - Deployments where reranking latency is unacceptable + - Debugging to isolate retrieval vs reranking issues + """ + + def __init__(self): + """Initialize RRF passthrough cross-encoder.""" + pass + + @property + def provider_name(self) -> str: + return "rrf" + + async def initialize(self) -> None: + """No initialization needed.""" + logger.info("Reranker: RRF passthrough provider initialized (neural reranking disabled)") + + async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: + """ + Return neutral scores - actual ranking uses RRF scores from retrieval. + + Args: + pairs: List of (query, document) tuples (ignored) + + Returns: + List of 0.5 scores (neutral, lets RRF scores dominate) + """ + # Return neutral scores so RRF ranking is preserved + return [0.5] * len(pairs) + + +class FlashRankCrossEncoder(CrossEncoderModel): + """ + FlashRank cross-encoder implementation. + + FlashRank is an ultra-lite reranking library that runs on CPU without + requiring PyTorch or Transformers. It's ideal for serverless deployments + with minimal cold-start overhead. + + Available models: + - ms-marco-TinyBERT-L-2-v2: Fastest, ~4MB + - ms-marco-MiniLM-L-12-v2: Best quality, ~34MB (default) + - rank-T5-flan: Best zero-shot, ~110MB + - ms-marco-MultiBERT-L-12: Multi-lingual, ~150MB + """ + + # Shared executor for CPU-bound reranking + _executor: ThreadPoolExecutor | None = None + _max_concurrent: int = 4 + + def __init__( + self, + model_name: str | None = None, + cache_dir: str | None = None, + max_length: int = 512, + max_concurrent: int = 4, + ): + """ + Initialize FlashRank cross-encoder. + + Args: + model_name: FlashRank model name. Default: ms-marco-MiniLM-L-12-v2 + cache_dir: Directory to cache downloaded models. Default: system cache + max_length: Maximum sequence length for reranking. Default: 512 + max_concurrent: Maximum concurrent reranking calls. Default: 4 + """ + self.model_name = model_name or DEFAULT_RERANKER_FLASHRANK_MODEL + self.cache_dir = cache_dir or DEFAULT_RERANKER_FLASHRANK_CACHE_DIR + self.max_length = max_length + self._ranker = None + FlashRankCrossEncoder._max_concurrent = max_concurrent + + @property + def provider_name(self) -> str: + return "flashrank" + + async def initialize(self) -> None: + """Load the FlashRank model.""" + if self._ranker is not None: + return + + try: + from flashrank import Ranker # type: ignore[import-untyped] + except ImportError: + raise ImportError("flashrank is required for FlashRankCrossEncoder. Install it with: pip install flashrank") + + logger.info(f"Reranker: initializing FlashRank provider with model {self.model_name}") + + # Initialize ranker with optional cache directory + ranker_kwargs = {"model_name": self.model_name, "max_length": self.max_length} + if self.cache_dir: + ranker_kwargs["cache_dir"] = self.cache_dir + + self._ranker = Ranker(**ranker_kwargs) + + # Initialize shared executor + if FlashRankCrossEncoder._executor is None: + FlashRankCrossEncoder._executor = ThreadPoolExecutor( + max_workers=FlashRankCrossEncoder._max_concurrent, + thread_name_prefix="flashrank", + ) + logger.info( + f"Reranker: FlashRank provider initialized (max_concurrent={FlashRankCrossEncoder._max_concurrent})" + ) + else: + logger.info("Reranker: FlashRank provider initialized (using existing executor)") + + def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]: + """Synchronous predict - processes each query group.""" + from flashrank import RerankRequest # type: ignore[import-untyped] + + if not pairs: + return [] + + # Group pairs by query + query_groups: dict[str, list[tuple[int, str]]] = {} + for idx, (query, text) in enumerate(pairs): + if query not in query_groups: + query_groups[query] = [] + query_groups[query].append((idx, text)) + + all_scores = [0.0] * len(pairs) + + for query, indexed_texts in query_groups.items(): + # Build passages list for FlashRank + passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)] + global_indices = [idx for idx, _ in indexed_texts] + + # Create rerank request + request = RerankRequest(query=query, passages=passages) + results = self._ranker.rerank(request) + + # Map scores back to original positions + for result in results: + local_idx = result["id"] + score = result["score"] + global_idx = global_indices[local_idx] + all_scores[global_idx] = score + + return all_scores + + async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: + """ + Score query-document pairs using FlashRank. + + Args: + pairs: List of (query, document) tuples to score + + Returns: + List of relevance scores (higher = more relevant) + """ + if self._ranker is None: + raise RuntimeError("Reranker not initialized. Call initialize() first.") + + # Run in thread pool to avoid blocking event loop + loop = asyncio.get_event_loop() + return await loop.run_in_executor(FlashRankCrossEncoder._executor, self._predict_sync, pairs) + + def create_cross_encoder_from_env() -> CrossEncoderModel: """ Create a CrossEncoderModel instance based on environment variables. @@ -489,5 +672,13 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'") model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL) return CohereCrossEncoder(api_key=api_key, model=model) + elif provider == "flashrank": + model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL) + cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR) + return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir) + elif provider == "rrf": + return RRFPassthroughCrossEncoder() else: - raise ValueError(f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere'") + raise ValueError( + f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'rrf'" + ) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 4319fa4c..b02addbb 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -409,10 +409,8 @@ class MemoryEngine(MemoryEngineInterface): self._task_backend = AsyncIOQueueBackend(batch_size=_task_batch_size, batch_interval=_task_batch_interval) # 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) + # Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50) + self._search_semaphore = asyncio.Semaphore(get_config().recall_max_concurrent) # Backpressure for put operations: limit concurrent puts to prevent database contention # Each put_batch holds a connection for the entire transaction, so we limit to 5 @@ -1418,7 +1416,9 @@ class MemoryEngine(MemoryEngineInterface): # Backpressure: limit concurrent recalls to prevent overwhelming the database result = None error_msg = None + semaphore_wait_start = time.time() async with self._search_semaphore: + semaphore_wait = time.time() - semaphore_wait_start # Retry loop for connection errors max_retries = 3 for attempt in range(max_retries + 1): @@ -1436,6 +1436,7 @@ class MemoryEngine(MemoryEngineInterface): include_chunks, max_chunk_tokens, request_context, + semaphore_wait=semaphore_wait, ) break # Success - exit retry loop except Exception as e: @@ -1553,6 +1554,7 @@ class MemoryEngine(MemoryEngineInterface): include_chunks: bool = False, max_chunk_tokens: int = 8192, request_context: "RequestContext" = None, + semaphore_wait: float = 0.0, ) -> RecallResultModel: """ Search implementation with modular retrieval and reranking. @@ -1612,20 +1614,16 @@ class MemoryEngine(MemoryEngineInterface): query_embedding_str = str(query_embedding) from .search.retrieval import get_default_graph_retriever, retrieve_parallel - from .search.temporal_extraction import extract_temporal_constraint # Track each retrieval start time retrieval_start = time.time() - # Pre-extract temporal constraint once (shared across all fact types) - tc_start = time.time() - temporal_constraint = extract_temporal_constraint( - query, reference_date=question_date, analyzer=self.query_analyzer - ) - tc_duration = time.time() - tc_start + # Temporal extraction now runs IN PARALLEL with other retrievals inside retrieve_parallel + # This prevents slow dateparser from blocking semantic/BM25/graph retrieval + tc_duration = 0.0 # Will be tracked inside temporal retrieval timing # Run retrieval for each fact type in parallel - # MPFP does lazy edge loading internally, no need to pre-load adjacency + # Each retrieve_parallel uses ~4 connections, so 3 fact types = ~12 concurrent connections retrieval_tasks = [ retrieve_parallel( pool, @@ -1636,7 +1634,7 @@ class MemoryEngine(MemoryEngineInterface): thinking_budget, question_date, self.query_analyzer, - temporal_constraint=temporal_constraint, + temporal_constraint=None, # Extracted in parallel inside retrieve_parallel ) for ft in fact_type ] @@ -1649,10 +1647,17 @@ class MemoryEngine(MemoryEngineInterface): bm25_results = [] graph_results = [] temporal_results = [] - aggregated_timings = {"semantic": 0.0, "bm25": 0.0, "graph": 0.0, "temporal": 0.0} + aggregated_timings = { + "semantic": 0.0, + "bm25": 0.0, + "graph": 0.0, + "temporal": 0.0, + "temporal_extraction": 0.0, + } all_mpfp_timings = [] detected_temporal_constraint = None + max_conn_wait = 0.0 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" @@ -1673,6 +1678,8 @@ class MemoryEngine(MemoryEngineInterface): detected_temporal_constraint = retrieval_result.temporal_constraint # Collect MPFP timings all_mpfp_timings.extend(retrieval_result.mpfp_timings) + # Track max connection wait + max_conn_wait = max(max_conn_wait, retrieval_result.max_conn_wait) # If no temporal results from any fact type, set to None if not temporal_results: @@ -1696,6 +1703,7 @@ class MemoryEngine(MemoryEngineInterface): f"semantic={len(semantic_results)}({aggregated_timings['semantic']:.3f}s)", f"bm25={len(bm25_results)}({aggregated_timings['bm25']:.3f}s)", f"graph={len(graph_results)}({aggregated_timings['graph']:.3f}s)", + f"temporal_extraction={aggregated_timings['temporal_extraction']:.3f}s", ] temporal_info = "" if detected_temporal_constraint: @@ -1709,8 +1717,9 @@ class MemoryEngine(MemoryEngineInterface): f" [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{setup_info}{temporal_info}" ) - # Log MPFP timing breakdown if available + # Log graph retriever timing breakdown if available if all_mpfp_timings: + retriever_name = get_default_graph_retriever().name.upper() mpfp_total = all_mpfp_timings[0] # Take first fact type's timing as representative mpfp_parts = [ f"db_queries={mpfp_total.db_queries}", @@ -1720,7 +1729,20 @@ class MemoryEngine(MemoryEngineInterface): ] if mpfp_total.seeds_time > 0.01: mpfp_parts.append(f"seeds={mpfp_total.seeds_time:.3f}s") - log_buffer.append(f" [MPFP] {', '.join(mpfp_parts)}") + if mpfp_total.fusion > 0.001: + mpfp_parts.append(f"fusion={mpfp_total.fusion:.3f}s") + if mpfp_total.fetch > 0.001: + mpfp_parts.append(f"fetch={mpfp_total.fetch:.3f}s") + log_buffer.append(f" [{retriever_name}] {', '.join(mpfp_parts)}") + # Log detailed hop timing for debugging slow queries + if mpfp_total.hop_details: + for hd in mpfp_total.hop_details: + log_buffer.append( + f" hop{hd['hop']}: exec={hd.get('exec_time', 0) * 1000:.0f}ms, " + f"uncached={hd.get('uncached_after_filter', 0)}, " + f"load={hd.get('load_time', 0) * 1000:.0f}ms, " + f"edges={hd.get('edges_loaded', 0)}" + ) # Record retrieval results for tracer - per fact type if tracer: @@ -1819,11 +1841,24 @@ class MemoryEngine(MemoryEngineInterface): # Ensure reranker is initialized (for lazy initialization mode) await reranker_instance.ensure_initialized() + # Pre-filter candidates to reduce reranking cost (RRF already provides good ranking) + # This is especially important for remote rerankers with network latency + reranker_max_candidates = get_config().reranker_max_candidates + pre_filtered_count = 0 + if len(merged_candidates) > reranker_max_candidates: + # Sort by RRF score and take top candidates + merged_candidates.sort(key=lambda mc: mc.rrf_score, reverse=True) + pre_filtered_count = len(merged_candidates) - reranker_max_candidates + merged_candidates = merged_candidates[:reranker_max_candidates] + # Rerank using cross-encoder scored_results = await reranker_instance.rerank(query, merged_candidates) step_duration = time.time() - step_start - log_buffer.append(f" [4] Reranking: {len(scored_results)} candidates scored in {step_duration:.3f}s") + pre_filter_note = f" (pre-filtered {pre_filtered_count})" if pre_filtered_count > 0 else "" + log_buffer.append( + f" [4] Reranking: {len(scored_results)} candidates scored in {step_duration:.3f}s{pre_filter_note}" + ) # Step 4.5: Combine cross-encoder score with retrieval signals # This preserves retrieval work (RRF, temporal, recency) instead of pure cross-encoder ranking @@ -2153,8 +2188,15 @@ class MemoryEngine(MemoryEngineInterface): 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 + # Include wait times in log if significant + wait_parts = [] + if semaphore_wait > 0.01: + wait_parts.append(f"sem={semaphore_wait:.3f}s") + if max_conn_wait > 0.01: + wait_parts.append(f"conn={max_conn_wait:.3f}s") + wait_info = f" | waits: {', '.join(wait_parts)}" if wait_parts else "" log_buffer.append( - f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok), {num_entities} entities ({total_entity_tokens} tok) | {fact_type_summary} | {total_time:.3f}s" + 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{wait_info}" ) logger.info("\n" + "\n".join(log_buffer)) diff --git a/hindsight-api/hindsight_api/engine/query_analyzer.py b/hindsight-api/hindsight_api/engine/query_analyzer.py index a6e6a303..5c9b568d 100644 --- a/hindsight-api/hindsight_api/engine/query_analyzer.py +++ b/hindsight-api/hindsight_api/engine/query_analyzer.py @@ -84,7 +84,7 @@ class DateparserQueryAnalyzer(QueryAnalyzer): Performance: - ~10-50ms per query - - No model loading required + - No model loading required (lazy import on first use) """ def __init__(self): @@ -112,8 +112,6 @@ class DateparserQueryAnalyzer(QueryAnalyzer): Returns: QueryAnalysis with temporal_constraint if found """ - self.load() - if reference_date is None: reference_date = datetime.now() @@ -123,6 +121,9 @@ class DateparserQueryAnalyzer(QueryAnalyzer): if period_result is not None: return QueryAnalysis(temporal_constraint=period_result) + # Lazy load dateparser (only imports on first call, then cached) + self.load() + # Use dateparser's search_dates to find temporal expressions settings = { "RELATIVE_BASE": reference_date, diff --git a/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py b/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py new file mode 100644 index 00000000..ccd88ef2 --- /dev/null +++ b/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py @@ -0,0 +1,230 @@ +""" +Link Expansion graph retrieval. + +A simple, fast graph retrieval that expands from seeds via: +1. Entity links: Find facts sharing entities with seeds (filtered by entity frequency) +2. Causal links: Find facts causally linked to seeds (top-k by weight) + +Characteristics: +- 2-3 DB queries (seed finding + parallel entity/causal expansion) +- Sublinear: only touches connected facts via indexes +- No iteration, no propagation, no normalization +- Target: <100ms +""" + +import logging +import time + +from ..db_utils import acquire_with_retry +from ..memory_engine import fq_table +from .graph_retrieval import GraphRetriever +from .types import MPFPTimings, RetrievalResult + +logger = logging.getLogger(__name__) + + +async def _find_semantic_seeds( + conn, + query_embedding_str: str, + bank_id: str, + fact_type: str, + limit: int = 20, + threshold: float = 0.3, +) -> list[RetrievalResult]: + """Find semantic seeds via embedding search.""" + rows = await conn.fetch( + f""" + SELECT id, text, context, event_date, occurred_start, occurred_end, + mentioned_at, access_count, embedding, fact_type, document_id, chunk_id, + 1 - (embedding <=> $1::vector) AS similarity + FROM {fq_table("memory_units")} + WHERE bank_id = $2 + AND embedding IS NOT NULL + AND fact_type = $3 + AND (1 - (embedding <=> $1::vector)) >= $4 + ORDER BY embedding <=> $1::vector + LIMIT $5 + """, + query_embedding_str, + bank_id, + fact_type, + threshold, + limit, + ) + return [RetrievalResult.from_db_row(dict(r)) for r in rows] + + +class LinkExpansionRetriever(GraphRetriever): + """ + Graph retrieval via direct link expansion from seeds. + + Expands through entity co-occurrence and causal links in a single query. + Fast and simple alternative to MPFP. + """ + + def __init__( + self, + max_entity_frequency: int = 500, + causal_weight_threshold: float = 0.3, + causal_limit_per_seed: int = 10, + ): + """ + Initialize link expansion retriever. + + Args: + max_entity_frequency: Skip entities appearing in more than this many facts + causal_weight_threshold: Minimum weight for causal links + causal_limit_per_seed: Max causal links to follow per seed + """ + self.max_entity_frequency = max_entity_frequency + self.causal_weight_threshold = causal_weight_threshold + self.causal_limit_per_seed = causal_limit_per_seed + + @property + def name(self) -> str: + return "link_expansion" + + async def retrieve( + self, + pool, + query_embedding_str: str, + bank_id: str, + fact_type: str, + budget: int, + query_text: str | None = None, + semantic_seeds: list[RetrievalResult] | None = None, + temporal_seeds: list[RetrievalResult] | None = None, + adjacency=None, + ) -> tuple[list[RetrievalResult], MPFPTimings | None]: + """ + Retrieve facts by expanding links from seeds. + + Args: + pool: Database connection pool + query_embedding_str: Query embedding (unused, kept for interface) + bank_id: Memory bank ID + fact_type: Fact type to filter + budget: Maximum results to return + query_text: Original query text (unused) + semantic_seeds: Pre-computed semantic entry points + temporal_seeds: Pre-computed temporal entry points + adjacency: Unused, kept for interface compatibility + + Returns: + Tuple of (results, timings) + """ + start_time = time.time() + timings = MPFPTimings(fact_type=fact_type) + + # Use single connection for all queries to reduce pool pressure + # (queries are fast ~50ms each, connection acquisition is the bottleneck) + async with acquire_with_retry(pool) as conn: + # Find seeds if not provided + if semantic_seeds: + all_seeds = list(semantic_seeds) + else: + seeds_start = time.time() + all_seeds = await _find_semantic_seeds( + conn, query_embedding_str, bank_id, fact_type, limit=20, threshold=0.3 + ) + timings.seeds_time = time.time() - seeds_start + + # Add temporal seeds if provided + if temporal_seeds: + all_seeds.extend(temporal_seeds) + + if not all_seeds: + return [], timings + + seed_ids = list({s.id for s in all_seeds}) + timings.pattern_count = len(seed_ids) + + # Run entity and causal expansion sequentially on same connection + query_start = time.time() + + entity_rows = await conn.fetch( + f""" + SELECT + mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, + mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, + mu.fact_type, mu.document_id, mu.chunk_id, + COUNT(*)::float AS score + FROM {fq_table("unit_entities")} seed_ue + JOIN {fq_table("entities")} e ON seed_ue.entity_id = e.id + JOIN {fq_table("unit_entities")} other_ue ON seed_ue.entity_id = other_ue.entity_id + JOIN {fq_table("memory_units")} mu ON other_ue.unit_id = mu.id + WHERE seed_ue.unit_id = ANY($1::uuid[]) + AND e.mention_count < $2 + AND mu.id != ALL($1::uuid[]) + AND mu.fact_type = $3 + GROUP BY mu.id + ORDER BY score DESC + LIMIT $4 + """, + seed_ids, + self.max_entity_frequency, + fact_type, + budget, + ) + + causal_rows = await conn.fetch( + f""" + SELECT DISTINCT ON (mu.id) + mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, + mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, + mu.fact_type, mu.document_id, mu.chunk_id, + ml.weight + 1.0 AS score + FROM {fq_table("memory_links")} ml + JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id + WHERE ml.from_unit_id = ANY($1::uuid[]) + AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents') + AND ml.weight >= $2 + AND mu.fact_type = $3 + ORDER BY mu.id, ml.weight DESC + LIMIT $4 + """, + seed_ids, + self.causal_weight_threshold, + fact_type, + budget, + ) + + timings.edge_load_time = time.time() - query_start + timings.db_queries = 2 + timings.edge_count = len(entity_rows) + len(causal_rows) + + # Merge results, taking max score per fact + score_map: dict[str, float] = {} + row_map: dict[str, dict] = {} + + for row in entity_rows: + fact_id = str(row["id"]) + score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"]) + row_map[fact_id] = dict(row) + + for row in causal_rows: + fact_id = str(row["id"]) + score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"]) + if fact_id not in row_map: + row_map[fact_id] = dict(row) + + # Sort by score and limit + sorted_ids = sorted(score_map.keys(), key=lambda x: score_map[x], reverse=True)[:budget] + rows = [row_map[fact_id] for fact_id in sorted_ids] + + # Convert to results + results = [] + for row in rows: + result = RetrievalResult.from_db_row(dict(row)) + result.activation = row["score"] + results.append(result) + + timings.result_count = len(results) + timings.traverse = time.time() - start_time + + logger.debug( + f"LinkExpansion: {len(results)} results from {len(seed_ids)} seeds " + f"in {timings.traverse * 1000:.1f}ms (query: {timings.edge_load_time * 1000:.1f}ms)" + ) + + return results, timings diff --git a/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py index 2f20810c..aa896930 100644 --- a/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py @@ -49,6 +49,7 @@ class EdgeCache: Grows per-hop as edges are loaded for frontier nodes. Shared across patterns to avoid redundant loads. Loads ALL edge types at once to minimize DB queries. + Thread-safe via asyncio lock to prevent redundant concurrent loads. """ # edge_type -> from_node_id -> list of EdgeTarget @@ -58,6 +59,10 @@ class EdgeCache: # Timing stats db_queries: int = 0 edge_load_time: float = 0.0 + # Detailed hop timing for debugging + hop_details: list[dict] = field(default_factory=list) + # Lock to prevent redundant concurrent loads + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) def get_neighbors(self, edge_type: str, node_id: str) -> list[EdgeTarget]: """Get neighbors for a node via a specific edge type.""" @@ -153,13 +158,20 @@ class SeedNode: async def load_all_edges_for_frontier( pool, node_ids: list[str], + top_k_per_type: int = 20, ) -> dict[str, dict[str, list[EdgeTarget]]]: """ - Load ALL edge types for frontier nodes in one query. + Load top-k edges per (node, edge_type) for frontier nodes. + + Uses a LATERAL join to efficiently fetch only the top-k edges per type, + avoiding loading hundreds of entity edges when only 20 are needed. + + Requires composite index: (from_unit_id, link_type, weight DESC) Args: pool: Database connection pool node_ids: Frontier node IDs to load edges for + top_k_per_type: Max edges to load per (node, link_type) pair Returns: Dict mapping edge_type -> from_node_id -> list of EdgeTarget @@ -168,15 +180,26 @@ async def load_all_edges_for_frontier( return {} async with acquire_with_retry(pool) as conn: + # Use LATERAL join to get top-k per (from_node, link_type) + # This leverages the composite index for efficient early termination rows = await conn.fetch( f""" - SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight - FROM {fq_table("memory_links")} ml - WHERE ml.from_unit_id = ANY($1::uuid[]) - AND ml.weight >= 0.1 - ORDER BY ml.from_unit_id, ml.link_type, ml.weight DESC + WITH frontier(node_id) AS (SELECT unnest($1::uuid[])) + SELECT f.node_id as from_unit_id, lt.link_type, edges.to_unit_id, edges.weight + FROM frontier f + CROSS JOIN (VALUES ('semantic'), ('temporal'), ('entity'), ('causes'), ('caused_by')) AS lt(link_type) + CROSS JOIN LATERAL ( + SELECT ml.to_unit_id, ml.weight + FROM {fq_table("memory_links")} ml + WHERE ml.from_unit_id = f.node_id + AND ml.link_type = lt.link_type + AND ml.weight >= 0.1 + ORDER BY ml.weight DESC + LIMIT $2 + ) edges """, node_ids, + top_k_per_type, ) # Group by edge_type -> from_node -> neighbors @@ -197,6 +220,162 @@ async def load_all_edges_for_frontier( # ----------------------------------------------------------------------------- +@dataclass +class PatternState: + """State for a pattern traversal between hops.""" + + pattern: list[str] + hop_index: int + scores: dict[str, float] + frontier: dict[str, float] + + +def _init_pattern_state(seeds: list[SeedNode], pattern: list[str]) -> PatternState: + """Initialize pattern state from seeds.""" + if not seeds: + return PatternState(pattern=pattern, hop_index=0, scores={}, frontier={}) + + total_seed_score = sum(s.score for s in seeds) + if total_seed_score == 0: + total_seed_score = len(seeds) + + frontier = {s.node_id: s.score / total_seed_score for s in seeds} + return PatternState(pattern=pattern, hop_index=0, scores={}, frontier=frontier) + + +def _execute_hop(state: PatternState, cache: EdgeCache, config: MPFPConfig) -> set[str]: + """ + Execute ONE hop of traversal, return frontier nodes for next hop. + + This is a pure function that uses cached edges (no DB access). + Returns set of uncached nodes needed for next hop. + """ + if state.hop_index >= len(state.pattern): + return set() + + edge_type = state.pattern[state.hop_index] + + # Collect active nodes above threshold + active_nodes = [node_id for node_id, mass in state.frontier.items() if mass >= config.threshold] + if not active_nodes: + state.frontier = {} + return set() + + # Propagate mass using cached edges + next_frontier: dict[str, float] = {} + uncached_for_next: set[str] = set() + + for node_id, mass in state.frontier.items(): + if mass < config.threshold: + continue + + # Keep α portion for this node + state.scores[node_id] = state.scores.get(node_id, 0) + config.alpha * mass + + # Push (1-α) to neighbors + push_mass = (1 - config.alpha) * mass + neighbors = cache.get_normalized_neighbors(edge_type, node_id, config.top_k_neighbors) + + for neighbor in neighbors: + next_frontier[neighbor.node_id] = next_frontier.get(neighbor.node_id, 0) + push_mass * neighbor.weight + # Track if we'll need edges for this node in the next hop + if not cache.is_fully_loaded(neighbor.node_id): + uncached_for_next.add(neighbor.node_id) + + state.frontier = next_frontier + state.hop_index += 1 + + return uncached_for_next + + +def _finalize_pattern(state: PatternState, config: MPFPConfig) -> PatternResult: + """Finalize pattern by adding remaining frontier mass to scores.""" + for node_id, mass in state.frontier.items(): + if mass >= config.threshold: + state.scores[node_id] = state.scores.get(node_id, 0) + mass + + return PatternResult(pattern=state.pattern, scores=state.scores) + + +async def mpfp_traverse_hop_synchronized( + pool, + pattern_jobs: list[tuple[list[SeedNode], list[str]]], + config: MPFPConfig, + cache: EdgeCache, +) -> list[PatternResult]: + """ + Execute ALL patterns with hop-synchronized edge loading. + + Instead of running each pattern independently (causing multiple DB queries), + this function: + 1. Runs hop 1 for ALL patterns (using pre-warmed seed edges) + 2. Collects ALL unique hop-2 frontier nodes across patterns + 3. Pre-warms hop-2 edges in ONE query + 4. Runs hop 2 for ALL patterns + + This reduces DB queries from O(patterns * hops) to O(hops). + + Args: + pool: Database connection pool + pattern_jobs: List of (seeds, pattern) tuples + config: Algorithm parameters + cache: Shared edge cache (should be pre-warmed with seed edges) + + Returns: + List of PatternResult for each pattern + """ + import time + + # Initialize all pattern states + states = [_init_pattern_state(seeds, pattern) for seeds, pattern in pattern_jobs] + + # Determine max hops (all patterns should be same length, but be safe) + max_hops = max((len(p) for _, p in pattern_jobs), default=0) + + # Detailed timing for debugging + hop_times: list[dict] = [] + + # Execute hop-by-hop across ALL patterns + for hop in range(max_hops): + hop_start = time.time() + hop_timing = {"hop": hop, "patterns_executed": 0, "uncached_count": 0, "load_time": 0.0} + + # Execute this hop for all patterns, collect uncached nodes for next hop + all_uncached: set[str] = set() + exec_start = time.time() + for state in states: + if state.hop_index < len(state.pattern): + uncached = _execute_hop(state, cache, config) + all_uncached.update(uncached) + hop_timing["patterns_executed"] += 1 + hop_timing["exec_time"] = time.time() - exec_start + + # Pre-warm edges for ALL uncached nodes before next hop + hop_timing["uncached_count"] = len(all_uncached) + if all_uncached: + uncached_list = list(all_uncached - cache._fully_loaded) + hop_timing["uncached_after_filter"] = len(uncached_list) + if uncached_list: + load_start = time.time() + edges_by_type = await load_all_edges_for_frontier(pool, uncached_list, config.top_k_neighbors) + hop_timing["load_time"] = time.time() - load_start + cache.edge_load_time += hop_timing["load_time"] + cache.db_queries += 1 + cache.add_all_edges(edges_by_type, uncached_list) + hop_timing["edges_loaded"] = sum( + len(neighbors) for edges in edges_by_type.values() for neighbors in edges.values() + ) + + hop_timing["total_time"] = time.time() - hop_start + hop_times.append(hop_timing) + + # Store hop timing details in cache for logging + cache.hop_details = hop_times + + # Finalize all patterns + return [_finalize_pattern(state, config) for state in states] + + async def mpfp_traverse_async( pool, seeds: list[SeedNode], @@ -207,76 +386,14 @@ async def mpfp_traverse_async( """ Async Forward Push traversal with lazy edge loading. - Loads ALL edge types per hop to minimize DB queries. - - Args: - pool: Database connection pool - seeds: Entry point nodes with initial scores - pattern: Sequence of edge types to follow - config: Algorithm parameters - cache: Shared edge cache (grows as edges are loaded) - - Returns: - PatternResult with accumulated scores per node + NOTE: For better performance with multiple patterns, use mpfp_traverse_hop_synchronized(). + This function is kept for single-pattern use cases. """ if not seeds: return PatternResult(pattern=pattern, scores={}) - scores: dict[str, float] = {} - - # Initialize frontier with seed masses (normalized) - total_seed_score = sum(s.score for s in seeds) - if total_seed_score == 0: - total_seed_score = len(seeds) # fallback to uniform - - frontier: dict[str, float] = {s.node_id: s.score / total_seed_score for s in seeds} - - # Follow pattern hop by hop - for edge_type in pattern: - # Collect frontier nodes above threshold - active_nodes = [node_id for node_id, mass in frontier.items() if mass >= config.threshold] - - if not active_nodes: - break - - # Find nodes that need edge loading (all edge types at once) - uncached = cache.get_uncached(active_nodes) - - # Batch load ALL edges for uncached nodes (one query for all edge types) - if uncached: - import time - - load_start = time.time() - edges_by_type = await load_all_edges_for_frontier(pool, uncached) - cache.edge_load_time += time.time() - load_start - cache.db_queries += 1 - cache.add_all_edges(edges_by_type, uncached) - - # Propagate mass - next_frontier: dict[str, float] = {} - - for node_id, mass in frontier.items(): - if mass < config.threshold: - continue - - # Keep α portion for this node - scores[node_id] = scores.get(node_id, 0) + config.alpha * mass - - # Push (1-α) to neighbors - push_mass = (1 - config.alpha) * mass - neighbors = cache.get_normalized_neighbors(edge_type, node_id, config.top_k_neighbors) - - for neighbor in neighbors: - next_frontier[neighbor.node_id] = next_frontier.get(neighbor.node_id, 0) + push_mass * neighbor.weight - - frontier = next_frontier - - # Final frontier nodes get their remaining mass - for node_id, mass in frontier.items(): - if mass >= config.threshold: - scores[node_id] = scores.get(node_id, 0) + mass - - return PatternResult(pattern=pattern, scores=scores) + results = await mpfp_traverse_hop_synchronized(pool, [(seeds, pattern)], config, cache) + return results[0] if results else PatternResult(pattern=pattern, scores={}) def rrf_fusion( @@ -363,7 +480,13 @@ class MPFPGraphRetriever(GraphRetriever): Args: config: Algorithm configuration (uses defaults if None) """ - self.config = config or MPFPConfig() + if config is None: + # Read top_k_neighbors from global config + from ...config import get_config + + global_config = get_config() + config = MPFPConfig(top_k_neighbors=global_config.mpfp_top_k_neighbors) + self.config = config @property def name(self) -> str: @@ -433,18 +556,30 @@ class MPFPGraphRetriever(GraphRetriever): # Shared edge cache across all patterns cache = EdgeCache() - # Run all patterns in parallel (each does lazy edge loading) + # Pre-warm cache with ALL seed node edges BEFORE running patterns + # This prevents redundant DB queries at hop 1 + all_seed_ids = list({s.node_id for seeds, _ in pattern_jobs for s in seeds}) + if all_seed_ids: + import time as time_module + + prewarm_start = time_module.time() + edges_by_type = await load_all_edges_for_frontier(pool, all_seed_ids, self.config.top_k_neighbors) + cache.edge_load_time += time_module.time() - prewarm_start + cache.db_queries += 1 + cache.add_all_edges(edges_by_type, all_seed_ids) + + # Run all patterns with HOP-SYNCHRONIZED edge loading + # This batches hop-2 edge loads across ALL patterns into ONE query + # Reduces DB queries from O(patterns * hops) to O(hops) step_start = time.time() - pattern_tasks = [ - mpfp_traverse_async(pool, seeds, pattern, self.config, cache) for seeds, pattern in pattern_jobs - ] - pattern_results = await asyncio.gather(*pattern_tasks) + pattern_results = await mpfp_traverse_hop_synchronized(pool, pattern_jobs, self.config, cache) timings.traverse = time.time() - step_start # Record edge loading stats from cache timings.edge_count = sum(len(neighbors) for g in cache.graphs.values() for neighbors in g.values()) timings.db_queries = cache.db_queries timings.edge_load_time = cache.edge_load_time + timings.hop_details = cache.hop_details # Fuse results step_start = time.time() diff --git a/hindsight-api/hindsight_api/engine/search/retrieval.py b/hindsight-api/hindsight_api/engine/search/retrieval.py index fc9619f4..dc379744 100644 --- a/hindsight-api/hindsight_api/engine/search/retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/retrieval.py @@ -18,6 +18,7 @@ from ...config import get_config from ..db_utils import acquire_with_retry from ..memory_engine import fq_table from .graph_retrieval import BFSGraphRetriever, GraphRetriever +from .link_expansion_retrieval import LinkExpansionRetriever from .mpfp_retrieval import MPFPGraphRetriever from .types import MPFPTimings, RetrievalResult @@ -35,6 +36,7 @@ class ParallelRetrievalResult: timings: dict[str, float] = field(default_factory=dict) temporal_constraint: tuple | None = None # (start_date, end_date) mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type + max_conn_wait: float = 0.0 # Maximum connection acquisition wait time across all methods # Default graph retriever instance (can be overridden) @@ -49,13 +51,18 @@ def get_default_graph_retriever() -> GraphRetriever: retriever_type = config.graph_retriever.lower() if retriever_type == "mpfp": _default_graph_retriever = MPFPGraphRetriever() - logger.info("Using MPFP graph retriever") + logger.info( + f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})" + ) elif retriever_type == "bfs": _default_graph_retriever = BFSGraphRetriever() logger.info("Using BFS graph retriever") + elif retriever_type == "link_expansion": + _default_graph_retriever = LinkExpansionRetriever() + logger.info("Using LinkExpansion graph retriever") else: - logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to MPFP") - _default_graph_retriever = MPFPGraphRetriever() + logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to link_expansion") + _default_graph_retriever = LinkExpansionRetriever() return _default_graph_retriever @@ -390,17 +397,11 @@ async def retrieve_parallel( Returns: ParallelRetrievalResult with semantic, bm25, graph, temporal results and timings """ - # Extract temporal constraint if not pre-provided - if temporal_constraint is None: - from .temporal_extraction import extract_temporal_constraint - - temporal_constraint = extract_temporal_constraint( - query_text, reference_date=question_date, analyzer=query_analyzer - ) - retriever = graph_retriever or get_default_graph_retriever() - if retriever.name == "mpfp": + # Use optimized parallel path for MPFP and LinkExpansion (runs all methods truly in parallel) + # BFS uses legacy path that extracts temporal constraint upfront + if retriever.name in ("mpfp", "link_expansion"): return await _retrieve_parallel_mpfp( pool, query_text, @@ -410,8 +411,17 @@ async def retrieve_parallel( thinking_budget, temporal_constraint, retriever, + question_date, + query_analyzer, ) else: + # For BFS, extract temporal constraint upfront (legacy path) + if temporal_constraint is None: + from .temporal_extraction import extract_temporal_constraint + + temporal_constraint = extract_temporal_constraint( + query_text, reference_date=question_date, analyzer=query_analyzer + ) return await _retrieve_parallel_bfs( pool, query_text, query_embedding_str, bank_id, fact_type, thinking_budget, temporal_constraint, retriever ) @@ -423,6 +433,7 @@ class _TimedResult: results: list[RetrievalResult] time: float + conn_wait: float = 0.0 # Connection acquisition wait time async def _retrieve_parallel_mpfp( @@ -434,6 +445,8 @@ async def _retrieve_parallel_mpfp( thinking_budget: int, temporal_constraint: tuple | None, retriever: GraphRetriever, + question_date: datetime | None = None, + query_analyzer=None, ) -> ParallelRetrievalResult: """ MPFP retrieval with true parallelization. @@ -442,40 +455,37 @@ async def _retrieve_parallel_mpfp( - Semantic: vector similarity search - BM25: keyword search - Graph: MPFP traversal (does its own semantic seeds internally) - - Temporal: date-range search (if constraint detected) + - Temporal: date extraction (if needed) + date-range search - Graph does its own semantic query for seeds, avoiding chain dependency. + Temporal extraction runs IN PARALLEL with other retrievals, so even if + dateparser is slow, it doesn't block semantic/BM25/graph. """ import time async def run_semantic() -> _TimedResult: """Independent semantic retrieval.""" start = time.time() + acquire_start = time.time() async with acquire_with_retry(pool) as conn: + conn_wait = time.time() - acquire_start results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget) - return _TimedResult(results, time.time() - start) + return _TimedResult(results, time.time() - start, conn_wait) async def run_bm25() -> _TimedResult: """Independent BM25 retrieval.""" start = time.time() + acquire_start = time.time() async with acquire_with_retry(pool) as conn: + conn_wait = time.time() - acquire_start results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget) - return _TimedResult(results, time.time() - start) + return _TimedResult(results, time.time() - start, conn_wait) async def run_graph() -> tuple[list[RetrievalResult], float, MPFPTimings | None]: """Independent graph retrieval - does its own semantic seeds.""" start = time.time() - # Get temporal seeds if needed (graph uses them for temporal patterns) - temporal_seeds = None - if temporal_constraint: - tc_start, tc_end = temporal_constraint - async with acquire_with_retry(pool) as conn: - temporal_seeds = await _get_temporal_entry_points( - conn, query_embedding_str, bank_id, fact_type, tc_start, tc_end, limit=20 - ) - # MPFP does its own semantic seeds via _find_semantic_seeds + # Note: temporal_seeds not used here to avoid dependency on temporal extraction results, mpfp_timing = await retriever.retrieve( pool=pool, query_embedding_str=query_embedding_str, @@ -484,14 +494,50 @@ async def _retrieve_parallel_mpfp( budget=thinking_budget, query_text=query_text, semantic_seeds=None, # Let MPFP find its own seeds - temporal_seeds=temporal_seeds, + temporal_seeds=None, # Don't wait for temporal extraction ) return results, time.time() - start, mpfp_timing - async def run_temporal(tc_start, tc_end) -> _TimedResult: - """Independent temporal retrieval.""" + @dataclass + class _TemporalWithConstraint: + """Temporal results with the extracted constraint.""" + + results: list[RetrievalResult] + time: float + constraint: tuple | None + extraction_time: float # Time spent in query analyzer (dateparser) + conn_wait: float = 0.0 # Connection acquisition wait time + + async def run_temporal_with_extraction() -> _TemporalWithConstraint: + """ + Extract temporal constraint AND run temporal retrieval. + + This runs in parallel with semantic/BM25/graph, so dateparser + latency doesn't block other retrievals. + """ start = time.time() + + # Use pre-provided constraint if available + tc = temporal_constraint + extraction_time = 0.0 + + # Otherwise extract from query (this is the potentially slow dateparser call) + if tc is None: + from .temporal_extraction import extract_temporal_constraint + + extraction_start = time.time() + tc = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer) + extraction_time = time.time() - extraction_start + + # If no temporal constraint found, return empty (but still report extraction time) + if tc is None: + return _TemporalWithConstraint([], time.time() - start, None, extraction_time, 0.0) + + # Run temporal retrieval with the extracted constraint + tc_start, tc_end = tc + acquire_start = time.time() async with acquire_with_retry(pool) as conn: + conn_wait = time.time() - acquire_start results = await retrieve_temporal( conn, query_embedding_str, @@ -502,52 +548,36 @@ async def _retrieve_parallel_mpfp( budget=thinking_budget, semantic_threshold=0.1, ) - return _TimedResult(results, time.time() - start) + return _TemporalWithConstraint(results, time.time() - start, tc, extraction_time, conn_wait) - # Run all methods in parallel (no chain dependencies) - if temporal_constraint: - tc_start, tc_end = temporal_constraint - semantic_result, bm25_result, graph_result, temporal_result = await asyncio.gather( - run_semantic(), - run_bm25(), - run_graph(), - run_temporal(tc_start, tc_end), - ) - graph_results, graph_time, mpfp_timing = graph_result - return ParallelRetrievalResult( - semantic=semantic_result.results, - bm25=bm25_result.results, - graph=graph_results, - temporal=temporal_result.results, - timings={ - "semantic": semantic_result.time, - "bm25": bm25_result.time, - "graph": graph_time, - "temporal": temporal_result.time, - }, - temporal_constraint=temporal_constraint, - mpfp_timings=[mpfp_timing] if mpfp_timing else [], - ) - else: - semantic_result, bm25_result, graph_result = await asyncio.gather( - run_semantic(), - run_bm25(), - run_graph(), - ) - graph_results, graph_time, mpfp_timing = graph_result - return ParallelRetrievalResult( - semantic=semantic_result.results, - bm25=bm25_result.results, - graph=graph_results, - temporal=None, - timings={ - "semantic": semantic_result.time, - "bm25": bm25_result.time, - "graph": graph_time, - }, - temporal_constraint=None, - mpfp_timings=[mpfp_timing] if mpfp_timing else [], - ) + # Run ALL methods in parallel (including temporal extraction!) + semantic_result, bm25_result, graph_result, temporal_result = await asyncio.gather( + run_semantic(), + run_bm25(), + run_graph(), + run_temporal_with_extraction(), + ) + graph_results, graph_time, mpfp_timing = graph_result + + # Compute max connection wait across all methods (graph handles its own connections) + max_conn_wait = max(semantic_result.conn_wait, bm25_result.conn_wait, temporal_result.conn_wait) + + return ParallelRetrievalResult( + semantic=semantic_result.results, + bm25=bm25_result.results, + graph=graph_results, + temporal=temporal_result.results if temporal_result.results else None, + timings={ + "semantic": semantic_result.time, + "bm25": bm25_result.time, + "graph": graph_time, + "temporal": temporal_result.time, + "temporal_extraction": temporal_result.extraction_time, + }, + temporal_constraint=temporal_result.constraint, + mpfp_timings=[mpfp_timing] if mpfp_timing else [], + max_conn_wait=max_conn_wait, + ) async def _get_temporal_entry_points( diff --git a/hindsight-api/hindsight_api/engine/search/types.py b/hindsight-api/hindsight_api/engine/search/types.py index eb7f48db..34af67c2 100644 --- a/hindsight-api/hindsight_api/engine/search/types.py +++ b/hindsight-api/hindsight_api/engine/search/types.py @@ -24,6 +24,8 @@ class MPFPTimings: fetch: float = 0.0 # Time to fetch memory unit details seeds_time: float = 0.0 # Time to find semantic seeds (if fallback used) result_count: int = 0 # Number of results returned + # Detailed per-hop timing: list of {hop, exec_time, uncached, load_time, edges_loaded, total_time} + hop_details: list[dict] = field(default_factory=list) @dataclass diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 858829b7..9b99b463 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -23,7 +23,7 @@ import uvicorn from . import MemoryEngine from .api import create_app from .banner import print_banner -from .config import HindsightConfig, get_config +from .config import DEFAULT_WORKERS, ENV_WORKERS, HindsightConfig, get_config from .daemon import ( DEFAULT_DAEMON_PORT, DEFAULT_IDLE_TIMEOUT, @@ -95,7 +95,12 @@ def main(): # Development options parser.add_argument("--reload", action="store_true", help="Enable auto-reload on code changes (development only)") - parser.add_argument("--workers", type=int, default=1, help="Number of worker processes (default: 1)") + parser.add_argument( + "--workers", + type=int, + default=int(os.getenv(ENV_WORKERS, str(DEFAULT_WORKERS))), + help=f"Number of worker processes (env: {ENV_WORKERS}, default: {DEFAULT_WORKERS})", + ) # Access log options parser.add_argument("--access-log", action="store_true", help="Enable access log") @@ -187,11 +192,14 @@ def main(): reranker_tei_url=config.reranker_tei_url, reranker_tei_batch_size=config.reranker_tei_batch_size, reranker_tei_max_concurrent=config.reranker_tei_max_concurrent, + reranker_max_candidates=config.reranker_max_candidates, host=args.host, port=args.port, log_level=args.log_level, mcp_enabled=config.mcp_enabled, graph_retriever=config.graph_retriever, + mpfp_top_k_neighbors=config.mpfp_top_k_neighbors, + recall_max_concurrent=config.recall_max_concurrent, observation_min_facts=config.observation_min_facts, observation_top_entities=config.observation_top_entities, retain_max_completion_tokens=config.retain_max_completion_tokens, @@ -265,14 +273,27 @@ def main(): app = idle_middleware # Prepare uvicorn config + # When using workers or reload, we must use import string so each worker can import the app + use_import_string = args.workers > 1 or args.reload + # Check for uvloop availability + try: + import uvloop # noqa: F401 + + loop_impl = "uvloop" + print("uvloop available, will use for event loop") + except ImportError: + loop_impl = "asyncio" + print("uvloop not installed, using default asyncio event loop") + uvicorn_config = { - "app": app, + "app": "hindsight_api.server:app" if use_import_string else app, "host": args.host, "port": args.port, "log_level": args.log_level, "access_log": args.access_log, "proxy_headers": args.proxy_headers, "ws": "wsproto", # Use wsproto instead of websockets to avoid deprecation warnings + "loop": loop_impl, # Explicitly set event loop implementation } # Add optional parameters if provided diff --git a/hindsight-api/hindsight_api/server.py b/hindsight-api/hindsight_api/server.py index 8a4f7496..38fe4238 100644 --- a/hindsight-api/hindsight_api/server.py +++ b/hindsight-api/hindsight_api/server.py @@ -27,10 +27,17 @@ config.configure_logging() # Create app at module level (required for uvicorn import string) # MemoryEngine reads configuration from environment variables automatically -_memory = MemoryEngine() +# Note: run_migrations=True by default, but migrations are idempotent so safe with workers +_memory = MemoryEngine(run_migrations=config.run_migrations_on_startup) # Create unified app with both HTTP and optionally MCP -app = create_app(memory=_memory, http_api_enabled=True, mcp_api_enabled=config.mcp_enabled, mcp_mount_path="/mcp") +app = create_app( + memory=_memory, + http_api_enabled=True, + mcp_api_enabled=config.mcp_enabled, + mcp_mount_path="/mcp", + initialize_memory=True, +) if __name__ == "__main__": diff --git a/hindsight-api/pyproject.toml b/hindsight-api/pyproject.toml index 1bf76ed1..a71c986e 100644 --- a/hindsight-api/pyproject.toml +++ b/hindsight-api/pyproject.toml @@ -37,10 +37,12 @@ dependencies = [ "anthropic>=0.40.0", "typer>=0.9.0", "cohere>=5.0.0", + "flashrank>=0.2.0", # Local ML models for embeddings/reranking - can be excluded in Docker with INCLUDE_LOCAL_MODELS=false "sentence-transformers>=3.0.0,<3.3.0", "transformers>=4.30.0,<4.46.0", "torch>=2.0.0", + "uvloop>=0.22.1", ] [project.optional-dependencies] diff --git a/hindsight-api/tests/test_mpfp_retrieval.py b/hindsight-api/tests/test_mpfp_retrieval.py index 5325303a..a36464db 100644 --- a/hindsight-api/tests/test_mpfp_retrieval.py +++ b/hindsight-api/tests/test_mpfp_retrieval.py @@ -247,17 +247,21 @@ class TestMPFPTraverseAsync: seeds = [SeedNode("seed-1", 1.0)] - # Mock edge loading (returns all edge types at once) - async def mock_load_all_edges(pool, node_ids): - if "seed-1" in node_ids: - return { - "semantic": { - "seed-1": [ - EdgeTarget("neighbor-1", 0.8), - EdgeTarget("neighbor-2", 0.4), - ] - } + # Pre-populate cache with seed edges (mimics pre-warming in retrieve()) + cache.add_all_edges( + { + "semantic": { + "seed-1": [ + EdgeTarget("neighbor-1", 0.8), + EdgeTarget("neighbor-2", 0.4), + ] } + }, + ["seed-1"], + ) + + # Mock for loading neighbor edges (after hop 0) + async def mock_load_all_edges(pool, node_ids, top_k=20): return {} with patch( @@ -291,11 +295,15 @@ class TestMPFPTraverseAsync: seeds = [SeedNode("seed-1", 1.0)] - # Mock edge loading for two hops (returns all edge types at once) - async def mock_load_all_edges(pool, node_ids): + # Pre-populate cache with seed edges (mimics pre-warming in retrieve()) + cache.add_all_edges( + {"semantic": {"seed-1": [EdgeTarget("hop1-node", 1.0)]}}, + ["seed-1"], + ) + + # Mock edge loading for hop 1 nodes + async def mock_load_all_edges(pool, node_ids, top_k=20): edges: dict[str, dict[str, list[EdgeTarget]]] = {"semantic": {}} - if "seed-1" in node_ids: - edges["semantic"]["seed-1"] = [EdgeTarget("hop1-node", 1.0)] if "hop1-node" in node_ids: edges["semantic"]["hop1-node"] = [EdgeTarget("hop2-node", 1.0)] return edges @@ -319,12 +327,17 @@ class TestMPFPTraverseAsync: @pytest.mark.asyncio async def test_cache_reuse(self): - """Cache should prevent redundant edge loading.""" + """Cache should prevent redundant edge loading for already-cached nodes.""" cache = EdgeCache() config = MPFPConfig(alpha=0.15, threshold=1e-6) - # Pre-load cache (marks seed-1 as fully loaded) - cache.add_all_edges({"semantic": {"seed-1": [EdgeTarget("neighbor-1", 1.0)]}}, ["seed-1"]) + # Pre-load cache (marks seed-1 AND neighbor-1 as fully loaded) + # neighbor-1 is also cached because after hop 0, the frontier contains neighbor-1 + # and the algorithm tries to pre-warm edges for the next hop + cache.add_all_edges( + {"semantic": {"seed-1": [EdgeTarget("neighbor-1", 1.0)], "neighbor-1": []}}, + ["seed-1", "neighbor-1"], + ) seeds = [SeedNode("seed-1", 1.0)] @@ -342,7 +355,7 @@ class TestMPFPTraverseAsync: cache=cache, ) - # Should not call load_all_edges_for_frontier since seed-1 is already cached + # Should not call load_all_edges_for_frontier since all nodes are already cached load_mock.assert_not_called() @@ -356,7 +369,9 @@ class TestMPFPGraphRetriever: def test_default_config(self): """Default config should have expected patterns.""" - retriever = MPFPGraphRetriever() + # Use explicit config to avoid global config dependency + config = MPFPConfig() + retriever = MPFPGraphRetriever(config=config) assert len(retriever.config.patterns_semantic) > 0 assert len(retriever.config.patterns_temporal) > 0 @@ -398,7 +413,9 @@ class TestMPFPGraphRetriever: @pytest.mark.asyncio async def test_retrieve_no_seeds_returns_empty(self): """Retrieve with no seeds should return empty results.""" - retriever = MPFPGraphRetriever() + # Use explicit config to avoid global config dependency + config = MPFPConfig() + retriever = MPFPGraphRetriever(config=config) # Mock _find_semantic_seeds to return empty with patch.object(retriever, "_find_semantic_seeds", new_callable=AsyncMock, return_value=[]): @@ -417,15 +434,18 @@ class TestMPFPGraphRetriever: @pytest.mark.asyncio async def test_retrieve_with_semantic_seeds(self): """Retrieve with semantic seeds should run patterns and return results.""" - retriever = MPFPGraphRetriever() + # Use explicit config to avoid global config dependency + config = MPFPConfig() + retriever = MPFPGraphRetriever(config=config) semantic_seeds = [ RetrievalResult(id="seed-1", text="seed text", fact_type="world", similarity=0.9), ] # Mock the internal functions + # mpfp_traverse_hop_synchronized returns a list of PatternResult (one per pattern) async def mock_traverse(*args, **kwargs): - return PatternResult(pattern=["semantic"], scores={"seed-1": 0.5, "result-1": 0.3}) + return [PatternResult(pattern=["semantic"], scores={"seed-1": 0.5, "result-1": 0.3})] async def mock_fetch(pool, node_ids, fact_type): return [ @@ -435,13 +455,18 @@ class TestMPFPGraphRetriever: with ( patch( - "hindsight_api.engine.search.mpfp_retrieval.mpfp_traverse_async", + "hindsight_api.engine.search.mpfp_retrieval.mpfp_traverse_hop_synchronized", side_effect=mock_traverse, ), patch( "hindsight_api.engine.search.mpfp_retrieval.fetch_memory_units_by_ids", side_effect=mock_fetch, ), + patch( + "hindsight_api.engine.search.mpfp_retrieval.load_all_edges_for_frontier", + new_callable=AsyncMock, + return_value={}, + ), ): results, timings = await retriever.retrieve( pool=MagicMock(), @@ -553,3 +578,242 @@ async def test_mpfp_lazy_loading_efficiency(memory, request_context): finally: await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================================ +# MPFP Performance Benchmark Tests +# ============================================================================ +# These tests require an external database with a large memory bank to be useful. +# Set EXTERNAL_DATABASE_URL and BENCHMARK_BANK_ID environment variables to run. +# Example: +# EXTERNAL_DATABASE_URL=postgresql://user:pass@host:port/db \ +# BENCHMARK_BANK_ID=load-test \ +# pytest tests/test_mpfp_retrieval.py::test_mpfp_edge_loading_performance -v -s + + +import os +import asyncpg + +EXTERNAL_DATABASE_URL = os.environ.get("EXTERNAL_DATABASE_URL") +BENCHMARK_BANK_ID = os.environ.get("BENCHMARK_BANK_ID", "load-test") + +requires_external_db = pytest.mark.skipif( + EXTERNAL_DATABASE_URL is None, + reason="EXTERNAL_DATABASE_URL not set - skipping external DB benchmark", +) + + +@requires_external_db +@pytest.mark.asyncio +async def test_mpfp_edge_loading_performance(): + """ + Benchmark MPFP edge loading performance. + + This test measures the performance of the LATERAL query optimization + for loading edges in the MPFP graph traversal algorithm. + + Set EXTERNAL_DATABASE_URL to point to a database with existing data. + Set BENCHMARK_BANK_ID to specify which bank to query (default: load-test). + + Example usage: + EXTERNAL_DATABASE_URL=postgresql://hindsight:hindsight@localhost:5435/hindsight \ + BENCHMARK_BANK_ID=load-test \ + pytest tests/test_mpfp_retrieval.py::test_mpfp_edge_loading_performance -v -s + """ + import time + + # Connect to external database + pool = await asyncpg.create_pool(EXTERNAL_DATABASE_URL, min_size=2, max_size=10) + + try: + # Get some sample node IDs from the database + async with pool.acquire() as conn: + # First check how many links exist + stats = await conn.fetchrow(""" + SELECT + count(*) as total_links, + count(DISTINCT from_unit_id) as unique_sources + FROM memory_links + """) + print(f"\n📊 Database Stats:") + print(f" Total links: {stats['total_links']:,}") + print(f" Unique sources: {stats['unique_sources']:,}") + + # Get edge distribution by type + type_stats = await conn.fetch(""" + SELECT link_type, count(*) as cnt, + round(avg(weight)::numeric, 3) as avg_weight + FROM memory_links + GROUP BY link_type + ORDER BY cnt DESC + """) + print(f"\n Edge distribution:") + for row in type_stats: + print(f" - {row['link_type']}: {row['cnt']:,} (avg_weight={row['avg_weight']})") + + # Get sample frontier nodes (from memory_units in the benchmark bank) + # bank_id is the text primary key in banks table + frontier_rows = await conn.fetch(""" + SELECT id FROM memory_units + WHERE bank_id = $1 + LIMIT 100 + """, BENCHMARK_BANK_ID) + + if not frontier_rows: + pytest.skip(f"No memory units found for bank '{BENCHMARK_BANK_ID}'") + + frontier_node_ids = [str(row['id']) for row in frontier_rows] + print(f"\n🎯 Testing with {len(frontier_node_ids)} frontier nodes from bank '{BENCHMARK_BANK_ID}'") + + # Test 1: Original query approach (all edges, no per-type limit) + async with pool.acquire() as conn: + start = time.time() + original_rows = await conn.fetch(""" + SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight + FROM memory_links ml + WHERE ml.from_unit_id = ANY($1::uuid[]) + AND ml.weight >= 0.1 + ORDER BY ml.from_unit_id, ml.link_type, ml.weight DESC + """, frontier_node_ids) + original_time = time.time() - start + original_count = len(original_rows) + + # Test 2: New LATERAL query approach (top-k per type) + async with pool.acquire() as conn: + start = time.time() + lateral_rows = await conn.fetch(""" + WITH frontier(node_id) AS (SELECT unnest($1::uuid[])) + SELECT f.node_id as from_unit_id, lt.link_type, edges.to_unit_id, edges.weight + FROM frontier f + CROSS JOIN (VALUES ('semantic'), ('temporal'), ('entity'), ('causes'), ('caused_by')) AS lt(link_type) + CROSS JOIN LATERAL ( + SELECT ml.to_unit_id, ml.weight + FROM memory_links ml + WHERE ml.from_unit_id = f.node_id + AND ml.link_type = lt.link_type + AND ml.weight >= 0.1 + ORDER BY ml.weight DESC + LIMIT 20 + ) edges + """, frontier_node_ids) + lateral_time = time.time() - start + lateral_count = len(lateral_rows) + + # Print results + print(f"\n⏱️ Performance Comparison ({len(frontier_node_ids)} nodes):") + print(f"\n Original (all edges):") + print(f" - Time: {original_time * 1000:.2f}ms") + print(f" - Rows: {original_count:,}") + print(f" - Rows/node: {original_count / len(frontier_node_ids):.1f}") + + print(f"\n LATERAL (top-20 per type):") + print(f" - Time: {lateral_time * 1000:.2f}ms") + print(f" - Rows: {lateral_count:,}") + print(f" - Rows/node: {lateral_count / len(frontier_node_ids):.1f}") + + speedup = original_time / lateral_time if lateral_time > 0 else float('inf') + reduction = (1 - lateral_count / original_count) * 100 if original_count > 0 else 0 + print(f"\n 📈 Improvement:") + print(f" - Speedup: {speedup:.2f}x faster") + print(f" - Data reduction: {reduction:.1f}% fewer rows") + + # Assert improvement (should be at least some improvement for large datasets) + if original_count > 1000: + # For large datasets, expect significant improvement + assert speedup >= 1.5, f"Expected at least 1.5x speedup, got {speedup:.2f}x" + assert reduction >= 30, f"Expected at least 30% data reduction, got {reduction:.1f}%" + print(f"\n✅ Performance test PASSED!") + else: + print(f"\n⚠️ Dataset too small ({original_count} rows) for meaningful performance comparison") + + finally: + await pool.close() + + +@requires_external_db +@pytest.mark.asyncio +async def test_mpfp_full_retrieval_performance(): + """ + Benchmark full MPFP retrieval including traversal and reranking. + + This test measures end-to-end MPFP retrieval performance. + """ + import time + + pool = await asyncpg.create_pool(EXTERNAL_DATABASE_URL, min_size=2, max_size=10) + + try: + # Get a sample query embedding from an existing memory unit + async with pool.acquire() as conn: + # Check if bank exists + bank_exists = await conn.fetchval(""" + SELECT 1 FROM banks WHERE bank_id = $1 + """, BENCHMARK_BANK_ID) + if not bank_exists: + pytest.skip(f"Bank '{BENCHMARK_BANK_ID}' not found") + + sample = await conn.fetchrow(""" + SELECT embedding::text as embedding_str + FROM memory_units + WHERE bank_id = $1 + AND embedding IS NOT NULL + LIMIT 1 + """, BENCHMARK_BANK_ID) + + if not sample: + pytest.skip("No memory units with embeddings found") + + query_embedding_str = sample['embedding_str'] + + # Run MPFP retrieval + retriever = MPFPGraphRetriever() + + print(f"\n🔍 Running MPFP retrieval benchmark on bank '{BENCHMARK_BANK_ID}'...") + + # Warm-up run + await retriever.retrieve( + pool=pool, + query_embedding_str=query_embedding_str, + bank_id=BENCHMARK_BANK_ID, + fact_type="world", + budget=100, + query_text="test query", + ) + + # Timed runs + timings_list = [] + for i in range(3): + start = time.time() + results, timings = await retriever.retrieve( + pool=pool, + query_embedding_str=query_embedding_str, + bank_id=BENCHMARK_BANK_ID, + fact_type="opinion", + budget=100, + query_text="What did I say about training models?", + ) + elapsed = time.time() - start + timings_list.append((elapsed, timings, len(results))) + + # Print results + print(f"\n⏱️ MPFP Retrieval Results (3 runs):") + for i, (elapsed, timings, count) in enumerate(timings_list): + print(f"\n Run {i + 1}:") + print(f" - Total: {elapsed * 1000:.2f}ms") + print(f" - Results: {count}") + if timings: + print(f" - Seeds: {timings.seeds_time * 1000:.2f}ms") + print(f" - Patterns: {timings.pattern_count}") + print(f" - Traverse: {timings.traverse * 1000:.2f}ms") + print(f" - Edge load: {timings.edge_load_time * 1000:.2f}ms") + print(f" - Edges: {timings.edge_count:,}") + print(f" - DB queries: {timings.db_queries}") + print(f" - Fusion: {timings.fusion * 1000:.2f}ms") + print(f" - Fetch: {timings.fetch * 1000:.2f}ms") + + avg_time = sum(t[0] for t in timings_list) / len(timings_list) + print(f"\n 📊 Average: {avg_time * 1000:.2f}ms") + print(f"\n✅ MPFP retrieval benchmark complete!") + + finally: + await pool.close() diff --git a/hindsight-api/tests/test_tei_cross_encoder.py b/hindsight-api/tests/test_tei_cross_encoder.py index 192daf4a..ea3b1b2e 100644 --- a/hindsight-api/tests/test_tei_cross_encoder.py +++ b/hindsight-api/tests/test_tei_cross_encoder.py @@ -544,3 +544,243 @@ class TestRemoteTEICrossEncoderConfig: assert encoder.base_url == "http://test:9000" assert encoder.batch_size == 256 assert encoder.max_concurrent == 16 + + +# ============================================================================ +# TEI Reranker Performance Benchmark Tests +# ============================================================================ +# These tests require a running TEI server to measure actual performance. +# Set TEI_RERANKER_URL environment variable to run. +# Example: +# TEI_RERANKER_URL=http://localhost:8000 \ +# pytest tests/test_tei_cross_encoder.py::test_tei_reranker_performance -v -s -n0 + +import os + +TEI_RERANKER_URL = os.environ.get("TEI_RERANKER_URL") + +requires_tei_server = pytest.mark.skipif( + TEI_RERANKER_URL is None, + reason="TEI_RERANKER_URL not set - skipping TEI performance benchmark", +) + + +@requires_tei_server +@pytest.mark.asyncio +async def test_tei_reranker_performance(): + """ + Benchmark TEI reranker performance with different configurations. + + This test measures latency for different batch sizes and concurrency levels + to find the optimal configuration for your TEI server. + + Example usage: + TEI_RERANKER_URL=http://localhost:8000 \ + pytest tests/test_tei_cross_encoder.py::test_tei_reranker_performance -v -s -n0 + """ + import httpx + + # Get server info + async with httpx.AsyncClient() as client: + response = await client.get(f"{TEI_RERANKER_URL}/info") + info = response.json() + print(f"\n📊 TEI Server Info:") + print(f" URL: {TEI_RERANKER_URL}") + print(f" Model: {info.get('model_id', 'unknown')}") + if "reranker_model" in info: + print(f" Reranker Model: {info['reranker_model']}") + + # Generate test data (800 pairs to simulate real workload) + num_pairs = 800 + query = "What did I say about training machine learning models and artificial intelligence?" + test_pairs = [ + (query, f"Document {i} about machine learning, neural networks, and AI training techniques.") + for i in range(num_pairs) + ] + + # Test configurations: (batch_size, max_concurrent) + configs = [ + (128, 8), # Default + (256, 4), # Larger batches, fewer concurrent + (256, 8), # Larger batches, same concurrent + (512, 2), # Very large batches, few concurrent + (512, 4), # Very large batches, moderate concurrent + (64, 16), # Smaller batches, more concurrent + (800, 1), # Single batch (all at once) + ] + + results = [] + print(f"\n⏱️ Benchmarking {num_pairs} pairs with different configurations:\n") + + for batch_size, max_concurrent in configs: + encoder = RemoteTEICrossEncoder( + base_url=TEI_RERANKER_URL, + batch_size=batch_size, + max_concurrent=max_concurrent, + timeout=60.0, + ) + await encoder.initialize() + + # Warm-up run + await encoder.predict(test_pairs[:100]) + + # Timed runs (3 iterations) + times = [] + for _ in range(3): + start = time.time() + scores = await encoder.predict(test_pairs) + elapsed = time.time() - start + times.append(elapsed) + assert len(scores) == num_pairs + + avg_time = sum(times) / len(times) + min_time = min(times) + results.append({ + "batch_size": batch_size, + "max_concurrent": max_concurrent, + "avg_ms": avg_time * 1000, + "min_ms": min_time * 1000, + "num_batches": (num_pairs + batch_size - 1) // batch_size, + }) + + print(f" batch_size={batch_size:4d}, max_concurrent={max_concurrent:2d}: " + f"avg={avg_time * 1000:6.1f}ms, min={min_time * 1000:6.1f}ms " + f"({results[-1]['num_batches']} batches)") + + # Find best configuration + best = min(results, key=lambda x: x["avg_ms"]) + print(f"\n🏆 Best Configuration:") + print(f" batch_size={best['batch_size']}, max_concurrent={best['max_concurrent']}") + print(f" Average: {best['avg_ms']:.1f}ms, Min: {best['min_ms']:.1f}ms") + + # Performance target check + target_ms = 100 + if best["avg_ms"] <= target_ms: + print(f"\n✅ Target met! Average {best['avg_ms']:.1f}ms <= {target_ms}ms") + else: + print(f"\n⚠️ Target NOT met. Average {best['avg_ms']:.1f}ms > {target_ms}ms") + print(f" Consider: larger batch size, GPU optimization, or faster network") + + +@requires_tei_server +@pytest.mark.asyncio +async def test_tei_reranker_concurrent_requests(): + """ + Test TEI reranker performance under concurrent request load. + + This simulates multiple parallel recall requests hitting the reranker + at the same time. + """ + # Smaller batches to simulate typical recall workload + num_pairs_per_request = 200 + num_concurrent_requests = 4 + + query = "Tell me about machine learning and AI training" + test_pairs = [ + (query, f"Document {i} about ML and training.") + for i in range(num_pairs_per_request) + ] + + # Test configurations + configs = [ + (128, 8), # Default + (256, 4), # Larger batches + (512, 2), # Very large batches + (200, 1), # Single batch per request + ] + + print(f"\n⏱️ Concurrent Load Test: {num_concurrent_requests} parallel requests, " + f"{num_pairs_per_request} pairs each:\n") + + for batch_size, max_concurrent in configs: + encoder = RemoteTEICrossEncoder( + base_url=TEI_RERANKER_URL, + batch_size=batch_size, + max_concurrent=max_concurrent, + timeout=60.0, + ) + await encoder.initialize() + + # Warm-up + await encoder.predict(test_pairs[:50]) + + async def run_single_request(): + start = time.time() + scores = await encoder.predict(test_pairs) + return time.time() - start, len(scores) + + # Run concurrent requests + times = [] + for _ in range(3): # 3 iterations + start = time.time() + results = await asyncio.gather(*[run_single_request() for _ in range(num_concurrent_requests)]) + total_time = time.time() - start + + individual_times = [r[0] for r in results] + times.append({ + "total": total_time, + "max_individual": max(individual_times), + "avg_individual": sum(individual_times) / len(individual_times), + }) + + avg_total = sum(t["total"] for t in times) / len(times) + avg_max_individual = sum(t["max_individual"] for t in times) / len(times) + + print(f" batch_size={batch_size:4d}, max_concurrent={max_concurrent:2d}: " + f"total={avg_total * 1000:6.1f}ms, slowest_req={avg_max_individual * 1000:6.1f}ms") + + +@requires_tei_server +@pytest.mark.asyncio +async def test_tei_reranker_latency_breakdown(): + """ + Measure latency breakdown for TEI reranker requests. + + This helps identify where time is spent: network vs processing. + """ + import httpx + + print(f"\n⏱️ Latency Breakdown Test:\n") + + # Test single document latency (network overhead) + async with httpx.AsyncClient(timeout=30.0) as client: + times = [] + for _ in range(10): + start = time.time() + await client.post( + f"{TEI_RERANKER_URL}/rerank", + json={ + "query": "test query", + "texts": ["test document"], + "return_text": False, + }, + ) + times.append((time.time() - start) * 1000) + + avg_single = sum(times) / len(times) + print(f" Single doc latency (raw HTTP): {avg_single:.2f}ms") + + # Test batch latencies + batch_sizes = [10, 50, 100, 200, 500] + for batch_size in batch_sizes: + texts = [f"Document {i} about machine learning" for i in range(batch_size)] + async with httpx.AsyncClient(timeout=30.0) as client: + times = [] + for _ in range(5): + start = time.time() + await client.post( + f"{TEI_RERANKER_URL}/rerank", + json={ + "query": "What about machine learning?", + "texts": texts, + "return_text": False, + }, + ) + times.append((time.time() - start) * 1000) + + avg = sum(times) / len(times) + per_doc = avg / batch_size + print(f" Batch size {batch_size:4d}: {avg:6.1f}ms total, {per_doc:.2f}ms/doc") + + print(f"\n 💡 Insight: Higher per-doc time at small batches = network overhead dominant") + print(f" 💡 Insight: Lower per-doc time at large batches = GPU efficiently utilized") diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 1a2c7fd5..02356daf 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -210,15 +210,6 @@ export HINDSIGHT_API_COHERE_API_KEY=your-api-key # shared with embeddings export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 ``` -### Server - -| Variable | Description | Default | -|----------|-------------|---------| -| `HINDSIGHT_API_HOST` | Bind address | `0.0.0.0` | -| `HINDSIGHT_API_PORT` | Server port | `8888` | -| `HINDSIGHT_API_LOG_LEVEL` | Log level: `debug`, `info`, `warning`, `error` | `info` | -| `HINDSIGHT_API_MCP_ENABLED` | Enable MCP server at `/mcp/{bank_id}/` | `true` | - ### Authentication By default, Hindsight runs without authentication. For production deployments, enable API key authentication using the built-in tenant extension: @@ -242,11 +233,29 @@ Requests without a valid API key receive a `401 Unauthorized` response. For advanced authentication (JWT, OAuth, multi-tenant schemas), implement a custom `TenantExtension`. See the [Extensions documentation](./extensions.md) for details. ::: +### Server + +| Variable | Description | Default | +|----------|-------------|---------| +| `HINDSIGHT_API_HOST` | Bind address | `0.0.0.0` | +| `HINDSIGHT_API_PORT` | Server port | `8888` | +| `HINDSIGHT_API_WORKERS` | Number of uvicorn worker processes | `1` | +| `HINDSIGHT_API_LOG_LEVEL` | Log level: `debug`, `info`, `warning`, `error` | `info` | +| `HINDSIGHT_API_MCP_ENABLED` | Enable MCP server at `/mcp/{bank_id}/` | `true` | + ### Retrieval | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_GRAPH_RETRIEVER` | Graph retrieval algorithm: `bfs` or `mpfp` | `bfs` | +| `HINDSIGHT_API_GRAPH_RETRIEVER` | Graph retrieval algorithm: `link_expansion`, `mpfp`, or `bfs` | `link_expansion` | +| `HINDSIGHT_API_RECALL_MAX_CONCURRENT` | Max concurrent recall operations per worker (backpressure) | `32` | +| `HINDSIGHT_API_RERANKER_MAX_CANDIDATES` | Max candidates to rerank per recall (RRF pre-filters the rest) | `300` | + +#### Graph Retrieval Algorithms + +- **`link_expansion`** (default): Fast, simple graph expansion from semantic seeds via entity co-occurrence and causal links. Target latency under 100ms. Recommended for most use cases. +- **`mpfp`**: Multi-Path Fact Propagation - iterative graph traversal with activation spreading. More thorough but slower. +- **`bfs`**: Breadth-first search from seed facts. Simple but less effective for large graphs. ### Entity Observations diff --git a/hindsight-docs/docs/developer/retrieval.md b/hindsight-docs/docs/developer/retrieval.md index 099fa5a3..4ccaed69 100644 --- a/hindsight-docs/docs/developer/retrieval.md +++ b/hindsight-docs/docs/developer/retrieval.md @@ -231,44 +231,9 @@ Budget and max_tokens control different aspects of recall: ## Graph Retrieval Algorithms -Hindsight supports two graph traversal algorithms, each optimized for different scenarios: +Hindsight supports multiple graph traversal algorithms. The default (`link_expansion`) is optimized for fast retrieval with target latency under 100ms. -| Algorithm | Default | Best For | Complexity | -|-----------|---------|----------|------------| -| **MPFP** | ✓ | Large graphs, production | O(P × H × F × K) | -| **BFS** | | Small graphs, debugging | O(V + E) | - -### MPFP (Meta-Path Forward Push) - -A sublinear graph traversal algorithm that follows predefined meta-paths (patterns of edge types) using lazy edge loading. - -**How it works:** -1. Starts from semantic entry points (top similar facts) -2. Follows multiple meta-path patterns in parallel: - - `semantic → semantic` (topic expansion) - - `entity → temporal` (entity timeline) - - `semantic → causes` (causal reasoning) - - `entity → semantic` (entity context) -3. Loads edges lazily per hop, only for active frontier nodes -4. Fuses results from all patterns via Reciprocal Rank Fusion (RRF) - -**Complexity:** O(P × H × F × K) where P = patterns (~7), H = hops (2), F = frontier size (~20-100), K = neighbors per node (20). - -**Use case:** Production workloads with large memory banks (10k+ facts). Only loads the edges it needs, avoiding full graph scans. - -### BFS (Breadth-First Spreading Activation) - -Classic spreading activation that propagates relevance scores through the graph using breadth-first traversal. - -**How it works:** -1. Starts from semantic entry points with initial activation scores -2. Spreads activation to neighbors with decay (α = 0.8 per hop) -3. Boosts causal links (causes, enables, prevents) -4. Continues until budget exhausted or activation below threshold - -**Complexity:** O(V + E) where V and E are visited nodes and edges, bounded by budget. - -**Use case:** Small memory banks, debugging, or when you need to understand exactly how results were found. +See [Configuration → Retrieval](./configuration#retrieval) for available algorithms and how to configure them. --- diff --git a/scripts/dev/start-api.sh b/scripts/dev/start-api.sh index f18f64a1..ef2301d6 100755 --- a/scripts/dev/start-api.sh +++ b/scripts/dev/start-api.sh @@ -17,4 +17,4 @@ set -a source "$ENV_FILE" set +a -uv run hindsight-api "${SERVER_ARGS[@]}" +uv run hindsight-api "$@" diff --git a/uv.lock b/uv.lock index d700d1d4..12150df2 100644 --- a/uv.lock +++ b/uv.lock @@ -621,6 +621,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 }, ] +[[package]] +name = "coloredlogs" +version = "15.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "humanfriendly" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934", size = 46018 }, +] + [[package]] name = "cryptography" version = "46.0.3" @@ -955,6 +967,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/76/91/7216b27286936c16f5b4d0c530087e4a54eead683e6b0b73dd0c64844af6/filelock-3.20.0-py3-none-any.whl", hash = "sha256:339b4732ffda5cd79b13f4e2711a31b0365ce445d95d243bb996273d072546a2", size = 16054 }, ] +[[package]] +name = "flashrank" +version = "0.2.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "onnxruntime" }, + { name = "requests" }, + { name = "tokenizers" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/1f/176cb4a857a70c3538f637e19389ab6aed21548a1ba1d1424fccc8bba108/FlashRank-0.2.10.tar.gz", hash = "sha256:f8f82a25c32fdfc668a09dc4089421d6aab8e7f71308424b541f40bb3f01d9db", size = 18905 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/99/72639cc1c9221c5bc77a2df1c2d352fe11965553bdf7d3e0856e7fcc8fd6/FlashRank-0.2.10-py3-none-any.whl", hash = "sha256:5d3272ae657d793c132d1e7917ed9e2adf49e0e1c60735583a67b051c6f0434a", size = 14511 }, +] + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661 }, +] + [[package]] name = "frozenlist" version = "1.8.0" @@ -1255,6 +1291,7 @@ dependencies = [ { name = "dateparser" }, { name = "fastapi", extra = ["standard"] }, { name = "fastmcp" }, + { name = "flashrank" }, { name = "google-genai" }, { name = "greenlet" }, { name = "httpx" }, @@ -1278,6 +1315,7 @@ dependencies = [ { name = "transformers" }, { name = "typer" }, { name = "uvicorn" }, + { name = "uvloop" }, { name = "wsproto" }, ] @@ -1312,6 +1350,7 @@ requires-dist = [ { name = "fastapi", extras = ["standard"], specifier = ">=0.120.3" }, { name = "fastmcp", specifier = ">=2.3.0" }, { name = "filelock", marker = "extra == 'test'", specifier = ">=3.0.0" }, + { name = "flashrank", specifier = ">=0.2.0" }, { name = "google-genai", specifier = ">=1.0.0" }, { name = "greenlet", specifier = ">=3.2.4" }, { name = "httpx", specifier = ">=0.27.0" }, @@ -1339,6 +1378,7 @@ requires-dist = [ { name = "transformers", specifier = ">=4.30.0,<4.46.0" }, { name = "typer", specifier = ">=0.9.0" }, { name = "uvicorn", specifier = ">=0.38.0" }, + { name = "uvloop", specifier = ">=0.22.1" }, { name = "wsproto", specifier = ">=1.0.0" }, ] provides-extras = ["test"] @@ -1539,6 +1579,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/bd/1a875e0d592d447cbc02805fd3fe0f497714d6a2583f59d14fa9ebad96eb/huggingface_hub-0.36.0-py3-none-any.whl", hash = "sha256:7bcc9ad17d5b3f07b57c78e79d527102d08313caa278a641993acddcb894548d", size = 566094 }, ] +[[package]] +name = "humanfriendly" +version = "10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyreadline3", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794 }, +] + [[package]] name = "idna" version = "3.11" @@ -2341,6 +2393,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/9c/92789c596b8df838baa98fa71844d84283302f7604ed565dafe5a6b5041a/oauthlib-3.3.1-py3-none-any.whl", hash = "sha256:88119c938d2b8fb88561af5f6ee0eec8cc8d552b7bb1f712743136eb7523b7a1", size = 160065 }, ] +[[package]] +name = "onnxruntime" +version = "1.23.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "coloredlogs" }, + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "sympy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/be/467b00f09061572f022ffd17e49e49e5a7a789056bad95b54dfd3bee73ff/onnxruntime-1.23.2-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:6f91d2c9b0965e86827a5ba01531d5b669770b01775b23199565d6c1f136616c", size = 17196113 }, + { url = "https://files.pythonhosted.org/packages/9f/a8/3c23a8f75f93122d2b3410bfb74d06d0f8da4ac663185f91866b03f7da1b/onnxruntime-1.23.2-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:87d8b6eaf0fbeb6835a60a4265fde7a3b60157cf1b2764773ac47237b4d48612", size = 19153857 }, + { url = "https://files.pythonhosted.org/packages/3f/d8/506eed9af03d86f8db4880a4c47cd0dffee973ef7e4f4cff9f1d4bcf7d22/onnxruntime-1.23.2-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bbfd2fca76c855317568c1b36a885ddea2272c13cb0e395002c402f2360429a6", size = 15220095 }, + { url = "https://files.pythonhosted.org/packages/e9/80/113381ba832d5e777accedc6cb41d10f9eca82321ae31ebb6bcede530cea/onnxruntime-1.23.2-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:da44b99206e77734c5819aa2142c69e64f3b46edc3bd314f6a45a932defc0b3e", size = 17372080 }, + { url = "https://files.pythonhosted.org/packages/3a/db/1b4a62e23183a0c3fe441782462c0ede9a2a65c6bbffb9582fab7c7a0d38/onnxruntime-1.23.2-cp311-cp311-win_amd64.whl", hash = "sha256:902c756d8b633ce0dedd889b7c08459433fbcf35e9c38d1c03ddc020f0648c6e", size = 13468349 }, + { url = "https://files.pythonhosted.org/packages/1b/9e/f748cd64161213adeef83d0cb16cb8ace1e62fa501033acdd9f9341fff57/onnxruntime-1.23.2-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:b8f029a6b98d3cf5be564d52802bb50a8489ab73409fa9db0bf583eabb7c2321", size = 17195929 }, + { url = "https://files.pythonhosted.org/packages/91/9d/a81aafd899b900101988ead7fb14974c8a58695338ab6a0f3d6b0100f30b/onnxruntime-1.23.2-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:218295a8acae83905f6f1aed8cacb8e3eb3bd7513a13fe4ba3b2664a19fc4a6b", size = 19157705 }, + { url = "https://files.pythonhosted.org/packages/3c/35/4e40f2fba272a6698d62be2cd21ddc3675edfc1a4b9ddefcc4648f115315/onnxruntime-1.23.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76ff670550dc23e58ea9bc53b5149b99a44e63b34b524f7b8547469aaa0dcb8c", size = 15226915 }, + { url = "https://files.pythonhosted.org/packages/ef/88/9cc25d2bafe6bc0d4d3c1db3ade98196d5b355c0b273e6a5dc09c5d5d0d5/onnxruntime-1.23.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f9b4ae77f8e3c9bee50c27bc1beede83f786fe1d52e99ac85aa8d65a01e9b77", size = 17382649 }, + { url = "https://files.pythonhosted.org/packages/c0/b4/569d298f9fc4d286c11c45e85d9ffa9e877af12ace98af8cab52396e8f46/onnxruntime-1.23.2-cp312-cp312-win_amd64.whl", hash = "sha256:25de5214923ce941a3523739d34a520aac30f21e631de53bba9174dc9c004435", size = 13470528 }, + { url = "https://files.pythonhosted.org/packages/3d/41/fba0cabccecefe4a1b5fc8020c44febb334637f133acefc7ec492029dd2c/onnxruntime-1.23.2-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:2ff531ad8496281b4297f32b83b01cdd719617e2351ffe0dba5684fb283afa1f", size = 17196337 }, + { url = "https://files.pythonhosted.org/packages/fe/f9/2d49ca491c6a986acce9f1d1d5fc2099108958cc1710c28e89a032c9cfe9/onnxruntime-1.23.2-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:162f4ca894ec3de1a6fd53589e511e06ecdc3ff646849b62a9da7489dee9ce95", size = 19157691 }, + { url = "https://files.pythonhosted.org/packages/1c/a1/428ee29c6eaf09a6f6be56f836213f104618fb35ac6cc586ff0f477263eb/onnxruntime-1.23.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:45d127d6e1e9b99d1ebeae9bcd8f98617a812f53f46699eafeb976275744826b", size = 15226898 }, + { url = "https://files.pythonhosted.org/packages/f2/2b/b57c8a2466a3126dbe0a792f56ad7290949b02f47b86216cd47d857e4b77/onnxruntime-1.23.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8bace4e0d46480fbeeb7bbe1ffe1f080e6663a42d1086ff95c1551f2d39e7872", size = 17382518 }, + { url = "https://files.pythonhosted.org/packages/4a/93/aba75358133b3a941d736816dd392f687e7eab77215a6e429879080b76b6/onnxruntime-1.23.2-cp313-cp313-win_amd64.whl", hash = "sha256:1f9cc0a55349c584f083c1c076e611a7c35d5b867d5d6e6d6c823bf821978088", size = 13470276 }, + { url = "https://files.pythonhosted.org/packages/7c/3d/6830fa61c69ca8e905f237001dbfc01689a4e4ab06147020a4518318881f/onnxruntime-1.23.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9d2385e774f46ac38f02b3a91a91e30263d41b2f1f4f26ae34805b2a9ddef466", size = 15229610 }, + { url = "https://files.pythonhosted.org/packages/b6/ca/862b1e7a639460f0ca25fd5b6135fb42cf9deea86d398a92e44dfda2279d/onnxruntime-1.23.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e2b9233c4947907fd1818d0e581c049c41ccc39b2856cc942ff6d26317cee145", size = 17394184 }, +] + [[package]] name = "openai" version = "2.7.2" @@ -3215,6 +3299,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063 }, ] +[[package]] +name = "pyreadline3" +version = "3.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/49/4cea918a08f02817aabae639e3d0ac046fef9f9180518a3ad394e22da148/pyreadline3-3.5.4.tar.gz", hash = "sha256:8d57d53039a1c75adba8e50dd3d992b28143480816187ea5efbd5c78e6c885b7", size = 99839 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl", hash = "sha256:eaf8e6cc3c49bcccf145fc6067ba8643d1df34d604a1ec0eccbf7a18e6d3fae6", size = 83178 }, +] + [[package]] name = "pytest" version = "9.0.0"