diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 87bcf7d7..cf42d6cf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -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() diff --git a/.gitignore b/.gitignore index b6d6e6bb..1c046c31 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/CLAUDE.md b/CLAUDE.md index 64a3be57..cc8340e5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 # Requires API server running + +# Results viewer ./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001 ``` diff --git a/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py b/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py new file mode 100644 index 00000000..11a143d2 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py @@ -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") diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 2d351d28..bf469221 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -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): diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index d2334278..a966232d 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -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( diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index a45e9613..cc564c3d 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -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, diff --git a/hindsight-api/hindsight_api/engine/operation_metadata.py b/hindsight-api/hindsight_api/engine/operation_metadata.py new file mode 100644 index 00000000..c1afcf5f --- /dev/null +++ b/hindsight-api/hindsight_api/engine/operation_metadata.py @@ -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) diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index b5387164..e29bf76a 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -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, diff --git a/hindsight-api/tests/test_async_batch_retain.py b/hindsight-api/tests/test_async_batch_retain.py new file mode 100644 index 00000000..d82d502d --- /dev/null +++ b/hindsight-api/tests/test_async_batch_retain.py @@ -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 diff --git a/hindsight-api/tests/test_async_retain_tags.py b/hindsight-api/tests/test_async_retain_tags.py index 24fd4380..6804cc33 100644 --- a/hindsight-api/tests/test_async_retain_tags.py +++ b/hindsight-api/tests/test_async_retain_tags.py @@ -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 = { diff --git a/hindsight-api/tests/test_http_api_integration.py b/hindsight-api/tests/test_http_api_integration.py index 2a05e98c..a2b2b543 100644 --- a/hindsight-api/tests/test_http_api_integration.py +++ b/hindsight-api/tests/test_http_api_integration.py @@ -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", }, ] }, diff --git a/hindsight-clients/python/.openapi-generator/FILES b/hindsight-clients/python/.openapi-generator/FILES index 596e70ad..01d35830 100644 --- a/hindsight-clients/python/.openapi-generator/FILES +++ b/hindsight-clients/python/.openapi-generator/FILES @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/__init__.py b/hindsight-clients/python/hindsight_client_api/__init__.py index 0c880bca..6fa1da7c 100644 --- a/hindsight-clients/python/hindsight_client_api/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/__init__.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 diff --git a/hindsight-clients/python/hindsight_client_api/models/__init__.py b/hindsight-clients/python/hindsight_client_api/models/__init__.py index be231f69..0b3b710a 100644 --- a/hindsight-clients/python/hindsight_client_api/models/__init__.py +++ b/hindsight-clients/python/hindsight_client_api/models/__init__.py @@ -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 diff --git a/hindsight-clients/python/hindsight_client_api/models/child_operation_status.py b/hindsight-clients/python/hindsight_client_api/models/child_operation_status.py new file mode 100644 index 00000000..42ea31b3 --- /dev/null +++ b/hindsight-clients/python/hindsight_client_api/models/child_operation_status.py @@ -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 + + diff --git a/hindsight-clients/python/hindsight_client_api/models/operation_status_response.py b/hindsight-clients/python/hindsight_client_api/models/operation_status_response.py index 98aa37e9..52bf4708 100644 --- a/hindsight-clients/python/hindsight_client_api/models/operation_status_response.py +++ b/hindsight-clients/python/hindsight_client_api/models/operation_status_response.py @@ -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 diff --git a/hindsight-clients/typescript/generated/types.gen.ts b/hindsight-clients/typescript/generated/types.gen.ts index 2312d9d7..5121dfe0 100644 --- a/hindsight-clients/typescript/generated/types.gen.ts +++ b/hindsight-clients/typescript/generated/types.gen.ts @@ -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 | null; }; /** diff --git a/hindsight-control-plane/src/components/bank-operations-view.tsx b/hindsight-control-plane/src/components/bank-operations-view.tsx index 7e055362..ffd20da7 100644 --- a/hindsight-control-plane/src/components/bank-operations-view.tsx +++ b/hindsight-control-plane/src/components/bank-operations-view.tsx @@ -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([]); @@ -33,6 +80,9 @@ export function BankOperationsView() { const [offset, setOffset] = useState(0); const [cancellingOpId, setCancellingOpId] = useState(null); const [loading, setLoading] = useState(false); + const [selectedOperation, setSelectedOperation] = useState(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() { {operations.map((op) => ( - + handleOperationClick(op.id)} + > {op.id.substring(0, 8)} @@ -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() {

)} + + {/* Operation Details Dialog */} + + + + Operation Details + + {selectedOperation?.operation_id && ( + {selectedOperation.operation_id} + )} + + + {loadingDetails ? ( +
+ +
+ ) : selectedOperation ? ( +
+ {selectedOperation.error ? ( +
{selectedOperation.error}
+ ) : ( + <> + {/* Basic Info */} +
+
+
Status
+
+ {selectedOperation.status === "pending" && ( + + + pending + + )} + {selectedOperation.status === "failed" && ( + + + failed + + )} + {selectedOperation.status === "completed" && ( + + + completed + + )} +
+
+
+
Type
+
+ {selectedOperation.operation_type} +
+
+
+
Created
+
+ {selectedOperation.created_at + ? new Date(selectedOperation.created_at).toLocaleString() + : "N/A"} +
+
+
+
Updated
+
+ {selectedOperation.updated_at + ? new Date(selectedOperation.updated_at).toLocaleString() + : "N/A"} +
+
+ {selectedOperation.completed_at && ( +
+
Completed
+
+ {new Date(selectedOperation.completed_at).toLocaleString()} +
+
+ )} + {selectedOperation.result_metadata?.items_count !== undefined && ( +
+
Total Items
+
+ {selectedOperation.result_metadata.items_count} +
+
+ )} +
+ + {/* Error Message */} + {selectedOperation.error_message && ( +
+
+ Error +
+
+ {selectedOperation.error_message} +
+
+ )} + + {/* Child Operations (for parent operations) */} + {selectedOperation.child_operations && + selectedOperation.child_operations.length > 0 && ( +
+
+ Sub-batches ( + {selectedOperation.result_metadata?.num_sub_batches || + selectedOperation.child_operations.length} + ) +
+
+ + + + Index + ID + Items + Status + + + + {selectedOperation.child_operations.map((child) => ( + + {child.sub_batch_index} + + {child.operation_id.substring(0, 8)} + + {child.items_count} + + {child.status === "pending" && ( + + + pending + + )} + {child.status === "failed" && ( + + + failed + + )} + {child.status === "completed" && ( + + + completed + + )} + + + ))} + +
+
+
+ )} + + )} +
+ ) : null} + + ); } diff --git a/hindsight-dev/benchmarks/README.md b/hindsight-dev/benchmarks/README.md index 0a9d4ce9..5b98cb67 100644 --- a/hindsight-dev/benchmarks/README.md +++ b/hindsight-dev/benchmarks/README.md @@ -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: diff --git a/hindsight-dev/benchmarks/perf/__init__.py b/hindsight-dev/benchmarks/perf/__init__.py new file mode 100644 index 00000000..7bdd257c --- /dev/null +++ b/hindsight-dev/benchmarks/perf/__init__.py @@ -0,0 +1 @@ +"""Performance benchmarks for Hindsight operations.""" diff --git a/hindsight-dev/benchmarks/perf/retain_perf.py b/hindsight-dev/benchmarks/perf/retain_perf.py new file mode 100644 index 00000000..ed8214cf --- /dev/null +++ b/hindsight-dev/benchmarks/perf/retain_perf.py @@ -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 [options] + + # Directory (batches all files) + uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document [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()) diff --git a/hindsight-dev/pyproject.toml b/hindsight-dev/pyproject.toml index a4b25761..9df0f48a 100644 --- a/hindsight-dev/pyproject.toml +++ b/hindsight-dev/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "openai>=1.0.0", "rich>=13.0.0", "pydantic>=2.0.0", + "httpx>=0.27.0", ] [project.optional-dependencies] diff --git a/hindsight-docs/docs/developer/performance.md b/hindsight-docs/docs/developer/performance.md index 5e52ca21..5896cc4e 100644 --- a/hindsight-docs/docs/developer/performance.md +++ b/hindsight-docs/docs/developer/performance.md @@ -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: diff --git a/hindsight-docs/examples/api/documents.mjs b/hindsight-docs/examples/api/documents.mjs index 3042f1b3..458f5613 100644 --- a/hindsight-docs/examples/api/documents.mjs +++ b/hindsight-docs/examples/api/documents.mjs @@ -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}`); diff --git a/hindsight-docs/examples/api/documents.py b/hindsight-docs/examples/api/documents.py index 1b008649..f51e0489 100644 --- a/hindsight-docs/examples/api/documents.py +++ b/hindsight-docs/examples/api/documents.py @@ -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] diff --git a/hindsight-docs/examples/api/retain.mjs b/hindsight-docs/examples/api/retain.mjs index e41322b3..3aa3e8d2 100644 --- a/hindsight-docs/examples/api/retain.mjs +++ b/hindsight-docs/examples/api/retain.mjs @@ -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] diff --git a/hindsight-docs/examples/api/retain.py b/hindsight-docs/examples/api/retain.py index b7c68c8b..bf127ed0 100644 --- a/hindsight-docs/examples/api/retain.py +++ b/hindsight-docs/examples/api/retain.py @@ -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] diff --git a/hindsight-docs/static/openapi.json b/hindsight-docs/static/openapi.json index 6111d1c3..0e9e7ce1 100644 --- a/hindsight-docs/static/openapi.json +++ b/hindsight-docs/static/openapi.json @@ -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", diff --git a/scripts/benchmarks/run-retain-perf.sh b/scripts/benchmarks/run-retain-perf.sh new file mode 100755 index 00000000..92c72ace --- /dev/null +++ b/scripts/benchmarks/run-retain-perf.sh @@ -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 [options] + +Required: + --document Path to document file to retain + +Options: + --bank-id Bank ID to use (default: perf-test) + --context Optional context for the retain operation + --api-url API base URL (default: http://localhost:8000) + --timeout Request timeout (default: 300) + --output 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 diff --git a/scripts/test-doc-examples.sh b/scripts/test-doc-examples.sh new file mode 100755 index 00000000..89044122 --- /dev/null +++ b/scripts/test-doc-examples.sh @@ -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 diff --git a/uv.lock b/uv.lock index 97a32062..7fb239dc 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },