Spaces:
Running
Running
File size: 5,244 Bytes
0039ed6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import os
import sys
# Add app to path
app_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "app"))
if app_dir not in sys.path:
sys.path.insert(0, app_dir)
from fastapi.testclient import TestClient
from app import app
from core.storage import storage
from core.persistent_db import persistent_db
client = TestClient(app)
def test_full_memory_and_multidoc_isolation():
print("\n--- 1. Testing Session Creation & Isolation ---")
res_a = client.post("/chat/new")
assert res_a.status_code == 200
chat_a = res_a.json()["chat_id"]
res_b = client.post("/chat/new")
assert res_b.status_code == 200
chat_b = res_b.json()["chat_b" if "chat_b" in res_b.json() else "chat_id"]
print(f"Chat A ID: {chat_a}")
print(f"Chat B ID: {chat_b}")
assert chat_a != chat_b
# Create dummy mock docs in storage
doc1_id = "mock-doc-1"
doc2_id = "mock-doc-2"
mock_tree_1 = [
{"title": "Section 1: Architecture", "page": 1, "text": "Atlas uses tree-based retrieval."},
{"title": "Section 2: Memory", "page": 2, "text": "Session memory isolates short and long context."},
]
mock_tree_2 = [
{"title": "Overview: Financials", "page": 1, "text": "Revenue grew by 24% in Q3."},
]
storage.set_document_ready(doc1_id, tree=mock_tree_1, filename="Architecture_Guide.pdf", page_count=2, section_count=2)
storage.set_document_ready(doc2_id, tree=mock_tree_2, filename="Financial_Report_Q3.pdf", page_count=1, section_count=1)
print("\n--- 2. Testing Multi-Document Attachment to Chat A ---")
attach_res = client.post(f"/chat/{chat_a}/attach-multiple", json={"doc_ids": [doc1_id, doc2_id]})
assert attach_res.status_code == 200
assert len(attach_res.json()["doc_ids"]) == 2
docs_a_res = client.get(f"/chat/{chat_a}/documents")
assert docs_a_res.status_code == 200
docs_a = docs_a_res.json()["documents"]
print(f"Chat A attached documents count: {len(docs_a)}")
assert len(docs_a) == 2
assert {d["doc_id"] for d in docs_a} == {doc1_id, doc2_id}
# Verify Chat B has zero attached documents (Strict Session Isolation)
docs_b_res = client.get(f"/chat/{chat_b}/documents")
assert docs_b_res.status_code == 200
docs_b = docs_b_res.json()["documents"]
print(f"Chat B attached documents count: {len(docs_b)}")
assert len(docs_b) == 0, "Violation of session isolation: Chat B should have 0 docs"
print("\n--- 3. Testing Message Isolation Between Sessions ---")
storage.add_message(chat_a, "user", "What is Atlas?")
storage.add_message(chat_a, "assistant", "Atlas is a document intelligence platform.")
storage.add_message(chat_b, "user", "Hello from user B")
storage.add_message(chat_b, "assistant", "Greetings, this is session B.")
hist_a = client.get(f"/chat/{chat_a}/history").json()["history"]
hist_b = client.get(f"/chat/{chat_b}/history").json()["history"]
assert len(hist_a) == 2
assert hist_a[0]["content"] == "What is Atlas?"
assert len(hist_b) == 2
assert hist_b[0]["content"] == "Hello from user B"
print("Session history isolation verified!")
print("\n--- 4. Testing Session Memory Diagnostic Endpoint ---")
mem_res = client.get(f"/memory/chat/{chat_a}")
assert mem_res.status_code == 200
mem_data = mem_res.json()
assert mem_data["chat_id"] == chat_a
assert mem_data["recent_turns_count"] == 2
print(f"Memory diagnostic for Chat A: {mem_data['recent_turns_count']} turns recorded.")
print("\n--- 5. Testing Detach Single Document ---")
detach_res = client.delete(f"/chat/{chat_a}/document/{doc1_id}")
assert detach_res.status_code == 200
remaining_docs = client.get(f"/chat/{chat_a}/documents").json()["documents"]
assert len(remaining_docs) == 1
assert remaining_docs[0]["doc_id"] == doc2_id
print(f"Remaining document in Chat A: {remaining_docs[0]['filename']}")
print("\n--- 6. Testing Global Knowledge & Dossier Storage ---")
# Save a test dossier
persistent_db.save_knowledge_dossier(
doc_id=doc1_id,
filename="Architecture_Guide.pdf",
executive_summary="Explains Atlas tree-based vectorless RAG architecture.",
key_entities=["Atlas", "PageIndex", "Groq", "FastAPI"],
key_findings=["Tree search outperforms embeddings on dense tabular texts."],
topics=["Architecture", "Document Intelligence", "FastAPI"]
)
persistent_db.save_global_knowledge_summary(
collective_summary="Atlas workspace contains technical documentation on tree retrieval and quarterly financial reports.",
total_documents=2,
key_insights=["Tree retrieval avoids embedding drift", "Multi-doc synthesis enables cross-sectional queries"]
)
g_res = client.get("/memory/knowledge")
assert g_res.status_code == 200
g_data = g_res.json()
print(f"Global Knowledge Summary: {g_data['collective_summary']}")
print(f"Dossiers count: {len(g_data['dossiers'])}")
assert len(g_data["dossiers"]) >= 1
assert "Atlas" in g_data["collective_summary"]
print("\n ALL BACKEND MULTI-DOC & MEMORY TESTS PASSED PERFECTLY!")
if __name__ == "__main__":
test_full_memory_and_multidoc_isolation()
|