feat(02-01): add two-layer memory system — Redis sliding window + pgvector long-term
- ConversationEmbedding ORM model with Vector(384) column (pgvector) - memory_short_key, escalation_status_key, pending_tool_confirm_key in redis_keys.py - orchestrator/memory/short_term.py: RPUSH/LTRIM sliding window (get_recent_messages, append_message) - orchestrator/memory/long_term.py: pgvector HNSW cosine search (retrieve_relevant, store_embedding) - Migration 002: conversation_embeddings table, HNSW index, RLS with FORCE, SELECT/INSERT only - 10 unit tests (fakeredis), 6 integration tests (pgvector) — all passing - Auto-fix [Rule 3]: postgres image updated to pgvector/pgvector:pg16 (extension required)
This commit is contained in:
259
tests/integration/test_memory_long_term.py
Normal file
259
tests/integration/test_memory_long_term.py
Normal file
@@ -0,0 +1,259 @@
|
||||
"""
|
||||
Integration tests for pgvector long-term memory.
|
||||
|
||||
Requires a live PostgreSQL instance with pgvector extension installed.
|
||||
Tests are automatically skipped if the database is not available
|
||||
(fixture from conftest.py handles that via pytest.skip).
|
||||
|
||||
Key scenarios tested:
|
||||
- store_embedding inserts with correct scoping
|
||||
- retrieve_relevant returns matching content above threshold
|
||||
- Cross-tenant isolation: tenant A's embeddings never returned for tenant B
|
||||
- High threshold returns empty list for dissimilar queries
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from orchestrator.memory.long_term import retrieve_relevant, store_embedding
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def agent_a_id() -> uuid.UUID:
|
||||
"""Return a stable agent UUID for tenant A tests."""
|
||||
return uuid.UUID("aaaaaaaa-0000-0000-0000-000000000001")
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def agent_b_id() -> uuid.UUID:
|
||||
"""Return a stable agent UUID for tenant B tests."""
|
||||
return uuid.UUID("bbbbbbbb-0000-0000-0000-000000000002")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_store_embedding_inserts_row(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
):
|
||||
"""store_embedding inserts a row into conversation_embeddings."""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
tenant_id = tenant_a["id"]
|
||||
user_id = "user-store-test"
|
||||
embedding = [0.1] * 384
|
||||
content = "I prefer concise answers."
|
||||
|
||||
token = current_tenant_id.set(tenant_id)
|
||||
try:
|
||||
await store_embedding(db_session, tenant_id, agent_a_id, user_id, content, "user", embedding)
|
||||
await db_session.commit()
|
||||
|
||||
result = await db_session.execute(
|
||||
text("SELECT content, role FROM conversation_embeddings WHERE tenant_id = :tid AND user_id = :uid"),
|
||||
{"tid": str(tenant_id), "uid": user_id},
|
||||
)
|
||||
rows = result.fetchall()
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0].content == content
|
||||
assert rows[0].role == "user"
|
||||
|
||||
|
||||
async def test_retrieve_relevant_returns_similar_content(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
):
|
||||
"""retrieve_relevant returns content above cosine similarity threshold."""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
tenant_id = tenant_a["id"]
|
||||
user_id = "user-retrieve-test"
|
||||
|
||||
# Store two embeddings: one very similar to the query, one dissimilar
|
||||
# We simulate similarity by using identical embeddings
|
||||
similar_embedding = [1.0] + [0.0] * 383
|
||||
dissimilar_embedding = [0.0] * 383 + [1.0]
|
||||
query_embedding = [1.0] + [0.0] * 383 # identical to similar_embedding
|
||||
|
||||
token = current_tenant_id.set(tenant_id)
|
||||
try:
|
||||
await store_embedding(
|
||||
db_session, tenant_id, agent_a_id, user_id,
|
||||
"The user likes Python programming.", "user", similar_embedding
|
||||
)
|
||||
await store_embedding(
|
||||
db_session, tenant_id, agent_a_id, user_id,
|
||||
"This is completely unrelated content.", "user", dissimilar_embedding
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
results = await retrieve_relevant(
|
||||
db_session, tenant_id, agent_a_id, user_id, query_embedding, top_k=3, threshold=0.5
|
||||
)
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
# Should return the similar content
|
||||
assert len(results) >= 1
|
||||
assert any("Python" in r for r in results)
|
||||
|
||||
|
||||
async def test_retrieve_relevant_high_threshold_returns_empty(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
):
|
||||
"""retrieve_relevant with threshold=0.99 and dissimilar query returns empty list."""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
tenant_id = tenant_a["id"]
|
||||
user_id = "user-threshold-test"
|
||||
|
||||
# Store an embedding pointing in one direction
|
||||
stored_embedding = [1.0] + [0.0] * 383
|
||||
# Query pointing in orthogonal direction — cosine distance ~= 1.0, similarity ~= 0.0
|
||||
query_embedding = [0.0] + [1.0] + [0.0] * 382
|
||||
|
||||
token = current_tenant_id.set(tenant_id)
|
||||
try:
|
||||
await store_embedding(
|
||||
db_session, tenant_id, agent_a_id, user_id,
|
||||
"Some stored content.", "user", stored_embedding
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
results = await retrieve_relevant(
|
||||
db_session, tenant_id, agent_a_id, user_id, query_embedding, top_k=3, threshold=0.99
|
||||
)
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
assert results == []
|
||||
|
||||
|
||||
async def test_cross_tenant_isolation(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
tenant_b: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
agent_b_id: uuid.UUID,
|
||||
):
|
||||
"""
|
||||
retrieve_relevant with tenant_id=A NEVER returns tenant_id=B embeddings.
|
||||
|
||||
This is the critical security test — cross-tenant contamination would be
|
||||
a catastrophic data leak in a multi-tenant system.
|
||||
"""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
user_id = "shared-user-id"
|
||||
tenant_a_id = tenant_a["id"]
|
||||
tenant_b_id = tenant_b["id"]
|
||||
|
||||
# Same query embedding for both tenants
|
||||
embedding = [1.0] + [0.0] * 383
|
||||
|
||||
# Store embedding for tenant B
|
||||
token = current_tenant_id.set(tenant_b_id)
|
||||
try:
|
||||
await store_embedding(
|
||||
db_session, tenant_b_id, agent_b_id, user_id,
|
||||
"Tenant B secret information.", "user", embedding
|
||||
)
|
||||
await db_session.commit()
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
# Query as tenant A — should NOT see tenant B's data
|
||||
token = current_tenant_id.set(tenant_a_id)
|
||||
try:
|
||||
results = await retrieve_relevant(
|
||||
db_session, tenant_a_id, agent_a_id, user_id, embedding, top_k=10, threshold=0.0
|
||||
)
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
# Tenant A should get nothing — it has no embeddings of its own
|
||||
# and it MUST NOT see tenant B's embeddings
|
||||
for result in results:
|
||||
assert "Tenant B" not in result, "Cross-tenant data leakage detected!"
|
||||
|
||||
|
||||
async def test_retrieve_relevant_user_isolation(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
):
|
||||
"""retrieve_relevant for user A never returns user B embeddings."""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
tenant_id = tenant_a["id"]
|
||||
embedding = [1.0] + [0.0] * 383
|
||||
|
||||
token = current_tenant_id.set(tenant_id)
|
||||
try:
|
||||
await store_embedding(
|
||||
db_session, tenant_id, agent_a_id, "user-A",
|
||||
"User A private information.", "user", embedding
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
# Query as user B — should not see user A's data
|
||||
results = await retrieve_relevant(
|
||||
db_session, tenant_id, agent_a_id, "user-B", embedding, top_k=10, threshold=0.0
|
||||
)
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
for result in results:
|
||||
assert "User A private" not in result
|
||||
|
||||
|
||||
async def test_retrieve_relevant_top_k_limits_results(
|
||||
db_session: AsyncSession,
|
||||
tenant_a: dict,
|
||||
agent_a_id: uuid.UUID,
|
||||
):
|
||||
"""retrieve_relevant respects top_k limit."""
|
||||
from shared.rls import current_tenant_id
|
||||
|
||||
tenant_id = tenant_a["id"]
|
||||
user_id = "user-topk-test"
|
||||
embedding = [1.0] + [0.0] * 383
|
||||
|
||||
token = current_tenant_id.set(tenant_id)
|
||||
try:
|
||||
# Store 5 very similar embeddings
|
||||
for i in range(5):
|
||||
await store_embedding(
|
||||
db_session, tenant_id, agent_a_id, user_id,
|
||||
f"Content item {i}", "user", embedding
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
results = await retrieve_relevant(
|
||||
db_session, tenant_id, agent_a_id, user_id, embedding, top_k=2, threshold=0.0
|
||||
)
|
||||
finally:
|
||||
current_tenant_id.reset(token)
|
||||
|
||||
assert len(results) <= 2
|
||||
139
tests/unit/test_memory_short_term.py
Normal file
139
tests/unit/test_memory_short_term.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""
|
||||
Unit tests for the Redis short-term memory sliding window.
|
||||
|
||||
Uses fakeredis to avoid requiring a real Redis connection.
|
||||
All tests verify tenant+agent+user namespacing and RPUSH/LTRIM correctness.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import fakeredis.aioredis
|
||||
import pytest
|
||||
|
||||
from orchestrator.memory.short_term import append_message, get_recent_messages
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def redis():
|
||||
"""Return a fakeredis async client for testing."""
|
||||
client = fakeredis.aioredis.FakeRedis()
|
||||
yield client
|
||||
await client.aclose()
|
||||
|
||||
|
||||
TENANT = "tenant-abc"
|
||||
AGENT = "agent-xyz"
|
||||
USER = "user-123"
|
||||
|
||||
|
||||
async def test_get_recent_messages_empty(redis):
|
||||
"""get_recent_messages on empty key returns empty list."""
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER)
|
||||
assert result == []
|
||||
|
||||
|
||||
async def test_append_and_get_single_message(redis):
|
||||
"""append_message stores a message; get_recent_messages retrieves it."""
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content="Hello!")
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER)
|
||||
assert len(result) == 1
|
||||
assert result[0] == {"role": "user", "content": "Hello!"}
|
||||
|
||||
|
||||
async def test_append_multiple_messages_ordering(redis):
|
||||
"""Messages are returned in insertion order (oldest first)."""
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content="First")
|
||||
await append_message(redis, TENANT, AGENT, USER, role="assistant", content="Second")
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content="Third")
|
||||
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER)
|
||||
assert len(result) == 3
|
||||
assert result[0]["content"] == "First"
|
||||
assert result[1]["content"] == "Second"
|
||||
assert result[2]["content"] == "Third"
|
||||
|
||||
|
||||
async def test_sliding_window_trims_to_window_size(redis):
|
||||
"""append_message with window=5 keeps only last 5 messages."""
|
||||
for i in range(10):
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content=f"msg-{i}", window=5)
|
||||
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER, n=20)
|
||||
assert len(result) == 5
|
||||
# Should have the last 5 messages: msg-5 through msg-9
|
||||
contents = [m["content"] for m in result]
|
||||
assert contents == ["msg-5", "msg-6", "msg-7", "msg-8", "msg-9"]
|
||||
|
||||
|
||||
async def test_default_window_20(redis):
|
||||
"""Default window is 20 — 21st message pushes out the first."""
|
||||
for i in range(21):
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content=f"msg-{i}")
|
||||
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER)
|
||||
assert len(result) == 20
|
||||
assert result[0]["content"] == "msg-1"
|
||||
assert result[-1]["content"] == "msg-20"
|
||||
|
||||
|
||||
async def test_get_recent_messages_n_parameter(redis):
|
||||
"""get_recent_messages n parameter limits results."""
|
||||
for i in range(10):
|
||||
await append_message(redis, TENANT, AGENT, USER, role="user", content=f"msg-{i}")
|
||||
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER, n=3)
|
||||
assert len(result) == 3
|
||||
# n=3 returns last 3: msg-7, msg-8, msg-9
|
||||
contents = [m["content"] for m in result]
|
||||
assert contents == ["msg-7", "msg-8", "msg-9"]
|
||||
|
||||
|
||||
async def test_key_namespacing_user_isolation(redis):
|
||||
"""Different users of the same agent have isolated memory."""
|
||||
await append_message(redis, TENANT, AGENT, "user-A", role="user", content="User A message")
|
||||
await append_message(redis, TENANT, AGENT, "user-B", role="user", content="User B message")
|
||||
|
||||
result_a = await get_recent_messages(redis, TENANT, AGENT, "user-A")
|
||||
result_b = await get_recent_messages(redis, TENANT, AGENT, "user-B")
|
||||
|
||||
assert len(result_a) == 1
|
||||
assert result_a[0]["content"] == "User A message"
|
||||
assert len(result_b) == 1
|
||||
assert result_b[0]["content"] == "User B message"
|
||||
|
||||
|
||||
async def test_key_namespacing_tenant_isolation(redis):
|
||||
"""Different tenants with same agent+user IDs have isolated memory."""
|
||||
await append_message(redis, "tenant-1", AGENT, USER, role="user", content="Tenant 1 message")
|
||||
await append_message(redis, "tenant-2", AGENT, USER, role="user", content="Tenant 2 message")
|
||||
|
||||
result_1 = await get_recent_messages(redis, "tenant-1", AGENT, USER)
|
||||
result_2 = await get_recent_messages(redis, "tenant-2", AGENT, USER)
|
||||
|
||||
assert result_1[0]["content"] == "Tenant 1 message"
|
||||
assert result_2[0]["content"] == "Tenant 2 message"
|
||||
|
||||
|
||||
async def test_key_namespacing_agent_isolation(redis):
|
||||
"""Different agents for the same tenant+user have isolated memory."""
|
||||
await append_message(redis, TENANT, "agent-1", USER, role="user", content="Agent 1 context")
|
||||
await append_message(redis, TENANT, "agent-2", USER, role="user", content="Agent 2 context")
|
||||
|
||||
result_1 = await get_recent_messages(redis, TENANT, "agent-1", USER)
|
||||
result_2 = await get_recent_messages(redis, TENANT, "agent-2", USER)
|
||||
|
||||
assert result_1[0]["content"] == "Agent 1 context"
|
||||
assert result_2[0]["content"] == "Agent 2 context"
|
||||
|
||||
|
||||
async def test_message_role_and_content_round_trip(redis):
|
||||
"""Messages store and retrieve role + content correctly."""
|
||||
await append_message(redis, TENANT, AGENT, USER, role="assistant", content="I can help you with that.")
|
||||
result = await get_recent_messages(redis, TENANT, AGENT, USER)
|
||||
msg = result[0]
|
||||
assert msg["role"] == "assistant"
|
||||
assert msg["content"] == "I can help you with that."
|
||||
# Verify it has exactly role and content keys
|
||||
assert set(msg.keys()) == {"role", "content"}
|
||||
Reference in New Issue
Block a user