Spaces:
Running
Running
AI Assistant
feat: add multi-document vectorless RAG, session isolation, and memory services
0039ed6 | 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() | |