Compare commits
20
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2b02fd13fb | ||
|
|
7df7afabfe | ||
|
|
62a8476678 | ||
|
|
7fc4404568 | ||
|
|
5676dda35c | ||
|
|
cc820fcd89 | ||
|
|
6458460040 | ||
|
|
be4b5b830b | ||
|
|
ded59240c9 | ||
|
|
431bd7f1f9 | ||
|
|
ac097d2f79 | ||
|
|
0ba16e21b7 | ||
|
|
a8d1d319d5 | ||
|
|
8aa61725c3 | ||
|
|
5c12950025 | ||
|
|
26731b383a | ||
|
|
2efb387da4 | ||
|
|
c5f1400ec5 | ||
|
|
ef8b603dcc | ||
|
|
f517bc172a |
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -57,8 +57,15 @@ cd hindsight-control-plane && npm run dev
|
||||
|
||||
### Benchmarks
|
||||
```bash
|
||||
# Accuracy benchmarks
|
||||
./scripts/benchmarks/run-longmemeval.sh
|
||||
./scripts/benchmarks/run-locomo.sh
|
||||
|
||||
# Performance benchmarks
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
|
||||
|
||||
# Results viewer
|
||||
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
|
||||
```
|
||||
|
||||
|
||||
+49
@@ -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")
|
||||
@@ -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):
|
||||
|
||||
@@ -250,6 +250,7 @@ ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
|
||||
ENV_RETAIN_EXTRACT_CAUSAL_LINKS = "HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS"
|
||||
ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
|
||||
ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
@@ -371,6 +372,7 @@ DEFAULT_RETAIN_EXTRACT_CAUSAL_LINKS = True # Extract causal links between facts
|
||||
DEFAULT_RETAIN_EXTRACTION_MODE = "concise" # Extraction mode: "concise", "verbose", or "custom"
|
||||
RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction modes
|
||||
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
@@ -590,6 +592,7 @@ class HindsightConfig:
|
||||
retain_extract_causal_links: bool
|
||||
retain_extraction_mode: str
|
||||
retain_custom_instructions: str | None
|
||||
retain_batch_tokens: int
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
@@ -939,6 +942,7 @@ class HindsightConfig:
|
||||
os.getenv(ENV_RETAIN_EXTRACTION_MODE, DEFAULT_RETAIN_EXTRACTION_MODE)
|
||||
),
|
||||
retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
|
||||
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
consolidation_batch_size=int(
|
||||
|
||||
@@ -18,11 +18,20 @@ import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import tiktoken
|
||||
|
||||
from ..config import get_config
|
||||
from ..metrics import get_metrics_collector
|
||||
from ..tracing import create_operation_span
|
||||
from ..utils import mask_network_location
|
||||
from .db_budget import budgeted_operation
|
||||
from .operation_metadata import (
|
||||
BatchRetainChildMetadata,
|
||||
BatchRetainParentMetadata,
|
||||
ConsolidationMetadata,
|
||||
RefreshMentalModelMetadata,
|
||||
RetainMetadata,
|
||||
)
|
||||
|
||||
# Context variable for current schema (async-safe, per-task isolation)
|
||||
# Note: default is None, actual default comes from config via get_current_schema()
|
||||
@@ -38,6 +47,15 @@ def get_current_schema() -> str:
|
||||
return schema
|
||||
|
||||
|
||||
# Initialize tiktoken encoder once at module level for efficiency
|
||||
_tiktoken_encoder = tiktoken.get_encoding("cl100k_base") # GPT-4/GPT-3.5-turbo encoding
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Count tokens in text using tiktoken (cl100k_base encoding for GPT-4/3.5)."""
|
||||
return len(_tiktoken_encoder.encode(text))
|
||||
|
||||
|
||||
def fq_table(table_name: str) -> str:
|
||||
"""
|
||||
Get fully-qualified table name with current schema.
|
||||
@@ -826,7 +844,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
logger.error(f"Failed to delete async operation record {operation_id}: {e}")
|
||||
|
||||
async def _mark_operation_failed(self, operation_id: str, error_message: str, error_traceback: str):
|
||||
"""Helper to mark an operation as failed in the database."""
|
||||
"""Helper to mark an operation as failed in the database.
|
||||
|
||||
Also checks if this is a child operation and updates the parent if all siblings are done.
|
||||
Uses a single transaction to avoid race conditions when multiple children fail simultaneously.
|
||||
"""
|
||||
try:
|
||||
pool = await self._get_pool()
|
||||
# Truncate error message to avoid extremely long strings
|
||||
@@ -834,35 +856,159 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
truncated_error = full_error[:5000] if len(full_error) > 5000 else full_error
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'failed', error_message = $2, updated_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
truncated_error,
|
||||
)
|
||||
logger.info(f"Marked async operation as failed: {operation_id}")
|
||||
async with conn.transaction():
|
||||
# Mark this operation as failed
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'failed', error_message = $2, updated_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
truncated_error,
|
||||
)
|
||||
logger.info(f"Marked async operation as failed: {operation_id}")
|
||||
|
||||
# Check if this is a child operation and update parent if all siblings are done
|
||||
# This happens in the same transaction after the child status is updated
|
||||
await self._maybe_update_parent_operation(operation_id, conn)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to mark operation as failed {operation_id}: {e}")
|
||||
|
||||
async def _mark_operation_completed(self, operation_id: str):
|
||||
"""Helper to mark an operation as completed in the database."""
|
||||
"""Helper to mark an operation as completed in the database.
|
||||
|
||||
Also checks if this is a child operation and updates the parent if all siblings are done.
|
||||
Uses a single transaction to avoid race conditions when multiple children complete simultaneously.
|
||||
"""
|
||||
try:
|
||||
pool = await self._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Mark this operation as completed
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
logger.info(f"Marked async operation as completed: {operation_id}")
|
||||
|
||||
# Check if this is a child operation and update parent if all siblings are done
|
||||
# This happens in the same transaction after the child status is updated
|
||||
await self._maybe_update_parent_operation(operation_id, conn)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to mark operation as completed {operation_id}: {e}")
|
||||
|
||||
async def _maybe_update_parent_operation(self, child_operation_id: str, conn):
|
||||
"""Check if this is a child operation and update parent status if all siblings are done.
|
||||
|
||||
Must be called within an active transaction that has already updated the child's status.
|
||||
Uses SELECT FOR UPDATE to lock the parent and prevent race conditions.
|
||||
|
||||
Args:
|
||||
child_operation_id: The operation ID that just completed or failed
|
||||
conn: Database connection with an active transaction
|
||||
"""
|
||||
try:
|
||||
# Get this operation's metadata to check if it has a parent
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT result_metadata, bank_id
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(child_operation_id),
|
||||
)
|
||||
|
||||
if not row:
|
||||
return
|
||||
|
||||
result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
|
||||
parent_operation_id = result_metadata.get("parent_operation_id")
|
||||
|
||||
if not parent_operation_id:
|
||||
# Not a child operation
|
||||
return
|
||||
|
||||
bank_id = row["bank_id"]
|
||||
|
||||
# Lock the parent operation to prevent concurrent updates from other children
|
||||
# Use FOR UPDATE to ensure only one child can update the parent at a time
|
||||
parent_row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT operation_id
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE operation_id = $1 AND bank_id = $2
|
||||
FOR UPDATE
|
||||
""",
|
||||
uuid.UUID(parent_operation_id),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not parent_row:
|
||||
# Parent doesn't exist (shouldn't happen)
|
||||
return
|
||||
|
||||
# Get all sibling operations (including this one)
|
||||
# This query runs in the same transaction, so it sees the current child's updated status
|
||||
siblings = await conn.fetch(
|
||||
f"""
|
||||
SELECT status
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE bank_id = $1
|
||||
AND result_metadata::jsonb @> $2::jsonb
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps({"parent_operation_id": parent_operation_id}),
|
||||
)
|
||||
|
||||
if not siblings:
|
||||
return
|
||||
|
||||
# Check if all siblings are done (completed or failed)
|
||||
all_completed = all(sib["status"] == "completed" for sib in siblings)
|
||||
any_failed = any(sib["status"] == "failed" for sib in siblings)
|
||||
all_done = all(sib["status"] in ("completed", "failed") for sib in siblings)
|
||||
|
||||
if not all_done:
|
||||
# Some siblings still pending/processing
|
||||
return
|
||||
|
||||
# All siblings are done - update parent status
|
||||
if any_failed:
|
||||
new_status = "failed"
|
||||
# Set parent error message to indicate child failure
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
|
||||
SET status = $2, error_message = $3, updated_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
uuid.UUID(parent_operation_id),
|
||||
new_status,
|
||||
"One or more sub-batches failed",
|
||||
)
|
||||
logger.info(f"Marked async operation as completed: {operation_id}")
|
||||
elif all_completed:
|
||||
new_status = "completed"
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = $2, updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(parent_operation_id),
|
||||
new_status,
|
||||
)
|
||||
|
||||
logger.info(f"Updated parent operation {parent_operation_id} to status '{new_status}' (all children done)")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to mark operation as completed {operation_id}: {e}")
|
||||
logger.error(f"Failed to update parent operation for child {child_operation_id}: {e}")
|
||||
# Re-raise to rollback the transaction
|
||||
raise
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize the connection pool, models, and background workers.
|
||||
@@ -1430,35 +1576,49 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if "document_id" not in item:
|
||||
item["document_id"] = document_id
|
||||
|
||||
# Auto-chunk large batches by character count to avoid timeouts and memory issues
|
||||
# Calculate total character count
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents)
|
||||
# Validate no duplicate document_ids in the batch
|
||||
# Having duplicate document_ids causes race conditions in document upserts during parallel processing
|
||||
doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
|
||||
if len(doc_ids) != len(set(doc_ids)):
|
||||
from collections import Counter
|
||||
|
||||
duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
|
||||
raise ValueError(
|
||||
f"Batch contains duplicate document_ids: {duplicates}. "
|
||||
f"Each content item in a batch must have a unique document_id to avoid race conditions."
|
||||
)
|
||||
|
||||
# Auto-chunk large batches by token count to avoid timeouts and memory issues
|
||||
# Calculate total token count
|
||||
total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
|
||||
total_usage = TokenUsage()
|
||||
|
||||
CHARS_PER_BATCH = 600_000
|
||||
# Get batch size threshold from config
|
||||
config = get_config()
|
||||
tokens_per_batch = config.retain_batch_tokens
|
||||
|
||||
if total_chars > CHARS_PER_BATCH:
|
||||
# Split into smaller batches based on character count
|
||||
if total_tokens > tokens_per_batch:
|
||||
# Split into smaller batches based on token count
|
||||
logger.info(
|
||||
f"Large batch detected ({total_chars:,} chars from {len(contents)} items). Splitting into sub-batches of ~{CHARS_PER_BATCH:,} chars each..."
|
||||
f"Large batch detected ({total_tokens:,} tokens from {len(contents)} items). Splitting into sub-batches of ~{tokens_per_batch:,} tokens each..."
|
||||
)
|
||||
|
||||
sub_batches = []
|
||||
current_batch = []
|
||||
current_batch_chars = 0
|
||||
current_batch_tokens = 0
|
||||
|
||||
for item in contents:
|
||||
item_chars = len(item.get("content", ""))
|
||||
item_tokens = count_tokens(item.get("content", ""))
|
||||
|
||||
# If adding this item would exceed the limit, start a new batch
|
||||
# (unless current batch is empty - then we must include it even if it's large)
|
||||
if current_batch and current_batch_chars + item_chars > CHARS_PER_BATCH:
|
||||
if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
|
||||
sub_batches.append(current_batch)
|
||||
current_batch = [item]
|
||||
current_batch_chars = item_chars
|
||||
current_batch_tokens = item_tokens
|
||||
else:
|
||||
current_batch.append(item)
|
||||
current_batch_chars += item_chars
|
||||
current_batch_tokens += item_tokens
|
||||
|
||||
# Add the last batch
|
||||
if current_batch:
|
||||
@@ -1469,9 +1629,9 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# Process each sub-batch
|
||||
all_results = []
|
||||
for i, sub_batch in enumerate(sub_batches, 1):
|
||||
sub_batch_chars = sum(len(item.get("content", "")) for item in sub_batch)
|
||||
sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
|
||||
logger.info(
|
||||
f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars"
|
||||
f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
|
||||
)
|
||||
|
||||
sub_results, sub_usage = await self._retain_batch_async_internal(
|
||||
@@ -5463,10 +5623,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
total = total_row["total"] if total_row else 0
|
||||
|
||||
# Get operations with pagination
|
||||
# Get operations with pagination (include result_metadata to check for parent operations)
|
||||
operations = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, created_at, status, error_message
|
||||
SELECT operation_id, operation_type, created_at, status, error_message, result_metadata
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY created_at DESC
|
||||
@@ -5477,21 +5637,29 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
offset,
|
||||
)
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"operations": [
|
||||
# Build operation list using status from database
|
||||
# Parent operations have their status updated when all children complete/fail
|
||||
operation_list = []
|
||||
for row in operations:
|
||||
# Map DB status to API status (pending includes processing)
|
||||
db_status = row["status"]
|
||||
api_status = "pending" if db_status in ("pending", "processing") else db_status
|
||||
|
||||
operation_list.append(
|
||||
{
|
||||
"id": str(row["operation_id"]),
|
||||
"task_type": row["operation_type"],
|
||||
"items_count": 0,
|
||||
"document_id": None,
|
||||
"created_at": row["created_at"].isoformat(),
|
||||
# Map DB status to API status (processing -> pending for simplicity)
|
||||
"status": "pending" if row["status"] in ("pending", "processing") else row["status"],
|
||||
"status": api_status,
|
||||
"error_message": row["error_message"],
|
||||
}
|
||||
for row in operations
|
||||
],
|
||||
)
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"operations": operation_list,
|
||||
}
|
||||
|
||||
async def get_operation_status(
|
||||
@@ -5503,10 +5671,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
) -> dict[str, Any]:
|
||||
"""Get the status of a specific async operation.
|
||||
|
||||
For parent operations, the status is automatically updated in the database when all children complete/fail.
|
||||
|
||||
Returns:
|
||||
- status: "pending", "completed", or "failed"
|
||||
- status: "pending", "completed", or "failed" (from database)
|
||||
- updated_at: last update timestamp
|
||||
- completed_at: completion timestamp (if completed)
|
||||
- child_operations: (for parent operations) list of child operation statuses
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
pool = await self._get_pool()
|
||||
@@ -5516,7 +5687,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message
|
||||
SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message, result_metadata
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE operation_id = $1 AND bank_id = $2
|
||||
""",
|
||||
@@ -5525,18 +5696,98 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
if row:
|
||||
# Map DB status to API status (processing -> pending for simplicity)
|
||||
# Check if this is a parent operation
|
||||
result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
|
||||
is_parent = result_metadata.get("is_parent", False)
|
||||
|
||||
# Use status from database (parent status is updated when all children complete/fail)
|
||||
db_status = row["status"]
|
||||
api_status = "pending" if db_status in ("pending", "processing") else db_status
|
||||
return {
|
||||
"operation_id": operation_id,
|
||||
"status": api_status,
|
||||
"operation_type": row["operation_type"],
|
||||
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
|
||||
"error_message": row["error_message"],
|
||||
}
|
||||
|
||||
# For parent operations, include child operations list
|
||||
if is_parent:
|
||||
# Query child operations
|
||||
child_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, error_message, result_metadata
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE bank_id = $1
|
||||
AND result_metadata::jsonb @> $2::jsonb
|
||||
ORDER BY (result_metadata->>'sub_batch_index')::int
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps({"parent_operation_id": operation_id}),
|
||||
)
|
||||
|
||||
# Build child operations list and check if parent status needs updating
|
||||
child_statuses = []
|
||||
all_done = True
|
||||
any_failed = False
|
||||
all_completed = True
|
||||
|
||||
for child_row in child_rows:
|
||||
child_metadata = (
|
||||
json.loads(child_row["result_metadata"]) if child_row["result_metadata"] else {}
|
||||
)
|
||||
child_status = child_row["status"]
|
||||
|
||||
child_statuses.append(
|
||||
{
|
||||
"operation_id": str(child_row["operation_id"]),
|
||||
"status": child_status,
|
||||
"sub_batch_index": child_metadata.get("sub_batch_index"),
|
||||
"items_count": child_metadata.get("items_count"),
|
||||
"error_message": child_row["error_message"],
|
||||
}
|
||||
)
|
||||
|
||||
if child_status not in ("completed", "failed"):
|
||||
all_done = False
|
||||
if child_status == "failed":
|
||||
any_failed = True
|
||||
if child_status != "completed":
|
||||
all_completed = False
|
||||
|
||||
# Self-healing: if parent status is out of sync with children, update it
|
||||
if all_done and api_status == "pending":
|
||||
correct_status = "failed" if any_failed else "completed"
|
||||
logger.warning(
|
||||
f"Parent operation {operation_id} status out of sync (DB: pending, should be: {correct_status}). Fixing."
|
||||
)
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = $2, updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
op_uuid,
|
||||
correct_status,
|
||||
)
|
||||
api_status = correct_status
|
||||
|
||||
return {
|
||||
"operation_id": operation_id,
|
||||
"status": api_status,
|
||||
"operation_type": row["operation_type"],
|
||||
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
|
||||
"error_message": row["error_message"],
|
||||
"result_metadata": result_metadata,
|
||||
"child_operations": child_statuses,
|
||||
}
|
||||
else:
|
||||
# Regular operation (not a parent)
|
||||
return {
|
||||
"operation_id": operation_id,
|
||||
"status": api_status,
|
||||
"operation_type": row["operation_type"],
|
||||
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
|
||||
"error_message": row["error_message"],
|
||||
"result_metadata": result_metadata,
|
||||
}
|
||||
else:
|
||||
# Operation not found
|
||||
return {
|
||||
@@ -5712,31 +5963,126 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
request_context: "RequestContext",
|
||||
document_tags: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Submit a batch retain operation to run asynchronously."""
|
||||
"""Submit a batch retain operation to run asynchronously.
|
||||
|
||||
For large batches (exceeding retain_batch_chars threshold), automatically splits
|
||||
into smaller sub-batches and creates a parent operation that tracks all children.
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
|
||||
task_payload: dict[str, Any] = {"contents": contents}
|
||||
if document_tags:
|
||||
task_payload["document_tags"] = document_tags
|
||||
# Pass tenant_id and api_key_id through task payload so the worker
|
||||
# can propagate request context to downstream operations (e.g.,
|
||||
# consolidation and mental model refreshes triggered after retain).
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
# Validate no duplicate document_ids in the batch
|
||||
# Having duplicate document_ids causes race conditions in document upserts during parallel processing
|
||||
doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
|
||||
if len(doc_ids) != len(set(doc_ids)):
|
||||
from collections import Counter
|
||||
|
||||
result = await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="retain",
|
||||
task_type="batch_retain",
|
||||
task_payload=task_payload,
|
||||
result_metadata={"items_count": len(contents)},
|
||||
dedupe_by_bank=False,
|
||||
duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
|
||||
raise ValueError(
|
||||
f"Batch contains duplicate document_ids: {duplicates}. "
|
||||
f"Each content item in a batch must have a unique document_id to avoid race conditions."
|
||||
)
|
||||
|
||||
# Calculate total token count and determine if we need to split
|
||||
total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
|
||||
config = get_config()
|
||||
tokens_per_batch = config.retain_batch_tokens
|
||||
|
||||
# Split into sub-batches based on token count
|
||||
sub_batches = []
|
||||
current_batch = []
|
||||
current_batch_tokens = 0
|
||||
|
||||
for item in contents:
|
||||
item_tokens = count_tokens(item.get("content", ""))
|
||||
|
||||
# If adding this item would exceed the limit, start a new batch
|
||||
# (unless current batch is empty - then we must include it even if it's large)
|
||||
if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
|
||||
sub_batches.append(current_batch)
|
||||
current_batch = [item]
|
||||
current_batch_tokens = item_tokens
|
||||
else:
|
||||
current_batch.append(item)
|
||||
current_batch_tokens += item_tokens
|
||||
|
||||
# Add the last batch
|
||||
if current_batch:
|
||||
sub_batches.append(current_batch)
|
||||
|
||||
# Log splitting info if we actually split
|
||||
if len(sub_batches) > 1:
|
||||
logger.info(
|
||||
f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). "
|
||||
f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each"
|
||||
)
|
||||
|
||||
# Always create parent operation (even for single batch - simpler, more reliable code path)
|
||||
import uuid
|
||||
|
||||
parent_operation_id = uuid.uuid4()
|
||||
pool = await self._get_pool()
|
||||
|
||||
# Create typed metadata for parent operation
|
||||
parent_metadata = BatchRetainParentMetadata(
|
||||
items_count=len(contents),
|
||||
total_tokens=total_tokens,
|
||||
num_sub_batches=len(sub_batches),
|
||||
)
|
||||
|
||||
result["items_count"] = len(contents)
|
||||
return result
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_operation_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps(parent_metadata.to_dict()),
|
||||
"pending", # Will be updated by status aggregation
|
||||
)
|
||||
|
||||
logger.info(f"Created parent operation {parent_operation_id} for {len(sub_batches)} sub-batch(es)")
|
||||
|
||||
# Submit child operations for each sub-batch
|
||||
for i, sub_batch in enumerate(sub_batches, 1):
|
||||
if len(sub_batches) > 1:
|
||||
sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
|
||||
logger.info(
|
||||
f"Submitting sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
|
||||
)
|
||||
|
||||
task_payload: dict[str, Any] = {"contents": sub_batch}
|
||||
if document_tags:
|
||||
task_payload["document_tags"] = document_tags
|
||||
# Pass tenant_id and api_key_id through task payload
|
||||
if request_context.tenant_id:
|
||||
task_payload["_tenant_id"] = request_context.tenant_id
|
||||
if request_context.api_key_id:
|
||||
task_payload["_api_key_id"] = request_context.api_key_id
|
||||
|
||||
# Create typed metadata for child operation
|
||||
child_metadata = BatchRetainChildMetadata(
|
||||
items_count=len(sub_batch),
|
||||
parent_operation_id=str(parent_operation_id),
|
||||
sub_batch_index=i,
|
||||
total_sub_batches=len(sub_batches),
|
||||
)
|
||||
|
||||
# Create child operation with reference to parent
|
||||
await self._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="retain",
|
||||
task_type="batch_retain",
|
||||
task_payload=task_payload,
|
||||
result_metadata=child_metadata.to_dict(),
|
||||
dedupe_by_bank=False,
|
||||
)
|
||||
|
||||
return {
|
||||
"operation_id": str(parent_operation_id),
|
||||
"items_count": len(contents),
|
||||
}
|
||||
|
||||
async def submit_async_consolidation(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Typed metadata models for async operations.
|
||||
|
||||
These dataclasses define the structure of result_metadata for different operation types.
|
||||
The metadata is exposed in the API for debugging purposes and may change without notice.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainParentMetadata:
|
||||
"""Metadata for parent batch_retain operations (when split into sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
total_tokens: int
|
||||
num_sub_batches: int
|
||||
is_parent: bool = True
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainChildMetadata:
|
||||
"""Metadata for child batch_retain operations (individual sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
parent_operation_id: str
|
||||
sub_batch_index: int
|
||||
total_sub_batches: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainMetadata:
|
||||
"""Metadata for regular retain operations (non-batched, deprecated async path)."""
|
||||
|
||||
items_count: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsolidationMetadata:
|
||||
"""Metadata for consolidation operations."""
|
||||
|
||||
# Currently empty, but structure for future fields
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefreshMentalModelMetadata:
|
||||
"""Metadata for mental model refresh operations."""
|
||||
|
||||
mental_model_id: str
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
@@ -245,6 +245,7 @@ def main():
|
||||
retain_extract_causal_links=config.retain_extract_causal_links,
|
||||
retain_extraction_mode=config.retain_extraction_mode,
|
||||
retain_custom_instructions=config.retain_custom_instructions,
|
||||
retain_batch_tokens=config.retain_batch_tokens,
|
||||
enable_observations=config.enable_observations,
|
||||
consolidation_batch_size=config.consolidation_batch_size,
|
||||
consolidation_max_tokens=config.consolidation_max_tokens,
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
"""Test async batch retain with smart batching and parent-child operations."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_async(memory, request_context):
|
||||
"""Test that async retain rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_async"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc2"},
|
||||
{"content": "Third item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_sync(memory, request_context):
|
||||
"""Test that sync retain also rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_sync"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_small_async_batch_no_splitting(memory, request_context):
|
||||
"""Test that small async batches create parent with single child (simplified code path)."""
|
||||
bank_id = "test_small_async"
|
||||
contents = [{"content": "Alice works at Google", "document_id": f"doc{i}"} for i in range(5)]
|
||||
|
||||
# Calculate total chars (should be well under threshold)
|
||||
total_chars = sum(len(item["content"]) for item in contents)
|
||||
assert total_chars < 10_000, "Test batch should be small"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 5
|
||||
|
||||
operation_id = result["operation_id"]
|
||||
|
||||
# Wait for task to complete (SyncTaskBackend executes immediately)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check operation status
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with single child (simplified code path)
|
||||
assert status["status"] == "completed"
|
||||
assert status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1 # Single sub-batch
|
||||
assert len(status["child_operations"]) == 1
|
||||
assert status["child_operations"][0]["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_async_batch_auto_splits(memory, request_context):
|
||||
"""Test that large async batches automatically split into sub-batches with parent operation."""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_large_async"
|
||||
|
||||
# Create a large batch that exceeds the threshold (10k tokens default)
|
||||
# Repeating "A"s gets heavily compressed by tokenizer, use varied content
|
||||
# Use ~22k chars per item = ~5.5k tokens per item, 2 items = ~11k tokens total (exceeds 10k)
|
||||
large_content = "The quick brown fox jumps over the lazy dog. " * 500 # ~22k chars = ~5.5k tokens
|
||||
contents = [{"content": large_content + f" item {i}", "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
# Calculate total tokens (should exceed threshold)
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
assert total_tokens > 10_000, "Test batch should exceed threshold"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 2
|
||||
|
||||
parent_operation_id = result["operation_id"]
|
||||
|
||||
# Wait for tasks to complete
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Check parent operation status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with children
|
||||
assert parent_status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in parent_status
|
||||
assert "num_sub_batches" in parent_status["result_metadata"]
|
||||
assert parent_status["result_metadata"]["num_sub_batches"] >= 2 # Should split into at least 2 batches
|
||||
assert parent_status["result_metadata"]["items_count"] == 2
|
||||
|
||||
# Verify child operations
|
||||
child_ops = parent_status["child_operations"]
|
||||
assert len(child_ops) >= 2, "Should have at least 2 child operations"
|
||||
|
||||
# All children should be completed (SyncTaskBackend executes immediately)
|
||||
for child in child_ops:
|
||||
assert child["status"] == "completed"
|
||||
assert child["sub_batch_index"] is not None
|
||||
assert child["items_count"] > 0
|
||||
|
||||
# Parent status should be aggregated as "completed"
|
||||
assert parent_status["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_pending(memory, request_context):
|
||||
"""Test that parent operation shows 'pending' when children are pending."""
|
||||
bank_id = "test_parent_pending"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one pending
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "pending" since one child is still pending
|
||||
assert parent_status["status"] == "pending"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_failed(memory, request_context):
|
||||
"""Test that parent operation shows 'failed' when any child fails."""
|
||||
bank_id = "test_parent_failed"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one failed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status, error_message)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"failed",
|
||||
"Test error",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "failed" since one child failed
|
||||
assert parent_status["status"] == "failed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
# Verify child with error is included
|
||||
failed_child = [c for c in parent_status["child_operations"] if c["status"] == "failed"][0]
|
||||
assert failed_child["error_message"] == "Test error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_completed(memory, request_context):
|
||||
"""Test that parent operation shows 'completed' when all children are completed."""
|
||||
bank_id = "test_parent_completed"
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - both completed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "completed" since all children are completed
|
||||
assert parent_status["status"] == "completed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
assert all(c["status"] == "completed" for c in parent_status["child_operations"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_retain_batch_tokens_respected(memory, request_context):
|
||||
"""Test that the retain_batch_tokens config setting is respected."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_config_batch_tokens"
|
||||
config = get_config()
|
||||
|
||||
# Check that config has the retain_batch_tokens setting
|
||||
assert hasattr(config, "retain_batch_tokens")
|
||||
assert config.retain_batch_tokens > 0
|
||||
|
||||
# Create a batch that's just under the threshold
|
||||
# Use content that produces roughly half the token limit per item
|
||||
content_size = config.retain_batch_tokens * 2 # chars (rough estimate: 1 token ~= 4 chars)
|
||||
contents = [{"content": "A" * content_size, "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
# Should be equal to threshold (boundary case, no splitting since we use > not >=)
|
||||
assert total_tokens <= config.retain_batch_tokens
|
||||
|
||||
# Submit - should NOT split
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for completion
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check status - should be a parent with single child (even for small batches)
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=result["operation_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Even small batches use parent-child pattern now (simpler code path)
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Unit tests for async retain tag propagation."""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -12,9 +12,23 @@ from hindsight_api.models import RequestContext
|
||||
async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
"""submit_async_retain should include document_tags in queued task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine._authenticate_tenant = AsyncMock()
|
||||
engine._submit_async_operation = AsyncMock(return_value={"operation_id": "op-1"})
|
||||
|
||||
# Mock the pool and connection for parent operation creation
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.execute = AsyncMock()
|
||||
mock_conn.transaction = MagicMock()
|
||||
mock_conn.transaction.return_value.__aenter__ = AsyncMock()
|
||||
mock_conn.transaction.return_value.__aexit__ = AsyncMock()
|
||||
|
||||
mock_pool = AsyncMock()
|
||||
mock_pool.acquire = AsyncMock(return_value=mock_conn)
|
||||
mock_pool.release = AsyncMock()
|
||||
|
||||
engine._get_pool = AsyncMock(return_value=mock_pool)
|
||||
|
||||
request_context = RequestContext(tenant_id="tenant-a", api_key_id="key-a")
|
||||
contents = [{"content": "Async retain payload test."}]
|
||||
document_tags = ["scope:tools", "user:alice"]
|
||||
@@ -27,10 +41,18 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result == {"operation_id": "op-1", "items_count": 1}
|
||||
# Check result structure
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 1
|
||||
|
||||
# Verify authentication was called
|
||||
engine._authenticate_tenant.assert_awaited_once_with(request_context)
|
||||
|
||||
# Verify child operation was submitted
|
||||
engine._submit_async_operation.assert_awaited_once()
|
||||
|
||||
# Verify child operation payload contains document_tags
|
||||
kwargs = engine._submit_async_operation.await_args.kwargs
|
||||
assert kwargs["bank_id"] == "bank-1"
|
||||
assert kwargs["operation_type"] == "retain"
|
||||
@@ -45,6 +67,7 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
async def test_handle_batch_retain_forwards_document_tags_to_retain_batch_async():
|
||||
"""Worker handler should forward document_tags from task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine.retain_batch_async = AsyncMock(return_value={"items_count": 1})
|
||||
|
||||
task_dict = {
|
||||
|
||||
@@ -528,7 +528,7 @@ async def test_delete_bank(api_client):
|
||||
{
|
||||
"content": "Bob is the CTO and leads the engineering team.",
|
||||
"context": "team info",
|
||||
"document_id": "team-doc-1",
|
||||
"document_id": "team-doc-2",
|
||||
},
|
||||
]
|
||||
},
|
||||
|
||||
@@ -24,6 +24,7 @@ hindsight_client_api/models/bank_profile_response.py
|
||||
hindsight_client_api/models/bank_stats_response.py
|
||||
hindsight_client_api/models/budget.py
|
||||
hindsight_client_api/models/cancel_operation_response.py
|
||||
hindsight_client_api/models/child_operation_status.py
|
||||
hindsight_client_api/models/chunk_data.py
|
||||
hindsight_client_api/models/chunk_include_options.py
|
||||
hindsight_client_api/models/chunk_response.py
|
||||
|
||||
@@ -49,6 +49,7 @@ from hindsight_client_api.models.bank_profile_response import BankProfileRespons
|
||||
from hindsight_client_api.models.bank_stats_response import BankStatsResponse
|
||||
from hindsight_client_api.models.budget import Budget
|
||||
from hindsight_client_api.models.cancel_operation_response import CancelOperationResponse
|
||||
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
|
||||
from hindsight_client_api.models.chunk_data import ChunkData
|
||||
from hindsight_client_api.models.chunk_include_options import ChunkIncludeOptions
|
||||
from hindsight_client_api.models.chunk_response import ChunkResponse
|
||||
|
||||
@@ -25,6 +25,7 @@ from hindsight_client_api.models.bank_profile_response import BankProfileRespons
|
||||
from hindsight_client_api.models.bank_stats_response import BankStatsResponse
|
||||
from hindsight_client_api.models.budget import Budget
|
||||
from hindsight_client_api.models.cancel_operation_response import CancelOperationResponse
|
||||
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
|
||||
from hindsight_client_api.models.chunk_data import ChunkData
|
||||
from hindsight_client_api.models.chunk_include_options import ChunkIncludeOptions
|
||||
from hindsight_client_api.models.chunk_response import ChunkResponse
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
The version of the OpenAPI document: 0.4.11
|
||||
Generated by OpenAPI Generator (https://openapi-generator.tech)
|
||||
|
||||
Do not edit the class manually.
|
||||
""" # noqa: E501
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
import pprint
|
||||
import re # noqa: F401
|
||||
import json
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, StrictInt, StrictStr
|
||||
from typing import Any, ClassVar, Dict, List, Optional
|
||||
from typing import Optional, Set
|
||||
from typing_extensions import Self
|
||||
|
||||
class ChildOperationStatus(BaseModel):
|
||||
"""
|
||||
Status of a child operation (for batch operations).
|
||||
""" # noqa: E501
|
||||
operation_id: StrictStr
|
||||
status: StrictStr
|
||||
sub_batch_index: Optional[StrictInt] = None
|
||||
items_count: Optional[StrictInt] = None
|
||||
error_message: Optional[StrictStr] = None
|
||||
__properties: ClassVar[List[str]] = ["operation_id", "status", "sub_batch_index", "items_count", "error_message"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
validate_assignment=True,
|
||||
protected_namespaces=(),
|
||||
)
|
||||
|
||||
|
||||
def to_str(self) -> str:
|
||||
"""Returns the string representation of the model using alias"""
|
||||
return pprint.pformat(self.model_dump(by_alias=True))
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""Returns the JSON representation of the model using alias"""
|
||||
# TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
|
||||
return json.dumps(self.to_dict())
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str: str) -> Optional[Self]:
|
||||
"""Create an instance of ChildOperationStatus from a JSON string"""
|
||||
return cls.from_dict(json.loads(json_str))
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Return the dictionary representation of the model using alias.
|
||||
|
||||
This has the following differences from calling pydantic's
|
||||
`self.model_dump(by_alias=True)`:
|
||||
|
||||
* `None` is only added to the output dict for nullable fields that
|
||||
were set at model initialization. Other fields with value `None`
|
||||
are ignored.
|
||||
"""
|
||||
excluded_fields: Set[str] = set([
|
||||
])
|
||||
|
||||
_dict = self.model_dump(
|
||||
by_alias=True,
|
||||
exclude=excluded_fields,
|
||||
exclude_none=True,
|
||||
)
|
||||
# set to None if sub_batch_index (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.sub_batch_index is None and "sub_batch_index" in self.model_fields_set:
|
||||
_dict['sub_batch_index'] = None
|
||||
|
||||
# set to None if items_count (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.items_count is None and "items_count" in self.model_fields_set:
|
||||
_dict['items_count'] = None
|
||||
|
||||
# set to None if error_message (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.error_message is None and "error_message" in self.model_fields_set:
|
||||
_dict['error_message'] = None
|
||||
|
||||
return _dict
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
|
||||
"""Create an instance of ChildOperationStatus from a dict"""
|
||||
if obj is None:
|
||||
return None
|
||||
|
||||
if not isinstance(obj, dict):
|
||||
return cls.model_validate(obj)
|
||||
|
||||
_obj = cls.model_validate({
|
||||
"operation_id": obj.get("operation_id"),
|
||||
"status": obj.get("status"),
|
||||
"sub_batch_index": obj.get("sub_batch_index"),
|
||||
"items_count": obj.get("items_count"),
|
||||
"error_message": obj.get("error_message")
|
||||
})
|
||||
return _obj
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ import json
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, StrictStr, field_validator
|
||||
from typing import Any, ClassVar, Dict, List, Optional
|
||||
from hindsight_client_api.models.child_operation_status import ChildOperationStatus
|
||||
from typing import Optional, Set
|
||||
from typing_extensions import Self
|
||||
|
||||
@@ -33,7 +34,9 @@ class OperationStatusResponse(BaseModel):
|
||||
updated_at: Optional[StrictStr] = None
|
||||
completed_at: Optional[StrictStr] = None
|
||||
error_message: Optional[StrictStr] = None
|
||||
__properties: ClassVar[List[str]] = ["operation_id", "status", "operation_type", "created_at", "updated_at", "completed_at", "error_message"]
|
||||
result_metadata: Optional[Dict[str, Any]] = None
|
||||
child_operations: Optional[List[ChildOperationStatus]] = None
|
||||
__properties: ClassVar[List[str]] = ["operation_id", "status", "operation_type", "created_at", "updated_at", "completed_at", "error_message", "result_metadata", "child_operations"]
|
||||
|
||||
@field_validator('status')
|
||||
def status_validate_enum(cls, value):
|
||||
@@ -81,6 +84,13 @@ class OperationStatusResponse(BaseModel):
|
||||
exclude=excluded_fields,
|
||||
exclude_none=True,
|
||||
)
|
||||
# override the default output from pydantic by calling `to_dict()` of each item in child_operations (list)
|
||||
_items = []
|
||||
if self.child_operations:
|
||||
for _item_child_operations in self.child_operations:
|
||||
if _item_child_operations:
|
||||
_items.append(_item_child_operations.to_dict())
|
||||
_dict['child_operations'] = _items
|
||||
# set to None if operation_type (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.operation_type is None and "operation_type" in self.model_fields_set:
|
||||
@@ -106,6 +116,16 @@ class OperationStatusResponse(BaseModel):
|
||||
if self.error_message is None and "error_message" in self.model_fields_set:
|
||||
_dict['error_message'] = None
|
||||
|
||||
# set to None if result_metadata (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.result_metadata is None and "result_metadata" in self.model_fields_set:
|
||||
_dict['result_metadata'] = None
|
||||
|
||||
# set to None if child_operations (nullable) is None
|
||||
# and model_fields_set contains the field
|
||||
if self.child_operations is None and "child_operations" in self.model_fields_set:
|
||||
_dict['child_operations'] = None
|
||||
|
||||
return _dict
|
||||
|
||||
@classmethod
|
||||
@@ -124,7 +144,9 @@ class OperationStatusResponse(BaseModel):
|
||||
"created_at": obj.get("created_at"),
|
||||
"updated_at": obj.get("updated_at"),
|
||||
"completed_at": obj.get("completed_at"),
|
||||
"error_message": obj.get("error_message")
|
||||
"error_message": obj.get("error_message"),
|
||||
"result_metadata": obj.get("result_metadata"),
|
||||
"child_operations": [ChildOperationStatus.from_dict(_item) for _item in obj["child_operations"]] if obj.get("child_operations") is not None else None
|
||||
})
|
||||
return _obj
|
||||
|
||||
|
||||
@@ -278,6 +278,34 @@ export type CancelOperationResponse = {
|
||||
operation_id: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* ChildOperationStatus
|
||||
*
|
||||
* Status of a child operation (for batch operations).
|
||||
*/
|
||||
export type ChildOperationStatus = {
|
||||
/**
|
||||
* Operation Id
|
||||
*/
|
||||
operation_id: string;
|
||||
/**
|
||||
* Status
|
||||
*/
|
||||
status: string;
|
||||
/**
|
||||
* Sub Batch Index
|
||||
*/
|
||||
sub_batch_index?: number | null;
|
||||
/**
|
||||
* Items Count
|
||||
*/
|
||||
items_count?: number | null;
|
||||
/**
|
||||
* Error Message
|
||||
*/
|
||||
error_message?: string | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* ChunkData
|
||||
*
|
||||
@@ -1201,6 +1229,20 @@ export type OperationStatusResponse = {
|
||||
* Error Message
|
||||
*/
|
||||
error_message?: string | null;
|
||||
/**
|
||||
* Result Metadata
|
||||
*
|
||||
* Internal metadata for debugging. Structure may change without notice. Not for production use.
|
||||
*/
|
||||
result_metadata?: {
|
||||
[key: string]: unknown;
|
||||
} | null;
|
||||
/**
|
||||
* Child Operations
|
||||
*
|
||||
* Child operations for batch operations (if applicable)
|
||||
*/
|
||||
child_operations?: Array<ChildOperationStatus> | null;
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -12,6 +12,13 @@ import {
|
||||
TableHeader,
|
||||
TableRow,
|
||||
} from "@/components/ui/table";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { RefreshCw, Clock, AlertCircle, CheckCircle, Loader2, X } from "lucide-react";
|
||||
|
||||
interface Operation {
|
||||
@@ -24,6 +31,46 @@ interface Operation {
|
||||
error_message: string | null;
|
||||
}
|
||||
|
||||
interface ChildOperationStatus {
|
||||
operation_id: string;
|
||||
status: string;
|
||||
sub_batch_index: number | null;
|
||||
items_count: number | null;
|
||||
error_message: string | null;
|
||||
}
|
||||
|
||||
type OperationDetails =
|
||||
| {
|
||||
operation_id: string;
|
||||
status: string;
|
||||
operation_type: string | null;
|
||||
created_at: string | null;
|
||||
updated_at: string | null;
|
||||
completed_at: string | null;
|
||||
error_message: string | null;
|
||||
result_metadata?: {
|
||||
items_count?: number;
|
||||
total_tokens?: number;
|
||||
num_sub_batches?: number;
|
||||
is_parent?: boolean;
|
||||
[key: string]: any;
|
||||
};
|
||||
child_operations?: ChildOperationStatus[];
|
||||
error?: never; // Not present in success case
|
||||
}
|
||||
| {
|
||||
error: string; // Error state when loading fails
|
||||
operation_id?: never;
|
||||
status?: never;
|
||||
operation_type?: never;
|
||||
created_at?: never;
|
||||
updated_at?: never;
|
||||
completed_at?: never;
|
||||
error_message?: never;
|
||||
result_metadata?: never;
|
||||
child_operations?: never;
|
||||
};
|
||||
|
||||
export function BankOperationsView() {
|
||||
const { currentBank } = useBank();
|
||||
const [operations, setOperations] = useState<Operation[]>([]);
|
||||
@@ -33,6 +80,9 @@ export function BankOperationsView() {
|
||||
const [offset, setOffset] = useState(0);
|
||||
const [cancellingOpId, setCancellingOpId] = useState<string | null>(null);
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [selectedOperation, setSelectedOperation] = useState<OperationDetails | null>(null);
|
||||
const [dialogOpen, setDialogOpen] = useState(false);
|
||||
const [loadingDetails, setLoadingDetails] = useState(false);
|
||||
|
||||
const loadOperations = async (
|
||||
newStatusFilter: string | null = statusFilter,
|
||||
@@ -82,6 +132,22 @@ export function BankOperationsView() {
|
||||
}
|
||||
};
|
||||
|
||||
const handleOperationClick = async (operationId: string) => {
|
||||
if (!currentBank) return;
|
||||
|
||||
setLoadingDetails(true);
|
||||
setDialogOpen(true);
|
||||
try {
|
||||
const details = await client.getOperationStatus(currentBank, operationId);
|
||||
setSelectedOperation(details);
|
||||
} catch (error) {
|
||||
console.error("Error loading operation details:", error);
|
||||
setSelectedOperation({ error: "Failed to load operation details" });
|
||||
} finally {
|
||||
setLoadingDetails(false);
|
||||
}
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
if (currentBank) {
|
||||
loadOperations();
|
||||
@@ -152,7 +218,11 @@ export function BankOperationsView() {
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{operations.map((op) => (
|
||||
<TableRow key={op.id} className={op.status === "failed" ? "bg-red-500/5" : ""}>
|
||||
<TableRow
|
||||
key={op.id}
|
||||
className={`cursor-pointer hover:bg-muted/50 ${op.status === "failed" ? "bg-red-500/5" : ""}`}
|
||||
onClick={() => handleOperationClick(op.id)}
|
||||
>
|
||||
<TableCell className="font-mono text-xs text-muted-foreground">
|
||||
{op.id.substring(0, 8)}
|
||||
</TableCell>
|
||||
@@ -189,7 +259,10 @@ export function BankOperationsView() {
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
className="h-7 text-xs text-muted-foreground hover:text-red-600 dark:hover:text-red-400"
|
||||
onClick={() => handleCancelOperation(op.id)}
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
handleCancelOperation(op.id);
|
||||
}}
|
||||
disabled={cancellingOpId === op.id}
|
||||
>
|
||||
{cancellingOpId === op.id ? (
|
||||
@@ -240,6 +313,169 @@ export function BankOperationsView() {
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Operation Details Dialog */}
|
||||
<Dialog open={dialogOpen} onOpenChange={setDialogOpen}>
|
||||
<DialogContent className="max-w-3xl max-h-[80vh] overflow-y-auto">
|
||||
<DialogHeader>
|
||||
<DialogTitle>Operation Details</DialogTitle>
|
||||
<DialogDescription>
|
||||
{selectedOperation?.operation_id && (
|
||||
<span className="font-mono text-xs">{selectedOperation.operation_id}</span>
|
||||
)}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
{loadingDetails ? (
|
||||
<div className="flex items-center justify-center py-8">
|
||||
<Loader2 className="w-6 h-6 animate-spin text-muted-foreground" />
|
||||
</div>
|
||||
) : selectedOperation ? (
|
||||
<div className="space-y-4">
|
||||
{selectedOperation.error ? (
|
||||
<div className="text-red-600 dark:text-red-400">{selectedOperation.error}</div>
|
||||
) : (
|
||||
<>
|
||||
{/* Basic Info */}
|
||||
<div className="grid grid-cols-2 gap-4">
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Status</div>
|
||||
<div className="mt-1">
|
||||
{selectedOperation.status === "pending" && (
|
||||
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-amber-500/10 text-amber-600 dark:text-amber-400 border border-amber-500/20">
|
||||
<Clock className="w-3 h-3" />
|
||||
pending
|
||||
</span>
|
||||
)}
|
||||
{selectedOperation.status === "failed" && (
|
||||
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-red-500/10 text-red-600 dark:text-red-400 border border-red-500/20">
|
||||
<AlertCircle className="w-3 h-3" />
|
||||
failed
|
||||
</span>
|
||||
)}
|
||||
{selectedOperation.status === "completed" && (
|
||||
<span className="inline-flex items-center gap-1.5 px-2 py-0.5 rounded-full text-xs font-medium bg-emerald-500/10 text-emerald-600 dark:text-emerald-400 border border-emerald-500/20">
|
||||
<CheckCircle className="w-3 h-3" />
|
||||
completed
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Type</div>
|
||||
<div className="mt-1 font-mono text-sm">
|
||||
{selectedOperation.operation_type}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Created</div>
|
||||
<div className="mt-1 text-sm">
|
||||
{selectedOperation.created_at
|
||||
? new Date(selectedOperation.created_at).toLocaleString()
|
||||
: "N/A"}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Updated</div>
|
||||
<div className="mt-1 text-sm">
|
||||
{selectedOperation.updated_at
|
||||
? new Date(selectedOperation.updated_at).toLocaleString()
|
||||
: "N/A"}
|
||||
</div>
|
||||
</div>
|
||||
{selectedOperation.completed_at && (
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Completed</div>
|
||||
<div className="mt-1 text-sm">
|
||||
{new Date(selectedOperation.completed_at).toLocaleString()}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{selectedOperation.result_metadata?.items_count !== undefined && (
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground">Total Items</div>
|
||||
<div className="mt-1 text-sm">
|
||||
{selectedOperation.result_metadata.items_count}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Error Message */}
|
||||
{selectedOperation.error_message && (
|
||||
<div className="rounded-lg border border-red-500/20 bg-red-500/5 p-3">
|
||||
<div className="text-sm font-medium text-red-600 dark:text-red-400 mb-1">
|
||||
Error
|
||||
</div>
|
||||
<div className="text-sm text-red-600/80 dark:text-red-400/80 font-mono">
|
||||
{selectedOperation.error_message}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Child Operations (for parent operations) */}
|
||||
{selectedOperation.child_operations &&
|
||||
selectedOperation.child_operations.length > 0 && (
|
||||
<div>
|
||||
<div className="text-sm font-medium text-muted-foreground mb-2">
|
||||
Sub-batches (
|
||||
{selectedOperation.result_metadata?.num_sub_batches ||
|
||||
selectedOperation.child_operations.length}
|
||||
)
|
||||
</div>
|
||||
<div className="rounded-lg border">
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead className="w-[60px]">Index</TableHead>
|
||||
<TableHead className="w-[100px]">ID</TableHead>
|
||||
<TableHead className="w-[80px]">Items</TableHead>
|
||||
<TableHead>Status</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{selectedOperation.child_operations.map((child) => (
|
||||
<TableRow key={child.operation_id}>
|
||||
<TableCell className="text-sm">{child.sub_batch_index}</TableCell>
|
||||
<TableCell className="font-mono text-xs text-muted-foreground">
|
||||
{child.operation_id.substring(0, 8)}
|
||||
</TableCell>
|
||||
<TableCell className="text-sm">{child.items_count}</TableCell>
|
||||
<TableCell>
|
||||
{child.status === "pending" && (
|
||||
<span className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-amber-500/10 text-amber-600 dark:text-amber-400">
|
||||
<Clock className="w-3 h-3" />
|
||||
pending
|
||||
</span>
|
||||
)}
|
||||
{child.status === "failed" && (
|
||||
<span
|
||||
className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-red-500/10 text-red-600 dark:text-red-400"
|
||||
title={child.error_message ?? undefined}
|
||||
>
|
||||
<AlertCircle className="w-3 h-3" />
|
||||
failed
|
||||
</span>
|
||||
)}
|
||||
{child.status === "completed" && (
|
||||
<span className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full text-xs font-medium bg-emerald-500/10 text-emerald-600 dark:text-emerald-400">
|
||||
<CheckCircle className="w-3 h-3" />
|
||||
completed
|
||||
</span>
|
||||
)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
) : null}
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -68,6 +68,45 @@ Tests long-term memory across different categories.
|
||||
- `--only-failed` - Retry failed questions
|
||||
- `--fill` - Resume interrupted runs
|
||||
|
||||
### Consolidation Performance
|
||||
|
||||
Tests consolidation throughput and identifies bottlenecks.
|
||||
|
||||
```bash
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
|
||||
# With custom memory count
|
||||
NUM_MEMORIES=200 ./scripts/benchmarks/run-consolidation.sh
|
||||
```
|
||||
|
||||
### Retain Performance
|
||||
|
||||
Measures retain operation performance (throughput and token usage).
|
||||
|
||||
**Prerequisites:** API server must be running (`./scripts/dev/start-api.sh`)
|
||||
|
||||
```bash
|
||||
# Basic usage
|
||||
./scripts/benchmarks/run-retain-perf.sh \
|
||||
--document hindsight-dev/benchmarks/perf/test_data/sample_document.txt
|
||||
|
||||
# Save results to JSON
|
||||
./scripts/benchmarks/run-retain-perf.sh \
|
||||
--document ./my_document.txt \
|
||||
--bank-id my-test-bank \
|
||||
--output results/retain_perf.json
|
||||
```
|
||||
|
||||
**Options:**
|
||||
- `--document PATH` - Document file to retain (required)
|
||||
- `--bank-id ID` - Bank ID to use (default: perf-test)
|
||||
- `--context TEXT` - Optional context
|
||||
- `--api-url URL` - API URL (default: http://localhost:8000)
|
||||
- `--timeout SECONDS` - Request timeout (default: 300)
|
||||
- `--output PATH` - Save results to JSON file
|
||||
|
||||
See [perf/README.md](perf/README.md) for detailed documentation.
|
||||
|
||||
## Visualizer
|
||||
|
||||
View benchmark results in a web UI:
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Performance benchmarks for Hindsight operations."""
|
||||
@@ -0,0 +1,456 @@
|
||||
"""
|
||||
Retain operation performance benchmark.
|
||||
|
||||
Measures retain operation performance by:
|
||||
1. Loading a document from a file or directory
|
||||
2. Sending it to the retain endpoint via HTTP (batched for directories)
|
||||
3. Measuring time taken and token usage
|
||||
4. Reporting performance metrics
|
||||
|
||||
Usage:
|
||||
# Single file
|
||||
uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document <file_path> [options]
|
||||
|
||||
# Directory (batches all files)
|
||||
uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document <dir_path> [options]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
async def retain_via_memory_engine(
|
||||
bank_id: str,
|
||||
items: list[dict[str, Any]],
|
||||
) -> tuple[float, dict[str, Any]]:
|
||||
"""
|
||||
Send retain request directly to MemoryEngine (in-memory, no HTTP).
|
||||
|
||||
Args:
|
||||
bank_id: Bank ID to retain into
|
||||
items: List of items to retain
|
||||
|
||||
Returns:
|
||||
Tuple of (duration_seconds, response_data)
|
||||
"""
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Initialize memory engine
|
||||
memory = MemoryEngine(
|
||||
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"),
|
||||
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
|
||||
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
|
||||
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-20b"),
|
||||
memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
|
||||
)
|
||||
await memory.initialize()
|
||||
|
||||
# Measure time
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Call retain_batch_async directly
|
||||
result, usage = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=items,
|
||||
request_context=RequestContext(),
|
||||
return_usage=True,
|
||||
)
|
||||
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Format response to match HTTP response structure
|
||||
response_data = {
|
||||
"success": True,
|
||||
"bank_id": bank_id,
|
||||
"items_count": len(items),
|
||||
"async": False,
|
||||
"usage": usage.model_dump() if usage else None,
|
||||
}
|
||||
|
||||
return duration, response_data
|
||||
finally:
|
||||
# Close memory engine connections
|
||||
pool = await memory._get_pool()
|
||||
await pool.close()
|
||||
|
||||
|
||||
async def retain_via_http(
|
||||
base_url: str,
|
||||
bank_id: str,
|
||||
items: list[dict[str, Any]],
|
||||
timeout: float = 300.0,
|
||||
) -> tuple[float, dict[str, Any]]:
|
||||
"""
|
||||
Send retain request via HTTP and measure performance.
|
||||
|
||||
Args:
|
||||
base_url: API base URL (e.g., http://localhost:8000)
|
||||
bank_id: Bank ID to retain into
|
||||
items: List of items to retain (each with 'content' and optional 'context', 'metadata')
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
Tuple of (duration_seconds, response_data)
|
||||
"""
|
||||
url = f"{base_url}/v1/default/banks/{bank_id}/memories"
|
||||
|
||||
payload = {"items": items}
|
||||
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
# Measure time
|
||||
start_time = time.time()
|
||||
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
response = await client.post(url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
duration = time.time() - start_time
|
||||
|
||||
return duration, result
|
||||
|
||||
|
||||
def load_documents(path: str) -> tuple[list[dict[str, Any]], int]:
|
||||
"""
|
||||
Load document(s) from file or directory.
|
||||
|
||||
For directories: loads all .json, .txt, and .md files
|
||||
For JSON files with 'content' field: extracts content
|
||||
For other files: reads entire file as content
|
||||
|
||||
Returns:
|
||||
Tuple of (items_list, total_content_length)
|
||||
items_list: List of dicts with 'content' and optional 'metadata'/'context'
|
||||
total_content_length: Total character count across all documents
|
||||
"""
|
||||
file_path = Path(path)
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"Path not found: {path}")
|
||||
|
||||
items = []
|
||||
total_length = 0
|
||||
|
||||
if file_path.is_file():
|
||||
# Single file
|
||||
content, metadata = _load_single_file(file_path)
|
||||
total_length = len(content)
|
||||
item = {"content": content}
|
||||
if metadata:
|
||||
item["metadata"] = metadata
|
||||
items.append(item)
|
||||
else:
|
||||
# Directory - load all supported files
|
||||
supported_extensions = {".json", ".txt", ".md"}
|
||||
files = [f for f in file_path.rglob("*") if f.is_file() and f.suffix in supported_extensions]
|
||||
|
||||
if not files:
|
||||
raise ValueError(f"No supported files (.json, .txt, .md) found in directory: {path}")
|
||||
|
||||
console.print(f"Found {len(files)} files in directory")
|
||||
|
||||
for file in sorted(files):
|
||||
try:
|
||||
content, metadata = _load_single_file(file)
|
||||
total_length += len(content)
|
||||
item = {"content": content}
|
||||
if metadata:
|
||||
item["metadata"] = metadata
|
||||
# Add filename as context for batch processing
|
||||
item["context"] = f"Source: {file.name}"
|
||||
items.append(item)
|
||||
except Exception as e:
|
||||
console.print(f"[yellow]Warning: Failed to load {file.name}: {e}[/yellow]")
|
||||
continue
|
||||
|
||||
return items, total_length
|
||||
|
||||
|
||||
def _load_single_file(file_path: Path) -> tuple[str, dict[str, Any] | None]:
|
||||
"""
|
||||
Load a single file and extract content.
|
||||
|
||||
Returns:
|
||||
Tuple of (content, metadata)
|
||||
"""
|
||||
if file_path.suffix == ".json":
|
||||
# Try to parse as JSON and extract 'content' field
|
||||
try:
|
||||
data = json.loads(file_path.read_text())
|
||||
if isinstance(data, dict) and "content" in data:
|
||||
# Extract metadata if present
|
||||
metadata = data.get("metadata", {})
|
||||
# Add doc_id to metadata if present
|
||||
if "doc_id" in data:
|
||||
metadata["doc_id"] = data["doc_id"]
|
||||
return data["content"], metadata if metadata else None
|
||||
else:
|
||||
# Fallback: use entire JSON as string
|
||||
return file_path.read_text(), None
|
||||
except json.JSONDecodeError:
|
||||
# Not valid JSON, read as text
|
||||
return file_path.read_text(), None
|
||||
else:
|
||||
# Read as plain text
|
||||
return file_path.read_text(), None
|
||||
|
||||
|
||||
def display_results(
|
||||
duration: float,
|
||||
usage: dict[str, int] | None,
|
||||
content_length: int,
|
||||
bank_id: str,
|
||||
num_documents: int,
|
||||
) -> None:
|
||||
"""Display benchmark results in a formatted table."""
|
||||
table = Table(title="Retain Performance Benchmark Results")
|
||||
table.add_column("Metric", style="cyan")
|
||||
table.add_column("Value", style="green")
|
||||
|
||||
table.add_row("Bank ID", bank_id)
|
||||
table.add_row("Documents", f"{num_documents:,}")
|
||||
table.add_row("Total Content Length", f"{content_length:,} chars")
|
||||
if num_documents > 1:
|
||||
table.add_row("Avg Content/Doc", f"{content_length / num_documents:,.0f} chars")
|
||||
table.add_row("", "") # Separator
|
||||
table.add_row("Duration", f"{duration:.3f}s")
|
||||
table.add_row("Throughput", f"{content_length / duration:,.0f} chars/sec")
|
||||
if num_documents > 1:
|
||||
table.add_row("Docs/Second", f"{num_documents / duration:.2f}")
|
||||
|
||||
if usage:
|
||||
table.add_row("", "") # Separator
|
||||
table.add_row("Input Tokens", f"{usage.get('input_tokens', 0):,}")
|
||||
table.add_row("Output Tokens", f"{usage.get('output_tokens', 0):,}")
|
||||
table.add_row("Total Tokens", f"{usage.get('total_tokens', 0):,}")
|
||||
table.add_row("Tokens/Second", f"{usage.get('total_tokens', 0) / duration:,.1f}")
|
||||
if num_documents > 1:
|
||||
table.add_row("Avg Tokens/Doc", f"{usage.get('total_tokens', 0) / num_documents:,.0f}")
|
||||
else:
|
||||
table.add_row("", "") # Separator
|
||||
table.add_row("Token Usage", "Not available (async mode or error)")
|
||||
|
||||
console.print("\n")
|
||||
console.print(table)
|
||||
|
||||
|
||||
def save_results(
|
||||
output_path: Path,
|
||||
duration: float,
|
||||
usage: dict[str, int] | None,
|
||||
content_length: int,
|
||||
bank_id: str,
|
||||
document_path: str,
|
||||
num_documents: int,
|
||||
) -> None:
|
||||
"""Save results to JSON file."""
|
||||
results = {
|
||||
"bank_id": bank_id,
|
||||
"document_path": document_path,
|
||||
"num_documents": num_documents,
|
||||
"content_length": content_length,
|
||||
"avg_content_per_doc": content_length / num_documents if num_documents > 0 else 0,
|
||||
"duration_seconds": duration,
|
||||
"chars_per_second": content_length / duration,
|
||||
"docs_per_second": num_documents / duration if num_documents > 0 else 0,
|
||||
"usage": usage,
|
||||
}
|
||||
|
||||
if usage:
|
||||
results["tokens_per_second"] = usage.get("total_tokens", 0) / duration
|
||||
results["avg_tokens_per_doc"] = usage.get("total_tokens", 0) / num_documents if num_documents > 0 else 0
|
||||
|
||||
with open(output_path, "w") as f:
|
||||
json.dump(results, f, indent=2)
|
||||
|
||||
console.print(f"\n[green]✓[/green] Results saved to {output_path}")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run the retain performance benchmark."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark retain operation performance",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Benchmark with a single document file
|
||||
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
|
||||
--document ./test_data/large_doc.txt \\
|
||||
--bank-id perf-test-001
|
||||
|
||||
# Benchmark with a directory (batches all files)
|
||||
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
|
||||
--document ~/Documents/my-docs/ \\
|
||||
--bank-id perf-test-batch \\
|
||||
--output results/batch_perf.json
|
||||
|
||||
# With custom API URL and save results
|
||||
uv run python hindsight-dev/benchmarks/perf/retain_perf.py \\
|
||||
--document ./test_data/ \\
|
||||
--bank-id perf-test-001 \\
|
||||
--api-url http://localhost:8000 \\
|
||||
--output results/retain_perf_001.json
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--document",
|
||||
required=True,
|
||||
help="Path to document file or directory (for directories, batches all .json/.txt/.md files)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bank-id",
|
||||
default="perf-test",
|
||||
help="Bank ID to use (default: perf-test)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context",
|
||||
help="Optional context for the retain operation (only used for single file mode)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
default="http://localhost:8000",
|
||||
help="API base URL (default: http://localhost:8000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=float,
|
||||
default=300.0,
|
||||
help="Request timeout in seconds (default: 300)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
help="Path to save results JSON (optional)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--in-memory",
|
||||
action="store_true",
|
||||
help="Use in-memory MemoryEngine instead of HTTP (bypasses API server, useful for isolating performance)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
console.print("\n[bold cyan]Retain Performance Benchmark[/bold cyan]")
|
||||
console.print("=" * 80)
|
||||
|
||||
# Check mode
|
||||
if args.in_memory:
|
||||
console.print("\n[cyan]Mode: IN-MEMORY (direct MemoryEngine, no HTTP)[/cyan]")
|
||||
else:
|
||||
console.print(f"\n[cyan]Mode: HTTP (via {args.api_url})[/cyan]")
|
||||
|
||||
# Check if server is running (skip for in-memory mode)
|
||||
if not args.in_memory:
|
||||
console.print(f"\n[1] Checking API server at {args.api_url}...")
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"{args.api_url}/health", timeout=5.0)
|
||||
response.raise_for_status()
|
||||
console.print(" [green]✓[/green] API server is running")
|
||||
except Exception as e:
|
||||
console.print(f" [red]✗[/red] API server is not accessible: {e}")
|
||||
console.print("\n[yellow]Please ensure the API server is running:[/yellow]")
|
||||
console.print(" ./scripts/dev/start-api.sh")
|
||||
sys.exit(1)
|
||||
|
||||
# Load document(s)
|
||||
doc_path = Path(args.document)
|
||||
if doc_path.is_dir():
|
||||
console.print(f"\n[2] Loading documents from directory {args.document}...")
|
||||
else:
|
||||
console.print(f"\n[2] Loading document from {args.document}...")
|
||||
|
||||
try:
|
||||
items, total_content_length = load_documents(args.document)
|
||||
num_docs = len(items)
|
||||
|
||||
# Add context to single file if provided
|
||||
if num_docs == 1 and args.context:
|
||||
items[0]["context"] = args.context
|
||||
|
||||
console.print(
|
||||
f" [green]✓[/green] Loaded {num_docs:,} document{'s' if num_docs > 1 else ''} ({total_content_length:,} characters)"
|
||||
)
|
||||
if num_docs > 1:
|
||||
console.print(
|
||||
f" [cyan]Average content per document: {total_content_length / num_docs:,.0f} chars[/cyan]"
|
||||
)
|
||||
except Exception as e:
|
||||
console.print(f" [red]✗[/red] Failed to load documents: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Run benchmark
|
||||
console.print(f"\n[3] {'Processing' if args.in_memory else 'Sending retain request to'} bank '{args.bank_id}'...")
|
||||
console.print(f" [cyan]Retaining {num_docs:,} document{'s' if num_docs > 1 else ''} in batch...[/cyan]")
|
||||
try:
|
||||
if args.in_memory:
|
||||
# In-memory mode: call MemoryEngine directly
|
||||
duration, result = await retain_via_memory_engine(
|
||||
bank_id=args.bank_id,
|
||||
items=items,
|
||||
)
|
||||
else:
|
||||
# HTTP mode: call API endpoint
|
||||
duration, result = await retain_via_http(
|
||||
base_url=args.api_url,
|
||||
bank_id=args.bank_id,
|
||||
items=items,
|
||||
timeout=args.timeout,
|
||||
)
|
||||
console.print(f" [green]✓[/green] Retain completed in {duration:.3f}s")
|
||||
|
||||
# Extract usage
|
||||
usage = result.get("usage")
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
console.print(f" [red]✗[/red] HTTP error: {e.response.status_code}")
|
||||
console.print(f" Response: {e.response.text}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
console.print(f" [red]✗[/red] Request failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Display results
|
||||
console.print("\n[4] Results:")
|
||||
display_results(
|
||||
duration=duration,
|
||||
usage=usage,
|
||||
content_length=total_content_length,
|
||||
bank_id=args.bank_id,
|
||||
num_documents=num_docs,
|
||||
)
|
||||
|
||||
# Save results if requested
|
||||
if args.output:
|
||||
console.print("\n[5] Saving results...")
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
save_results(
|
||||
output_path=args.output,
|
||||
duration=duration,
|
||||
usage=usage,
|
||||
content_length=total_content_length,
|
||||
bank_id=args.bank_id,
|
||||
document_path=args.document,
|
||||
num_documents=num_docs,
|
||||
)
|
||||
|
||||
console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -14,6 +14,7 @@ dependencies = [
|
||||
"openai>=1.0.0",
|
||||
"rich>=13.0.0",
|
||||
"pydantic>=2.0.0",
|
||||
"httpx>=0.27.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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}`);
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
Executable
+134
@@ -0,0 +1,134 @@
|
||||
#!/bin/bash
|
||||
# Run retain performance benchmark
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
||||
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
# Default values
|
||||
DOCUMENT="${DOCUMENT:-}"
|
||||
BANK_ID="${BANK_ID:-perf-test}"
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
TIMEOUT="${TIMEOUT:-300}"
|
||||
OUTPUT="${OUTPUT:-}"
|
||||
|
||||
# Help message
|
||||
show_help() {
|
||||
cat << EOF
|
||||
Run retain performance benchmark
|
||||
|
||||
Usage: $0 --document <path> [options]
|
||||
|
||||
Required:
|
||||
--document <path> Path to document file to retain
|
||||
|
||||
Options:
|
||||
--bank-id <id> Bank ID to use (default: perf-test)
|
||||
--context <text> Optional context for the retain operation
|
||||
--api-url <url> API base URL (default: http://localhost:8000)
|
||||
--timeout <seconds> Request timeout (default: 300)
|
||||
--output <path> Path to save results JSON (optional)
|
||||
--in-memory Use in-memory MemoryEngine instead of HTTP
|
||||
-h, --help Show this help message
|
||||
|
||||
Environment Variables:
|
||||
DOCUMENT Document path (can be used instead of --document)
|
||||
BANK_ID Bank ID (default: perf-test)
|
||||
API_URL API URL (default: http://localhost:8000)
|
||||
TIMEOUT Timeout in seconds (default: 300)
|
||||
OUTPUT Output path for results JSON
|
||||
|
||||
Examples:
|
||||
# Basic usage
|
||||
$0 --document ./test_data/large_doc.txt
|
||||
|
||||
# With custom bank ID and save results
|
||||
$0 --document ./test_data/large_doc.txt \\
|
||||
--bank-id my-test-bank \\
|
||||
--output results/retain_perf.json
|
||||
|
||||
# Using environment variables
|
||||
DOCUMENT=./test_data/large_doc.txt \\
|
||||
BANK_ID=my-test-bank \\
|
||||
$0
|
||||
EOF
|
||||
}
|
||||
|
||||
# Parse arguments
|
||||
CONTEXT=""
|
||||
IN_MEMORY=""
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--document)
|
||||
DOCUMENT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--bank-id)
|
||||
BANK_ID="$2"
|
||||
shift 2
|
||||
;;
|
||||
--context)
|
||||
CONTEXT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--api-url)
|
||||
API_URL="$2"
|
||||
shift 2
|
||||
;;
|
||||
--timeout)
|
||||
TIMEOUT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--output)
|
||||
OUTPUT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--in-memory)
|
||||
IN_MEMORY="--in-memory"
|
||||
shift 1
|
||||
;;
|
||||
-h|--help)
|
||||
show_help
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
echo "Use --help for usage information"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Validate required arguments
|
||||
if [ -z "$DOCUMENT" ]; then
|
||||
echo "Error: --document is required"
|
||||
echo "Use --help for usage information"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build command
|
||||
CMD="uv run python hindsight-dev/benchmarks/perf/retain_perf.py --document \"$DOCUMENT\" --bank-id \"$BANK_ID\" --api-url \"$API_URL\" --timeout $TIMEOUT"
|
||||
|
||||
if [ -n "$CONTEXT" ]; then
|
||||
CMD="$CMD --context \"$CONTEXT\""
|
||||
fi
|
||||
|
||||
if [ -n "$OUTPUT" ]; then
|
||||
CMD="$CMD --output \"$OUTPUT\""
|
||||
fi
|
||||
|
||||
if [ -n "$IN_MEMORY" ]; then
|
||||
CMD="$CMD --in-memory"
|
||||
fi
|
||||
|
||||
# Run benchmark
|
||||
echo "Running retain performance benchmark..."
|
||||
echo "Document: $DOCUMENT"
|
||||
echo "Bank ID: $BANK_ID"
|
||||
echo "API URL: $API_URL"
|
||||
echo ""
|
||||
|
||||
eval $CMD
|
||||
Executable
+111
@@ -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
|
||||
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user