Compare commits

...
20 Commits
Author SHA1 Message Date
Nicolò Boschi 2b02fd13fb fix: update tests for duplicate document_id validation
- test_async_retain_tags: verify operation structure instead of exact UUID
- test_delete_bank: use unique document_ids (team-doc-1, team-doc-2)
2026-02-16 12:14:32 +01:00
Nicolò Boschi 7df7afabfe fix: doc examples - duplicate document_ids and error handling
- retain.py: move document_id to item level to avoid duplicates
- documents.mjs: add error handling for getDocument to show clear error message
2026-02-16 12:12:58 +01:00
Nicolò Boschi 62a8476678 refactor: extract doc example testing to standalone script
- Create scripts/test-doc-examples.sh to run all examples
- Collects logs of failed examples separately
- Shows full error logs only for failures at the end
- Clean summary with pass/fail counts
- Proper exit codes
- Replaces inline bash in CI workflow
2026-02-16 12:01:35 +01:00
Nicolò Boschi 7fc4404568 ci: collect all doc example failures and show summary
- Run all Python/Node.js/CLI examples regardless of individual failures
- Collect failure list and display summary at the end
- Show pass/fail count and list of failed files
- Exit with failure only after running all examples
2026-02-16 11:58:00 +01:00
Nicolò Boschi 5676dda35c fix: properly mock async pool/connection and fix more duplicate document_ids
- Use AsyncMock for pool.acquire() to fix 'can't be used in await' error
- Fix duplicate document_ids in retain-async examples (retain.py and retain.mjs)
- Remove batch-level document_id parameter that caused duplicates
2026-02-16 11:55:13 +01:00
Nicolò Boschi cc820fcd89 fix: test mocking and duplicate document_ids in examples
- Mock _get_pool() in test_async_retain_tags.py to avoid _initialized error
- Set _initialized = True on mocked MemoryEngine instances
- Fix duplicate document_ids in retain.py and retain.mjs examples
2026-02-16 11:38:02 +01:00
Nicolò Boschi 6458460040 chore: trigger CI 2026-02-16 10:56:57 +01:00
Nicolò Boschi be4b5b830b fix: use unique document_ids in batch retain examples
- Each item in a batch must have unique document_id
- Update both Python and JavaScript examples
- Fixes test-doc-examples CI failure
2026-02-16 10:48:30 +01:00
Nicolò Boschi ded59240c9 fix: use discriminated union for OperationDetails type
- Support both success and error states properly
- Fixes TypeScript error when setting error state
2026-02-16 10:48:30 +01:00
Nicolò Boschi 431bd7f1f9 fix: correct OperationDetails type to match API response
- Change optional fields to use | null instead of ?
- Fixes TypeScript compilation error in control plane build
2026-02-16 10:48:30 +01:00
Nicolò Boschi ac097d2f79 docs: remove emojis and code example from performance page 2026-02-16 10:48:30 +01:00
Nicolò Boschi 0ba16e21b7 docs: explain automatic batch optimization for async retain
- Add section explaining Hindsight automatically handles batch sizing
- Users don't need to manually tune batch sizes with async mode
- Hindsight splits large batches (>10k tokens) into optimized sub-batches
- Include example showing best practices
2026-02-16 10:48:30 +01:00
Nicolò Boschi a8d1d319d5 Clean up perf benchmark: keep only Python files
- Remove README.md and PERFORMANCE_FINDINGS.md
- Remove results/ JSON files (gitignored)
- Remove test_data/ directory
- Keep only __init__.py and retain_perf.py
2026-02-16 10:48:30 +01:00
Nicolò Boschi 8aa61725c3 api 2026-02-16 10:48:30 +01:00
Nicolò Boschi 5c12950025 api 2026-02-16 10:48:30 +01:00
Nicolò Boschi 26731b383a api 2026-02-16 10:48:30 +01:00
Nicolò Boschi 2efb387da4 api 2026-02-16 10:48:30 +01:00
Nicolò Boschi c5f1400ec5 api 2026-02-16 10:48:30 +01:00
Nicolò Boschi ef8b603dcc fix: improve async batch retain with large payloads 2026-02-16 10:48:30 +01:00
Nicolò Boschi f517bc172a fix: improve async batch retain with large payloads 2026-02-16 10:48:30 +01:00
32 changed files with 2307 additions and 131 deletions
+2 -21
View File
@@ -941,30 +941,11 @@ jobs:
sleep 1
done
- name: Run Python doc examples
working-directory: ./hindsight-clients/python
run: |
for f in ../../hindsight-docs/examples/api/*.py; do
echo "Running $f..."
uv run python "$f"
done
- name: Run Node.js doc examples
run: |
for f in hindsight-docs/examples/api/*.mjs; do
echo "Running $f..."
node "$f"
done
- name: Configure CLI
run: hindsight configure --api-url http://localhost:8888
- name: Run CLI doc examples
run: |
for f in hindsight-docs/examples/api/*.sh; do
echo "Running $f..."
bash "$f"
done
- name: Run all doc examples
run: ./scripts/test-doc-examples.sh
- name: Show API server logs
if: always()
+1
View File
@@ -46,6 +46,7 @@ hindsight-docs/static/llms-full.txt
hindsight-dev/benchmarks/locomo/results/
hindsight-dev/benchmarks/longmemeval/results/
hindsight-dev/benchmarks/consolidation/results/
hindsight-dev/benchmarks/perf/results/
benchmarks/results/
hindsight-cli/target
hindsight-clients/rust/target
+7
View File
@@ -57,8 +57,15 @@ cd hindsight-control-plane && npm run dev
### Benchmarks
```bash
# Accuracy benchmarks
./scripts/benchmarks/run-longmemeval.sh
./scripts/benchmarks/run-locomo.sh
# Performance benchmarks
./scripts/benchmarks/run-consolidation.sh
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
# Results viewer
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
```
@@ -0,0 +1,49 @@
"""Add GIN index on async_operations.result_metadata for parent_operation_id queries
Revision ID: y0t1u2v3w4x5
Revises: x9s0t1u2v3w4
Create Date: 2026-02-13
This migration adds a GIN index on the result_metadata JSONB column in the
async_operations table to support efficient queries for child operations by
parent_operation_id.
The index enables fast lookups when querying for child operations:
SELECT * FROM async_operations
WHERE result_metadata::jsonb @> '{"parent_operation_id": "uuid"}'::jsonb
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "y0t1u2v3w4x5"
down_revision: str | Sequence[str] | None = "x9s0t1u2v3w4"
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:
"""Add GIN index on result_metadata for efficient parent_operation_id queries."""
schema = _get_schema_prefix()
# Add GIN index for JSONB containment queries (@> operator)
op.execute(f"""
CREATE INDEX idx_async_operations_result_metadata
ON {schema}async_operations
USING gin(result_metadata)
""")
def downgrade() -> None:
"""Remove GIN index on result_metadata."""
schema = _get_schema_prefix()
# Drop index
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_result_metadata")
+17
View File
@@ -1357,6 +1357,16 @@ class CancelOperationResponse(BaseModel):
operation_id: str
class ChildOperationStatus(BaseModel):
"""Status of a child operation (for batch operations)."""
operation_id: str
status: str
sub_batch_index: int | None = None
items_count: int | None = None
error_message: str | None = None
class OperationStatusResponse(BaseModel):
"""Response model for getting a single operation status."""
@@ -1381,6 +1391,13 @@ class OperationStatusResponse(BaseModel):
updated_at: str | None = None
completed_at: str | None = None
error_message: str | None = None
result_metadata: dict[str, Any] | None = Field(
default=None,
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
)
child_operations: list[ChildOperationStatus] | None = Field(
default=None, description="Child operations for batch operations (if applicable)"
)
class AsyncOperationSubmitResponse(BaseModel):
+4
View File
@@ -250,6 +250,7 @@ ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
ENV_RETAIN_EXTRACT_CAUSAL_LINKS = "HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS"
ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
# Observations settings (consolidated knowledge from facts)
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
@@ -371,6 +372,7 @@ DEFAULT_RETAIN_EXTRACT_CAUSAL_LINKS = True # Extract causal links between facts
DEFAULT_RETAIN_EXTRACTION_MODE = "concise" # Extraction mode: "concise", "verbose", or "custom"
RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction modes
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
# Observations defaults (consolidated knowledge from facts)
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
@@ -590,6 +592,7 @@ class HindsightConfig:
retain_extract_causal_links: bool
retain_extraction_mode: str
retain_custom_instructions: str | None
retain_batch_tokens: int
# Observations settings (consolidated knowledge from facts)
enable_observations: bool
@@ -939,6 +942,7 @@ class HindsightConfig:
os.getenv(ENV_RETAIN_EXTRACTION_MODE, DEFAULT_RETAIN_EXTRACTION_MODE)
),
retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
# Observations settings (consolidated knowledge from facts)
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
consolidation_batch_size=int(
@@ -18,11 +18,20 @@ import uuid
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any
import tiktoken
from ..config import get_config
from ..metrics import get_metrics_collector
from ..tracing import create_operation_span
from ..utils import mask_network_location
from .db_budget import budgeted_operation
from .operation_metadata import (
BatchRetainChildMetadata,
BatchRetainParentMetadata,
ConsolidationMetadata,
RefreshMentalModelMetadata,
RetainMetadata,
)
# Context variable for current schema (async-safe, per-task isolation)
# Note: default is None, actual default comes from config via get_current_schema()
@@ -38,6 +47,15 @@ def get_current_schema() -> str:
return schema
# Initialize tiktoken encoder once at module level for efficiency
_tiktoken_encoder = tiktoken.get_encoding("cl100k_base") # GPT-4/GPT-3.5-turbo encoding
def count_tokens(text: str) -> int:
"""Count tokens in text using tiktoken (cl100k_base encoding for GPT-4/3.5)."""
return len(_tiktoken_encoder.encode(text))
def fq_table(table_name: str) -> str:
"""
Get fully-qualified table name with current schema.
@@ -826,7 +844,11 @@ class MemoryEngine(MemoryEngineInterface):
logger.error(f"Failed to delete async operation record {operation_id}: {e}")
async def _mark_operation_failed(self, operation_id: str, error_message: str, error_traceback: str):
"""Helper to mark an operation as failed in the database."""
"""Helper to mark an operation as failed in the database.
Also checks if this is a child operation and updates the parent if all siblings are done.
Uses a single transaction to avoid race conditions when multiple children fail simultaneously.
"""
try:
pool = await self._get_pool()
# Truncate error message to avoid extremely long strings
@@ -834,35 +856,159 @@ class MemoryEngine(MemoryEngineInterface):
truncated_error = full_error[:5000] if len(full_error) > 5000 else full_error
async with acquire_with_retry(pool) as conn:
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = 'failed', error_message = $2, updated_at = NOW()
WHERE operation_id = $1
""",
uuid.UUID(operation_id),
truncated_error,
)
logger.info(f"Marked async operation as failed: {operation_id}")
async with conn.transaction():
# Mark this operation as failed
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = 'failed', error_message = $2, updated_at = NOW()
WHERE operation_id = $1
""",
uuid.UUID(operation_id),
truncated_error,
)
logger.info(f"Marked async operation as failed: {operation_id}")
# Check if this is a child operation and update parent if all siblings are done
# This happens in the same transaction after the child status is updated
await self._maybe_update_parent_operation(operation_id, conn)
except Exception as e:
logger.error(f"Failed to mark operation as failed {operation_id}: {e}")
async def _mark_operation_completed(self, operation_id: str):
"""Helper to mark an operation as completed in the database."""
"""Helper to mark an operation as completed in the database.
Also checks if this is a child operation and updates the parent if all siblings are done.
Uses a single transaction to avoid race conditions when multiple children complete simultaneously.
"""
try:
pool = await self._get_pool()
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# Mark this operation as completed
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
WHERE operation_id = $1
""",
uuid.UUID(operation_id),
)
logger.info(f"Marked async operation as completed: {operation_id}")
# Check if this is a child operation and update parent if all siblings are done
# This happens in the same transaction after the child status is updated
await self._maybe_update_parent_operation(operation_id, conn)
except Exception as e:
logger.error(f"Failed to mark operation as completed {operation_id}: {e}")
async def _maybe_update_parent_operation(self, child_operation_id: str, conn):
"""Check if this is a child operation and update parent status if all siblings are done.
Must be called within an active transaction that has already updated the child's status.
Uses SELECT FOR UPDATE to lock the parent and prevent race conditions.
Args:
child_operation_id: The operation ID that just completed or failed
conn: Database connection with an active transaction
"""
try:
# Get this operation's metadata to check if it has a parent
row = await conn.fetchrow(
f"""
SELECT result_metadata, bank_id
FROM {fq_table("async_operations")}
WHERE operation_id = $1
""",
uuid.UUID(child_operation_id),
)
if not row:
return
result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
parent_operation_id = result_metadata.get("parent_operation_id")
if not parent_operation_id:
# Not a child operation
return
bank_id = row["bank_id"]
# Lock the parent operation to prevent concurrent updates from other children
# Use FOR UPDATE to ensure only one child can update the parent at a time
parent_row = await conn.fetchrow(
f"""
SELECT operation_id
FROM {fq_table("async_operations")}
WHERE operation_id = $1 AND bank_id = $2
FOR UPDATE
""",
uuid.UUID(parent_operation_id),
bank_id,
)
if not parent_row:
# Parent doesn't exist (shouldn't happen)
return
# Get all sibling operations (including this one)
# This query runs in the same transaction, so it sees the current child's updated status
siblings = await conn.fetch(
f"""
SELECT status
FROM {fq_table("async_operations")}
WHERE bank_id = $1
AND result_metadata::jsonb @> $2::jsonb
""",
bank_id,
json.dumps({"parent_operation_id": parent_operation_id}),
)
if not siblings:
return
# Check if all siblings are done (completed or failed)
all_completed = all(sib["status"] == "completed" for sib in siblings)
any_failed = any(sib["status"] == "failed" for sib in siblings)
all_done = all(sib["status"] in ("completed", "failed") for sib in siblings)
if not all_done:
# Some siblings still pending/processing
return
# All siblings are done - update parent status
if any_failed:
new_status = "failed"
# Set parent error message to indicate child failure
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
SET status = $2, error_message = $3, updated_at = NOW()
WHERE operation_id = $1
""",
uuid.UUID(operation_id),
uuid.UUID(parent_operation_id),
new_status,
"One or more sub-batches failed",
)
logger.info(f"Marked async operation as completed: {operation_id}")
elif all_completed:
new_status = "completed"
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = $2, updated_at = NOW(), completed_at = NOW()
WHERE operation_id = $1
""",
uuid.UUID(parent_operation_id),
new_status,
)
logger.info(f"Updated parent operation {parent_operation_id} to status '{new_status}' (all children done)")
except Exception as e:
logger.error(f"Failed to mark operation as completed {operation_id}: {e}")
logger.error(f"Failed to update parent operation for child {child_operation_id}: {e}")
# Re-raise to rollback the transaction
raise
async def initialize(self):
"""Initialize the connection pool, models, and background workers.
@@ -1430,35 +1576,49 @@ class MemoryEngine(MemoryEngineInterface):
if "document_id" not in item:
item["document_id"] = document_id
# Auto-chunk large batches by character count to avoid timeouts and memory issues
# Calculate total character count
total_chars = sum(len(item.get("content", "")) for item in contents)
# Validate no duplicate document_ids in the batch
# Having duplicate document_ids causes race conditions in document upserts during parallel processing
doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
if len(doc_ids) != len(set(doc_ids)):
from collections import Counter
duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
raise ValueError(
f"Batch contains duplicate document_ids: {duplicates}. "
f"Each content item in a batch must have a unique document_id to avoid race conditions."
)
# Auto-chunk large batches by token count to avoid timeouts and memory issues
# Calculate total token count
total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
total_usage = TokenUsage()
CHARS_PER_BATCH = 600_000
# Get batch size threshold from config
config = get_config()
tokens_per_batch = config.retain_batch_tokens
if total_chars > CHARS_PER_BATCH:
# Split into smaller batches based on character count
if total_tokens > tokens_per_batch:
# Split into smaller batches based on token count
logger.info(
f"Large batch detected ({total_chars:,} chars from {len(contents)} items). Splitting into sub-batches of ~{CHARS_PER_BATCH:,} chars each..."
f"Large batch detected ({total_tokens:,} tokens from {len(contents)} items). Splitting into sub-batches of ~{tokens_per_batch:,} tokens each..."
)
sub_batches = []
current_batch = []
current_batch_chars = 0
current_batch_tokens = 0
for item in contents:
item_chars = len(item.get("content", ""))
item_tokens = count_tokens(item.get("content", ""))
# If adding this item would exceed the limit, start a new batch
# (unless current batch is empty - then we must include it even if it's large)
if current_batch and current_batch_chars + item_chars > CHARS_PER_BATCH:
if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
sub_batches.append(current_batch)
current_batch = [item]
current_batch_chars = item_chars
current_batch_tokens = item_tokens
else:
current_batch.append(item)
current_batch_chars += item_chars
current_batch_tokens += item_tokens
# Add the last batch
if current_batch:
@@ -1469,9 +1629,9 @@ class MemoryEngine(MemoryEngineInterface):
# Process each sub-batch
all_results = []
for i, sub_batch in enumerate(sub_batches, 1):
sub_batch_chars = sum(len(item.get("content", "")) for item in sub_batch)
sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
logger.info(
f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars"
f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
)
sub_results, sub_usage = await self._retain_batch_async_internal(
@@ -5463,10 +5623,10 @@ class MemoryEngine(MemoryEngineInterface):
)
total = total_row["total"] if total_row else 0
# Get operations with pagination
# Get operations with pagination (include result_metadata to check for parent operations)
operations = await conn.fetch(
f"""
SELECT operation_id, operation_type, created_at, status, error_message
SELECT operation_id, operation_type, created_at, status, error_message, result_metadata
FROM {fq_table("async_operations")}
WHERE {where_clause}
ORDER BY created_at DESC
@@ -5477,21 +5637,29 @@ class MemoryEngine(MemoryEngineInterface):
offset,
)
return {
"total": total,
"operations": [
# Build operation list using status from database
# Parent operations have their status updated when all children complete/fail
operation_list = []
for row in operations:
# Map DB status to API status (pending includes processing)
db_status = row["status"]
api_status = "pending" if db_status in ("pending", "processing") else db_status
operation_list.append(
{
"id": str(row["operation_id"]),
"task_type": row["operation_type"],
"items_count": 0,
"document_id": None,
"created_at": row["created_at"].isoformat(),
# Map DB status to API status (processing -> pending for simplicity)
"status": "pending" if row["status"] in ("pending", "processing") else row["status"],
"status": api_status,
"error_message": row["error_message"],
}
for row in operations
],
)
return {
"total": total,
"operations": operation_list,
}
async def get_operation_status(
@@ -5503,10 +5671,13 @@ class MemoryEngine(MemoryEngineInterface):
) -> dict[str, Any]:
"""Get the status of a specific async operation.
For parent operations, the status is automatically updated in the database when all children complete/fail.
Returns:
- status: "pending", "completed", or "failed"
- status: "pending", "completed", or "failed" (from database)
- updated_at: last update timestamp
- completed_at: completion timestamp (if completed)
- child_operations: (for parent operations) list of child operation statuses
"""
await self._authenticate_tenant(request_context)
pool = await self._get_pool()
@@ -5516,7 +5687,7 @@ class MemoryEngine(MemoryEngineInterface):
async with acquire_with_retry(pool) as conn:
row = await conn.fetchrow(
f"""
SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message
SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message, result_metadata
FROM {fq_table("async_operations")}
WHERE operation_id = $1 AND bank_id = $2
""",
@@ -5525,18 +5696,98 @@ class MemoryEngine(MemoryEngineInterface):
)
if row:
# Map DB status to API status (processing -> pending for simplicity)
# Check if this is a parent operation
result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
is_parent = result_metadata.get("is_parent", False)
# Use status from database (parent status is updated when all children complete/fail)
db_status = row["status"]
api_status = "pending" if db_status in ("pending", "processing") else db_status
return {
"operation_id": operation_id,
"status": api_status,
"operation_type": row["operation_type"],
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
"error_message": row["error_message"],
}
# For parent operations, include child operations list
if is_parent:
# Query child operations
child_rows = await conn.fetch(
f"""
SELECT operation_id, status, error_message, result_metadata
FROM {fq_table("async_operations")}
WHERE bank_id = $1
AND result_metadata::jsonb @> $2::jsonb
ORDER BY (result_metadata->>'sub_batch_index')::int
""",
bank_id,
json.dumps({"parent_operation_id": operation_id}),
)
# Build child operations list and check if parent status needs updating
child_statuses = []
all_done = True
any_failed = False
all_completed = True
for child_row in child_rows:
child_metadata = (
json.loads(child_row["result_metadata"]) if child_row["result_metadata"] else {}
)
child_status = child_row["status"]
child_statuses.append(
{
"operation_id": str(child_row["operation_id"]),
"status": child_status,
"sub_batch_index": child_metadata.get("sub_batch_index"),
"items_count": child_metadata.get("items_count"),
"error_message": child_row["error_message"],
}
)
if child_status not in ("completed", "failed"):
all_done = False
if child_status == "failed":
any_failed = True
if child_status != "completed":
all_completed = False
# Self-healing: if parent status is out of sync with children, update it
if all_done and api_status == "pending":
correct_status = "failed" if any_failed else "completed"
logger.warning(
f"Parent operation {operation_id} status out of sync (DB: pending, should be: {correct_status}). Fixing."
)
await conn.execute(
f"""
UPDATE {fq_table("async_operations")}
SET status = $2, updated_at = NOW(), completed_at = NOW()
WHERE operation_id = $1
""",
op_uuid,
correct_status,
)
api_status = correct_status
return {
"operation_id": operation_id,
"status": api_status,
"operation_type": row["operation_type"],
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
"error_message": row["error_message"],
"result_metadata": result_metadata,
"child_operations": child_statuses,
}
else:
# Regular operation (not a parent)
return {
"operation_id": operation_id,
"status": api_status,
"operation_type": row["operation_type"],
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
"error_message": row["error_message"],
"result_metadata": result_metadata,
}
else:
# Operation not found
return {
@@ -5712,31 +5963,126 @@ class MemoryEngine(MemoryEngineInterface):
request_context: "RequestContext",
document_tags: list[str] | None = None,
) -> dict[str, Any]:
"""Submit a batch retain operation to run asynchronously."""
"""Submit a batch retain operation to run asynchronously.
For large batches (exceeding retain_batch_chars threshold), automatically splits
into smaller sub-batches and creates a parent operation that tracks all children.
"""
await self._authenticate_tenant(request_context)
task_payload: dict[str, Any] = {"contents": contents}
if document_tags:
task_payload["document_tags"] = document_tags
# Pass tenant_id and api_key_id through task payload so the worker
# can propagate request context to downstream operations (e.g.,
# consolidation and mental model refreshes triggered after retain).
if request_context.tenant_id:
task_payload["_tenant_id"] = request_context.tenant_id
if request_context.api_key_id:
task_payload["_api_key_id"] = request_context.api_key_id
# Validate no duplicate document_ids in the batch
# Having duplicate document_ids causes race conditions in document upserts during parallel processing
doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
if len(doc_ids) != len(set(doc_ids)):
from collections import Counter
result = await self._submit_async_operation(
bank_id=bank_id,
operation_type="retain",
task_type="batch_retain",
task_payload=task_payload,
result_metadata={"items_count": len(contents)},
dedupe_by_bank=False,
duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
raise ValueError(
f"Batch contains duplicate document_ids: {duplicates}. "
f"Each content item in a batch must have a unique document_id to avoid race conditions."
)
# Calculate total token count and determine if we need to split
total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
config = get_config()
tokens_per_batch = config.retain_batch_tokens
# Split into sub-batches based on token count
sub_batches = []
current_batch = []
current_batch_tokens = 0
for item in contents:
item_tokens = count_tokens(item.get("content", ""))
# If adding this item would exceed the limit, start a new batch
# (unless current batch is empty - then we must include it even if it's large)
if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
sub_batches.append(current_batch)
current_batch = [item]
current_batch_tokens = item_tokens
else:
current_batch.append(item)
current_batch_tokens += item_tokens
# Add the last batch
if current_batch:
sub_batches.append(current_batch)
# Log splitting info if we actually split
if len(sub_batches) > 1:
logger.info(
f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). "
f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each"
)
# Always create parent operation (even for single batch - simpler, more reliable code path)
import uuid
parent_operation_id = uuid.uuid4()
pool = await self._get_pool()
# Create typed metadata for parent operation
parent_metadata = BatchRetainParentMetadata(
items_count=len(contents),
total_tokens=total_tokens,
num_sub_batches=len(sub_batches),
)
result["items_count"] = len(contents)
return result
async with acquire_with_retry(pool) as conn:
await conn.execute(
f"""
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
parent_operation_id,
bank_id,
"batch_retain",
json.dumps(parent_metadata.to_dict()),
"pending", # Will be updated by status aggregation
)
logger.info(f"Created parent operation {parent_operation_id} for {len(sub_batches)} sub-batch(es)")
# Submit child operations for each sub-batch
for i, sub_batch in enumerate(sub_batches, 1):
if len(sub_batches) > 1:
sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
logger.info(
f"Submitting sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
)
task_payload: dict[str, Any] = {"contents": sub_batch}
if document_tags:
task_payload["document_tags"] = document_tags
# Pass tenant_id and api_key_id through task payload
if request_context.tenant_id:
task_payload["_tenant_id"] = request_context.tenant_id
if request_context.api_key_id:
task_payload["_api_key_id"] = request_context.api_key_id
# Create typed metadata for child operation
child_metadata = BatchRetainChildMetadata(
items_count=len(sub_batch),
parent_operation_id=str(parent_operation_id),
sub_batch_index=i,
total_sub_batches=len(sub_batches),
)
# Create child operation with reference to parent
await self._submit_async_operation(
bank_id=bank_id,
operation_type="retain",
task_type="batch_retain",
task_payload=task_payload,
result_metadata=child_metadata.to_dict(),
dedupe_by_bank=False,
)
return {
"operation_id": str(parent_operation_id),
"items_count": len(contents),
}
async def submit_async_consolidation(
self,
@@ -0,0 +1,69 @@
"""
Typed metadata models for async operations.
These dataclasses define the structure of result_metadata for different operation types.
The metadata is exposed in the API for debugging purposes and may change without notice.
"""
from dataclasses import asdict, dataclass
from typing import Any
@dataclass
class BatchRetainParentMetadata:
"""Metadata for parent batch_retain operations (when split into sub-batches)."""
items_count: int
total_tokens: int
num_sub_batches: int
is_parent: bool = True
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
class BatchRetainChildMetadata:
"""Metadata for child batch_retain operations (individual sub-batches)."""
items_count: int
parent_operation_id: str
sub_batch_index: int
total_sub_batches: int
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
class RetainMetadata:
"""Metadata for regular retain operations (non-batched, deprecated async path)."""
items_count: int
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
class ConsolidationMetadata:
"""Metadata for consolidation operations."""
# Currently empty, but structure for future fields
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
class RefreshMentalModelMetadata:
"""Metadata for mental model refresh operations."""
mental_model_id: str
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
+1
View File
@@ -245,6 +245,7 @@ def main():
retain_extract_causal_links=config.retain_extract_causal_links,
retain_extraction_mode=config.retain_extraction_mode,
retain_custom_instructions=config.retain_custom_instructions,
retain_batch_tokens=config.retain_batch_tokens,
enable_observations=config.enable_observations,
consolidation_batch_size=config.consolidation_batch_size,
consolidation_max_tokens=config.consolidation_max_tokens,
@@ -0,0 +1,423 @@
"""Test async batch retain with smart batching and parent-child operations."""
import asyncio
import json
import uuid
import pytest
from hindsight_api.extensions import RequestContext
@pytest.mark.asyncio
async def test_duplicate_document_ids_rejected_async(memory, request_context):
"""Test that async retain rejects batches with duplicate document_ids."""
bank_id = "test_duplicate_async"
contents = [
{"content": "First item", "document_id": "doc1"},
{"content": "Second item", "document_id": "doc2"},
{"content": "Third item", "document_id": "doc1"}, # Duplicate!
]
# Should raise ValueError due to duplicate document_ids
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
await memory.submit_async_retain(
bank_id=bank_id,
contents=contents,
request_context=request_context,
)
@pytest.mark.asyncio
async def test_duplicate_document_ids_rejected_sync(memory, request_context):
"""Test that sync retain also rejects batches with duplicate document_ids."""
bank_id = "test_duplicate_sync"
contents = [
{"content": "First item", "document_id": "doc1"},
{"content": "Second item", "document_id": "doc1"}, # Duplicate!
]
# Should raise ValueError due to duplicate document_ids
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
await memory.retain_batch_async(
bank_id=bank_id,
contents=contents,
request_context=request_context,
)
@pytest.mark.asyncio
async def test_small_async_batch_no_splitting(memory, request_context):
"""Test that small async batches create parent with single child (simplified code path)."""
bank_id = "test_small_async"
contents = [{"content": "Alice works at Google", "document_id": f"doc{i}"} for i in range(5)]
# Calculate total chars (should be well under threshold)
total_chars = sum(len(item["content"]) for item in contents)
assert total_chars < 10_000, "Test batch should be small"
# Submit async retain
result = await memory.submit_async_retain(
bank_id=bank_id,
contents=contents,
request_context=request_context,
)
# Verify we got an operation_id back
assert "operation_id" in result
assert "items_count" in result
assert result["items_count"] == 5
operation_id = result["operation_id"]
# Wait for task to complete (SyncTaskBackend executes immediately)
await asyncio.sleep(0.1)
# Check operation status
status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=operation_id,
request_context=request_context,
)
# Should be a parent operation with single child (simplified code path)
assert status["status"] == "completed"
assert status["operation_type"] == "batch_retain"
assert "child_operations" in status
assert status["result_metadata"]["num_sub_batches"] == 1 # Single sub-batch
assert len(status["child_operations"]) == 1
assert status["child_operations"][0]["status"] == "completed"
@pytest.mark.asyncio
async def test_large_async_batch_auto_splits(memory, request_context):
"""Test that large async batches automatically split into sub-batches with parent operation."""
from hindsight_api.engine.memory_engine import count_tokens
bank_id = "test_large_async"
# Create a large batch that exceeds the threshold (10k tokens default)
# Repeating "A"s gets heavily compressed by tokenizer, use varied content
# Use ~22k chars per item = ~5.5k tokens per item, 2 items = ~11k tokens total (exceeds 10k)
large_content = "The quick brown fox jumps over the lazy dog. " * 500 # ~22k chars = ~5.5k tokens
contents = [{"content": large_content + f" item {i}", "document_id": f"doc{i}"} for i in range(2)]
# Calculate total tokens (should exceed threshold)
total_tokens = sum(count_tokens(item["content"]) for item in contents)
assert total_tokens > 10_000, "Test batch should exceed threshold"
# Submit async retain
result = await memory.submit_async_retain(
bank_id=bank_id,
contents=contents,
request_context=request_context,
)
# Verify we got an operation_id back
assert "operation_id" in result
assert "items_count" in result
assert result["items_count"] == 2
parent_operation_id = result["operation_id"]
# Wait for tasks to complete
await asyncio.sleep(0.5)
# Check parent operation status
parent_status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=parent_operation_id,
request_context=request_context,
)
# Should be a parent operation with children
assert parent_status["operation_type"] == "batch_retain"
assert "child_operations" in parent_status
assert "num_sub_batches" in parent_status["result_metadata"]
assert parent_status["result_metadata"]["num_sub_batches"] >= 2 # Should split into at least 2 batches
assert parent_status["result_metadata"]["items_count"] == 2
# Verify child operations
child_ops = parent_status["child_operations"]
assert len(child_ops) >= 2, "Should have at least 2 child operations"
# All children should be completed (SyncTaskBackend executes immediately)
for child in child_ops:
assert child["status"] == "completed"
assert child["sub_batch_index"] is not None
assert child["items_count"] > 0
# Parent status should be aggregated as "completed"
assert parent_status["status"] == "completed"
@pytest.mark.asyncio
async def test_parent_operation_status_aggregation_pending(memory, request_context):
"""Test that parent operation shows 'pending' when children are pending."""
bank_id = "test_parent_pending"
pool = await memory._get_pool()
# Manually create a parent operation
parent_id = uuid.uuid4()
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
parent_id,
bank_id,
"batch_retain",
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
"pending",
)
# Create 2 child operations - one completed, one pending
child1_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
child1_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 1,
"total_sub_batches": 2,
}
),
"completed",
)
child2_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
child2_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 2,
"total_sub_batches": 2,
}
),
"pending",
)
# Check parent status
parent_status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=str(parent_id),
request_context=request_context,
)
# Parent should aggregate as "pending" since one child is still pending
assert parent_status["status"] == "pending"
assert len(parent_status["child_operations"]) == 2
@pytest.mark.asyncio
async def test_parent_operation_status_aggregation_failed(memory, request_context):
"""Test that parent operation shows 'failed' when any child fails."""
bank_id = "test_parent_failed"
pool = await memory._get_pool()
# Manually create a parent operation
parent_id = uuid.uuid4()
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
parent_id,
bank_id,
"batch_retain",
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
"pending",
)
# Create 2 child operations - one completed, one failed
child1_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
child1_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 1,
"total_sub_batches": 2,
}
),
"completed",
)
child2_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status, error_message)
VALUES ($1, $2, $3, $4, $5, $6)
""",
child2_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 2,
"total_sub_batches": 2,
}
),
"failed",
"Test error",
)
# Check parent status
parent_status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=str(parent_id),
request_context=request_context,
)
# Parent should aggregate as "failed" since one child failed
assert parent_status["status"] == "failed"
assert len(parent_status["child_operations"]) == 2
# Verify child with error is included
failed_child = [c for c in parent_status["child_operations"] if c["status"] == "failed"][0]
assert failed_child["error_message"] == "Test error"
@pytest.mark.asyncio
async def test_parent_operation_status_aggregation_completed(memory, request_context):
"""Test that parent operation shows 'completed' when all children are completed."""
bank_id = "test_parent_completed"
pool = await memory._get_pool()
# Manually create a parent operation
parent_id = uuid.uuid4()
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
parent_id,
bank_id,
"batch_retain",
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
"pending",
)
# Create 2 child operations - both completed
child1_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
child1_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 1,
"total_sub_batches": 2,
}
),
"completed",
)
child2_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
VALUES ($1, $2, $3, $4, $5)
""",
child2_id,
bank_id,
"retain",
json.dumps(
{
"items_count": 10,
"parent_operation_id": str(parent_id),
"sub_batch_index": 2,
"total_sub_batches": 2,
}
),
"completed",
)
# Check parent status
parent_status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=str(parent_id),
request_context=request_context,
)
# Parent should aggregate as "completed" since all children are completed
assert parent_status["status"] == "completed"
assert len(parent_status["child_operations"]) == 2
assert all(c["status"] == "completed" for c in parent_status["child_operations"])
@pytest.mark.asyncio
async def test_config_retain_batch_tokens_respected(memory, request_context):
"""Test that the retain_batch_tokens config setting is respected."""
from hindsight_api.config import get_config
from hindsight_api.engine.memory_engine import count_tokens
bank_id = "test_config_batch_tokens"
config = get_config()
# Check that config has the retain_batch_tokens setting
assert hasattr(config, "retain_batch_tokens")
assert config.retain_batch_tokens > 0
# Create a batch that's just under the threshold
# Use content that produces roughly half the token limit per item
content_size = config.retain_batch_tokens * 2 # chars (rough estimate: 1 token ~= 4 chars)
contents = [{"content": "A" * content_size, "document_id": f"doc{i}"} for i in range(2)]
total_tokens = sum(count_tokens(item["content"]) for item in contents)
# Should be equal to threshold (boundary case, no splitting since we use > not >=)
assert total_tokens <= config.retain_batch_tokens
# Submit - should NOT split
result = await memory.submit_async_retain(
bank_id=bank_id,
contents=contents,
request_context=request_context,
)
# Wait for completion
await asyncio.sleep(0.1)
# Check status - should be a parent with single child (even for small batches)
status = await memory.get_operation_status(
bank_id=bank_id,
operation_id=result["operation_id"],
request_context=request_context,
)
# Even small batches use parent-child pattern now (simpler code path)
assert "child_operations" in status
assert status["result_metadata"]["num_sub_batches"] == 1
+25 -2
View File
@@ -1,6 +1,6 @@
"""Unit tests for async retain tag propagation."""
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -12,9 +12,23 @@ from hindsight_api.models import RequestContext
async def test_submit_async_retain_includes_document_tags_in_task_payload():
"""submit_async_retain should include document_tags in queued task payload."""
engine = MemoryEngine.__new__(MemoryEngine)
engine._initialized = True
engine._authenticate_tenant = AsyncMock()
engine._submit_async_operation = AsyncMock(return_value={"operation_id": "op-1"})
# Mock the pool and connection for parent operation creation
mock_conn = AsyncMock()
mock_conn.execute = AsyncMock()
mock_conn.transaction = MagicMock()
mock_conn.transaction.return_value.__aenter__ = AsyncMock()
mock_conn.transaction.return_value.__aexit__ = AsyncMock()
mock_pool = AsyncMock()
mock_pool.acquire = AsyncMock(return_value=mock_conn)
mock_pool.release = AsyncMock()
engine._get_pool = AsyncMock(return_value=mock_pool)
request_context = RequestContext(tenant_id="tenant-a", api_key_id="key-a")
contents = [{"content": "Async retain payload test."}]
document_tags = ["scope:tools", "user:alice"]
@@ -27,10 +41,18 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
request_context=request_context,
)
assert result == {"operation_id": "op-1", "items_count": 1}
# Check result structure
assert "operation_id" in result
assert "items_count" in result
assert result["items_count"] == 1
# Verify authentication was called
engine._authenticate_tenant.assert_awaited_once_with(request_context)
# Verify child operation was submitted
engine._submit_async_operation.assert_awaited_once()
# Verify child operation payload contains document_tags
kwargs = engine._submit_async_operation.await_args.kwargs
assert kwargs["bank_id"] == "bank-1"
assert kwargs["operation_type"] == "retain"
@@ -45,6 +67,7 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
async def test_handle_batch_retain_forwards_document_tags_to_retain_batch_async():
"""Worker handler should forward document_tags from task payload."""
engine = MemoryEngine.__new__(MemoryEngine)
engine._initialized = True
engine.retain_batch_async = AsyncMock(return_value={"items_count": 1})
task_dict = {
@@ -528,7 +528,7 @@ async def test_delete_bank(api_client):
{
"content": "Bob is the CTO and leads the engineering team.",
"context": "team info",
"document_id": "team-doc-1",
"document_id": "team-doc-2",
},
]
},
@@ -24,6 +24,7 @@ hindsight_client_api/models/bank_profile_response.py
hindsight_client_api/models/bank_stats_response.py
hindsight_client_api/models/budget.py
hindsight_client_api/models/cancel_operation_response.py
hindsight_client_api/models/child_operation_status.py
hindsight_client_api/models/chunk_data.py
hindsight_client_api/models/chunk_include_options.py
hindsight_client_api/models/chunk_response.py
@@ -49,6 +49,7 @@ from hindsight_client_api.models.bank_profile_response import BankProfileRespons
from hindsight_client_api.models.bank_stats_response import BankStatsResponse
from hindsight_client_api.models.budget import Budget
from hindsight_client_api.models.cancel_operation_response import CancelOperationResponse
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
from hindsight_client_api.models.chunk_data import ChunkData
from hindsight_client_api.models.chunk_include_options import ChunkIncludeOptions
from hindsight_client_api.models.chunk_response import ChunkResponse
@@ -25,6 +25,7 @@ from hindsight_client_api.models.bank_profile_response import BankProfileRespons
from hindsight_client_api.models.bank_stats_response import BankStatsResponse
from hindsight_client_api.models.budget import Budget
from hindsight_client_api.models.cancel_operation_response import CancelOperationResponse
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
from hindsight_client_api.models.chunk_data import ChunkData
from hindsight_client_api.models.chunk_include_options import ChunkIncludeOptions
from hindsight_client_api.models.chunk_response import ChunkResponse
@@ -0,0 +1,110 @@
# coding: utf-8
"""
Hindsight HTTP API
HTTP API for Hindsight
The version of the OpenAPI document: 0.4.11
Generated by OpenAPI Generator (https://openapi-generator.tech)
Do not edit the class manually.
""" # noqa: E501
from __future__ import annotations
import pprint
import re # noqa: F401
import json
from pydantic import BaseModel, ConfigDict, StrictInt, StrictStr
from typing import Any, ClassVar, Dict, List, Optional
from typing import Optional, Set
from typing_extensions import Self
class ChildOperationStatus(BaseModel):
"""
Status of a child operation (for batch operations).
""" # noqa: E501
operation_id: StrictStr
status: StrictStr
sub_batch_index: Optional[StrictInt] = None
items_count: Optional[StrictInt] = None
error_message: Optional[StrictStr] = None
__properties: ClassVar[List[str]] = ["operation_id", "status", "sub_batch_index", "items_count", "error_message"]
model_config = ConfigDict(
populate_by_name=True,
validate_assignment=True,
protected_namespaces=(),
)
def to_str(self) -> str:
"""Returns the string representation of the model using alias"""
return pprint.pformat(self.model_dump(by_alias=True))
def to_json(self) -> str:
"""Returns the JSON representation of the model using alias"""
# TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
return json.dumps(self.to_dict())
@classmethod
def from_json(cls, json_str: str) -> Optional[Self]:
"""Create an instance of ChildOperationStatus from a JSON string"""
return cls.from_dict(json.loads(json_str))
def to_dict(self) -> Dict[str, Any]:
"""Return the dictionary representation of the model using alias.
This has the following differences from calling pydantic's
`self.model_dump(by_alias=True)`:
* `None` is only added to the output dict for nullable fields that
were set at model initialization. Other fields with value `None`
are ignored.
"""
excluded_fields: Set[str] = set([
])
_dict = self.model_dump(
by_alias=True,
exclude=excluded_fields,
exclude_none=True,
)
# set to None if sub_batch_index (nullable) is None
# and model_fields_set contains the field
if self.sub_batch_index is None and "sub_batch_index" in self.model_fields_set:
_dict['sub_batch_index'] = None
# set to None if items_count (nullable) is None
# and model_fields_set contains the field
if self.items_count is None and "items_count" in self.model_fields_set:
_dict['items_count'] = None
# set to None if error_message (nullable) is None
# and model_fields_set contains the field
if self.error_message is None and "error_message" in self.model_fields_set:
_dict['error_message'] = None
return _dict
@classmethod
def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
"""Create an instance of ChildOperationStatus from a dict"""
if obj is None:
return None
if not isinstance(obj, dict):
return cls.model_validate(obj)
_obj = cls.model_validate({
"operation_id": obj.get("operation_id"),
"status": obj.get("status"),
"sub_batch_index": obj.get("sub_batch_index"),
"items_count": obj.get("items_count"),
"error_message": obj.get("error_message")
})
return _obj
@@ -19,6 +19,7 @@ import json
from pydantic import BaseModel, ConfigDict, StrictStr, field_validator
from typing import Any, ClassVar, Dict, List, Optional
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
from typing import Optional, Set
from typing_extensions import Self
@@ -33,7 +34,9 @@ class OperationStatusResponse(BaseModel):
updated_at: Optional[StrictStr] = None
completed_at: Optional[StrictStr] = None
error_message: Optional[StrictStr] = None
__properties: ClassVar[List[str]] = ["operation_id", "status", "operation_type", "created_at", "updated_at", "completed_at", "error_message"]
result_metadata: Optional[Dict[str, Any]] = None
child_operations: Optional[List[ChildOperationStatus]] = None
__properties: ClassVar[List[str]] = ["operation_id", "status", "operation_type", "created_at", "updated_at", "completed_at", "error_message", "result_metadata", "child_operations"]
@field_validator('status')
def status_validate_enum(cls, value):
@@ -81,6 +84,13 @@ class OperationStatusResponse(BaseModel):
exclude=excluded_fields,
exclude_none=True,
)
# override the default output from pydantic by calling `to_dict()` of each item in child_operations (list)
_items = []
if self.child_operations:
for _item_child_operations in self.child_operations:
if _item_child_operations:
_items.append(_item_child_operations.to_dict())
_dict['child_operations'] = _items
# set to None if operation_type (nullable) is None
# and model_fields_set contains the field
if self.operation_type is None and "operation_type" in self.model_fields_set:
@@ -106,6 +116,16 @@ class OperationStatusResponse(BaseModel):
if self.error_message is None and "error_message" in self.model_fields_set:
_dict['error_message'] = None
# set to None if result_metadata (nullable) is None
# and model_fields_set contains the field
if self.result_metadata is None and "result_metadata" in self.model_fields_set:
_dict['result_metadata'] = None
# set to None if child_operations (nullable) is None
# and model_fields_set contains the field
if self.child_operations is None and "child_operations" in self.model_fields_set:
_dict['child_operations'] = None
return _dict
@classmethod
@@ -124,7 +144,9 @@ class OperationStatusResponse(BaseModel):
"created_at": obj.get("created_at"),
"updated_at": obj.get("updated_at"),
"completed_at": obj.get("completed_at"),
"error_message": obj.get("error_message")
"error_message": obj.get("error_message"),
"result_metadata": obj.get("result_metadata"),
"child_operations": [ChildOperationStatus.from_dict(_item) for _item in obj["child_operations"]] if obj.get("child_operations") is not None else None
})
return _obj
@@ -278,6 +278,34 @@ export type CancelOperationResponse = {
operation_id: string;
};
/**
* ChildOperationStatus
*
* Status of a child operation (for batch operations).
*/
export type ChildOperationStatus = {
/**
* Operation Id
*/
operation_id: string;
/**
* Status
*/
status: string;
/**
* Sub Batch Index
*/
sub_batch_index?: number | null;
/**
* Items Count
*/
items_count?: number | null;
/**
* Error Message
*/
error_message?: string | null;
};
/**
* ChunkData
*
@@ -1201,6 +1229,20 @@ export type OperationStatusResponse = {
* Error Message
*/
error_message?: string | null;
/**
* Result Metadata
*
* Internal metadata for debugging. Structure may change without notice. Not for production use.
*/
result_metadata?: {
[key: string]: unknown;
} | null;
/**
* Child Operations
*
* Child operations for batch operations (if applicable)
*/
child_operations?: Array<ChildOperationStatus> | null;
};
/**
@@ -12,6 +12,13 @@ import {
TableHeader,
TableRow,
} from "@/components/ui/table";
import {
Dialog,
DialogContent,
DialogDescription,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { RefreshCw, Clock, AlertCircle, CheckCircle, Loader2, X } from "lucide-react";
interface Operation {
@@ -24,6 +31,46 @@ interface Operation {
error_message: string | null;
}
interface ChildOperationStatus {
operation_id: string;
status: string;
sub_batch_index: number | null;
items_count: number | null;
error_message: string | null;
}
type OperationDetails =
| {
operation_id: string;
status: string;
operation_type: string | null;
created_at: string | null;
updated_at: string | null;
completed_at: string | null;
error_message: string | null;
result_metadata?: {
items_count?: number;
total_tokens?: number;
num_sub_batches?: number;
is_parent?: boolean;
[key: string]: any;
};
child_operations?: ChildOperationStatus[];
error?: never; // Not present in success case
}
| {
error: string; // Error state when loading fails
operation_id?: never;
status?: never;
operation_type?: never;
created_at?: never;
updated_at?: never;
completed_at?: never;
error_message?: never;
result_metadata?: never;
child_operations?: never;
};
export function BankOperationsView() {
const { currentBank } = useBank();
const [operations, setOperations] = useState<Operation[]>([]);
@@ -33,6 +80,9 @@ export function BankOperationsView() {
const [offset, setOffset] = useState(0);
const [cancellingOpId, setCancellingOpId] = useState<string | null>(null);
const [loading, setLoading] = useState(false);
const [selectedOperation, setSelectedOperation] = useState<OperationDetails | null>(null);
const [dialogOpen, setDialogOpen] = useState(false);
const [loadingDetails, setLoadingDetails] = useState(false);
const loadOperations = async (
newStatusFilter: string | null = statusFilter,
@@ -82,6 +132,22 @@ export function BankOperationsView() {
}
};
const handleOperationClick = async (operationId: string) => {
if (!currentBank) return;
setLoadingDetails(true);
setDialogOpen(true);
try {
const details = await client.getOperationStatus(currentBank, operationId);
setSelectedOperation(details);
} catch (error) {
console.error("Error loading operation details:", error);
setSelectedOperation({ error: "Failed to load operation details" });
} finally {
setLoadingDetails(false);
}
};
useEffect(() => {
if (currentBank) {
loadOperations();
@@ -152,7 +218,11 @@ export function BankOperationsView() {
</TableHeader>
<TableBody>
{operations.map((op) => (
<TableRow key={op.id} className={op.status === "failed" ? "bg-red-500/5" : ""}>
<TableRow
key={op.id}
className={`cursor-pointer hover:bg-muted/50 ${op.status === "failed" ? "bg-red-500/5" : ""}`}
onClick={() => handleOperationClick(op.id)}
>
<TableCell className="font-mono text-xs text-muted-foreground">
{op.id.substring(0, 8)}
</TableCell>
@@ -189,7 +259,10 @@ export function BankOperationsView() {
variant="ghost"
size="sm"
className="h-7 text-xs text-muted-foreground hover:text-red-600 dark:hover:text-red-400"
onClick={() => handleCancelOperation(op.id)}
onClick={(e) => {
e.stopPropagation();
handleCancelOperation(op.id);
}}
disabled={cancellingOpId === op.id}
>
{cancellingOpId === op.id ? (
@@ -240,6 +313,169 @@ export function BankOperationsView() {
</p>
)}
</div>
{/* Operation Details Dialog */}
<Dialog open={dialogOpen} onOpenChange={setDialogOpen}>
<DialogContent className="max-w-3xl max-h-[80vh] overflow-y-auto">
<DialogHeader>
<DialogTitle>Operation Details</DialogTitle>
<DialogDescription>
{selectedOperation?.operation_id && (
<span className="font-mono text-xs">{selectedOperation.operation_id}</span>
)}
</DialogDescription>
</DialogHeader>
{loadingDetails ? (
<div className="flex items-center justify-center py-8">
<Loader2 className="w-6 h-6 animate-spin text-muted-foreground" />
</div>
) : selectedOperation ? (
<div className="space-y-4">
{selectedOperation.error ? (
<div className="text-red-600 dark:text-red-400">{selectedOperation.error}</div>
) : (
<>
{/* Basic Info */}
<div className="grid grid-cols-2 gap-4">
<div>
<div className="text-sm font-medium text-muted-foreground">Status</div>
<div className="mt-1">
{selectedOperation.status === "pending" && (
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-amber-500/10 text-amber-600 dark:text-amber-400 border border-amber-500/20">
<Clock className="w-3 h-3" />
pending
</span>
)}
{selectedOperation.status === "failed" && (
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-red-500/10 text-red-600 dark:text-red-400 border border-red-500/20">
<AlertCircle className="w-3 h-3" />
failed
</span>
)}
{selectedOperation.status === "completed" && (
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-emerald-500/10 text-emerald-600 dark:text-emerald-400 border border-emerald-500/20">
<CheckCircle className="w-3 h-3" />
completed
</span>
)}
</div>
</div>
<div>
<div className="text-sm font-medium text-muted-foreground">Type</div>
<div className="mt-1 font-mono text-sm">
{selectedOperation.operation_type}
</div>
</div>
<div>
<div className="text-sm font-medium text-muted-foreground">Created</div>
<div className="mt-1 text-sm">
{selectedOperation.created_at
? new Date(selectedOperation.created_at).toLocaleString()
: "N/A"}
</div>
</div>
<div>
<div className="text-sm font-medium text-muted-foreground">Updated</div>
<div className="mt-1 text-sm">
{selectedOperation.updated_at
? new Date(selectedOperation.updated_at).toLocaleString()
: "N/A"}
</div>
</div>
{selectedOperation.completed_at && (
<div>
<div className="text-sm font-medium text-muted-foreground">Completed</div>
<div className="mt-1 text-sm">
{new Date(selectedOperation.completed_at).toLocaleString()}
</div>
</div>
)}
{selectedOperation.result_metadata?.items_count !== undefined && (
<div>
<div className="text-sm font-medium text-muted-foreground">Total Items</div>
<div className="mt-1 text-sm">
{selectedOperation.result_metadata.items_count}
</div>
</div>
)}
</div>
{/* Error Message */}
{selectedOperation.error_message && (
<div className="rounded-lg border border-red-500/20 bg-red-500/5 p-3">
<div className="text-sm font-medium text-red-600 dark:text-red-400 mb-1">
Error
</div>
<div className="text-sm text-red-600/80 dark:text-red-400/80 font-mono">
{selectedOperation.error_message}
</div>
</div>
)}
{/* Child Operations (for parent operations) */}
{selectedOperation.child_operations &&
selectedOperation.child_operations.length > 0 && (
<div>
<div className="text-sm font-medium text-muted-foreground mb-2">
Sub-batches (
{selectedOperation.result_metadata?.num_sub_batches ||
selectedOperation.child_operations.length}
)
</div>
<div className="rounded-lg border">
<Table>
<TableHeader>
<TableRow>
<TableHead className="w-[60px]">Index</TableHead>
<TableHead className="w-[100px]">ID</TableHead>
<TableHead className="w-[80px]">Items</TableHead>
<TableHead>Status</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{selectedOperation.child_operations.map((child) => (
<TableRow key={child.operation_id}>
<TableCell className="text-sm">{child.sub_batch_index}</TableCell>
<TableCell className="font-mono text-xs text-muted-foreground">
{child.operation_id.substring(0, 8)}
</TableCell>
<TableCell className="text-sm">{child.items_count}</TableCell>
<TableCell>
{child.status === "pending" && (
<span className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-amber-500/10 text-amber-600 dark:text-amber-400">
<Clock className="w-3 h-3" />
pending
</span>
)}
{child.status === "failed" && (
<span
className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-red-500/10 text-red-600 dark:text-red-400"
title={child.error_message ?? undefined}
>
<AlertCircle className="w-3 h-3" />
failed
</span>
)}
{child.status === "completed" && (
<span className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-emerald-500/10 text-emerald-600 dark:text-emerald-400">
<CheckCircle className="w-3 h-3" />
completed
</span>
)}
</TableCell>
</TableRow>
))}
</TableBody>
</Table>
</div>
</div>
)}
</>
)}
</div>
) : null}
</DialogContent>
</Dialog>
</div>
);
}
+39
View File
@@ -68,6 +68,45 @@ Tests long-term memory across different categories.
- `--only-failed` - Retry failed questions
- `--fill` - Resume interrupted runs
### Consolidation Performance
Tests consolidation throughput and identifies bottlenecks.
```bash
./scripts/benchmarks/run-consolidation.sh
# With custom memory count
NUM_MEMORIES=200 ./scripts/benchmarks/run-consolidation.sh
```
### Retain Performance
Measures retain operation performance (throughput and token usage).
**Prerequisites:** API server must be running (`./scripts/dev/start-api.sh`)
```bash
# Basic usage
./scripts/benchmarks/run-retain-perf.sh \
--document hindsight-dev/benchmarks/perf/test_data/sample_document.txt
# Save results to JSON
./scripts/benchmarks/run-retain-perf.sh \
--document ./my_document.txt \
--bank-id my-test-bank \
--output results/retain_perf.json
```
**Options:**
- `--document PATH` - Document file to retain (required)
- `--bank-id ID` - Bank ID to use (default: perf-test)
- `--context TEXT` - Optional context
- `--api-url URL` - API URL (default: http://localhost:8000)
- `--timeout SECONDS` - Request timeout (default: 300)
- `--output PATH` - Save results to JSON file
See [perf/README.md](perf/README.md) for detailed documentation.
## Visualizer
View benchmark results in a web UI:
@@ -0,0 +1 @@
"""Performance benchmarks for Hindsight operations."""
@@ -0,0 +1,456 @@
"""
Retain operation performance benchmark.
Measures retain operation performance by:
1. Loading a document from a file or directory
2. Sending it to the retain endpoint via HTTP (batched for directories)
3. Measuring time taken and token usage
4. Reporting performance metrics
Usage:
# Single file
uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document <file_path> [options]
# Directory (batches all files)
uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document <dir_path> [options]
"""
import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import Any
import httpx
from rich.console import Console
from rich.table import Table
console = Console()
async def retain_via_memory_engine(
bank_id: str,
items: list[dict[str, Any]],
) -> tuple[float, dict[str, Any]]:
"""
Send retain request directly to MemoryEngine (in-memory, no HTTP).
Args:
bank_id: Bank ID to retain into
items: List of items to retain
Returns:
Tuple of (duration_seconds, response_data)
"""
from hindsight_api import MemoryEngine
from hindsight_api.models import RequestContext
# Initialize memory engine
memory = MemoryEngine(
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"),
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-20b"),
memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
)
await memory.initialize()
# Measure time
start_time = time.time()
try:
# Call retain_batch_async directly
result, usage = await memory.retain_batch_async(
bank_id=bank_id,
contents=items,
request_context=RequestContext(),
return_usage=True,
)
duration = time.time() - start_time
# Format response to match HTTP response structure
response_data = {
"success": True,
"bank_id": bank_id,
"items_count": len(items),
"async": False,
"usage": usage.model_dump() if usage else None,
}
return duration, response_data
finally:
# Close memory engine connections
pool = await memory._get_pool()
await pool.close()
async def retain_via_http(
base_url: str,
bank_id: str,
items: list[dict[str, Any]],
timeout: float = 300.0,
) -> tuple[float, dict[str, Any]]:
"""
Send retain request via HTTP and measure performance.
Args:
base_url: API base URL (e.g., http://localhost:8000)
bank_id: Bank ID to retain into
items: List of items to retain (each with 'content' and optional 'context', 'metadata')
timeout: Request timeout in seconds
Returns:
Tuple of (duration_seconds, response_data)
"""
url = f"{base_url}/v1/default/banks/{bank_id}/memories"
payload = {"items": items}
headers = {"Content-Type": "application/json"}
# Measure time
start_time = time.time()
async with httpx.AsyncClient(timeout=timeout) as client:
response = await client.post(url, json=payload, headers=headers)
response.raise_for_status()
result = response.json()
duration = time.time() - start_time
return duration, result
def load_documents(path: str) -> tuple[list[dict[str, Any]], int]:
"""
Load document(s) from file or directory.
For directories: loads all .json, .txt, and .md files
For JSON files with 'content' field: extracts content
For other files: reads entire file as content
Returns:
Tuple of (items_list, total_content_length)
items_list: List of dicts with 'content' and optional 'metadata'/'context'
total_content_length: Total character count across all documents
"""
file_path = Path(path)
if not file_path.exists():
raise FileNotFoundError(f"Path not found: {path}")
items = []
total_length = 0
if file_path.is_file():
# Single file
content, metadata = _load_single_file(file_path)
total_length = len(content)
item = {"content": content}
if metadata:
item["metadata"] = metadata
items.append(item)
else:
# Directory - load all supported files
supported_extensions = {".json", ".txt", ".md"}
files = [f for f in file_path.rglob("*") if f.is_file() and f.suffix in supported_extensions]
if not files:
raise ValueError(f"No supported files (.json, .txt, .md) found in directory: {path}")
console.print(f"Found {len(files)} files in directory")
for file in sorted(files):
try:
content, metadata = _load_single_file(file)
total_length += len(content)
item = {"content": content}
if metadata:
item["metadata"] = metadata
# Add filename as context for batch processing
item["context"] = f"Source: {file.name}"
items.append(item)
except Exception as e:
console.print(f"[yellow]Warning: Failed to load {file.name}: {e}[/yellow]")
continue
return items, total_length
def _load_single_file(file_path: Path) -> tuple[str, dict[str, Any] | None]:
"""
Load a single file and extract content.
Returns:
Tuple of (content, metadata)
"""
if file_path.suffix == ".json":
# Try to parse as JSON and extract 'content' field
try:
data = json.loads(file_path.read_text())
if isinstance(data, dict) and "content" in data:
# Extract metadata if present
metadata = data.get("metadata", {})
# Add doc_id to metadata if present
if "doc_id" in data:
metadata["doc_id"] = data["doc_id"]
return data["content"], metadata if metadata else None
else:
# Fallback: use entire JSON as string
return file_path.read_text(), None
except json.JSONDecodeError:
# Not valid JSON, read as text
return file_path.read_text(), None
else:
# Read as plain text
return file_path.read_text(), None
def display_results(
duration: float,
usage: dict[str, int] | None,
content_length: int,
bank_id: str,
num_documents: int,
) -> None:
"""Display benchmark results in a formatted table."""
table = Table(title="Retain Performance Benchmark Results")
table.add_column("Metric", style="cyan")
table.add_column("Value", style="green")
table.add_row("Bank ID", bank_id)
table.add_row("Documents", f"{num_documents:,}")
table.add_row("Total Content Length", f"{content_length:,} chars")
if num_documents > 1:
table.add_row("Avg Content/Doc", f"{content_length / num_documents:,.0f} chars")
table.add_row("", "") # Separator
table.add_row("Duration", f"{duration:.3f}s")
table.add_row("Throughput", f"{content_length / duration:,.0f} chars/sec")
if num_documents > 1:
table.add_row("Docs/Second", f"{num_documents / duration:.2f}")
if usage:
table.add_row("", "") # Separator
table.add_row("Input Tokens", f"{usage.get('input_tokens', 0):,}")
table.add_row("Output Tokens", f"{usage.get('output_tokens', 0):,}")
table.add_row("Total Tokens", f"{usage.get('total_tokens', 0):,}")
table.add_row("Tokens/Second", f"{usage.get('total_tokens', 0) / duration:,.1f}")
if num_documents > 1:
table.add_row("Avg Tokens/Doc", f"{usage.get('total_tokens', 0) / num_documents:,.0f}")
else:
table.add_row("", "") # Separator
table.add_row("Token Usage", "Not available (async mode or error)")
console.print("\n")
console.print(table)
def save_results(
output_path: Path,
duration: float,
usage: dict[str, int] | None,
content_length: int,
bank_id: str,
document_path: str,
num_documents: int,
) -> None:
"""Save results to JSON file."""
results = {
"bank_id": bank_id,
"document_path": document_path,
"num_documents": num_documents,
"content_length": content_length,
"avg_content_per_doc": content_length / num_documents if num_documents > 0 else 0,
"duration_seconds": duration,
"chars_per_second": content_length / duration,
"docs_per_second": num_documents / duration if num_documents > 0 else 0,
"usage": usage,
}
if usage:
results["tokens_per_second"] = usage.get("total_tokens", 0) / duration
results["avg_tokens_per_doc"] = usage.get("total_tokens", 0) / num_documents if num_documents > 0 else 0
with open(output_path, "w") as f:
json.dump(results, f, indent=2)
console.print(f"\n[green]✓[/green] Results saved to {output_path}")
async def main():
"""Run the retain performance benchmark."""
parser = argparse.ArgumentParser(
description="Benchmark retain operation performance",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Benchmark with a single document file
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
--document ./test_data/large_doc.txt \\
--bank-id perf-test-001
# Benchmark with a directory (batches all files)
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
--document ~/Documents/my-docs/ \\
--bank-id perf-test-batch \\
--output results/batch_perf.json
# With custom API URL and save results
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
--document ./test_data/ \\
--bank-id perf-test-001 \\
--api-url http://localhost:8000 \\
--output results/retain_perf_001.json
""",
)
parser.add_argument(
"--document",
required=True,
help="Path to document file or directory (for directories, batches all .json/.txt/.md files)",
)
parser.add_argument(
"--bank-id",
default="perf-test",
help="Bank ID to use (default: perf-test)",
)
parser.add_argument(
"--context",
help="Optional context for the retain operation (only used for single file mode)",
)
parser.add_argument(
"--api-url",
default="http://localhost:8000",
help="API base URL (default: http://localhost:8000)",
)
parser.add_argument(
"--timeout",
type=float,
default=300.0,
help="Request timeout in seconds (default: 300)",
)
parser.add_argument(
"--output",
type=Path,
help="Path to save results JSON (optional)",
)
parser.add_argument(
"--in-memory",
action="store_true",
help="Use in-memory MemoryEngine instead of HTTP (bypasses API server, useful for isolating performance)",
)
args = parser.parse_args()
console.print("\n[bold cyan]Retain Performance Benchmark[/bold cyan]")
console.print("=" * 80)
# Check mode
if args.in_memory:
console.print("\n[cyan]Mode: IN-MEMORY (direct MemoryEngine, no HTTP)[/cyan]")
else:
console.print(f"\n[cyan]Mode: HTTP (via {args.api_url})[/cyan]")
# Check if server is running (skip for in-memory mode)
if not args.in_memory:
console.print(f"\n[1] Checking API server at {args.api_url}...")
try:
async with httpx.AsyncClient() as client:
response = await client.get(f"{args.api_url}/health", timeout=5.0)
response.raise_for_status()
console.print(" [green]✓[/green] API server is running")
except Exception as e:
console.print(f" [red]✗[/red] API server is not accessible: {e}")
console.print("\n[yellow]Please ensure the API server is running:[/yellow]")
console.print(" ./scripts/dev/start-api.sh")
sys.exit(1)
# Load document(s)
doc_path = Path(args.document)
if doc_path.is_dir():
console.print(f"\n[2] Loading documents from directory {args.document}...")
else:
console.print(f"\n[2] Loading document from {args.document}...")
try:
items, total_content_length = load_documents(args.document)
num_docs = len(items)
# Add context to single file if provided
if num_docs == 1 and args.context:
items[0]["context"] = args.context
console.print(
f" [green]✓[/green] Loaded {num_docs:,} document{'s' if num_docs > 1 else ''} ({total_content_length:,} characters)"
)
if num_docs > 1:
console.print(
f" [cyan]Average content per document: {total_content_length / num_docs:,.0f} chars[/cyan]"
)
except Exception as e:
console.print(f" [red]✗[/red] Failed to load documents: {e}")
sys.exit(1)
# Run benchmark
console.print(f"\n[3] {'Processing' if args.in_memory else 'Sending retain request to'} bank '{args.bank_id}'...")
console.print(f" [cyan]Retaining {num_docs:,} document{'s' if num_docs > 1 else ''} in batch...[/cyan]")
try:
if args.in_memory:
# In-memory mode: call MemoryEngine directly
duration, result = await retain_via_memory_engine(
bank_id=args.bank_id,
items=items,
)
else:
# HTTP mode: call API endpoint
duration, result = await retain_via_http(
base_url=args.api_url,
bank_id=args.bank_id,
items=items,
timeout=args.timeout,
)
console.print(f" [green]✓[/green] Retain completed in {duration:.3f}s")
# Extract usage
usage = result.get("usage")
except httpx.HTTPStatusError as e:
console.print(f" [red]✗[/red] HTTP error: {e.response.status_code}")
console.print(f" Response: {e.response.text}")
sys.exit(1)
except Exception as e:
console.print(f" [red]✗[/red] Request failed: {e}")
sys.exit(1)
# Display results
console.print("\n[4] Results:")
display_results(
duration=duration,
usage=usage,
content_length=total_content_length,
bank_id=args.bank_id,
num_documents=num_docs,
)
# Save results if requested
if args.output:
console.print("\n[5] Saving results...")
args.output.parent.mkdir(parents=True, exist_ok=True)
save_results(
output_path=args.output,
duration=duration,
usage=usage,
content_length=total_content_length,
bank_id=args.bank_id,
document_path=args.document,
num_documents=num_docs,
)
console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")
if __name__ == "__main__":
asyncio.run(main())
+1
View File
@@ -14,6 +14,7 @@ dependencies = [
"openai>=1.0.0",
"rich>=13.0.0",
"pydantic>=2.0.0",
"httpx>=0.27.0",
]
[project.optional-dependencies]
+18 -1
View File
@@ -53,12 +53,29 @@ To maximize retention throughput:
- **Fast**: [Groq](https://groq.com) with `gpt-oss-20b` or other openai-oss models, self-hosted models on GPU clusters (vLLM, TGI)
- **Slow**: Standard cloud LLM providers with rate limits
2. **Batch your operations**: Group related content into batch requests. The only limit is the HTTP payload size — Hindsight automatically splits large batches into smaller, optimized chunks under the hood, so you don't have to worry about it.
2. **Batch your operations**: Group related content into batch requests. Send as much data as you want in a single request — the only limit is the HTTP payload size.
3. **Use async mode for large datasets**: Queue operations in the background
4. **Parallel processing**: For very large datasets, use multiple concurrent retention requests with different `document_id` values
### Automatic Batch Optimization
**When using async retain, Hindsight automatically handles batch sizing for you.** You don't need to manually tune batch sizes or worry about optimal chunking.
How it works:
- **Send large batches**: Submit hundreds or thousands of items in a single async retain request
- **Automatic splitting**: Hindsight automatically splits large batches (>10,000 tokens) into optimized sub-batches
- **Parallel processing**: Sub-batches are processed concurrently in the background
- **Status tracking**: Parent operation aggregates status from all sub-batches
- **Token-based**: Batching uses tiktoken for accurate token counting, not character counts
Benefits:
- Send entire documents or datasets in one API call
- Let Hindsight optimize the processing strategy
- Track overall progress via the parent operation status
- No need to manually split data into small batches
### Throughput
Factors affecting throughput:
+10 -6
View File
@@ -22,12 +22,12 @@ await client.retain('my-bank', 'Alice presented the Q4 roadmap...', {
document_id: 'meeting-2024-03-15'
});
// Batch retain
// Batch retain for a document with different sections
await client.retainBatch('my-bank', [
{ content: 'Item 1: Product launch delayed to Q2' },
{ content: 'Item 2: New hiring targets announced' },
{ content: 'Item 3: Budget approved for ML team' }
], { documentId: 'meeting-2024-03-15' });
{ content: 'Item 1: Product launch delayed to Q2', document_id: 'meeting-2024-03-15-section-1' },
{ content: 'Item 2: New hiring targets announced', document_id: 'meeting-2024-03-15-section-2' },
{ content: 'Item 3: Budget approved for ML team', document_id: 'meeting-2024-03-15-section-3' }
]);
// [/docs:document-retain]
@@ -48,11 +48,15 @@ await client.retain('my-bank', 'Project deadline: April 15 (extended)', {
const apiClient = createClient(createConfig({ baseUrl: 'http://localhost:8888' }));
// Get document to expand context from recall results
const { data: doc } = await sdk.getDocument({
const { data: doc, error } = await sdk.getDocument({
client: apiClient,
path: { bank_id: 'my-bank', document_id: 'meeting-2024-03-15' }
});
if (error) {
throw new Error(`Failed to get document: ${JSON.stringify(error)}`);
}
console.log(`Document: ${doc.id}`);
console.log(`Original text: ${doc.original_text}`);
console.log(`Memory count: ${doc.memory_unit_count}`);
+5 -6
View File
@@ -27,15 +27,14 @@ client.retain(
document_id="meeting-2024-03-15"
)
# Batch retain for a document
# Batch retain for a document with different sections
client.retain_batch(
bank_id="my-bank",
items=[
{"content": "Item 1: Product launch delayed to Q2"},
{"content": "Item 2: New hiring targets announced"},
{"content": "Item 3: Budget approved for ML team"}
],
document_id="meeting-2024-03-15"
{"content": "Item 1: Product launch delayed to Q2", "document_id": "meeting-2024-03-15-section-1"},
{"content": "Item 2: New hiring targets announced", "document_id": "meeting-2024-03-15-section-2"},
{"content": "Item 3: Budget approved for ML team", "document_id": "meeting-2024-03-15-section-3"}
]
)
# [/docs:document-retain]
+6 -7
View File
@@ -31,20 +31,19 @@ await client.retain('my-bank', 'Alice got promoted to senior engineer', {
// [docs:retain-batch]
await client.retainBatch('my-bank', [
{ content: 'Alice works at Google', context: 'career' },
{ content: 'Bob is a data scientist at Meta', context: 'career' },
{ content: 'Alice and Bob are friends', context: 'relationship' }
], { documentId: 'conversation_001' });
{ content: 'Alice works at Google', context: 'career', document_id: 'conversation_001_msg_1' },
{ content: 'Bob is a data scientist at Meta', context: 'career', document_id: 'conversation_001_msg_2' },
{ content: 'Alice and Bob are friends', context: 'relationship', document_id: 'conversation_001_msg_3' }
]);
// [/docs:retain-batch]
// [docs:retain-async]
// Start async ingestion (returns immediately)
await client.retainBatch('my-bank', [
{ content: 'Large batch item 1' },
{ content: 'Large batch item 2' },
{ content: 'Large batch item 1', document_id: 'large-doc-1' },
{ content: 'Large batch item 2', document_id: 'large-doc-2' },
], {
documentId: 'large-doc',
async: true
});
// [/docs:retain-async]
+11 -12
View File
@@ -41,11 +41,10 @@ client.retain(
client.retain_batch(
bank_id="my-bank",
items=[
{"content": "Alice works at Google", "context": "career"},
{"content": "Bob is a data scientist at Meta", "context": "career"},
{"content": "Alice and Bob are friends", "context": "relationship"}
],
document_id="conversation_001"
{"content": "Alice works at Google", "context": "career", "document_id": "conversation_001_msg_1"},
{"content": "Bob is a data scientist at Meta", "context": "career", "document_id": "conversation_001_msg_2"},
{"content": "Alice and Bob are friends", "context": "relationship", "document_id": "conversation_001_msg_3"}
]
)
# [/docs:retain-batch]
@@ -55,10 +54,9 @@ client.retain_batch(
result = client.retain_batch(
bank_id="my-bank",
items=[
{"content": "Large batch item 1"},
{"content": "Large batch item 2"},
{"content": "Large batch item 1", "document_id": "large-doc-1"},
{"content": "Large batch item 2", "document_id": "large-doc-2"},
],
document_id="large-doc",
retain_async=True
)
@@ -74,14 +72,15 @@ client.retain_batch(
items=[
{
"content": "User Alice said she loves the new dashboard",
"tags": ["user:alice", "feedback"]
"tags": ["user:alice", "feedback"],
"document_id": "user_feedback_001"
},
{
"content": "User Bob reported a bug in the search feature",
"tags": ["user:bob", "bug-report"]
"tags": ["user:bob", "bug-report"],
"document_id": "user_feedback_002"
}
],
document_id="user_feedback_001"
]
)
# [/docs:retain-with-tags]
+80
View File
@@ -3605,6 +3605,58 @@
"success": true
}
},
"ChildOperationStatus": {
"properties": {
"operation_id": {
"type": "string",
"title": "Operation Id"
},
"status": {
"type": "string",
"title": "Status"
},
"sub_batch_index": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Sub Batch Index"
},
"items_count": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"title": "Items Count"
},
"error_message": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Error Message"
}
},
"type": "object",
"required": [
"operation_id",
"status"
],
"title": "ChildOperationStatus",
"description": "Status of a child operation (for batch operations)."
},
"ChunkData": {
"properties": {
"id": {
@@ -5150,6 +5202,34 @@
}
],
"title": "Error Message"
},
"result_metadata": {
"anyOf": [
{
"additionalProperties": true,
"type": "object"
},
{
"type": "null"
}
],
"title": "Result Metadata",
"description": "Internal metadata for debugging. Structure may change without notice. Not for production use."
},
"child_operations": {
"anyOf": [
{
"items": {
"$ref": "#/components/schemas/ChildOperationStatus"
},
"type": "array"
},
{
"type": "null"
}
],
"title": "Child Operations",
"description": "Child operations for batch operations (if applicable)"
}
},
"type": "object",
+134
View File
@@ -0,0 +1,134 @@
#!/bin/bash
# Run retain performance benchmark
set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
cd "$PROJECT_ROOT"
# Default values
DOCUMENT="${DOCUMENT:-}"
BANK_ID="${BANK_ID:-perf-test}"
API_URL="${API_URL:-http://localhost:8000}"
TIMEOUT="${TIMEOUT:-300}"
OUTPUT="${OUTPUT:-}"
# Help message
show_help() {
cat << EOF
Run retain performance benchmark
Usage: $0 --document <path> [options]
Required:
--document <path> Path to document file to retain
Options:
--bank-id <id> Bank ID to use (default: perf-test)
--context <text> Optional context for the retain operation
--api-url <url> API base URL (default: http://localhost:8000)
--timeout <seconds> Request timeout (default: 300)
--output <path> Path to save results JSON (optional)
--in-memory Use in-memory MemoryEngine instead of HTTP
-h, --help Show this help message
Environment Variables:
DOCUMENT Document path (can be used instead of --document)
BANK_ID Bank ID (default: perf-test)
API_URL API URL (default: http://localhost:8000)
TIMEOUT Timeout in seconds (default: 300)
OUTPUT Output path for results JSON
Examples:
# Basic usage
$0 --document ./test_data/large_doc.txt
# With custom bank ID and save results
$0 --document ./test_data/large_doc.txt \\
--bank-id my-test-bank \\
--output results/retain_perf.json
# Using environment variables
DOCUMENT=./test_data/large_doc.txt \\
BANK_ID=my-test-bank \\
$0
EOF
}
# Parse arguments
CONTEXT=""
IN_MEMORY=""
while [[ $# -gt 0 ]]; do
case $1 in
--document)
DOCUMENT="$2"
shift 2
;;
--bank-id)
BANK_ID="$2"
shift 2
;;
--context)
CONTEXT="$2"
shift 2
;;
--api-url)
API_URL="$2"
shift 2
;;
--timeout)
TIMEOUT="$2"
shift 2
;;
--output)
OUTPUT="$2"
shift 2
;;
--in-memory)
IN_MEMORY="--in-memory"
shift 1
;;
-h|--help)
show_help
exit 0
;;
*)
echo "Unknown option: $1"
echo "Use --help for usage information"
exit 1
;;
esac
done
# Validate required arguments
if [ -z "$DOCUMENT" ]; then
echo "Error: --document is required"
echo "Use --help for usage information"
exit 1
fi
# Build command
CMD="uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document \"$DOCUMENT\" --bank-id \"$BANK_ID\" --api-url \"$API_URL\" --timeout $TIMEOUT"
if [ -n "$CONTEXT" ]; then
CMD="$CMD --context \"$CONTEXT\""
fi
if [ -n "$OUTPUT" ]; then
CMD="$CMD --output \"$OUTPUT\""
fi
if [ -n "$IN_MEMORY" ]; then
CMD="$CMD --in-memory"
fi
# Run benchmark
echo "Running retain performance benchmark..."
echo "Document: $DOCUMENT"
echo "Bank ID: $BANK_ID"
echo "API URL: $API_URL"
echo ""
eval $CMD
+111
View File
@@ -0,0 +1,111 @@
#!/bin/bash
set +e # Don't exit on errors - we want to collect all failures
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
EXAMPLES_DIR="$PROJECT_ROOT/hindsight-docs/examples/api"
LOG_DIR="/tmp/doc-example-logs"
mkdir -p "$LOG_DIR"
TOTAL_PASSED=0
TOTAL_FAILED=0
FAILED_EXAMPLES=()
echo "======================================"
echo "Running Documentation Examples"
echo "======================================"
echo ""
# Function to run a single example
run_example() {
local file="$1"
local runner="$2"
local workdir="${3:-$PROJECT_ROOT}"
local basename=$(basename "$file")
local logfile="$LOG_DIR/$basename.log"
echo -n "Running $basename... "
pushd "$workdir" > /dev/null 2>&1
if $runner "$file" > "$logfile" 2>&1; then
echo -e "${GREEN}✓ PASS${NC}"
TOTAL_PASSED=$((TOTAL_PASSED + 1))
rm -f "$logfile" # Clean up successful test logs
popd > /dev/null 2>&1
return 0
else
echo -e "${RED}✗ FAIL${NC}"
TOTAL_FAILED=$((TOTAL_FAILED + 1))
FAILED_EXAMPLES+=("$basename:$logfile")
popd > /dev/null 2>&1
return 1
fi
}
# Run Python examples
echo "======================================"
echo "Python Examples"
echo "======================================"
cd "$PROJECT_ROOT/hindsight-clients/python"
for f in "$EXAMPLES_DIR"/*.py; do
[ -e "$f" ] || continue # Skip if no files match
run_example "$f" "uv run python" "$PROJECT_ROOT/hindsight-clients/python"
done
echo ""
# Run Node.js examples
echo "======================================"
echo "Node.js Examples"
echo "======================================"
cd "$PROJECT_ROOT"
for f in "$EXAMPLES_DIR"/*.mjs; do
[ -e "$f" ] || continue # Skip if no files match
run_example "$f" "node" "$PROJECT_ROOT"
done
echo ""
# Run CLI examples
echo "======================================"
echo "CLI Examples"
echo "======================================"
cd "$PROJECT_ROOT"
for f in "$EXAMPLES_DIR"/*.sh; do
[ -e "$f" ] || continue # Skip if no files match
run_example "$f" "bash" "$PROJECT_ROOT"
done
echo ""
# Print summary
echo "======================================"
echo "Summary"
echo "======================================"
echo -e "${GREEN}Passed: $TOTAL_PASSED${NC}"
echo -e "${RED}Failed: $TOTAL_FAILED${NC}"
echo ""
# If there are failures, show the logs
if [ $TOTAL_FAILED -gt 0 ]; then
echo "======================================"
echo "Failed Example Logs"
echo "======================================"
for entry in "${FAILED_EXAMPLES[@]}"; do
IFS=':' read -r name logfile <<< "$entry"
echo ""
echo -e "${YELLOW}=== $name ===${NC}"
cat "$logfile"
done
echo ""
echo -e "${RED}$TOTAL_FAILED example(s) failed${NC}"
exit 1
fi
echo -e "${GREEN}All examples passed!${NC}"
exit 0
Generated
+2
View File
@@ -1579,6 +1579,7 @@ version = "0.4.11"
source = { editable = "hindsight-dev" }
dependencies = [
{ name = "hindsight-api" },
{ name = "httpx" },
{ name = "openai" },
{ name = "pydantic" },
{ name = "python-fasthtml" },
@@ -1602,6 +1603,7 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "hindsight-api", editable = "hindsight-api" },
{ name = "httpx", specifier = ">=0.27.0" },
{ name = "httpx", marker = "extra == 'test'", specifier = ">=0.27.0" },
{ name = "openai", specifier = ">=1.0.0" },
{ name = "pydantic", specifier = ">=2.0.0" },