Human-in-the-Loop через middleware: agent.py

solutions/69a86305c46fd26feae6bcaa_Human-in-the-Loop_через_middleware/agent.py

@@ -0,0 +1,136 @@
+import os
+import logging
+from typing import Any, Dict, List
+
+from langchain_ollama import OllamaEmbeddings, ChatOllama
+from langchain_qdrant import QdrantStore
+from langchain_core.documents import Document
+from langchain_core.runnables import RunnableLambda
+
+# Настройка логирования
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+# Конфигурация подключения к Qdrant
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
+COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "knowledge_base")
+
+# Инициализация эмбеддингов и модели Ollama
+try:
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+except Exception as e:
+    logger.exception("Не удалось подключиться к Ollama Embeddings: %s", e)
+    raise
+
+try:
+    llm = ChatOllama(model="llama3")
+except Exception as e:
+    logger.exception("Не удалось подключиться к Ollama LLM: %s", e)
+    raise
+
+# Инициализация хранилища Qdrant
+try:
+    qdrant_store = QdrantStore(
+        url=QDRANT_URL,
+        collection_name=COLLECTION_NAME,
+        embeddings=embeddings,
+    )
+except Exception as e:
+    logger.exception("Не удалось подключиться к Qdrant: %s", e)
+    raise
+
+
+# ==========================
+# Middleware для обработки ошибок и логирования
+# ==========================
+def middleware(func):
+    """
+    Декоратор, который оборачивает функцию в try/except,
+    логирует входные параметры и результат.
+    """
+
+    def wrapper(*args, **kwargs):
+        func_name = func.__name__
+        logger.debug("Вызов %s с args=%s kwargs=%s", func_name, args, kwargs)
+        try:
+            result = func(*args, **kwargs)
+            logger.debug("Результат %s: %s", func_name, result)
+            return result
+        except Exception as exc:
+            logger.exception("Ошибка в %s: %s", func_name, exc)
+            raise
+
+    return wrapper
+
+
+# ==========================
+# RAG инструменты
+# ==========================
+@middleware
+def search_knowledge_base(query: str) -> List[Document]:
+    """
+    Поиск релевантных документов в Qdrant.
+    Возвращает список Document с полями content и metadata.
+    """
+    if not query:
+        raise ValueError("Запрос не может быть пустым")
+
+    try:
+        docs = qdrant_store.similarity_search(query, k=5)
+        logger.info("Найдено %d документов по запросу '%s'", len(docs), query)
+        return docs
+    except Exception as exc:
+        logger.exception("Ошибка при поиске: %s", exc)
+        raise
+
+
+@middleware
+def add_to_knowledge_base(content: str, metadata: Dict[str, Any] | None = None) -> None:
+    """
+    Добавляет новый документ в Qdrant.
+    content – текст документа,
+    metadata – произвольные метаданные (может быть None).
+    """
+    if not content:
+        raise ValueError("Содержимое не может быть пустым")
+
+    try:
+        doc = Document(page_content=content, metadata=metadata or {})
+        qdrant_store.add_documents([doc])
+        logger.info("Документ добавлен в коллекцию '%s'", COLLECTION_NAME)
+    except Exception as exc:
+        logger.exception("Ошибка при добавлении документа: %s", exc)
+        raise
+
+
+# ==========================
+# Пример использования (можно использовать как тесты)
+# ==========================
+if __name__ == "__main__":
+    # Минимальный набор unit‑тестов
+    import unittest
+
+    class TestAgent(unittest.TestCase):
+        def test_add_and_search(self):
+            test_content = "Тестовый документ для проверки RAG."
+            test_metadata = {"source": "unit_test"}
+
+            # Добавляем документ
+            add_to_knowledge_base(test_content, test_metadata)
+
+            # Ищем по ключевому слову
+            results = search_knowledge_base("тестовый")
+            self.assertTrue(any(test_content in doc.page_content for doc in results))
+
+        def test_empty_query(self):
+            with self.assertRaises(ValueError):
+                search_knowledge_base("")
+
+        def test_empty_content(self):
+            with self.assertRaises(ValueError):
+                add_to_knowledge_base("")
+
+    unittest.main(argv=["first-arg-is-ignored"], exit=False)