Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Integration tests for Memory module. | |
| Tests make use of LLM models, so ensure that they are accessible. | |
| These tests verify the full end-to-end functionality of the memory system | |
| including Milvus vector store, SQLite metadata, and tips extraction. | |
| """ | |
| import unittest | |
| import time | |
| import os | |
| os.environ['DYNACONF_ADVANCED_FEATURES__ENABLE_MEMORY'] = 'true' | |
| from cuga.backend.memory.agentic_memory.main import app | |
| from cuga.backend.memory.memory import Memory | |
| from cuga.backend.memory.agentic_memory.schema import Namespace, RecordedFact, Run | |
| from cuga.backend.memory.agentic_memory.client.exceptions import NamespaceNotFoundException | |
| from cuga.config import DBS_DIR | |
| from fastapi.testclient import TestClient | |
| from pathlib import Path | |
| from unittest import mock | |
| test_env = { | |
| "WXO_MILVUS_URI": str(Path(DBS_DIR) / "milvus.db"), | |
| "MEM0_HISTORY_DB_PATH": str(Path(DBS_DIR) / "mem0_history.db"), | |
| "MILVUS_DB": str(Path(DBS_DIR) / "test_memory.db"), | |
| "AGENTIC_DB": str(Path(DBS_DIR) / "agentic.db"), | |
| "MODEL_NAME": "Azure/gpt-4o", | |
| } | |
| class TestMemoryOperations(unittest.TestCase): | |
| """Test memory CRUD operations against a test service.""" | |
| def setUp(self): | |
| """Set up test client""" | |
| Memory._instance = None | |
| Memory._initialized = False | |
| self.client = Memory() | |
| self.client.memory_client.client = TestClient(app) | |
| def tearDown(self): | |
| """Clean up created databases.""" | |
| Memory._instance = None | |
| Memory._initialized = False | |
| self.client.memory_client.client.close() | |
| Path(test_env["WXO_MILVUS_URI"]).unlink(missing_ok=True) | |
| Path(test_env["MEM0_HISTORY_DB_PATH"]).unlink(missing_ok=True) | |
| Path(test_env["MILVUS_DB"]).unlink(missing_ok=True) | |
| Path(test_env["AGENTIC_DB"]).unlink(missing_ok=True) | |
| def test_full(self): | |
| """End-to-end test of memory CRUD operations.""" | |
| namespace_id = f"test_namespace_{int(time.time() * 1000)}" | |
| # Create Namespace | |
| namespace = self.client.create_namespace( | |
| namespace_id=namespace_id, user_id="test_user", agent_id="cuga", app_id="test_app" | |
| ) | |
| # Verify namespace was created correctly | |
| self.assertIsInstance(namespace, Namespace) | |
| self.assertEqual(namespace.id, namespace_id) | |
| self.assertEqual(namespace.user_id, "test_user") | |
| self.assertEqual(namespace.agent_id, "cuga") | |
| self.assertEqual(namespace.app_id, "test_app") | |
| self.assertIsNotNone(namespace.created_at) | |
| # Get namespace details another way | |
| namespace = self.client.get_namespace_details(namespace_id) | |
| self.assertEqual(namespace.id, namespace_id) | |
| self.assertEqual(namespace.user_id, "test_user") | |
| self.assertEqual(namespace.agent_id, "cuga") | |
| self.assertEqual(namespace.app_id, "test_app") | |
| # Add fact | |
| self.client.create_and_store_fact( | |
| namespace_id, "Python is a programming language", {"topic": "programming", "user_id": "test_user"} | |
| ) | |
| self.client.create_and_store_fact( | |
| namespace_id, | |
| "JavaScript is used for web development", | |
| {"topic": "programming", "user_id": "test_user"}, | |
| ) | |
| self.client.create_and_store_fact( | |
| namespace_id, "Machine learning uses neural networks", {"topic": "AI", "user_id": "test_user"} | |
| ) | |
| facts = self.client.search_for_facts( | |
| namespace_id=namespace_id, | |
| query="programming languages", | |
| filters={"user_id": "test_user"}, | |
| limit=5, | |
| ) | |
| self.assertIsInstance(facts, list) | |
| self.assertGreater(len(facts), 0) | |
| self.assertIsInstance(facts[0], RecordedFact) | |
| # Create Run | |
| run_id = f"run_{int(time.time() * 1000)}" | |
| run = self.client.create_run(namespace_id, run_id=run_id) | |
| self.assertIsInstance(run, Run) | |
| self.assertEqual(run.id, run_id) | |
| self.assertFalse(run.ended) | |
| self.assertEqual(len(run.steps), 0) | |
| step_1 = self.client.add_step( | |
| namespace_id=namespace_id, | |
| run_id=run.id, | |
| step={"name": "TaskAnalyzerAgent", "status": "success", "action": "analyzed task"}, | |
| prompt="Output a JSON object without any extra thoughts or commentary. Add a summary field.", | |
| ) | |
| step_2 = self.client.add_step( | |
| namespace_id=namespace_id, | |
| run_id=run.id, | |
| step={"agent": "APIPlannerAgent", "status": "success", "action": "selected APIs"}, | |
| prompt="Output a JSON object without any extra thoughts or commentary. Add a summary field.", | |
| ) | |
| self.assertIsNotNone(step_1) | |
| self.assertIsNotNone(step_2) | |
| run = self.client.get_run(namespace_id, run.id) | |
| self.assertGreaterEqual(len(run.steps), 2) | |
| # End the run | |
| self.client.end_run(namespace_id, run.id) | |
| retrieved_run = self.client.get_run(namespace_id, run.id) | |
| self.assertTrue(retrieved_run.ended) | |
| # Delete run | |
| self.client.delete_run(namespace_id, run.id) | |
| with self.assertRaises(Exception): | |
| self.client.get_run(namespace_id, run.id) | |
| # Delete it | |
| self.client.delete_namespace(namespace.id) | |
| # Verify it's gone | |
| with self.assertRaises(NamespaceNotFoundException): | |
| self.client.get_namespace_details(namespace_id) | |
| if __name__ == "__main__": | |
| unittest.main() | |