7990381f6a
* fix(ci): resolve all CI failures — unversioned integrations, test retries - Move integration docs to separate unversioned docs plugin (docs-integrations/) so new integrations don't need to be duplicated across versioned_docs - Remove integration pages from versioned_docs (v0.3, v0.4) — sidebar entries now use links instead of doc refs - Add missing title/description SEO frontmatter to autogen.md - Add retry logic (2 attempts) to test-doc-examples.sh for transient LLM timeouts - Add pytest-rerunfailures to test-api with --reruns 2 for flaky Gemini-dependent integration tests * ci: retrigger * fix: graph entity inheritance, SyncTaskBackend error propagation, fact_type test regressions - Fix observation entity inheritance in get_graph_data: the unit_entities query only fetched entities for visible observation IDs, not their source memory IDs, so the inheritance loop always found an empty entity_map - Remove error swallowing in SyncTaskBackend._execute_task so test failures surface instead of being silently logged - Wrap remaining consolidation submission call sites with try/except since consolidation is non-critical for those operations - Fix test_sync_backend test to expect errors to propagate - Remove fact_type=["world"] filter from test_document_upsert_behavior and test_mentioned_at_from_context_string (same PR #848 regression) - Remove flaky marker from consolidation test (now deterministic)
3334 lines
132 KiB
Python
3334 lines
132 KiB
Python
"""
|
||
Test retain function and chunk storage.
|
||
"""
|
||
import logging
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
import pytest
|
||
|
||
from hindsight_api import RequestContext
|
||
from hindsight_api.engine.memory_engine import Budget
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retain_with_chunks(memory, request_context):
|
||
"""
|
||
Test that retain function:
|
||
1. Stores facts with associated chunks
|
||
2. Recall returns chunk_id for each fact
|
||
"""
|
||
bank_id = f"test_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "test_doc_123"
|
||
|
||
try:
|
||
# Store content that will be chunked (long enough to create multiple facts)
|
||
long_content = """
|
||
Alice is a senior software engineer at TechCorp. She has been working there for 5 years.
|
||
Alice specializes in distributed systems and has led the development of the company's
|
||
microservices architecture. She is known for writing clean, well-documented code.
|
||
|
||
Bob joined the team last month as a junior developer. He is learning React and Node.js.
|
||
Bob is enthusiastic and asks great questions during code reviews. He recently completed
|
||
his first feature, which was a user authentication flow.
|
||
|
||
The team uses Kubernetes for container orchestration and deploys to AWS. They follow
|
||
agile methodologies with two-week sprints. Code reviews are mandatory before merging.
|
||
"""
|
||
|
||
# Retain with document_id to enable chunk storage
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=long_content,
|
||
context="team overview",
|
||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
print(f"\n=== Retained {len(unit_ids)} facts ===")
|
||
assert len(unit_ids) > 0, "Should have extracted and stored facts"
|
||
|
||
# Test 1: Recall with chunks enabled
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about Alice",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world"], # Search for world facts
|
||
include_chunks=True, # Enable chunks
|
||
max_chunk_tokens=8192,
|
||
request_context=request_context,
|
||
)
|
||
|
||
print("\n=== Recall Results (with chunks) ===")
|
||
print(f"Found {len(result.results)} results")
|
||
|
||
assert len(result.results) > 0, "Should find facts about Alice"
|
||
|
||
# Verify that chunks are returned
|
||
assert result.chunks is not None, "Chunks should be included in the response"
|
||
assert len(result.chunks) > 0, "Should have at least one chunk"
|
||
|
||
print(f"Number of chunks returned: {len(result.chunks)}")
|
||
|
||
# Verify chunk structure
|
||
for chunk_id, chunk_info in result.chunks.items():
|
||
print(f"\nChunk {chunk_id}:")
|
||
print(f" - chunk_index: {chunk_info.chunk_index}")
|
||
print(f" - chunk_text length: {len(chunk_info.chunk_text)} chars")
|
||
print(f" - truncated: {chunk_info.truncated}")
|
||
print(f" - text preview: {chunk_info.chunk_text[:100]}...")
|
||
|
||
# Verify chunk structure
|
||
assert isinstance(chunk_info.chunk_index, int), "Chunk index should be an integer"
|
||
assert chunk_info.chunk_index >= 0, "Chunk index should be non-negative"
|
||
assert len(chunk_info.chunk_text) > 0, "Chunk text should not be empty"
|
||
assert isinstance(chunk_info.truncated, bool), "Truncated should be boolean"
|
||
|
||
print("\n=== Test passed: Chunks are stored and retrieved correctly ===")
|
||
|
||
finally:
|
||
# Cleanup - delete the test bank
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_chunks_and_entities_follow_fact_order(memory, request_context):
|
||
"""
|
||
Test that chunks and entities in recall results follow the same order as facts.
|
||
This is critical because token limits may truncate later items.
|
||
|
||
The most relevant fact's chunk/entity should always be first in the returned data.
|
||
"""
|
||
bank_id = f"test_ordering_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store multiple distinct pieces of content as separate documents
|
||
# This ensures different chunks that we can identify
|
||
contents = [
|
||
{
|
||
"content": "Alice works at Google as a software engineer. She loves Python and has 10 years of experience.",
|
||
"document_id": "doc_alice",
|
||
"context": "Alice's profile"
|
||
},
|
||
{
|
||
"content": "Bob works at Meta as a data scientist. He specializes in machine learning and has published papers.",
|
||
"document_id": "doc_bob",
|
||
"context": "Bob's profile"
|
||
},
|
||
{
|
||
"content": "Charlie works at Amazon as a product manager. He leads a team of 15 people and ships features weekly.",
|
||
"document_id": "doc_charlie",
|
||
"context": "Charlie's profile"
|
||
},
|
||
]
|
||
|
||
# Store each content piece
|
||
for item in contents:
|
||
await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=item["content"],
|
||
context=item["context"],
|
||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||
document_id=item["document_id"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
print("\n=== Stored 3 separate documents ===")
|
||
|
||
# Recall with a query that matches all three, but Alice most closely
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about Alice's work at Google",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
fact_type=["world"],
|
||
include_chunks=True,
|
||
max_chunk_tokens=8192,
|
||
request_context=request_context,
|
||
)
|
||
|
||
print("\n=== Recall Results ===")
|
||
print(f"Found {len(result.results)} facts")
|
||
|
||
# Extract the order of entities mentioned in facts
|
||
fact_chunk_ids = []
|
||
fact_entities = []
|
||
|
||
for i, fact in enumerate(result.results):
|
||
print(f"\nFact {i}: {fact.text[:80]}...")
|
||
print(f" chunk_id: {fact.chunk_id}")
|
||
|
||
# Track chunk_id order
|
||
if fact.chunk_id:
|
||
fact_chunk_ids.append(fact.chunk_id)
|
||
|
||
# Track entities mentioned in this fact
|
||
if fact.entities:
|
||
for entity in fact.entities:
|
||
if entity not in fact_entities:
|
||
fact_entities.append(entity)
|
||
|
||
print(f"\n=== Fact chunk_ids in order: {fact_chunk_ids} ===")
|
||
print(f"=== Fact entities in order: {fact_entities} ===")
|
||
|
||
# Test 1: Verify chunks follow fact order
|
||
if result.chunks:
|
||
chunks_order = list(result.chunks.keys())
|
||
print(f"\n=== Chunks dict order: {chunks_order} ===")
|
||
|
||
# The chunks dict should contain chunks in the order they appear in facts
|
||
# (may be fewer chunks than facts due to deduplication)
|
||
chunk_positions = []
|
||
for chunk_id in chunks_order:
|
||
if chunk_id in fact_chunk_ids:
|
||
chunk_positions.append(fact_chunk_ids.index(chunk_id))
|
||
|
||
print(f"=== Chunk positions in fact order: {chunk_positions} ===")
|
||
|
||
# Verify chunks are in increasing order (following fact order)
|
||
assert chunk_positions == sorted(chunk_positions), \
|
||
f"Chunks should follow fact order! Got positions {chunk_positions} but expected {sorted(chunk_positions)}"
|
||
|
||
print("✓ Chunks follow fact order correctly")
|
||
|
||
# Test 2: Verify entities follow fact order
|
||
if result.entities:
|
||
entities_order = list(result.entities.keys())
|
||
print(f"\n=== Entities dict order: {entities_order} ===")
|
||
|
||
# The entities dict should contain entities in the order they first appear in facts
|
||
entity_positions = []
|
||
for entity_name in entities_order:
|
||
if entity_name in fact_entities:
|
||
entity_positions.append(fact_entities.index(entity_name))
|
||
|
||
print(f"=== Entity positions in fact order: {entity_positions} ===")
|
||
|
||
# Verify entities are in increasing order (following fact order)
|
||
assert entity_positions == sorted(entity_positions), \
|
||
f"Entities should follow fact order! Got positions {entity_positions} but expected {sorted(entity_positions)}"
|
||
|
||
print("✓ Entities follow fact order correctly")
|
||
|
||
print("\n=== Test passed: Chunks and entities follow fact relevance order ===")
|
||
|
||
finally:
|
||
# Cleanup
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_event_date_storage(memory, request_context):
|
||
"""
|
||
Test that event_date is correctly stored as occurred_start.
|
||
Verifies that we can track when events actually happened vs when they were stored.
|
||
"""
|
||
bank_id = f"test_temporal_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Event that occurred in the past
|
||
past_event_date = datetime(2023, 6, 15, 14, 30, tzinfo=timezone.utc)
|
||
|
||
# Store a fact about a past event
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice completed the Q2 product launch on June 15th, 2023.",
|
||
context="project history",
|
||
event_date=past_event_date,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should have created at least one memory unit"
|
||
|
||
# Recall the fact (no fact_type filter — LLM may classify as world or experience)
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="When did Alice complete the product launch?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the stored fact"
|
||
|
||
# Verify the occurred_start matches our event_date
|
||
fact = result.results[0]
|
||
assert fact.occurred_start is not None, "occurred_start should be set"
|
||
|
||
# Parse the occurred_start (it comes back as ISO string)
|
||
if isinstance(fact.occurred_start, str):
|
||
occurred_dt = datetime.fromisoformat(fact.occurred_start.replace('Z', '+00:00'))
|
||
else:
|
||
occurred_dt = fact.occurred_start
|
||
|
||
# Verify it matches our past event date (allowing for small time differences in extraction)
|
||
assert occurred_dt.year == past_event_date.year, f"Year should match: {occurred_dt.year} vs {past_event_date.year}"
|
||
assert occurred_dt.month == past_event_date.month, f"Month should match: {occurred_dt.month} vs {past_event_date.month}"
|
||
|
||
print(f"\n✓ Event date correctly stored: {occurred_dt}")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.xfail(reason="LLM date extraction from content is non-deterministic", strict=False)
|
||
async def test_temporal_ordering(memory, request_context):
|
||
"""
|
||
Test that facts can be stored and retrieved with correct temporal ordering.
|
||
Stores facts with different event_dates and verifies temporal relationships.
|
||
"""
|
||
bank_id = f"test_temporal_order_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store events in non-chronological order with different dates
|
||
events = [
|
||
{
|
||
"content": "Alice joined the team in January 2023.",
|
||
"event_date": datetime(2023, 1, 10, tzinfo=timezone.utc),
|
||
"context": "team history"
|
||
},
|
||
{
|
||
"content": "Alice got promoted to senior engineer in June 2023.",
|
||
"event_date": datetime(2023, 6, 15, tzinfo=timezone.utc),
|
||
"context": "team history"
|
||
},
|
||
{
|
||
"content": "Alice started as an intern in July 2022.",
|
||
"event_date": datetime(2022, 7, 1, tzinfo=timezone.utc),
|
||
"context": "team history"
|
||
},
|
||
]
|
||
|
||
# Store all events
|
||
for event in events:
|
||
await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=event["content"],
|
||
context=event["context"],
|
||
event_date=event["event_date"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
print("\n=== Stored 3 events with different temporal dates ===")
|
||
|
||
# Recall facts about Alice
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about Alice's career progression",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) >= 2, f"Should recall at least 2 events, got {len(result.results)}"
|
||
|
||
# Collect occurred dates
|
||
occurred_dates = []
|
||
for fact in result.results:
|
||
if fact.occurred_start:
|
||
if isinstance(fact.occurred_start, str):
|
||
dt = datetime.fromisoformat(fact.occurred_start.replace('Z', '+00:00'))
|
||
else:
|
||
dt = fact.occurred_start
|
||
occurred_dates.append((dt, fact.text[:50]))
|
||
print(f" - {dt.date()}: {fact.text[:60]}...")
|
||
|
||
# Verify we have temporal data for most facts (LLM may occasionally miss one)
|
||
assert len(occurred_dates) >= 2, "At least 2 facts should have temporal data"
|
||
|
||
# The dates should span the expected range (2022-2023)
|
||
min_date = min(dt for dt, _ in occurred_dates)
|
||
max_date = max(dt for dt, _ in occurred_dates)
|
||
|
||
assert min_date.year == 2022, f"Earliest event should be in 2022, got {min_date.year}"
|
||
assert max_date.year == 2023, f"Latest event should be in 2023, got {max_date.year}"
|
||
|
||
print(f"\n✓ Temporal ordering preserved: {min_date.date()} to {max_date.date()}")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mentioned_at_vs_occurred(memory, request_context):
|
||
"""
|
||
Test distinction between when fact occurred vs when it was mentioned.
|
||
|
||
Scenario: Ingesting a historical conversation from 2020
|
||
- event_date: When the conversation happened (2020-03-15)
|
||
- mentioned_at: When the conversation happened (same as event_date = 2020-03-15)
|
||
- occurred_start/end: When the event in the conversation happened (extracted by LLM, or falls back to mentioned_at)
|
||
"""
|
||
bank_id = f"test_mentioned_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Ingesting a conversation that happened in the past
|
||
conversation_date = datetime(2020, 3, 15, tzinfo=timezone.utc)
|
||
|
||
# Store a fact from a historical conversation
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice graduated from MIT in March 2020.",
|
||
context="education history",
|
||
event_date=conversation_date, # When this conversation happened
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create memory unit"
|
||
|
||
# Recall and check temporal fields
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Where did Alice go to school?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the fact"
|
||
fact = result.results[0]
|
||
|
||
# Parse occurred_start
|
||
if fact.occurred_start:
|
||
if isinstance(fact.occurred_start, str):
|
||
occurred_dt = datetime.fromisoformat(fact.occurred_start.replace('Z', '+00:00'))
|
||
else:
|
||
occurred_dt = fact.occurred_start
|
||
|
||
# Should be close to the conversation date (falls back to mentioned_at if LLM doesn't extract)
|
||
assert occurred_dt.year == 2020, f"occurred_start should be 2020, got {occurred_dt.year}"
|
||
print(f"✓ occurred_start (when event happened): {occurred_dt}")
|
||
|
||
# Parse mentioned_at
|
||
if fact.mentioned_at:
|
||
if isinstance(fact.mentioned_at, str):
|
||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace('Z', '+00:00'))
|
||
else:
|
||
mentioned_dt = fact.mentioned_at
|
||
|
||
# mentioned_at should match the conversation date (event_date)
|
||
time_diff = abs((conversation_date - mentioned_dt).total_seconds())
|
||
assert time_diff < 60, f"mentioned_at should match event_date (2020-03-15), but diff is {time_diff}s"
|
||
print(f"✓ mentioned_at (when conversation happened): {mentioned_dt}")
|
||
|
||
# Verify it's the historical date, not today
|
||
assert mentioned_dt.year == 2020, f"mentioned_at should be 2020, got {mentioned_dt.year}"
|
||
|
||
print("✓ Test passed: Historical conversation correctly ingested with event_date=2020")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_occurred_dates_not_defaulted(memory, request_context):
|
||
"""
|
||
Test that occurred_start and occurred_end are NOT defaulted to mentioned_at.
|
||
|
||
This is a regression test for a bug where occurred dates were incorrectly
|
||
defaulting to mentioned_at when the LLM didn't provide them.
|
||
|
||
Scenario: Store a fact where occurred dates are not applicable (current observation)
|
||
- mentioned_at should be set (to event_date or now())
|
||
- occurred_start and occurred_end should be None (not defaulted to mentioned_at)
|
||
"""
|
||
bank_id = f"test_occurred_not_defaulted_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store a current observation where occurred dates don't make sense
|
||
# Use present tense to avoid LLM extracting past dates
|
||
# Content needs to be substantial enough to not be filtered as trivial
|
||
event_date = datetime(2024, 2, 10, 15, 30, tzinfo=timezone.utc)
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice is a software engineer who specializes in Python and machine learning. She prefers dark roast coffee and works remotely from Seattle.",
|
||
context="current observations about Alice",
|
||
event_date=event_date,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create memory unit"
|
||
|
||
# Recall and check that occurred dates are None
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about Alice",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world", "experience"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the fact"
|
||
fact = result.results[0]
|
||
|
||
# mentioned_at should be set
|
||
assert fact.mentioned_at is not None, "mentioned_at should be set"
|
||
|
||
# Parse mentioned_at
|
||
if isinstance(fact.mentioned_at, str):
|
||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace('Z', '+00:00'))
|
||
else:
|
||
mentioned_dt = fact.mentioned_at
|
||
|
||
# Verify it matches event_date
|
||
time_diff = abs((event_date - mentioned_dt).total_seconds())
|
||
assert time_diff < 60, f"mentioned_at should match event_date, but diff is {time_diff}s"
|
||
|
||
# CRITICAL: occurred_start and occurred_end should be None
|
||
# They should NOT default to mentioned_at
|
||
if fact.occurred_start is not None:
|
||
# If occurred_start is set, it means the LLM extracted it
|
||
# In this case, log it but don't fail (LLM behavior can vary)
|
||
print(f"⚠ LLM extracted occurred_start: {fact.occurred_start}")
|
||
print(" This test expects None for present-tense observations")
|
||
else:
|
||
print("✓ occurred_start is correctly None (not defaulted to mentioned_at)")
|
||
|
||
if fact.occurred_end is not None:
|
||
print(f"⚠ LLM extracted occurred_end: {fact.occurred_end}")
|
||
print(" This test expects None for present-tense observations")
|
||
else:
|
||
print("✓ occurred_end is correctly None (not defaulted to mentioned_at)")
|
||
|
||
# At least verify they're not equal to mentioned_at if they are set
|
||
if fact.occurred_start is not None:
|
||
if isinstance(fact.occurred_start, str):
|
||
occurred_start_dt = datetime.fromisoformat(fact.occurred_start.replace('Z', '+00:00'))
|
||
else:
|
||
occurred_start_dt = fact.occurred_start
|
||
|
||
# If they're equal, it suggests the old defaulting bug
|
||
if occurred_start_dt == mentioned_dt:
|
||
raise AssertionError(
|
||
f"occurred_start should NOT be defaulted to mentioned_at! "
|
||
f"occurred_start={occurred_start_dt}, mentioned_at={mentioned_dt}"
|
||
)
|
||
|
||
print("✓ Test passed: occurred dates are not incorrectly defaulted to mentioned_at")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mentioned_at_from_context_string(memory, request_context):
|
||
"""
|
||
Test that mentioned_at is extracted from context string by LLM.
|
||
|
||
Scenario: User provides date in context like "happened on 2023-05-10 14:30:00 UTC"
|
||
- LLM should extract mentioned_at from this context
|
||
- If LLM fails to extract, should fall back to event_date (which defaults to now())
|
||
- mentioned_at should NEVER be None
|
||
"""
|
||
bank_id = f"test_context_date_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Test case 1: Date in context string (like longmemeval benchmark)
|
||
session_date = datetime(2023, 5, 10, 14, 30, 0, tzinfo=timezone.utc)
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice mentioned she loves hiking in the mountains.",
|
||
context=f"Session ABC123 - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||
event_date=None, # Not providing event_date - should default to now() if LLM doesn't extract
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create memory unit"
|
||
|
||
# Recall and verify mentioned_at is set (no fact_type filter — LLM may classify as world or experience)
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="What does Alice like?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the fact"
|
||
fact = result.results[0]
|
||
|
||
# mentioned_at must ALWAYS be set
|
||
assert fact.mentioned_at is not None, "mentioned_at should NEVER be None"
|
||
|
||
# Parse mentioned_at
|
||
if isinstance(fact.mentioned_at, str):
|
||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace('Z', '+00:00'))
|
||
else:
|
||
mentioned_dt = fact.mentioned_at
|
||
|
||
# Check if LLM extracted the date from context (ideal case)
|
||
# Or if it fell back to now() (acceptable fallback)
|
||
time_diff_from_context = abs((session_date - mentioned_dt).total_seconds())
|
||
time_diff_from_now = abs((datetime.now(timezone.utc) - mentioned_dt).total_seconds())
|
||
|
||
# Should either match the context date OR be recent (now)
|
||
is_from_context = time_diff_from_context < 60
|
||
is_from_now = time_diff_from_now < 60
|
||
|
||
assert is_from_context or is_from_now, \
|
||
f"mentioned_at should be either from context ({session_date}) or now(), but got {mentioned_dt}"
|
||
|
||
if is_from_context:
|
||
print(f"✓ LLM successfully extracted mentioned_at from context: {mentioned_dt}")
|
||
assert mentioned_dt.year == 2023
|
||
else:
|
||
print(f"⚠ LLM did not extract date from context, fell back to now(): {mentioned_dt}")
|
||
|
||
print("✓ mentioned_at is always set (never None)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# No Timestamp Tests
|
||
# ============================================================
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retain_no_timestamp(memory, request_context):
|
||
"""
|
||
Test retaining content with explicit "no timestamp" sentinel.
|
||
|
||
When event_date=None is passed explicitly in the dict (i.e. caller opted into
|
||
no timestamp), mentioned_at should be NULL in the DB rather than defaulting to now().
|
||
"""
|
||
bank_id = f"test_no_timestamp_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Use retain_batch_async with explicit event_date=None key to signal "no timestamp"
|
||
unit_ids_list = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "The capital of France is Paris. The Eiffel Tower is located in Paris.",
|
||
"context": "general knowledge",
|
||
"event_date": None, # Explicit sentinel: no timestamp
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_list) > 0, "Should create at least one batch result"
|
||
unit_ids = unit_ids_list[0]
|
||
assert len(unit_ids) > 0, "Should have extracted and stored facts"
|
||
|
||
# Recall the facts
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Where is the Eiffel Tower?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the stored fact"
|
||
|
||
# All temporal fields should be None for temporally agnostic content
|
||
for fact in result.results:
|
||
assert fact.mentioned_at is None, (
|
||
f"mentioned_at should be None for no-timestamp content, got {fact.mentioned_at}"
|
||
)
|
||
|
||
print(f"\n✓ Test passed: mentioned_at is None for {len(result.results)} fact(s)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retain_omit_timestamp_defaults_to_now(memory, request_context):
|
||
"""
|
||
Backward-compatibility regression test: omitting event_date still stores a real datetime.
|
||
|
||
When event_date is absent from the content dict (key not present), the orchestrator
|
||
should default to utcnow() — preserving existing behavior.
|
||
"""
|
||
bank_id = f"test_default_timestamp_{datetime.now(timezone.utc).timestamp()}"
|
||
before = datetime.now(timezone.utc)
|
||
|
||
try:
|
||
# Omit event_date entirely — should default to now()
|
||
unit_ids_list = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "Alice is a software engineer who loves Python.",
|
||
"context": "profile",
|
||
# event_date intentionally omitted
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
after = datetime.now(timezone.utc)
|
||
|
||
assert len(unit_ids_list) > 0
|
||
unit_ids = unit_ids_list[0]
|
||
assert len(unit_ids) > 0, "Should have extracted and stored facts"
|
||
|
||
# Recall and verify mentioned_at is a real datetime close to now
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Who is Alice?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the fact"
|
||
fact = result.results[0]
|
||
|
||
assert fact.mentioned_at is not None, "mentioned_at should be set when event_date is omitted"
|
||
|
||
if isinstance(fact.mentioned_at, str):
|
||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace("Z", "+00:00"))
|
||
else:
|
||
mentioned_dt = fact.mentioned_at
|
||
|
||
# Should be within 60s of when we ran the test
|
||
assert before <= mentioned_dt <= after + timedelta(seconds=60), (
|
||
f"mentioned_at {mentioned_dt} should be close to now ({before} – {after})"
|
||
)
|
||
|
||
print(f"\n✓ Test passed: mentioned_at={mentioned_dt} is a real datetime (backward compat)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Context Tracking Tests
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_context_preservation(memory, request_context):
|
||
"""
|
||
Test that context is preserved and retrievable.
|
||
Context helps understand why/how memory was formed.
|
||
"""
|
||
bank_id = f"test_context_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store content with specific context
|
||
specific_context = "team meeting notes from Q4 planning session"
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="The team decided to prioritize mobile development for next quarter.",
|
||
context=specific_context,
|
||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create at least one memory unit"
|
||
|
||
# Recall and verify context is returned (no fact_type filter — LLM may classify as world or experience)
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="What did the team decide?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall the stored fact"
|
||
|
||
# Verify context is preserved (context is stored in the database)
|
||
# Note: context might not be returned in the API response by default
|
||
# but it should be stored in the database
|
||
print(f"✓ Successfully stored fact with context: '{specific_context}'")
|
||
print(f" Retrieved {len(result.results)} facts")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_context_with_batch(memory, request_context):
|
||
"""
|
||
Test that each item in a batch can have different contexts.
|
||
|
||
Note: LLM fact extraction is non-deterministic. Simple sentences may
|
||
not always produce exactly 1 fact each. We verify the batch was
|
||
processed and at least some facts were extracted.
|
||
"""
|
||
bank_id = f"test_batch_context_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store batch with different contexts
|
||
unit_ids = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "Alice completed the authentication module.",
|
||
"context": "sprint 1 standup",
|
||
"event_date": datetime(2024, 1, 10, tzinfo=timezone.utc)
|
||
},
|
||
{
|
||
"content": "Bob started working on the database schema.",
|
||
"context": "sprint 1 planning",
|
||
"event_date": datetime(2024, 1, 11, tzinfo=timezone.utc)
|
||
},
|
||
{
|
||
"content": "Charlie fixed critical bugs in the payment flow.",
|
||
"context": "incident response",
|
||
"event_date": datetime(2024, 1, 12, tzinfo=timezone.utc)
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Should have created facts from at least some items
|
||
# LLM extraction is non-deterministic, so we allow some flexibility
|
||
total_units = sum(len(ids) for ids in unit_ids)
|
||
assert total_units >= 2, f"Should create at least 2 units from 3 batch items, got {total_units}"
|
||
|
||
print(f"✓ Stored {len(unit_ids)} batch items with different contexts")
|
||
print(f" Created {total_units} total memory units")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Metadata Storage Tests
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_metadata_storage_and_retrieval(memory, request_context):
|
||
"""
|
||
Test that user-defined metadata passed during retain is returned on recall.
|
||
Metadata allows arbitrary key-value data to be stored with facts.
|
||
"""
|
||
bank_id = f"test_metadata_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
custom_metadata = {
|
||
"source": "slack",
|
||
"channel": "engineering",
|
||
"importance": "high",
|
||
}
|
||
|
||
# Use retain_batch_async which supports the metadata parameter
|
||
unit_ids_list = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "The product launch is scheduled for March 1st.",
|
||
"context": "planning meeting",
|
||
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||
"metadata": custom_metadata,
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_list) > 0, "Should create memory units"
|
||
assert len(unit_ids_list[0]) > 0, "Should have at least one unit ID"
|
||
|
||
# Recall and verify metadata is returned
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="When is the product launch?",
|
||
budget=Budget.LOW,
|
||
max_tokens=500,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall stored facts"
|
||
|
||
# Verify metadata is present on recalled facts
|
||
fact = result.results[0]
|
||
assert fact.metadata is not None, "Metadata should not be null on recall"
|
||
assert fact.metadata.get("source") == "slack"
|
||
assert fact.metadata.get("channel") == "engineering"
|
||
assert fact.metadata.get("importance") == "high"
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Batch Processing Edge Cases
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_empty_batch(memory, request_context):
|
||
"""
|
||
Test that empty batch is handled gracefully without errors.
|
||
"""
|
||
bank_id = f"test_empty_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Attempt to store empty batch
|
||
unit_ids = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[],
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Should return empty list or handle gracefully
|
||
assert isinstance(unit_ids, list), "Should return a list"
|
||
assert len(unit_ids) == 0, "Empty batch should create no units"
|
||
|
||
print("✓ Empty batch handled gracefully")
|
||
|
||
finally:
|
||
# Clean up (though nothing should be stored)
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_single_item_batch(memory, request_context):
|
||
"""
|
||
Test that batch with one item works correctly.
|
||
"""
|
||
bank_id = f"test_single_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store batch with single item
|
||
unit_ids = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "Alice shipped the new feature to production.",
|
||
"context": "deployment log",
|
||
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) == 1, "Should return one list of unit IDs"
|
||
assert len(unit_ids[0]) > 0, "Should create at least one memory unit"
|
||
|
||
print(f"✓ Single-item batch created {len(unit_ids[0])} units")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mixed_content_batch(memory, request_context):
|
||
"""
|
||
Test batch with varying content sizes (short and long).
|
||
"""
|
||
bank_id = f"test_mixed_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Mix short and long content
|
||
short_content = "Alice joined the team."
|
||
long_content = """
|
||
Bob has been working on the authentication system for the past three months.
|
||
He implemented OAuth 2.0 integration, set up JWT token management, and built
|
||
a comprehensive role-based access control system. The system supports multiple
|
||
identity providers including Google, GitHub, and Microsoft. Bob also wrote
|
||
extensive documentation and unit tests covering over 90% of the codebase.
|
||
The team recognized his work with an excellence award at the quarterly meeting.
|
||
"""
|
||
|
||
unit_ids = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{"content": short_content, "context": "onboarding"},
|
||
{"content": long_content, "context": "performance review"},
|
||
{"content": "Charlie is on vacation this week.", "context": "team status"}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
# All items should be processed
|
||
assert len(unit_ids) == 3, "Should process all 3 items"
|
||
|
||
# Long content should create more facts
|
||
short_units = len(unit_ids[0])
|
||
long_units = len(unit_ids[1])
|
||
|
||
print("✓ Mixed batch processed successfully")
|
||
print(f" Short content: {short_units} units")
|
||
print(f" Long content: {long_units} units")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_batch_with_missing_optional_fields(memory, request_context):
|
||
"""
|
||
Test that batch handles items with missing optional fields.
|
||
"""
|
||
bank_id = f"test_optional_fields_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Some items have all fields, some have minimal fields
|
||
unit_ids = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[
|
||
{
|
||
"content": "Alice finished the project.",
|
||
"context": "complete record",
|
||
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||
},
|
||
{
|
||
"content": "Bob started a new task.",
|
||
# No context or event_date
|
||
},
|
||
{
|
||
"content": "Charlie reviewed code.",
|
||
"context": "code review",
|
||
# No event_date
|
||
}
|
||
],
|
||
request_context=request_context,
|
||
)
|
||
|
||
# All items should be processed successfully
|
||
assert len(unit_ids) == 3, "Should process all items even with missing optional fields"
|
||
|
||
total_units = sum(len(ids) for ids in unit_ids)
|
||
print(f"✓ Batch with mixed optional fields created {total_units} total units")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Multi-Document Batch Tests
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_single_batch_multiple_documents(memory, request_context):
|
||
"""
|
||
Test storing multiple distinct documents in a single batch call.
|
||
Each should be tracked separately.
|
||
"""
|
||
bank_id = f"test_multi_docs_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store single batch where each item could be a different document
|
||
# (In practice, document_id is a batch-level parameter, so we test
|
||
# that multiple retain_async calls work correctly)
|
||
|
||
doc1_units = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice's resume: 10 years Python experience, worked at Google.",
|
||
context="resume review",
|
||
document_id="resume_alice",
|
||
request_context=request_context,
|
||
)
|
||
|
||
doc2_units = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Bob's resume: 5 years JavaScript experience, worked at Meta.",
|
||
context="resume review",
|
||
document_id="resume_bob",
|
||
request_context=request_context,
|
||
)
|
||
|
||
doc3_units = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Charlie's resume: 8 years Go experience, worked at Amazon.",
|
||
context="resume review",
|
||
document_id="resume_charlie",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# All documents should be stored
|
||
assert len(doc1_units) > 0, "Should create units for doc1"
|
||
assert len(doc2_units) > 0, "Should create units for doc2"
|
||
assert len(doc3_units) > 0, "Should create units for doc3"
|
||
|
||
total_units = len(doc1_units) + len(doc2_units) + len(doc3_units)
|
||
print(f"✓ Stored 3 separate documents with {total_units} total units")
|
||
|
||
# Verify we can recall from any document
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Who worked at Google?",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
fact_type=["world"],
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should find facts about Alice"
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_document_upsert_behavior(memory, request_context):
|
||
"""
|
||
Test that upserting a document replaces the old content.
|
||
"""
|
||
bank_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "project_status"
|
||
|
||
try:
|
||
# Store initial version
|
||
v1_units = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Project is in planning phase. Alice is the lead.",
|
||
context="status update v1",
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(v1_units) > 0, "Should create units for v1"
|
||
|
||
# Update with new version (upsert)
|
||
v2_units = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Project is in development phase. Bob has joined as co-lead.",
|
||
context="status update v2",
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(v2_units) > 0, "Should create units for v2"
|
||
|
||
# Recall should return the updated information (no fact_type filter — LLM may classify as world or experience)
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="What is the project status?",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall facts"
|
||
|
||
print(f"✓ Document upsert created v1: {len(v1_units)} units, v2: {len(v2_units)} units")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Chunk Storage Advanced Tests
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_chunk_fact_mapping(memory, request_context):
|
||
"""
|
||
Test that facts correctly reference their source chunks via chunk_id.
|
||
"""
|
||
bank_id = f"test_chunk_mapping_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "technical_doc"
|
||
|
||
try:
|
||
# Store content that will be chunked
|
||
content = """
|
||
The authentication system uses JWT tokens for session management.
|
||
Tokens expire after 24 hours and must be refreshed using the refresh endpoint.
|
||
The system supports OAuth 2.0 integration with Google and GitHub.
|
||
|
||
The database layer uses PostgreSQL with connection pooling.
|
||
We maintain separate read and write connection pools for performance.
|
||
All queries use prepared statements to prevent SQL injection.
|
||
"""
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=content,
|
||
context="technical documentation",
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create memory units"
|
||
|
||
# Recall with chunks enabled
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="How does authentication work?",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
fact_type=["world"],
|
||
include_chunks=True,
|
||
max_chunk_tokens=8192,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result.results) > 0, "Should recall facts"
|
||
|
||
# Verify facts have chunk_id references
|
||
facts_with_chunks = [f for f in result.results if f.chunk_id]
|
||
|
||
print(f"✓ Created {len(unit_ids)} units from chunked document")
|
||
print(f" {len(facts_with_chunks)}/{len(result.results)} facts have chunk_id references")
|
||
|
||
# If chunks are returned, verify they match the chunk_ids in facts
|
||
if result.chunks:
|
||
fact_chunk_ids = {f.chunk_id for f in facts_with_chunks}
|
||
returned_chunk_ids = set(result.chunks.keys())
|
||
|
||
# All chunk_ids in facts should have corresponding chunk data
|
||
assert fact_chunk_ids.issubset(returned_chunk_ids) or len(fact_chunk_ids) == 0, \
|
||
"Fact chunk_ids should have corresponding chunk data"
|
||
|
||
print(f" Returned {len(result.chunks)} chunks matching fact references")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_chunk_ordering_preservation(memory, request_context):
|
||
"""
|
||
Test that chunk_index reflects the correct order within a document.
|
||
"""
|
||
bank_id = f"test_chunk_order_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "ordered_doc"
|
||
|
||
try:
|
||
# Store long content that will create multiple chunks with meaningful content
|
||
sections = []
|
||
sections.append("""
|
||
Alice is the team lead for the authentication project. She has 10 years of experience
|
||
with security systems and previously worked at Google on identity management.
|
||
She is responsible for architecture decisions and code review.
|
||
""")
|
||
sections.append("""
|
||
Bob is a backend engineer focusing on the API layer. He specializes in Python
|
||
and has built several microservices for the company. He joined the team in 2023.
|
||
""")
|
||
sections.append("""
|
||
Charlie is the DevOps engineer managing the deployment pipeline. He set up
|
||
our Kubernetes infrastructure and maintains the CI/CD system using GitHub Actions.
|
||
""")
|
||
sections.append("""
|
||
The project uses PostgreSQL as the main database with Redis for caching.
|
||
We deploy to AWS using Docker containers orchestrated by Kubernetes.
|
||
The team follows agile methodology with two-week sprints.
|
||
""")
|
||
sections.append("""
|
||
Security is a top priority. All API endpoints require JWT authentication.
|
||
We use OAuth 2.0 for third-party integrations and maintain strict access controls.
|
||
Regular security audits are conducted quarterly.
|
||
""")
|
||
|
||
content = "\n\n".join(sections)
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=content,
|
||
context="multi-section document",
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create units"
|
||
|
||
# Recall with chunks
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about the sections",
|
||
budget=Budget.MID,
|
||
max_tokens=2000,
|
||
fact_type=["world"],
|
||
include_chunks=True,
|
||
max_chunk_tokens=8192,
|
||
request_context=request_context,
|
||
)
|
||
|
||
if result.chunks:
|
||
# Verify chunk_index values are sequential and start from 0
|
||
chunk_indices = [chunk.chunk_index for chunk in result.chunks.values()]
|
||
chunk_indices_sorted = sorted(chunk_indices)
|
||
|
||
print(f"✓ Document created {len(result.chunks)} chunks")
|
||
print(f" Chunk indices: {chunk_indices}")
|
||
|
||
# Indices should start from 0 and be sequential
|
||
if len(chunk_indices) > 0:
|
||
assert min(chunk_indices) == 0, "Chunk indices should start from 0"
|
||
assert chunk_indices_sorted == list(range(len(chunk_indices))), \
|
||
"Chunk indices should be sequential"
|
||
else:
|
||
print("✓ Content stored (may have created single chunk or no chunks returned)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.timeout(180) # Allow up to 3 minutes for this test
|
||
async def test_chunks_truncation_behavior(memory, request_context):
|
||
"""
|
||
Test that when chunks exceed max_chunk_tokens, truncation is indicated.
|
||
|
||
Note: This test processes larger content and may take longer than typical tests.
|
||
"""
|
||
bank_id = f"test_chunk_truncation_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "large_doc"
|
||
|
||
try:
|
||
# Create a moderately large document with meaningful content
|
||
# Reduced from * 5 to * 2 for faster execution while still testing truncation
|
||
large_content = """
|
||
The company's product roadmap for 2024 includes several major initiatives.
|
||
The engineering team is expanding to support these efforts.
|
||
|
||
Alice leads the authentication team, which is implementing OAuth 2.0 and JWT tokens.
|
||
The team has been working on this for six months and expects to launch in Q2.
|
||
Security is the top priority, with regular penetration testing scheduled.
|
||
|
||
Bob manages the API development team. They are building RESTful endpoints
|
||
for all major features including user management, billing, and analytics.
|
||
The team uses Python with FastAPI and deploys to AWS Lambda.
|
||
|
||
Charlie oversees the infrastructure team. They maintain Kubernetes clusters
|
||
across three AWS regions for high availability. The team also manages
|
||
the CI/CD pipeline using GitHub Actions and ArgoCD.
|
||
|
||
The data engineering team, led by Diana, processes millions of events daily.
|
||
They use Apache Kafka for streaming and Snowflake for analytics.
|
||
Real-time dashboards are built with Grafana and Prometheus.
|
||
|
||
The mobile team is building iOS and Android apps using React Native.
|
||
They are targeting a beta launch in Q3 with select customers.
|
||
Push notifications and offline support are key features.
|
||
|
||
The design team has created a new design system that will be rolled out
|
||
across all products. The system includes components for accessibility
|
||
and internationalization support for 12 languages.
|
||
|
||
Customer support is being enhanced with AI-powered chatbots.
|
||
The system can handle common queries and escalate complex issues to humans.
|
||
Average response time has improved by 40% since implementation.
|
||
|
||
The marketing team is planning a major campaign for the product launch.
|
||
They are working with influencers and planning webinars for enterprise customers.
|
||
Early feedback from beta users has been very positive.
|
||
|
||
Sales operations are being streamlined with new CRM integrations.
|
||
The team can now track leads more effectively and automate follow-ups.
|
||
Conversion rates have increased by 25% in the pilot program.
|
||
|
||
The finance team is implementing new budgeting tools for better forecasting.
|
||
They are also working on automated expense reporting and approval workflows.
|
||
This will save approximately 100 hours per month in manual work.
|
||
""" * 2 # Repeat to create enough content for truncation testing
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=large_content,
|
||
context="large document test",
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should create units"
|
||
|
||
# Recall with very small chunk token limit to force truncation
|
||
result = await memory.recall_async(
|
||
bank_id=bank_id,
|
||
query="Tell me about the document",
|
||
budget=Budget.MID,
|
||
max_tokens=1000,
|
||
fact_type=["world"],
|
||
include_chunks=True,
|
||
max_chunk_tokens=500, # Small limit to test truncation
|
||
request_context=request_context,
|
||
)
|
||
|
||
if result.chunks:
|
||
# Check if any chunks show truncation
|
||
truncated_chunks = [
|
||
chunk_id for chunk_id, chunk_info in result.chunks.items()
|
||
if chunk_info.truncated
|
||
]
|
||
|
||
print(f"✓ Retrieved {len(result.chunks)} chunks")
|
||
if truncated_chunks:
|
||
print(f" {len(truncated_chunks)} chunks were truncated due to token limit")
|
||
else:
|
||
print(" No chunks were truncated (content within limit)")
|
||
|
||
else:
|
||
print("✓ No chunks returned (may be under token limit)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ============================================================
|
||
# Memory Links Tests
|
||
# ============================================================
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_temporal_links_creation(memory, request_context):
|
||
"""
|
||
Test that temporal links are created between facts with nearby event dates.
|
||
|
||
Temporal links connect facts that occurred close in time (within 24 hours).
|
||
"""
|
||
bank_id = f"test_temporal_links_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store facts with nearby timestamps (within 24 hours)
|
||
base_date = datetime(2024, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
|
||
|
||
# Fact 1 at 10:00 AM
|
||
unit_ids_1 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice started working on the authentication module.",
|
||
context="daily standup",
|
||
event_date=base_date,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Fact 2 at 2:00 PM same day (4 hours later)
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Bob reviewed the API design document.",
|
||
context="daily standup",
|
||
event_date=base_date.replace(hour=14),
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Fact 3 at 9:00 AM next day (23 hours later)
|
||
unit_ids_3 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Charlie deployed the new database schema.",
|
||
context="daily standup",
|
||
event_date=base_date.replace(day=16, hour=9),
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||
|
||
logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3)} facts")
|
||
|
||
# Query the memory_links table to verify temporal links exist
|
||
async with memory._pool.acquire() as conn:
|
||
# Get all temporal links for these units
|
||
all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3
|
||
|
||
temporal_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, link_type, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type = 'temporal'
|
||
ORDER BY weight DESC
|
||
""",
|
||
all_unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(temporal_links)} temporal links")
|
||
|
||
# Should have temporal links between the facts
|
||
assert len(temporal_links) > 0, "Should have created temporal links between facts with nearby dates"
|
||
|
||
# Verify link properties
|
||
for link in temporal_links:
|
||
from_id = str(link['from_unit_id'])
|
||
to_id = str(link['to_unit_id'])
|
||
logger.info(f" Link: {from_id[:8]}... -> {to_id[:8]}... (weight: {link['weight']:.2f})")
|
||
assert link['link_type'] == 'temporal', "Link type should be 'temporal'"
|
||
assert 0.0 <= link['weight'] <= 1.0, "Weight should be between 0 and 1"
|
||
|
||
logger.info("Temporal links created successfully with proper weights")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_semantic_links_creation(memory, request_context):
|
||
"""
|
||
Test that semantic links are created between facts with similar content.
|
||
|
||
Semantic links connect facts that are semantically similar based on embeddings.
|
||
"""
|
||
bank_id = f"test_semantic_links_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store facts with similar semantic content
|
||
unit_ids_1 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice is an expert in Python programming and has built many web applications.",
|
||
context="team skills",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Similar content - should create semantic link
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Bob is proficient in Python development and specializes in building APIs.",
|
||
context="team skills",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Different content - less likely to create strong semantic link
|
||
unit_ids_3 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="The quarterly sales meeting is scheduled for next Tuesday at 3 PM.",
|
||
context="calendar events",
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||
|
||
logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3)} facts")
|
||
|
||
# Query the memory_links table to verify semantic links exist
|
||
async with memory._pool.acquire() as conn:
|
||
all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3
|
||
|
||
semantic_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, link_type, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type = 'semantic'
|
||
ORDER BY weight DESC
|
||
""",
|
||
all_unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(semantic_links)} semantic links")
|
||
|
||
# Should have semantic links between similar facts
|
||
assert len(semantic_links) > 0, "Should have created semantic links between similar facts"
|
||
|
||
# Verify link properties
|
||
for link in semantic_links:
|
||
from_id = str(link['from_unit_id'])
|
||
to_id = str(link['to_unit_id'])
|
||
logger.info(f" Link: {from_id[:8]}... -> {to_id[:8]}... (weight: {link['weight']:.3f})")
|
||
assert link['link_type'] == 'semantic', "Link type should be 'semantic'"
|
||
assert 0.0 <= link['weight'] <= 1.0, "Weight should be between 0 and 1"
|
||
# Semantic links typically have weight >= 0.7 (threshold)
|
||
assert link['weight'] >= 0.7, f"Semantic links should have weight >= 0.7, got {link['weight']}"
|
||
|
||
logger.info("Semantic links created successfully between similar content")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_entity_links_creation(memory, request_context):
|
||
"""
|
||
Test that entity links are created between facts that mention the same entities.
|
||
|
||
Entity links connect facts that reference the same person, place, or concept.
|
||
This is core functionality and should work consistently.
|
||
"""
|
||
bank_id = f"test_entity_links_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store facts that mention the same entities
|
||
unit_ids_1 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice joined Google as a software engineer in 2020.",
|
||
context="career history",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Mentions same entity (Alice) - should create entity link
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice led the development of the new authentication system.",
|
||
context="project updates",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Mentions same entity (Google) - should create entity link
|
||
unit_ids_3 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Google announced new cloud services at their annual conference.",
|
||
context="tech news",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Different entities - no entity link expected
|
||
unit_ids_4 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Bob works at Meta on machine learning infrastructure.",
|
||
context="career history",
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0 and len(unit_ids_4) > 0
|
||
|
||
logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3) + len(unit_ids_4)} facts")
|
||
|
||
# Query the memory_links table to verify entity links exist
|
||
async with memory._pool.acquire() as conn:
|
||
all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3 + unit_ids_4
|
||
|
||
entity_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type = 'entity'
|
||
ORDER BY from_unit_id, to_unit_id
|
||
""",
|
||
all_unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(entity_links)} entity links")
|
||
|
||
# Entity extraction is core functionality and should work
|
||
assert len(entity_links) > 0, "Should have created entity links between facts with shared entities (Alice, Google)"
|
||
|
||
# Verify link properties
|
||
entities_seen = set()
|
||
for link in entity_links:
|
||
entity_id = link['entity_id']
|
||
entities_seen.add(str(entity_id))
|
||
from_id = str(link['from_unit_id'])
|
||
to_id = str(link['to_unit_id'])
|
||
logger.info(f" Link: {from_id[:8]}... -> {to_id[:8]}... via entity {str(entity_id)[:8]}...")
|
||
assert link['link_type'] == 'entity', "Link type should be 'entity'"
|
||
assert link['weight'] == 1.0, "Entity links should have weight 1.0"
|
||
assert entity_id is not None, "Entity links must reference an entity_id"
|
||
|
||
logger.info(f"Entity links created successfully for {len(entities_seen)} unique entities")
|
||
|
||
# Verify bidirectional links (entity links should be bidirectional)
|
||
link_pairs = set()
|
||
for link in entity_links:
|
||
from_id = str(link['from_unit_id'])
|
||
to_id = str(link['to_unit_id'])
|
||
entity_id = str(link['entity_id'])
|
||
link_pairs.add((from_id, to_id, entity_id))
|
||
|
||
# Check that for each (A -> B) link, there's a (B -> A) link with same entity
|
||
for from_id, to_id, entity_id in link_pairs:
|
||
reverse_exists = (to_id, from_id, entity_id) in link_pairs
|
||
assert reverse_exists, f"Entity links should be bidirectional: missing reverse link for {from_id[:8]} -> {to_id[:8]}"
|
||
|
||
logger.info("Entity links are properly bidirectional")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_people_name_extraction(memory, request_context):
|
||
"""
|
||
Test that people names are correctly extracted as entities.
|
||
|
||
This verifies that the entity resolver properly identifies and extracts
|
||
person names from content.
|
||
"""
|
||
bank_id = f"test_people_names_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store content with various people names
|
||
contents = [
|
||
"John Smith is a software engineer at Google.",
|
||
"Dr. Sarah Johnson presented her research at the conference.",
|
||
"Bob Williams and Alice Chen collaborated on the project.",
|
||
"Professor Michael Brown teaches computer science at MIT.",
|
||
]
|
||
|
||
for content in contents:
|
||
await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=content,
|
||
context="people info",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Query entities to verify people names were extracted
|
||
async with memory._pool.acquire() as conn:
|
||
entities = await conn.fetch(
|
||
"""
|
||
SELECT canonical_name, mention_count
|
||
FROM entities
|
||
WHERE bank_id = $1
|
||
ORDER BY mention_count DESC, canonical_name
|
||
""",
|
||
bank_id
|
||
)
|
||
|
||
logger.info(f"Extracted {len(entities)} entities")
|
||
for entity in entities:
|
||
logger.info(f" - {entity['canonical_name']} (mentions: {entity['mention_count']})")
|
||
|
||
# Verify we extracted the expected people names
|
||
entity_names = {e['canonical_name'].lower() for e in entities}
|
||
|
||
# Check for expected people (names may vary slightly based on LLM extraction)
|
||
expected_people = ["john", "sarah", "bob", "alice", "michael"]
|
||
found_people = []
|
||
for person in expected_people:
|
||
matching = [name for name in entity_names if person in name]
|
||
if matching:
|
||
found_people.append(person)
|
||
logger.info(f" Found '{person}' as: {matching}")
|
||
|
||
assert len(found_people) >= 3, \
|
||
f"Should extract at least 3 people names, found: {found_people}. All entities: {entity_names}"
|
||
|
||
logger.info(f"Successfully extracted {len(found_people)} people names: {found_people}")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mention_count_accuracy(memory, request_context):
|
||
"""
|
||
Test that mention_count is accurately tracked across retain calls.
|
||
|
||
Verifies that when an entity is mentioned multiple times across different
|
||
retain calls, the mention_count reflects the total number of mentions.
|
||
"""
|
||
bank_id = f"test_mention_count_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store content mentioning "Alice" multiple times across separate retain calls
|
||
contents = [
|
||
"Alice is a data scientist at Netflix.",
|
||
"Alice presented her research on recommendation algorithms.",
|
||
"Alice leads a team of 5 engineers.",
|
||
"Alice graduated from Stanford with honors.",
|
||
"Alice published a paper on machine learning.",
|
||
]
|
||
|
||
for content in contents:
|
||
await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=content,
|
||
context="career info",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Check Alice's mention count
|
||
async with memory._pool.acquire() as conn:
|
||
alice_entity = await conn.fetchrow(
|
||
"""
|
||
SELECT canonical_name, mention_count
|
||
FROM entities
|
||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%alice%'
|
||
""",
|
||
bank_id
|
||
)
|
||
|
||
assert alice_entity is not None, "Alice entity should exist"
|
||
logger.info(f"Alice mention_count after 5 separate retains: {alice_entity['mention_count']}")
|
||
|
||
# Alice should have mention_count >= 5 (one per content item)
|
||
assert alice_entity['mention_count'] >= 5, \
|
||
f"Alice should have at least 5 mentions, got {alice_entity['mention_count']}"
|
||
|
||
logger.info(f"Mention count accuracy verified: {alice_entity['mention_count']} mentions")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mention_count_batch_retain(memory, request_context):
|
||
"""
|
||
Test that mention_count is accurate when using batch retain with multiple items.
|
||
|
||
This specifically tests the scenario where multiple content items are retained
|
||
in a single batch call, ensuring mention_count is correctly aggregated.
|
||
"""
|
||
bank_id = f"test_mention_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Batch retain with multiple items mentioning "Bob"
|
||
batch_contents = [
|
||
{"content": "Bob is a frontend developer at Microsoft.", "context": "work"},
|
||
{"content": "Bob specializes in React and TypeScript.", "context": "skills"},
|
||
{"content": "Bob has 10 years of experience.", "context": "experience"},
|
||
{"content": "Bob mentors junior developers.", "context": "mentoring"},
|
||
{"content": "Bob presented at ReactConf 2024.", "context": "conferences"},
|
||
{"content": "Bob wrote a popular open-source library.", "context": "projects"},
|
||
]
|
||
|
||
# Use retain_batch_async for batch processing
|
||
await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=batch_contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Check Bob's mention count after batch retain
|
||
async with memory._pool.acquire() as conn:
|
||
bob_entity = await conn.fetchrow(
|
||
"""
|
||
SELECT canonical_name, mention_count
|
||
FROM entities
|
||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%bob%'
|
||
""",
|
||
bank_id
|
||
)
|
||
|
||
assert bob_entity is not None, "Bob entity should exist after batch retain"
|
||
logger.info(f"Bob mention_count after batch retain of 6 items: {bob_entity['mention_count']}")
|
||
|
||
# Bob should have mention_count >= 6 (mentioned in each batch item)
|
||
assert bob_entity['mention_count'] >= 6, \
|
||
f"Bob should have at least 6 mentions from batch retain, got {bob_entity['mention_count']}"
|
||
|
||
# Now do another batch retain with more Bob mentions
|
||
more_contents = [
|
||
{"content": "Bob loves hiking on weekends.", "context": "hobbies"},
|
||
{"content": "Bob has a dog named Max.", "context": "personal"},
|
||
]
|
||
|
||
await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=more_contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Check updated mention count
|
||
async with memory._pool.acquire() as conn:
|
||
bob_entity_updated = await conn.fetchrow(
|
||
"""
|
||
SELECT canonical_name, mention_count
|
||
FROM entities
|
||
WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%bob%'
|
||
""",
|
||
bank_id
|
||
)
|
||
|
||
logger.info(f"Bob mention_count after second batch: {bob_entity_updated['mention_count']}")
|
||
|
||
# Bob should now have mention_count >= 8 (6 + 2)
|
||
assert bob_entity_updated['mention_count'] >= 8, \
|
||
f"Bob should have at least 8 mentions after second batch, got {bob_entity_updated['mention_count']}"
|
||
|
||
# Verify the increment is correct
|
||
increment = bob_entity_updated['mention_count'] - bob_entity['mention_count']
|
||
assert increment >= 2, \
|
||
f"Mention count should have increased by at least 2, but increased by {increment}"
|
||
|
||
logger.info(f"Batch retain mention count verified: {bob_entity['mention_count']} -> {bob_entity_updated['mention_count']}")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_causal_links_creation(memory, request_context):
|
||
"""
|
||
Test that causal links are created between facts with causal relationships.
|
||
|
||
Causal links connect facts where one causes, enables, or prevents another.
|
||
Note: This depends on LLM extracting causal relationships, which may be non-deterministic.
|
||
"""
|
||
bank_id = f"test_causal_links_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store content with explicit causal relationships
|
||
# Using clear cause-and-effect language to maximize LLM detection
|
||
content = """
|
||
Alice completed the authentication module on Monday. Because Alice finished the auth module,
|
||
Bob was able to start integrating it with the API on Tuesday. Bob's API integration enabled
|
||
Charlie to begin testing the complete user flow on Wednesday. The successful testing caused
|
||
the team to schedule the production deployment for Friday.
|
||
"""
|
||
|
||
unit_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=content,
|
||
context="project timeline",
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids) > 0, "Should have created facts"
|
||
logger.info(f"Created {len(unit_ids)} facts from causal content")
|
||
|
||
# Query the memory_links table to check for causal links
|
||
async with memory._pool.acquire() as conn:
|
||
causal_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, link_type, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||
ORDER BY link_type, weight DESC
|
||
""",
|
||
unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(causal_links)} causal links")
|
||
|
||
if len(causal_links) > 0:
|
||
# Verify link properties
|
||
causal_types = {}
|
||
for link in causal_links:
|
||
link_type = link['link_type']
|
||
causal_types[link_type] = causal_types.get(link_type, 0) + 1
|
||
from_id = str(link['from_unit_id'])
|
||
to_id = str(link['to_unit_id'])
|
||
logger.info(f" Link: {from_id[:8]}... -> {to_id[:8]}... ({link_type}, weight: {link['weight']:.2f})")
|
||
assert link['link_type'] in ['causes', 'caused_by', 'enables', 'prevents'], \
|
||
f"Causal link type must be valid, got '{link['link_type']}'"
|
||
assert 0.0 <= link['weight'] <= 1.0, "Weight should be between 0 and 1"
|
||
|
||
logger.info("Causal links created successfully:")
|
||
for link_type, count in causal_types.items():
|
||
logger.info(f" - {link_type}: {count} links")
|
||
else:
|
||
logger.warning("No causal links detected (LLM may not have extracted causal relationships)")
|
||
logger.info(" This is expected as causal extraction depends on LLM interpretation")
|
||
|
||
# This test passes even if no causal links are found, since causal extraction
|
||
# is non-deterministic and depends on LLM behavior
|
||
logger.info("Test completed (causal link extraction is LLM-dependent)")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_all_link_types_together(memory, request_context):
|
||
"""
|
||
Integration test: Verify all link types can be created in a single retain operation.
|
||
|
||
Tests that temporal, semantic, entity, and potentially causal links are all
|
||
created when appropriate conditions are met.
|
||
"""
|
||
bank_id = f"test_all_links_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store multiple related facts that should trigger all link types
|
||
base_date = datetime(2024, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
|
||
|
||
# Fact 1: Alice at time T
|
||
unit_ids_1 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice completed the Python backend service for the authentication system.",
|
||
context="sprint review",
|
||
event_date=base_date,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Fact 2: Related to Alice, similar topic (Python), close in time
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice optimized the Python code and improved the authentication performance by 40%.",
|
||
context="sprint review",
|
||
event_date=base_date.replace(hour=14), # Same day, 4 hours later
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Fact 3: Related to Alice, different topic but same entity
|
||
unit_ids_3 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Alice presented the security architecture at the team meeting.",
|
||
context="team meeting",
|
||
event_date=base_date.replace(day=16), # Next day
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||
|
||
logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3)} facts")
|
||
|
||
# Query for all link types
|
||
async with memory._pool.acquire() as conn:
|
||
all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3
|
||
|
||
all_links = await conn.fetch(
|
||
"""
|
||
SELECT link_type, COUNT(*) as count
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
GROUP BY link_type
|
||
ORDER BY link_type
|
||
""",
|
||
all_unit_ids
|
||
)
|
||
|
||
logger.info("Link types created:")
|
||
link_types_found = {}
|
||
for row in all_links:
|
||
link_type = row['link_type']
|
||
count = row['count']
|
||
link_types_found[link_type] = count
|
||
logger.info(f" - {link_type}: {count} links")
|
||
|
||
# Should have temporal, semantic, and entity links
|
||
assert 'temporal' in link_types_found, "Should have temporal links (facts with nearby dates)"
|
||
assert 'semantic' in link_types_found, "Should have semantic links (similar content about Python/auth)"
|
||
assert 'entity' in link_types_found, "Should have entity links (all mention Alice)"
|
||
|
||
logger.info(f"Successfully created {len(link_types_found)} different link types")
|
||
logger.info("All major link types (temporal, semantic, entity) are working correctly")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_semantic_links_within_same_batch(memory, request_context):
|
||
"""
|
||
Test that semantic links are created between facts retained in the SAME batch.
|
||
|
||
This is a regression test - semantic links should connect similar facts
|
||
even when they are retained together in a single call.
|
||
"""
|
||
bank_id = f"test_semantic_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Retain multiple semantically similar facts in ONE batch
|
||
contents = [
|
||
{"content": "Alice is an expert in Python programming and machine learning.", "context": "team skills"},
|
||
{"content": "Bob specializes in Python development and data science.", "context": "team skills"},
|
||
{"content": "Charlie works with Python for backend API development.", "context": "team skills"},
|
||
]
|
||
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Flatten the list of lists
|
||
unit_ids = [uid for sublist in result for uid in sublist]
|
||
|
||
assert len(unit_ids) >= 3, f"Should have created at least 3 facts, got {len(unit_ids)}"
|
||
logger.info(f"Created {len(unit_ids)} facts in single batch")
|
||
|
||
# Query semantic links between these units
|
||
async with memory._pool.acquire() as conn:
|
||
semantic_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND to_unit_id::text = ANY($1)
|
||
AND link_type = 'semantic'
|
||
""",
|
||
unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(semantic_links)} semantic links within the batch")
|
||
|
||
# All three facts mention Python - they should be linked to each other
|
||
assert len(semantic_links) > 0, (
|
||
"REGRESSION: Semantic links should be created between similar facts "
|
||
"retained in the same batch, but none were found"
|
||
)
|
||
|
||
# Log the links for debugging
|
||
for link in semantic_links:
|
||
logger.info(f" Semantic link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_semantic_links_phase1_ann_cross_batch(memory, request_context):
|
||
"""
|
||
Test that Phase 1 ANN search creates semantic links between facts from
|
||
DIFFERENT retain batches.
|
||
|
||
The semantic ANN search runs in Phase 1 on a separate connection (outside
|
||
the write transaction) using placeholder unit IDs to avoid TimeoutErrors
|
||
from HNSW index contention under concurrent load. This test verifies that:
|
||
1. Phase 1 ANN with placeholder IDs works correctly
|
||
2. Placeholder IDs are remapped to real unit IDs before insertion
|
||
3. Cross-batch semantic links are created between similar facts
|
||
"""
|
||
bank_id = f"test_semantic_phase1_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# First batch: store some world facts about a topic
|
||
# Use clearly "world" content (general knowledge, not personal experience)
|
||
# to ensure consistent fact_type classification across batches,
|
||
# since ANN search filters by fact_type.
|
||
await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Python is a high-level programming language widely used for web development with frameworks like FastAPI.",
|
||
context="programming languages",
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Second batch: store similar world facts — Phase 1 ANN should find the first batch's
|
||
# facts via HNSW index and create cross-batch semantic links
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="FastAPI is a modern Python web framework known for its high performance and automatic API documentation.",
|
||
context="programming languages",
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(unit_ids_2) > 0
|
||
|
||
# Verify cross-batch semantic links exist
|
||
async with memory._pool.acquire() as conn:
|
||
cross_batch_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type = 'semantic'
|
||
AND to_unit_id::text != ALL($1)
|
||
""",
|
||
unit_ids_2,
|
||
)
|
||
|
||
logger.info(f"Cross-batch semantic links from batch 2: {len(cross_batch_links)}")
|
||
for link in cross_batch_links:
|
||
logger.info(
|
||
f" {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... "
|
||
f"(weight: {link['weight']:.3f})"
|
||
)
|
||
|
||
# Phase 1 ANN should have found similar facts from batch 1
|
||
assert len(cross_batch_links) > 0, (
|
||
"Phase 1 ANN search should create semantic links between similar facts "
|
||
"from different retain batches. This tests that placeholder unit IDs are "
|
||
"correctly remapped to real IDs after insert_facts_batch."
|
||
)
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_temporal_links_within_same_batch(memory, request_context):
|
||
"""
|
||
Test that temporal links are created between facts retained in the SAME batch.
|
||
|
||
This is a regression test - temporal links should connect facts with nearby
|
||
event dates even when they are retained together in a single call.
|
||
"""
|
||
bank_id = f"test_temporal_batch_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Retain multiple facts with nearby timestamps in ONE batch
|
||
base_date = datetime(2024, 6, 15, 10, 0, 0, tzinfo=timezone.utc)
|
||
|
||
contents = [
|
||
{
|
||
"content": "Morning standup: Alice presented the sprint goals.",
|
||
"context": "daily meeting",
|
||
"event_date": base_date
|
||
},
|
||
{
|
||
"content": "Bob demoed the new feature after standup.",
|
||
"context": "daily meeting",
|
||
"event_date": base_date + timedelta(hours=1) # 1 hour later
|
||
},
|
||
{
|
||
"content": "Charlie reviewed the pull requests in the afternoon.",
|
||
"context": "daily meeting",
|
||
"event_date": base_date + timedelta(hours=4) # 4 hours later
|
||
},
|
||
]
|
||
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Flatten the list of lists
|
||
unit_ids = [uid for sublist in result for uid in sublist]
|
||
|
||
assert len(unit_ids) >= 3, f"Should have created at least 3 facts, got {len(unit_ids)}"
|
||
logger.info(f"Created {len(unit_ids)} facts in single batch")
|
||
|
||
# Query temporal links between these units
|
||
async with memory._pool.acquire() as conn:
|
||
temporal_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND to_unit_id::text = ANY($1)
|
||
AND link_type = 'temporal'
|
||
""",
|
||
unit_ids
|
||
)
|
||
|
||
logger.info(f"Found {len(temporal_links)} temporal links within the batch")
|
||
|
||
# All three facts are within 24 hours - they should be linked to each other
|
||
assert len(temporal_links) > 0, (
|
||
"REGRESSION: Temporal links should be created between facts with nearby dates "
|
||
"retained in the same batch, but none were found"
|
||
)
|
||
|
||
# Log the links for debugging
|
||
for link in temporal_links:
|
||
logger.info(f" Temporal link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_user_provided_entities(memory, request_context):
|
||
"""
|
||
Test that user-provided entities are merged with auto-extracted entities.
|
||
|
||
This tests the feature added in PR #91 where users can provide entities
|
||
via the 'entities' field in the retain request. These should be combined
|
||
with LLM-extracted entities, with case-insensitive deduplication.
|
||
"""
|
||
bank_id = f"test_user_entities_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Store content with user-provided entities
|
||
# The content mentions "Alice" which LLM might extract,
|
||
# but we also provide "ProjectX" and "ACME Corp" which may not be in the text
|
||
contents = [
|
||
{
|
||
"content": "Alice completed the quarterly report.",
|
||
"context": "work update",
|
||
"entities": [
|
||
{"text": "ProjectX", "type": "PROJECT"},
|
||
{"text": "ACME Corp", "type": "ORG"},
|
||
{"text": "Alice"}, # May also be extracted by LLM (dedup test)
|
||
],
|
||
}
|
||
]
|
||
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Flatten the list of lists
|
||
unit_ids = [uid for sublist in result for uid in sublist]
|
||
assert len(unit_ids) > 0, "Should have created at least one fact"
|
||
|
||
logger.info(f"Created {len(unit_ids)} facts with user-provided entities")
|
||
|
||
# Query entity links to verify user-provided entities were stored
|
||
async with memory._pool.acquire() as conn:
|
||
# Get all entities linked to our facts via the unit_entities junction table
|
||
entity_rows = await conn.fetch(
|
||
"""
|
||
SELECT DISTINCT e.canonical_name
|
||
FROM entities e
|
||
JOIN unit_entities ue ON e.id = ue.entity_id
|
||
WHERE ue.unit_id::text = ANY($1)
|
||
""",
|
||
unit_ids
|
||
)
|
||
|
||
entity_names = {row['canonical_name'].lower() for row in entity_rows}
|
||
logger.info(f"Found entities linked to facts: {[row['canonical_name'] for row in entity_rows]}")
|
||
|
||
# Verify user-provided entities are present
|
||
assert "projectx" in entity_names, "User-provided entity 'ProjectX' should be linked"
|
||
assert "acme corp" in entity_names, "User-provided entity 'ACME Corp' should be linked"
|
||
|
||
# Alice should be present (either from LLM extraction or user-provided)
|
||
assert "alice" in entity_names, "Entity 'Alice' should be linked"
|
||
|
||
logger.info("✓ User-provided entities successfully merged with extracted entities")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
def test_recall_result_model_empty_construction():
|
||
"""
|
||
Test that RecallResultModel can be constructed with empty results.
|
||
|
||
This is a regression test for the bug where constructing an empty RecallResultModel
|
||
would cause an UnboundLocalError because RecallResult was imported as RecallResultModel
|
||
but the code mistakenly used the wrong name.
|
||
|
||
The fix ensures RecallResultModel is used consistently throughout memory_engine.py.
|
||
"""
|
||
from hindsight_api.engine.response_models import RecallResult
|
||
|
||
# This should not raise any errors
|
||
result = RecallResult(results=[], entities={}, chunks={})
|
||
|
||
assert result is not None, "Should create a result object"
|
||
assert result.results == [], "Should have empty results"
|
||
assert result.entities == {}, "Should have empty entities"
|
||
assert result.chunks == {}, "Should have empty chunks"
|
||
|
||
logger.info("✓ RecallResult empty construction works correctly")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_custom_extraction_mode():
|
||
"""
|
||
Test that custom extraction mode uses custom guidelines from env variable.
|
||
|
||
This test verifies that when HINDSIGHT_API_RETAIN_EXTRACTION_MODE=custom and
|
||
HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS is set, the fact extraction uses the
|
||
custom guidelines while keeping structural parts intact.
|
||
"""
|
||
import os
|
||
|
||
from hindsight_api import LLMConfig
|
||
from hindsight_api.config import _get_raw_config, clear_config_cache
|
||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||
|
||
# Save original env vars
|
||
original_mode = os.getenv("HINDSIGHT_API_RETAIN_EXTRACTION_MODE")
|
||
original_instructions = os.getenv("HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS")
|
||
|
||
try:
|
||
# Set custom extraction mode with challenging language-specific guidelines
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = "custom"
|
||
os.environ["HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"] = """ONLY extract facts that are in ITALIAN language.
|
||
|
||
DO NOT extract:
|
||
❌ Facts in English
|
||
❌ Facts in any other language besides Italian
|
||
|
||
If the text contains both Italian and English content, extract ONLY the Italian facts."""
|
||
|
||
# Clear config cache to pick up new env vars
|
||
clear_config_cache()
|
||
|
||
# Test content with BOTH Italian (should extract) and English (should NOT extract) facts
|
||
# This is a much harder test than filtering greetings
|
||
text = """
|
||
The team discussed the new architecture. We will use microservices.
|
||
|
||
Il database PostgreSQL ha ridotto la latenza delle query del 60%.
|
||
Alice ha suggerito di usare il connection pooling per migliorare le prestazioni.
|
||
|
||
Bob mentioned that the API endpoint is ready for testing.
|
||
The deployment pipeline has been updated to use Kubernetes.
|
||
|
||
Marco ha completato la revisione del codice e ha approvato le modifiche.
|
||
Il sistema di autenticazione è stato migrato a OAuth 2.0.
|
||
"""
|
||
|
||
llm_config = LLMConfig.from_env()
|
||
|
||
facts, _, _ = await extract_facts_from_text(
|
||
text=text,
|
||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||
context="team meeting notes",
|
||
llm_config=llm_config,
|
||
agent_name="TestUser",
|
||
config=_get_raw_config(),
|
||
)
|
||
|
||
logger.info(f"\nExtracted {len(facts)} facts with custom mode (Italian only):")
|
||
for i, fact in enumerate(facts):
|
||
logger.info(f" {i+1}. {fact.fact}")
|
||
|
||
assert len(facts) > 0, "Should extract at least one Italian fact"
|
||
|
||
# All facts text
|
||
all_facts_text = " ".join([f.fact for f in facts])
|
||
|
||
# Should HAVE Italian content
|
||
italian_keywords = ["postgresql", "latenza", "query", "alice", "connection pooling", "prestazioni",
|
||
"marco", "revisione", "codice", "autenticazione", "oauth"]
|
||
has_italian = any(keyword in all_facts_text.lower() for keyword in italian_keywords)
|
||
assert has_italian, f"Should extract Italian facts. Got: {all_facts_text}"
|
||
|
||
# Should NOT have English-only content
|
||
# These are facts that appear ONLY in English sections
|
||
english_only_keywords = ["microservices", "bob", "api endpoint", "testing", "deployment pipeline", "kubernetes"]
|
||
|
||
# Check if facts contain English-only content (this would be wrong)
|
||
facts_lower = all_facts_text.lower()
|
||
found_english_only = [kw for kw in english_only_keywords if kw in facts_lower]
|
||
|
||
if found_english_only:
|
||
logger.warning(f"⚠ Found English-only keywords in facts: {found_english_only}")
|
||
logger.warning(f" Facts: {all_facts_text}")
|
||
logger.warning(" This may indicate the LLM is not strictly following language-specific custom guidelines")
|
||
# Log but don't fail - LLM behavior can vary
|
||
else:
|
||
logger.info("✓ Successfully extracted only Italian facts, ignored English facts")
|
||
|
||
# At least verify we have some Italian indicators
|
||
italian_indicators = ["latenza", "prestazioni", "revisione", "codice", "autenticazione"]
|
||
italian_count = sum(1 for ind in italian_indicators if ind in facts_lower)
|
||
|
||
assert italian_count >= 1, \
|
||
f"Should extract facts with Italian words. Found {italian_count} Italian indicators in: {all_facts_text}"
|
||
|
||
logger.info("✓ Custom extraction mode works with language-specific guidelines")
|
||
logger.info(f"✓ Extracted {len(facts)} Italian facts, found {italian_count} Italian indicators")
|
||
|
||
finally:
|
||
# Restore original env vars
|
||
if original_mode is not None:
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = original_mode
|
||
else:
|
||
os.environ.pop("HINDSIGHT_API_RETAIN_EXTRACTION_MODE", None)
|
||
|
||
if original_instructions is not None:
|
||
os.environ["HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"] = original_instructions
|
||
else:
|
||
os.environ.pop("HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS", None)
|
||
|
||
# Clear cache again to restore original config
|
||
clear_config_cache()
|
||
|
||
|
||
def test_apply_strategy():
|
||
"""
|
||
Unit test for apply_strategy:
|
||
- Known strategy applies overrides on top of resolved config
|
||
- Unknown strategy returns config unchanged with a warning
|
||
- Non-hierarchical fields in a strategy are silently ignored
|
||
- entity_labels and entities_allow_free_form are overridable
|
||
"""
|
||
from hindsight_api.config import _get_raw_config, clear_config_cache
|
||
from hindsight_api.config_resolver import apply_strategy
|
||
|
||
clear_config_cache()
|
||
base_config = _get_raw_config()
|
||
|
||
strategies = {
|
||
"documents": {
|
||
"retain_extraction_mode": "chunks",
|
||
"retain_chunk_size": 800,
|
||
"entities_allow_free_form": False,
|
||
},
|
||
"bad_field": {
|
||
"database_url": "should-be-ignored", # static field, not hierarchical
|
||
"retain_extraction_mode": "verbose",
|
||
},
|
||
}
|
||
config_with_strategies = base_config.__class__(
|
||
**{**base_config.__dict__, "retain_strategies": strategies}
|
||
)
|
||
|
||
# Known strategy: overrides applied
|
||
result = apply_strategy(config_with_strategies, "documents")
|
||
assert result.retain_extraction_mode == "chunks"
|
||
assert result.retain_chunk_size == 800
|
||
assert result.entities_allow_free_form is False
|
||
|
||
# Non-hierarchical field silently ignored, hierarchical one applied
|
||
result2 = apply_strategy(config_with_strategies, "bad_field")
|
||
assert result2.retain_extraction_mode == "verbose"
|
||
assert result2.database_url == base_config.database_url # unchanged
|
||
|
||
# Unknown strategy: config returned unchanged
|
||
result3 = apply_strategy(config_with_strategies, "nonexistent")
|
||
assert result3.retain_extraction_mode == base_config.retain_extraction_mode
|
||
|
||
|
||
def test_collapse_to_verbatim_single_fact_per_chunk():
|
||
"""
|
||
Unit test for _collapse_to_verbatim:
|
||
- One fact per chunk → text overridden with original chunk text
|
||
- Two facts from same chunk → collapsed to one, entities merged
|
||
"""
|
||
from hindsight_api.engine.retain.fact_extraction import _collapse_to_verbatim
|
||
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact
|
||
|
||
chunks = [
|
||
ChunkMetadata(chunk_text="Alice went to Paris.", fact_count=1, content_index=0, chunk_index=0),
|
||
ChunkMetadata(chunk_text="Bob fixed the bug yesterday.", fact_count=2, content_index=0, chunk_index=1),
|
||
]
|
||
|
||
facts = [
|
||
ExtractedFact(fact_text="LLM paraphrase of Alice in Paris", fact_type="world", entities=["Alice", "Paris"], chunk_index=0, content_index=0),
|
||
ExtractedFact(fact_text="LLM first fact about Bob", fact_type="world", entities=["Bob"], chunk_index=1, content_index=0),
|
||
ExtractedFact(fact_text="LLM second fact about bug", fact_type="world", entities=["bug"], chunk_index=1, content_index=0),
|
||
]
|
||
|
||
result = _collapse_to_verbatim(facts, chunks)
|
||
|
||
assert len(result) == 2, "Should produce exactly one fact per chunk"
|
||
|
||
# Chunk 0: text overridden with original chunk text
|
||
assert result[0].fact_text == "Alice went to Paris.", "Text must be the raw chunk text"
|
||
assert result[0].entities == ["Alice", "Paris"]
|
||
|
||
# Chunk 1: collapsed to one fact, entities merged from both LLM facts
|
||
assert result[1].fact_text == "Bob fixed the bug yesterday.", "Text must be the raw chunk text"
|
||
assert "Bob" in result[1].entities
|
||
assert "bug" in result[1].entities
|
||
|
||
|
||
def test_chunks_extraction_mode():
|
||
"""
|
||
Unit test for chunks mode: no LLM, chunks stored as-is, zero token usage.
|
||
"""
|
||
import asyncio
|
||
import os
|
||
|
||
from hindsight_api.config import _get_raw_config, clear_config_cache
|
||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_contents
|
||
from hindsight_api.engine.retain.types import RetainContent
|
||
|
||
original_mode = os.getenv("HINDSIGHT_API_RETAIN_EXTRACTION_MODE")
|
||
|
||
try:
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = "chunks"
|
||
clear_config_cache()
|
||
|
||
contents = [
|
||
RetainContent(
|
||
content="Alice joined the infrastructure team on March 5, 2024.",
|
||
event_date=datetime(2024, 3, 10, tzinfo=timezone.utc),
|
||
entities=[{"text": "Alice"}, {"text": "infrastructure team"}],
|
||
),
|
||
RetainContent(content="Bob fixed the critical bug in the payment service."),
|
||
]
|
||
|
||
facts, chunks, usage = asyncio.get_event_loop().run_until_complete(
|
||
extract_facts_from_contents(
|
||
contents=contents,
|
||
llm_config=None, # Must not be called
|
||
agent_name="TestAgent",
|
||
config=_get_raw_config(),
|
||
)
|
||
)
|
||
|
||
# One fact per chunk (both contents fit in one chunk each)
|
||
assert len(facts) == len(chunks) == 2
|
||
|
||
# Text preserved exactly
|
||
assert facts[0].fact_text == contents[0].content
|
||
assert facts[1].fact_text == contents[1].content
|
||
|
||
# No LLM-extracted entities (user-provided entities handled downstream)
|
||
assert facts[0].entities == []
|
||
assert facts[1].entities == []
|
||
|
||
# Zero token usage
|
||
assert usage.total_tokens == 0
|
||
|
||
logger.info("✓ chunks mode: no LLM call, chunks stored as-is, zero token usage")
|
||
|
||
finally:
|
||
if original_mode is not None:
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = original_mode
|
||
else:
|
||
os.environ.pop("HINDSIGHT_API_RETAIN_EXTRACTION_MODE", None)
|
||
clear_config_cache()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_verbatim_extraction_mode():
|
||
"""
|
||
Integration test for verbatim extraction mode.
|
||
|
||
Verifies that:
|
||
1. Each chunk produces exactly one fact
|
||
2. The fact text is the original chunk text, not a paraphrase
|
||
3. Entities are still extracted by the LLM
|
||
4. Temporal info (occurred_start) is still extracted
|
||
"""
|
||
import os
|
||
|
||
from hindsight_api import LLMConfig
|
||
from hindsight_api.config import _get_raw_config, clear_config_cache
|
||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_contents
|
||
from hindsight_api.engine.retain.types import RetainContent
|
||
|
||
original_mode = os.getenv("HINDSIGHT_API_RETAIN_EXTRACTION_MODE")
|
||
|
||
try:
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = "verbatim"
|
||
clear_config_cache()
|
||
|
||
text = (
|
||
"Alice joined the infrastructure team on March 5, 2024. "
|
||
"She holds a CKA certification and has 5 years of Kubernetes experience."
|
||
)
|
||
|
||
llm_config = LLMConfig.from_env()
|
||
contents = [RetainContent(content=text, event_date=datetime(2024, 3, 10, tzinfo=timezone.utc), context="onboarding notes")]
|
||
facts, chunks, _ = await extract_facts_from_contents(
|
||
contents=contents,
|
||
llm_config=llm_config,
|
||
agent_name="TestAgent",
|
||
config=_get_raw_config(),
|
||
)
|
||
|
||
logger.info(f"Verbatim mode extracted {len(facts)} facts from {len(chunks)} chunks")
|
||
for i, f in enumerate(facts):
|
||
logger.info(f" fact[{i}]: {f.fact_text!r} entities={f.entities}")
|
||
|
||
# One fact per chunk
|
||
assert len(facts) == len(chunks), "Verbatim mode must produce exactly one fact per chunk"
|
||
|
||
# Text must match the original chunk exactly
|
||
for fact, chunk in zip(facts, chunks):
|
||
assert fact.fact_text == chunk.chunk_text, (
|
||
f"fact_text must equal original chunk text.\n"
|
||
f" expected: {chunk.chunk_text!r}\n"
|
||
f" got: {fact.fact_text!r}"
|
||
)
|
||
|
||
# Entities should still be extracted
|
||
all_entities = [e for f in facts for e in f.entities]
|
||
assert any("alice" in e.lower() for e in all_entities), (
|
||
f"Expected entity 'Alice' to be extracted. Entities: {all_entities}"
|
||
)
|
||
|
||
logger.info("✓ Verbatim mode preserves chunk text and still extracts entities")
|
||
|
||
finally:
|
||
if original_mode is not None:
|
||
os.environ["HINDSIGHT_API_RETAIN_EXTRACTION_MODE"] = original_mode
|
||
else:
|
||
os.environ.pop("HINDSIGHT_API_RETAIN_EXTRACTION_MODE", None)
|
||
clear_config_cache()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_retain_batch_with_per_item_tags_on_document(memory, request_context):
|
||
"""
|
||
Test that per-item tags are correctly stored on documents.
|
||
|
||
This test verifies the fix for a bug where per-item tags in content dictionaries
|
||
were not being merged and passed to document tracking, causing tags to be lost
|
||
even though they were correctly sent through the API.
|
||
|
||
Without the fix, this test would fail because:
|
||
- Tags are correctly passed in the content dict
|
||
- Tags are correctly stored on memory_units (facts)
|
||
- BUT tags were NOT stored on the document record itself
|
||
"""
|
||
bank_id = f"test_doc_tags_{datetime.now(timezone.utc).timestamp()}"
|
||
document_id = "app-state-testuser"
|
||
|
||
try:
|
||
# Retain content with per-item tags (simulating the TasteAI use case)
|
||
contents = [
|
||
{
|
||
"content": '{"username":"testuser","meals":[],"preferences":{"nickname":"testuser"}}',
|
||
"document_id": document_id,
|
||
"tags": ["user:testuser", "app-type:taste-ai"],
|
||
}
|
||
]
|
||
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=contents,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert len(result) > 0, "Should have retained content"
|
||
print("\n=== Retained content with tags ===")
|
||
|
||
# Retrieve the document
|
||
doc = await memory.get_document(
|
||
document_id=document_id,
|
||
bank_id=bank_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
assert doc is not None, "Document should exist"
|
||
assert "tags" in doc, "Document should have tags field"
|
||
|
||
# This is the critical assertion - tags should be stored on the document
|
||
doc_tags = doc["tags"] or []
|
||
print(f"Document tags: {doc_tags}")
|
||
|
||
assert "user:testuser" in doc_tags, \
|
||
f"Document should have 'user:testuser' tag, but got: {doc_tags}"
|
||
assert "app-type:taste-ai" in doc_tags, \
|
||
f"Document should have 'app-type:taste-ai' tag, but got: {doc_tags}"
|
||
|
||
print("✓ Per-item tags correctly stored on document")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||
|
||
|
||
def test_retain_mission_injected_into_prompt():
|
||
"""Test that retain_mission is injected as a FOCUS section into any extraction mode."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema
|
||
|
||
spec = "Focus on technical decisions and architecture choices only."
|
||
|
||
# Test with concise mode
|
||
config = MagicMock()
|
||
config.retain_extraction_mode = "concise"
|
||
config.retain_mission = spec
|
||
config.retain_custom_instructions = None
|
||
config.retain_extract_causal_links = False
|
||
|
||
prompt, _ = _build_extraction_prompt_and_schema(config)
|
||
assert spec in prompt
|
||
assert "FOCUS" in prompt
|
||
|
||
# retain_mission is injected into verbose mode as well
|
||
config.retain_extraction_mode = "verbose"
|
||
prompt_verbose, _ = _build_extraction_prompt_and_schema(config)
|
||
assert spec in prompt_verbose
|
||
assert "FOCUS" in prompt_verbose
|
||
|
||
|
||
def test_retain_mission_absent_when_not_set():
|
||
"""Test that no FOCUS section appears when retain_mission is not set."""
|
||
from unittest.mock import MagicMock
|
||
|
||
from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema
|
||
|
||
config = MagicMock()
|
||
config.retain_extraction_mode = "concise"
|
||
config.retain_mission = None
|
||
config.retain_custom_instructions = None
|
||
config.retain_extract_causal_links = False
|
||
|
||
prompt, _ = _build_extraction_prompt_and_schema(config)
|
||
assert "FOCUS" not in prompt
|
||
assert "retain_mission_section" not in prompt
|
||
|
||
|
||
def test_retain_mission_config_loaded_from_env():
|
||
"""Test that retain_mission is loaded from env and is a configurable field."""
|
||
import os
|
||
|
||
from hindsight_api.config import HindsightConfig, _get_raw_config, clear_config_cache
|
||
|
||
original = os.getenv("HINDSIGHT_API_RETAIN_MISSION")
|
||
try:
|
||
os.environ["HINDSIGHT_API_RETAIN_MISSION"] = "Only technical decisions."
|
||
clear_config_cache()
|
||
config = _get_raw_config()
|
||
assert config.retain_mission == "Only technical decisions."
|
||
assert "retain_mission" in HindsightConfig.get_configurable_fields()
|
||
finally:
|
||
if original is None:
|
||
os.environ.pop("HINDSIGHT_API_RETAIN_MISSION", None)
|
||
else:
|
||
os.environ["HINDSIGHT_API_RETAIN_MISSION"] = original
|
||
clear_config_cache()
|
||
|
||
|
||
def test_strategy_overrides_extraction_mode_for_chunks():
|
||
"""
|
||
Unit test: a named strategy with retain_extraction_mode=chunks causes
|
||
extract_facts_from_contents to skip the LLM and return verbatim chunks.
|
||
"""
|
||
import asyncio
|
||
|
||
from hindsight_api.config import _get_raw_config, clear_config_cache
|
||
from hindsight_api.config_resolver import apply_strategy
|
||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_contents
|
||
from hindsight_api.engine.retain.types import RetainContent
|
||
|
||
clear_config_cache()
|
||
base_config = _get_raw_config()
|
||
|
||
# Build a config that has a strategy overriding to chunks
|
||
strategies = {"fast": {"retain_extraction_mode": "chunks"}}
|
||
config_with_strategies = base_config.__class__(
|
||
**{**base_config.__dict__, "retain_strategies": strategies}
|
||
)
|
||
strategy_config = apply_strategy(config_with_strategies, "fast")
|
||
assert strategy_config.retain_extraction_mode == "chunks"
|
||
|
||
contents = [
|
||
RetainContent(content="Alice deployed the new API on Monday."),
|
||
RetainContent(content="Bob reviewed the pull request."),
|
||
]
|
||
|
||
facts, chunks, usage = asyncio.get_event_loop().run_until_complete(
|
||
extract_facts_from_contents(
|
||
contents=contents,
|
||
llm_config=None, # chunks must not call the LLM
|
||
agent_name="TestAgent",
|
||
config=strategy_config,
|
||
)
|
||
)
|
||
|
||
assert len(facts) == 2
|
||
assert facts[0].fact_text == contents[0].content
|
||
assert facts[1].fact_text == contents[1].content
|
||
assert usage.total_tokens == 0
|
||
logger.info("✓ strategy with chunks mode: no LLM, verbatim chunks, zero tokens")
|
||
|
||
|
||
def test_retain_request_per_item_strategy_field():
|
||
"""
|
||
Unit test: MemoryItem accepts a strategy field; items with different strategies
|
||
are grouped correctly by per-item strategy.
|
||
"""
|
||
from hindsight_api.api.http import RetainRequest
|
||
|
||
request = RetainRequest.model_validate(
|
||
{
|
||
"items": [
|
||
{"content": "Alice joined.", "strategy": "fast"},
|
||
{"content": "Bob left.", "strategy": "detailed"},
|
||
{"content": "Carol arrived."}, # no strategy — falls back to bank default
|
||
],
|
||
}
|
||
)
|
||
|
||
assert request.items[0].strategy == "fast"
|
||
assert request.items[1].strategy == "detailed"
|
||
assert request.items[2].strategy is None
|
||
|
||
# Simulate grouping logic from api_retain handler
|
||
strategy_groups: dict = {}
|
||
for item in request.items:
|
||
strategy_groups.setdefault(item.strategy, []).append(item.content)
|
||
|
||
assert set(strategy_groups.keys()) == {"fast", "detailed", None}
|
||
assert strategy_groups["fast"] == ["Alice joined."]
|
||
assert strategy_groups["detailed"] == ["Bob left."]
|
||
assert strategy_groups[None] == ["Carol arrived."]
|
||
logger.info("✓ per-item strategy grouping works correctly")
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_named_strategy_applied_end_to_end(memory, request_context):
|
||
"""
|
||
Integration test: a named strategy stored in bank config is actually applied
|
||
during retain_batch_async.
|
||
|
||
Regression test for the bug where strategy was passed through the HTTP layer
|
||
but the extraction mode override was silently ignored, always using the bank
|
||
default (e.g. 'concise') instead of the strategy's override (e.g. 'chunks').
|
||
"""
|
||
from hindsight_api.config_resolver import ConfigResolver
|
||
|
||
bank_id = f"test_strategy_e2e_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Seed the bank so the row exists before we write config to it
|
||
# (update_bank_config is a plain UPDATE — it silently no-ops on missing rows)
|
||
await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[{"content": "seed"}],
|
||
request_context=request_context,
|
||
)
|
||
|
||
# Now configure the bank with a named strategy that overrides to chunks
|
||
await memory._config_resolver.update_bank_config(
|
||
bank_id,
|
||
{
|
||
"retain_extraction_mode": "concise", # bank default
|
||
"retain_strategies": {
|
||
"chunks": {"retain_extraction_mode": "chunks"},
|
||
},
|
||
},
|
||
request_context,
|
||
)
|
||
|
||
contents = [{"content": "Alice deployed the new API on Monday."}]
|
||
|
||
# Retain using the named strategy
|
||
unit_ids_by_content, usage = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=contents,
|
||
strategy="chunks",
|
||
request_context=request_context,
|
||
return_usage=True,
|
||
)
|
||
|
||
# chunks produces exactly one fact per chunk (verbatim) and calls no LLM
|
||
assert usage.total_tokens == 0, f"chunks should use zero LLM tokens, got {usage.total_tokens}"
|
||
assert len(unit_ids_by_content) == 1
|
||
assert len(unit_ids_by_content[0]) == 1, "chunks should produce exactly one fact per content item"
|
||
|
||
# Verify the stored fact is the verbatim content
|
||
facts = await memory.recall_async(bank_id, "Alice", request_context=request_context)
|
||
assert any("Alice" in f.text for f in facts.results), "Verbatim content should be retrievable"
|
||
|
||
logger.info("✓ named strategy 'chunks' with chunks applied end-to-end: no LLM, verbatim storage")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_semantic_ann_uses_hnsw_index(memory, request_context):
|
||
"""
|
||
Test that Phase 1 ANN semantic search creates links between similar world
|
||
facts across batches. This exercises the per-fact_type partial HNSW index
|
||
and the placeholder-ID remap logic.
|
||
"""
|
||
bank_id = f"test_sem_ann_hnsw_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
# Batch 1: world facts about machine learning
|
||
unit_ids_1 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=(
|
||
"Deep learning models require large amounts of training data. "
|
||
"Gradient descent is the primary optimization algorithm used in neural networks."
|
||
),
|
||
context="ML knowledge base",
|
||
event_date=datetime(2024, 3, 1, tzinfo=timezone.utc),
|
||
request_context=request_context,
|
||
)
|
||
assert len(unit_ids_1) > 0, "Batch 1 should produce facts"
|
||
|
||
# Batch 2: similar ML world facts — Phase 1 ANN should link to batch 1
|
||
unit_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content=(
|
||
"Neural networks learn by adjusting weights through backpropagation. "
|
||
"Training deep learning models requires GPUs for fast gradient computation."
|
||
),
|
||
context="ML knowledge base",
|
||
event_date=datetime(2024, 3, 2, tzinfo=timezone.utc),
|
||
request_context=request_context,
|
||
)
|
||
assert len(unit_ids_2) > 0, "Batch 2 should produce facts"
|
||
|
||
logger.info(f"Batch 1: {len(unit_ids_1)} facts, Batch 2: {len(unit_ids_2)} facts")
|
||
|
||
# Verify cross-batch semantic links exist
|
||
async with memory._pool.acquire() as conn:
|
||
cross_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND link_type = 'semantic'
|
||
AND to_unit_id::text = ANY($2)
|
||
""",
|
||
unit_ids_2,
|
||
unit_ids_1,
|
||
)
|
||
|
||
logger.info(f"Cross-batch semantic links (batch2 -> batch1): {len(cross_links)}")
|
||
for link in cross_links:
|
||
logger.info(
|
||
f" {str(link['from_unit_id'])[:8]}... -> "
|
||
f"{str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})"
|
||
)
|
||
|
||
assert len(cross_links) > 0, (
|
||
"Phase 1 ANN should create semantic links between similar world facts "
|
||
"from different batches via the HNSW index with placeholder-ID remap."
|
||
)
|
||
|
||
# All weights must meet the similarity threshold
|
||
for link in cross_links:
|
||
assert link["weight"] >= 0.7, (
|
||
f"Semantic link weight {link['weight']:.3f} below threshold 0.7"
|
||
)
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_temporal_links_scoped_by_fact_type(memory, request_context):
|
||
"""
|
||
Test that temporal links only connect facts of the SAME fact_type.
|
||
|
||
World facts should not get temporal links to experience facts even when
|
||
their event dates fall within the time window.
|
||
"""
|
||
bank_id = f"test_temporal_scope_{datetime.now(timezone.utc).timestamp()}"
|
||
|
||
try:
|
||
base_date = datetime(2024, 5, 10, 12, 0, 0, tzinfo=timezone.utc)
|
||
|
||
# Store a world fact
|
||
world_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="Python 3.12 was released with significant performance improvements for the interpreter.",
|
||
context="tech news",
|
||
event_date=base_date,
|
||
fact_type_override="world",
|
||
request_context=request_context,
|
||
)
|
||
assert len(world_ids) > 0, "Should create world fact(s)"
|
||
|
||
# Store an experience fact at a nearby timestamp (same hour)
|
||
experience_ids = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="I upgraded all my projects to Python 3.12 and benchmarked the speed improvements.",
|
||
context="personal log",
|
||
event_date=base_date + timedelta(hours=1),
|
||
fact_type_override="experience",
|
||
request_context=request_context,
|
||
)
|
||
assert len(experience_ids) > 0, "Should create experience fact(s)"
|
||
|
||
# Store another world fact at a nearby timestamp so we can confirm
|
||
# same-type temporal links ARE created
|
||
world_ids_2 = await memory.retain_async(
|
||
bank_id=bank_id,
|
||
content="The Python Software Foundation announced long-term support plans for Python 3.12.",
|
||
context="tech news",
|
||
event_date=base_date + timedelta(hours=2),
|
||
fact_type_override="world",
|
||
request_context=request_context,
|
||
)
|
||
assert len(world_ids_2) > 0, "Should create second world fact(s)"
|
||
|
||
logger.info(
|
||
f"World1: {world_ids}, Experience: {experience_ids}, World2: {world_ids_2}"
|
||
)
|
||
|
||
async with memory._pool.acquire() as conn:
|
||
# Check that world facts DO have temporal links to each other
|
||
world_all = world_ids + world_ids_2
|
||
world_temporal = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE from_unit_id::text = ANY($1)
|
||
AND to_unit_id::text = ANY($1)
|
||
AND link_type = 'temporal'
|
||
""",
|
||
world_all,
|
||
)
|
||
logger.info(f"World-to-world temporal links: {len(world_temporal)}")
|
||
assert len(world_temporal) > 0, (
|
||
"World facts with nearby dates should have temporal links to each other"
|
||
)
|
||
|
||
# Check that world facts do NOT have temporal links to experience facts
|
||
cross_type_links = await conn.fetch(
|
||
"""
|
||
SELECT from_unit_id, to_unit_id, weight
|
||
FROM memory_links
|
||
WHERE (
|
||
(from_unit_id::text = ANY($1) AND to_unit_id::text = ANY($2))
|
||
OR
|
||
(from_unit_id::text = ANY($2) AND to_unit_id::text = ANY($1))
|
||
)
|
||
AND link_type = 'temporal'
|
||
""",
|
||
world_all,
|
||
experience_ids,
|
||
)
|
||
logger.info(f"Cross-type temporal links (world<->experience): {len(cross_type_links)}")
|
||
assert len(cross_type_links) == 0, (
|
||
f"Temporal links should NOT cross fact types, but found {len(cross_type_links)} "
|
||
f"world<->experience links"
|
||
)
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Streaming chunk batching tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
import json
|
||
import uuid
|
||
from unittest.mock import patch
|
||
|
||
import pytest_asyncio
|
||
|
||
from hindsight_api.engine.llm_wrapper import TokenUsage
|
||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||
from hindsight_api.engine.task_backend import SyncTaskBackend
|
||
|
||
|
||
def _make_mock_llm_call():
|
||
"""Create a mock LLM call function that returns deterministic facts."""
|
||
|
||
async def mock_llm_call(*args, **kwargs):
|
||
from hindsight_api.engine.consolidation.consolidator import _ConsolidationBatchResponse
|
||
|
||
if kwargs.get("scope") == "consolidation":
|
||
return_usage = kwargs.get("return_usage", False)
|
||
if return_usage:
|
||
return _ConsolidationBatchResponse(), TokenUsage(input_tokens=0, output_tokens=0)
|
||
return _ConsolidationBatchResponse()
|
||
|
||
messages = kwargs.get("messages", args[0] if args else [])
|
||
user_msg = messages[-1]["content"] if messages else ""
|
||
|
||
# Extract sentences from the content to generate one fact per sentence
|
||
sentences = [s.strip() for s in user_msg.split(".") if s.strip() and len(s.strip()) > 10]
|
||
num_facts = max(1, min(len(sentences), 10))
|
||
|
||
facts = []
|
||
for i in range(num_facts):
|
||
sentence = sentences[i] if i < len(sentences) else f"Fact {i}"
|
||
facts.append({
|
||
"what": sentence[:200],
|
||
"when": "2024-06-15",
|
||
"where": "N/A",
|
||
"who": "N/A",
|
||
"why": "N/A",
|
||
"fact_type": "world",
|
||
"entities": [{"text": f"Entity{i}"}],
|
||
"causal_relations": [],
|
||
})
|
||
|
||
response_dict = {"facts": facts}
|
||
return_usage = kwargs.get("return_usage", False)
|
||
if return_usage:
|
||
usage = TokenUsage(
|
||
input_tokens=len(user_msg) // 4,
|
||
output_tokens=len(json.dumps(response_dict)) // 4,
|
||
)
|
||
return response_dict, usage
|
||
return response_dict
|
||
|
||
return mock_llm_call
|
||
|
||
|
||
@pytest_asyncio.fixture(scope="function")
|
||
async def memory_mock_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||
"""MemoryEngine with mock LLM for streaming tests."""
|
||
mem = MemoryEngine(
|
||
db_url=pg0_db_url,
|
||
memory_llm_provider="openai",
|
||
memory_llm_api_key="mock-key",
|
||
memory_llm_model="gpt-4",
|
||
embeddings=embeddings,
|
||
cross_encoder=cross_encoder,
|
||
query_analyzer=query_analyzer,
|
||
pool_min_size=1,
|
||
pool_max_size=5,
|
||
run_migrations=False,
|
||
skip_llm_verification=True,
|
||
task_backend=SyncTaskBackend(),
|
||
)
|
||
await mem.initialize()
|
||
yield mem
|
||
try:
|
||
if mem._pool and not mem._pool._closing:
|
||
await mem.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _generate_chunky_content(num_chunks: int, chunk_size: int = 3000) -> str:
|
||
"""Generate content that will produce approximately num_chunks chunks.
|
||
|
||
Each chunk is chunk_size characters, separated by double newlines.
|
||
"""
|
||
base_sentences = [
|
||
"Alice works as a senior engineer at TechCorp in San Francisco.",
|
||
"Bob joined the marketing team last month from Chicago.",
|
||
"The project deadline was extended to December 15th.",
|
||
"Sarah mentioned she is planning a trip to Tokyo next month.",
|
||
"The quarterly budget review showed a 15% increase in revenue.",
|
||
"Mike suggested exploring alternative cloud providers.",
|
||
"The client feedback from beta testing was positive overall.",
|
||
"Emily started learning Rust programming language last week.",
|
||
"The new office will be located in the financial district.",
|
||
"David presented the annual technology roadmap to stakeholders.",
|
||
]
|
||
|
||
chunks = []
|
||
for chunk_idx in range(num_chunks):
|
||
# Generate enough text for one chunk
|
||
lines = []
|
||
chars = 0
|
||
line_idx = 0
|
||
while chars < chunk_size - 100:
|
||
sentence = f"[Chunk {chunk_idx}, Line {line_idx}] {base_sentences[line_idx % len(base_sentences)]}"
|
||
lines.append(sentence)
|
||
chars += len(sentence) + 1
|
||
line_idx += 1
|
||
chunks.append("\n".join(lines))
|
||
|
||
return "\n\n".join(chunks)
|
||
|
||
|
||
def _set_chunk_batch_size(memory: MemoryEngine, batch_size: int) -> None:
|
||
"""Set retain_chunk_batch_size on the config resolver's global config."""
|
||
memory._config_resolver._global_config.retain_chunk_batch_size = batch_size
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_chunk_batching_produces_same_facts(memory_mock_llm, request_context):
|
||
"""
|
||
Retain a medium document (~10 chunks) with batch_size=3.
|
||
Verify all facts are extracted (streaming should not lose facts).
|
||
"""
|
||
memory = memory_mock_llm
|
||
_set_chunk_batch_size(memory, 3)
|
||
bank_id = f"test_streaming_{uuid.uuid4().hex[:8]}"
|
||
document_id = f"streaming_doc_{uuid.uuid4().hex[:8]}"
|
||
|
||
# Generate content that produces ~10 chunks at default chunk_size (3000 chars)
|
||
content = _generate_chunky_content(num_chunks=10, chunk_size=3000)
|
||
|
||
mock_llm_call = _make_mock_llm_call()
|
||
|
||
try:
|
||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call):
|
||
# Retain with streaming enabled (batch_size=3, so 10 chunks -> 4 mini-batches)
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[{
|
||
"content": content,
|
||
"context": "streaming test",
|
||
"event_date": datetime(2024, 6, 15, tzinfo=timezone.utc),
|
||
}],
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
streaming_unit_ids = result[0] if result else []
|
||
logger.info(f"Streaming produced {len(streaming_unit_ids)} facts")
|
||
assert len(streaming_unit_ids) > 0, "Streaming should produce facts"
|
||
|
||
# Verify facts are in the DB
|
||
async with memory._pool.acquire() as conn:
|
||
fact_count = await conn.fetchval(
|
||
"SELECT COUNT(*) FROM memory_units WHERE bank_id = $1",
|
||
bank_id,
|
||
)
|
||
assert fact_count == len(streaming_unit_ids), (
|
||
f"DB has {fact_count} facts, but streaming returned {len(streaming_unit_ids)} unit_ids"
|
||
)
|
||
|
||
# Verify the document was tracked
|
||
doc = await conn.fetchrow(
|
||
"SELECT id FROM documents WHERE bank_id = $1 AND id = $2",
|
||
bank_id, document_id,
|
||
)
|
||
assert doc is not None, "Document should be tracked in DB"
|
||
|
||
# Verify chunks were stored with correct indices
|
||
chunk_count = await conn.fetchval(
|
||
"SELECT COUNT(*) FROM chunks WHERE bank_id = $1 AND document_id = $2",
|
||
bank_id, document_id,
|
||
)
|
||
assert chunk_count > 0, "Chunks should be stored in DB"
|
||
logger.info(f"Stored {chunk_count} chunks for document {document_id}")
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_chunk_batching_recovery(memory_mock_llm, request_context):
|
||
"""
|
||
Test recovery: retain a document with streaming, then retain the same
|
||
document again. Delta retain should detect existing chunks and skip
|
||
re-extraction. Fact count should be unchanged (no duplicates).
|
||
"""
|
||
memory = memory_mock_llm
|
||
_set_chunk_batch_size(memory, 3)
|
||
bank_id = f"test_streaming_recovery_{uuid.uuid4().hex[:8]}"
|
||
document_id = f"recovery_doc_{uuid.uuid4().hex[:8]}"
|
||
|
||
content = _generate_chunky_content(num_chunks=9, chunk_size=3000)
|
||
|
||
mock_llm_call = _make_mock_llm_call()
|
||
|
||
try:
|
||
# First retain — streaming mode
|
||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call):
|
||
result1 = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[{
|
||
"content": content,
|
||
"context": "recovery test",
|
||
"event_date": datetime(2024, 6, 15, tzinfo=timezone.utc),
|
||
}],
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
first_unit_ids = result1[0] if result1 else []
|
||
assert len(first_unit_ids) > 0, "First retain should produce facts"
|
||
|
||
async with memory._pool.acquire() as conn:
|
||
first_fact_count = await conn.fetchval(
|
||
"SELECT COUNT(*) FROM memory_units WHERE bank_id = $1",
|
||
bank_id,
|
||
)
|
||
|
||
logger.info(f"First retain: {first_fact_count} facts")
|
||
|
||
# Second retain — same document, same content (should be a no-op via delta retain)
|
||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call):
|
||
result2 = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[{
|
||
"content": content,
|
||
"context": "recovery test",
|
||
"event_date": datetime(2024, 6, 15, tzinfo=timezone.utc),
|
||
}],
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
async with memory._pool.acquire() as conn:
|
||
second_fact_count = await conn.fetchval(
|
||
"SELECT COUNT(*) FROM memory_units WHERE bank_id = $1",
|
||
bank_id,
|
||
)
|
||
|
||
logger.info(f"Second retain: {second_fact_count} facts")
|
||
|
||
# Fact count should be the same (delta retain skipped unchanged chunks)
|
||
assert second_fact_count == first_fact_count, (
|
||
f"Second retain should not create duplicates: first={first_fact_count}, second={second_fact_count}"
|
||
)
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_disabled_for_small_docs(memory_mock_llm, request_context):
|
||
"""
|
||
Retain a small document (2 chunks) with batch_size=500.
|
||
Verify it uses the non-streaming path (no batching overhead).
|
||
"""
|
||
memory = memory_mock_llm
|
||
_set_chunk_batch_size(memory, 500)
|
||
bank_id = f"test_streaming_small_{uuid.uuid4().hex[:8]}"
|
||
document_id = f"small_doc_{uuid.uuid4().hex[:8]}"
|
||
|
||
# Generate content that produces ~2 chunks
|
||
content = _generate_chunky_content(num_chunks=2, chunk_size=3000)
|
||
|
||
mock_llm_call = _make_mock_llm_call()
|
||
|
||
try:
|
||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call):
|
||
# batch_size=500 >> 2 chunks, so non-streaming path should be used
|
||
result = await memory.retain_batch_async(
|
||
bank_id=bank_id,
|
||
contents=[{
|
||
"content": content,
|
||
"context": "small doc test",
|
||
"event_date": datetime(2024, 6, 15, tzinfo=timezone.utc),
|
||
}],
|
||
document_id=document_id,
|
||
request_context=request_context,
|
||
)
|
||
|
||
unit_ids = result[0] if result else []
|
||
logger.info(f"Small doc produced {len(unit_ids)} facts")
|
||
assert len(unit_ids) > 0, "Should produce facts even through non-streaming path"
|
||
|
||
# Verify the document was tracked
|
||
async with memory._pool.acquire() as conn:
|
||
doc = await conn.fetchrow(
|
||
"SELECT id FROM documents WHERE bank_id = $1 AND id = $2",
|
||
bank_id, document_id,
|
||
)
|
||
assert doc is not None, "Document should be tracked"
|
||
|
||
finally:
|
||
await memory.delete_bank(bank_id, request_context=request_context)
|