Compare commits
6
Commits
workers
...
improve-perf
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
93774c20ba | ||
|
|
0076c1de36 | ||
|
|
905ab98235 | ||
|
|
983b70c087 | ||
|
|
a058e9e6c7 | ||
|
|
bd99fc0d74 |
@@ -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: `<revision_id>_<description>.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: <previous_revision_id>
|
||||
Create Date: YYYY-MM-DD
|
||||
"""
|
||||
from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "f1a2b3c4d5e6"
|
||||
down_revision: str | Sequence[str] | None = "<previous_revision_id>"
|
||||
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
|
||||
|
||||
+44
@@ -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")
|
||||
@@ -1190,6 +1190,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:
|
||||
@@ -1215,10 +1218,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,
|
||||
@@ -1277,9 +1282,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:
|
||||
@@ -1289,8 +1306,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(
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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'"
|
||||
)
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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__":
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -17,4 +17,4 @@ set -a
|
||||
source "$ENV_FILE"
|
||||
set +a
|
||||
|
||||
uv run hindsight-api "${SERVER_ARGS[@]}"
|
||||
uv run hindsight-api "$@"
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user