7e339e1677
* feat: ai sdk integration * more fixes * fix(security): mental model refresh tag-based security - Mental model refresh now passes tags with all_strict matching - Consolidation only triggers refresh for mental models with matching tags - Consolidation filters related observations by tags (all_strict) - Added tests to verify tag-based security boundaries - Updated OpenAPI spec to include tags and text_preview in list_documents - Added tags column to documents UI table * chore: regenerate OpenAPI spec after rebase * fix: improve consolidation prompt for contradiction handling and mental model refresh security - Enhanced consolidation prompt to be more explicit about capturing temporal changes in contradictions - Fixed mental model refresh security: tagged memories now only trigger refresh of mental models with matching tags - Added stricter tag filtering to prevent cross-scope mental model refreshes Fixes test_consolidation_merges_contradictions by improving LLM instructions to use temporal markers like "used to X, now Y" when merging contradictory facts. Note: test_refresh_with_tags_only_accesses_same_tagged_models still needs investigation - REFLECT operation may need additional tag filtering. * fix: mental model refresh security - proper tag filtering in search Fixed tool_search_mental_models to properly handle all_strict tag matching mode by using the centralized build_tags_where_clause function. Previously, the function only handled "all" vs "any" modes and always included untagged mental models when using non-"all" modes. This ensures that when a tagged mental model is refreshed with all_strict matching, it cannot access untagged mental models, preventing cross-scope information leakage. Fixes test_refresh_with_tags_only_accesses_same_tagged_models. Note: test_sensory_dimension_preservation is failing but this is a pre-existing issue on main branch - the LLM model (gpt-oss-20b) is not extracting facts from sensory text. Not related to security changes. * chore: apply formatting from pre-commit hook * fix: allow untagged mental models to be refreshed by any consolidation Untagged mental models are considered "global" and should be refreshed by any consolidation, regardless of whether tagged or untagged memories were consolidated. This maintains security boundaries while allowing global mental models to stay fresh. When tagged memories are consolidated: - Refresh mental models with matching tags (security boundary) - Also refresh untagged mental models (they're global) - DO NOT refresh mental models with different tags When untagged memories are consolidated: - Only refresh untagged mental models - DO NOT refresh tagged mental models (security boundary) Fixes test_consolidation_only_refreshes_matching_tagged_models.
2257 lines
90 KiB
Python
2257 lines
90 KiB
Python
"""
|
|
Test retain function and chunk storage.
|
|
"""
|
|
import pytest
|
|
import logging
|
|
from datetime import datetime, timezone, timedelta
|
|
from hindsight_api.engine.memory_engine import Budget
|
|
from hindsight_api import RequestContext
|
|
|
|
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(f"\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(f"\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
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="When did Alice complete the product launch?",
|
|
budget=Budget.LOW,
|
|
max_tokens=500,
|
|
fact_type=["world"],
|
|
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(f"✓ 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(f" This test expects None for present-tense observations")
|
|
else:
|
|
print(f"✓ 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(f" This test expects None for present-tense observations")
|
|
else:
|
|
print(f"✓ 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(f"✓ 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
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="What does Alice like?",
|
|
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]
|
|
|
|
# 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(f"✓ mentioned_at is always set (never None)")
|
|
|
|
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
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="What did the team decide?",
|
|
budget=Budget.LOW,
|
|
max_tokens=500,
|
|
fact_type=["world"],
|
|
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 is preserved.
|
|
Metadata allows arbitrary key-value data to be stored with facts.
|
|
"""
|
|
bank_id = f"test_metadata_{datetime.now(timezone.utc).timestamp()}"
|
|
|
|
try:
|
|
# Store content with custom metadata
|
|
custom_metadata = {
|
|
"source": "slack",
|
|
"channel": "engineering",
|
|
"importance": "high",
|
|
"tags": "product,launch"
|
|
}
|
|
|
|
# Note: retain_async doesn't directly support metadata parameter
|
|
# Metadata would need to be supported in the API layer
|
|
# For now, we test that the system handles content without errors
|
|
unit_ids = await memory.retain_async(
|
|
bank_id=bank_id,
|
|
content="The product launch is scheduled for March 1st.",
|
|
context="planning meeting",
|
|
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
|
request_context=request_context,
|
|
)
|
|
|
|
assert len(unit_ids) > 0, "Should create memory units"
|
|
|
|
# Recall to verify storage worked
|
|
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"
|
|
|
|
print(f"✓ Successfully stored and retrieved facts")
|
|
print(f" (Note: Metadata support depends on API implementation)")
|
|
|
|
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(f"✓ 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
|
|
result = await memory.recall_async(
|
|
bank_id=bank_id,
|
|
query="What is the project status?",
|
|
budget=Budget.MID,
|
|
max_tokens=1000,
|
|
fact_type=["world"],
|
|
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(f" 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_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.engine.retain.fact_extraction import extract_facts_from_text
|
|
from hindsight_api.config import clear_config_cache
|
|
|
|
# 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.for_memory()
|
|
|
|
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"
|
|
)
|
|
|
|
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(f" 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()
|
|
|
|
|
|
@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(f"\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} ===")
|