2319 lines
93 KiB
Python
2319 lines
93 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, _get_raw_config
|
|
|
|
# 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",
|
|
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(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} ===")
|
|
|
|
|
|
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 present regardless of extraction mode (verbose has its own template, no spec injection)
|
|
config.retain_extraction_mode = "verbose"
|
|
prompt_verbose, _ = _build_extraction_prompt_and_schema(config)
|
|
# verbose uses its own template - spec not injected there
|
|
assert spec not 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()
|