feat: solution for 'Агент с RAG-памятью'

tests/test_agent.py

@@ -1,33 +1,99 @@
+"""
+Unit tests for the RAG agent and auto-check graph.
+"""
+
 import os
-import unittest
+import json
+import tempfile
+import shutil
+import pytest
 
-from src.agent import RAGAgent
+from src.index import RAGAgent, auto_check_graph
 
-class TestRAGAgent(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        # Ensure data directory exists with at least one document
-        data_dir = "data"
-        os.makedirs(data_dir, exist_ok=True)
-        sample_path = os.path.join(data_dir, "sample.txt")
-        with open(sample_path, "w", encoding="utf-8") as f:
-            f.write("Python is a versatile programming language used for web development, data science, and automation.")
-        cls.agent = RAGAgent(config_path="src/config.yaml")
+# Helper to create a temporary agent with fake embeddings/LLM
+def create_temp_agent(tmp_path):
+    # Ensure no OpenAI key
+    os.environ.pop("OPENAI_API_KEY", None)
+    agent = RAGAgent(
+        vector_store_path=tmp_path / "vector_store.faiss",
+        documents_dir=tmp_path / "docs",
+    )
+    return agent
 
-    def test_retrieve_non_empty(self):
-        passages = self.agent.kb.retrieve("Python programming", top_k=2)
-        self.assertTrue(len(passages) > 0)
-        self.assertIn("Python is a versatile programming language", passages[0][0])
+def test_add_and_query():
+    tmp_dir = tempfile.mkdtemp()
+    try:
+        agent = create_temp_agent(tmp_dir)
+        docs = [
+            "The capital of France is Paris.",
+            "William Shakespeare wrote Hamlet.",
+        ]
+        agent.add_documents(docs)
+        # Query for first doc
+        answer = agent.query("What is the capital of France?")
+        assert "Paris" in answer
+        # Query for second doc
+        answer2 = agent.query("Who wrote Hamlet?")
+        assert "Shakespeare" in answer2
+    finally:
+        shutil.rmtree(tmp_dir)
 
-    def test_generate_response(self):
-        answer = self.agent.generate_response("What is Python?")
-        self.assertIsInstance(answer, str)
-        self.assertTrue(len(answer) > 0)
+def test_auto_check_pass():
+    tmp_dir = tempfile.mkdtemp()
+    try:
+        agent = create_temp_agent(tmp_dir)
+        docs = [
+            "The capital of France is Paris.",
+            "William Shakespeare wrote Hamlet.",
+        ]
+        agent.add_documents(docs)
+        ground_truth = {
+            "What is the capital of France?": "Paris",
+            "Who wrote Hamlet?": "William Shakespeare",
+        }
+        result = auto_check_graph(
+            "What is the capital of France?", agent, ground_truth
+        )
+        assert result["verdict_row"] == "PASS"
+        assert "Paris" in result["answer"]
+    finally:
+        shutil.rmtree(tmp_dir)
 
-    def test_empty_query(self):
-        answer = self.agent.generate_response("")
-        self.assertIsInstance(answer, str)
-        self.assertIn("No relevant information found", answer)
+def test_auto_check_fail():
+    tmp_dir = tempfile.mkdtemp()
+    try:
+        agent = create_temp_agent(tmp_dir)
+        docs = [
+            "The capital of France is Paris.",
+        ]
+        agent.add_documents(docs)
+        ground_truth = {
+            "What is the capital of France?": "Berlin",
+        }
+        result = auto_check_graph(
+            "What is the capital of France?", agent, ground_truth
+        )
+        assert result["verdict_row"] == "FAIL"
+    finally:
+        shutil.rmtree(tmp_dir)
+
+def test_auto_check_unknown():
+    tmp_dir = tempfile.mkdtemp()
+    try:
+        agent = create_temp_agent(tmp_dir)
+        docs = [
+            "The capital of France is Paris.",
+        ]
+        agent.add_documents(docs)
+        ground_truth = {
+            "What is the capital of Germany?": "Berlin",
+        }
+        result = auto_check_graph(
+            "What is the capital of Germany?", agent, ground_truth
+        )
+        assert result["verdict_row"] == "UNKNOWN"
+    finally:
+        shutil.rmtree(tmp_dir)
 
 if __name__ == "__main__":
-    unittest.main()
+    pytest.main([__file__])