15ea23d5d6
* feat: introduce hindsight-api-slim and hindsight-all-slim packages Closes #552 - Move all source code from hindsight-api/ to new hindsight-api-slim/ - hindsight-api-slim has heavy ML deps (torch, sentence-transformers, transformers, einops, flashrank, mlx, mlx-lm, safetensors) and pg0-embedded as optional extras: [local-ml], [embedded-db], [all] - hindsight-api becomes a zero-code meta-package depending on hindsight-api-slim[all] for full backward compatibility - Add hindsight-all-slim meta-package: hindsight-api-slim + client + embed - hindsight-all updated to depend on hindsight-api-slim[all] - pg0.py: lazy-import pg0 with clear ImportError pointing to [embedded-db] - Dockerfile: replace sed hack with proper uv sync --extra flags - Update release.yml, test.yml, lint.sh, release.sh, CLAUDE.md and all path references throughout the repo * refactor: rename hindsight/ directory to hindsight-all/ * docs: document hindsight-api-slim and hindsight-all-slim package variants Add package variants table and extras explanation to installation.md * docs: remove emojis from installation.md, use professional tone * docs: link Docker slim variant to pip package variants section * docs: consolidate Docker image variants into single table * ci: fix working-directory paths after package restructure - Replace all hindsight-api → hindsight-api-slim in test.yml - Replace hindsight → hindsight-all in test.yml - Add --extra embedded-db to test-embed API install step * ci: add local-ml and embedded-db extras to API sync steps These extras were previously implicit in the old hindsight-api package (which bundled everything). Now that hindsight-api-slim uses optional extras, we must explicitly request local-ml and embedded-db in CI. * ci: add API install step with embedded-db to test-embed smoke test The smoke test starts hindsight-api as a daemon, which requires pg0-embedded. Add a dedicated install step for hindsight-api-slim with embedded-db extra so the daemon can start successfully. * ci: remove --no-install-project when using optional extras When --no-install-project is combined with --extra, the optional deps are not installed because extras require the project to be active. Remove --no-install-project from steps that need local-ml or embedded-db. * ci: fix ordering of uv sync steps to preserve optional extras When uv sync runs for a different workspace member, it removes optional extras installed for other members. Fix by always running extra-requiring API sync last, after other workspace member syncs. Also remove --no-install-project from embedded-db sync in test-embed, as --no-install-project prevents optional extras from being active. * ci: add local-ml extra to test-embed API install for smoke test The smoke test starts the full API server which needs sentence-transformers for local embeddings (default provider). Add local-ml extra to the install. * ci: simplify extras with --all-extras and add slim pip smoke test - Replace explicit --extra local-ml --extra embedded-db with --all-extras for cleaner, more maintainable sync steps - Add test-pip-slim job: tests hindsight-api-slim[embedded-db] without local ML models, using Cohere for embeddings/reranking (mirrors Docker slim smoke test approach) * ci: simplify slim smoke test to health check only (mirrors Docker test)
435 lines
15 KiB
Python
435 lines
15 KiB
Python
"""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
|
|
|
|
|
|
async def _ensure_bank(pool, bank_id: str) -> None:
|
|
"""Upsert a minimal bank row so FK on async_operations passes."""
|
|
await pool.execute(
|
|
"INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
|
|
bank_id,
|
|
bank_id,
|
|
)
|
|
|
|
|
|
@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()
|
|
await _ensure_bank(pool, bank_id)
|
|
|
|
# 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()
|
|
await _ensure_bank(pool, bank_id)
|
|
|
|
# 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()
|
|
await _ensure_bank(pool, bank_id)
|
|
|
|
# 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
|