0560f6260d
* fix: cancel async ops on bank delete via CASCADE FK + heartbeat checkpoints - Add migration e5f6g7h8i9j0: FK ON DELETE CASCADE from async_operations and webhooks to banks, so deleting a bank auto-removes all its ops/webhooks - Add _check_op_alive() helper: returns False if op row was deleted (cascade) - Add consolidation checkpoint: after each LLM batch commit, abort early if op was deleted mid-run (returns status='cancelled') - Add retain checkpoint: between sub-batches, abort early if op was deleted - _mark_operation_completed/failed/completed_and_fire_webhook: gracefully handle missing row (UPDATE 0) with log instead of silent error - Thread operation_id into run_consolidation_job() for checkpoint access - Fix y0t1u2v3w4x5 and a1b2c3d4e5f6 migrations: add IF NOT EXISTS to prevent failure on idempotent re-runs - Add 10 tests covering cascade delete, _check_op_alive, graceful mark methods, consolidation checkpoint, and retain checkpoint * refactor: use RETURNING + fetchrow instead of execute + string comparison * fix: add bank upsert before async_operations FK inserts and update tests - memory_engine.py: upsert bank in submit_async_retain before async_operations INSERT - http.py: upsert bank in api_create_webhook before webhooks INSERT - test_worker.py, test_async_batch_retain.py, test_webhooks.py: add _ensure_bank helper calls before direct async_operations/webhooks inserts to satisfy FK constraint * fix: mock bank_utils.get_bank_profile in unit test with mocked pool
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
|