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

2026-05-27 10:49:57 +00:00
parent 3d822df232
commit bdca1f348d
1 changed files with 165 additions and 0 deletions
@@ -0,0 +1,165 @@
+import sys
+import logging
+from typing import List
+
+# LangChain imports
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain.schema import Document
+
+# Настройка логирования
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+logger = logging.getLogger(__name__)
+
+# Конфигурация Ollama и Qdrant
+EMBEDDING_MODEL = "nomic-embed-text"
+LLM_MODEL = "llama3"
+
+# Создаём объект для получения эмбеддингов
+embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+
+# Подключаемся к локальному экземпляру Qdrant (по умолчанию localhost:6333)
+vector_store = QdrantVectorStore(
+    embedding_function=embeddings,
+    url="http://localhost:6333",
+    collection_name="documents",
+)
+
+
+def middleware(func):
+    """
+    Декоратор для обработки ошибок и логирования.
+    """
+
+    def wrapper(*args, **kwargs):
+        try:
+            logger.info(f"Выполняется команда: {func.__name__}")
+            result = func(*args, **kwargs)
+            logger.info("Команда выполнена успешно")
+            return result
+        except Exception as e:
+            logger.exception(f"Ошибка в команде {func.__name__}: {e}")
+            print(f"[ERROR] {e}")
+
+    return wrapper
+
+
+@middleware
+def add_document(text: str) -> None:
+    """
+    Добавляет документ в Qdrant.
+    """
+    doc = Document(page_content=text)
+    vector_store.add_documents([doc])
+    print("Документ добавлен.")
+
+
+@middleware
+def search(query: str, k: int = 3) -> List[Document]:
+    """
+    Выполняет поиск по запросу и выводит результаты.
+    """
+    results = vector_store.similarity_search_with_score(query, k=k)
+    if not results:
+        print("Ничего не найдено.")
+        return []
+
+    for idx, (doc, score) in enumerate(results, start=1):
+        print(f"\nРезультат {idx} (score: {score:.4f})")
+        print(doc.page_content[:500] + ("..." if len(doc.page_content) > 500 else ""))
+    return [doc for doc, _ in results]
+
+
+def parse_command(line: str):
+    """
+    Разбирает строку команды и вызывает соответствующую функцию.
+    """
+    if not line.strip():
+        return
+
+    parts = line.split(maxsplit=1)
+    cmd = parts[0].lower()
+
+    if cmd == "/add":
+        if len(parts) < 2:
+            print("[ERROR] Необходимо указать текст для добавления.")
+            return
+        add_document(parts[1])
+    elif cmd == "/search":
+        if len(parts) < 2:
+            print("[ERROR] Необходимо указать запрос для поиска.")
+            return
+        search(parts[1])
+    elif cmd == "/quit":
+        print("Выход из программы.")
+        sys.exit(0)
+    else:
+        print(f"[ERROR] Нераспознанная команда: {cmd}")
+
+
+def main():
+    """
+    Основной цикл CLI.
+    """
+    logger.info("Запуск клиента Human-in-the-Loop")
+    print("Добро пожаловать! Используйте команды /add, /search и /quit.")
+    while True:
+        try:
+            line = input("> ")
+            parse_command(line)
+        except KeyboardInterrupt:
+            print("\nПрервано пользователем. Выход.")
+            break
+        except EOFError:
+            print("\nEOF. Выход.")
+            break
+
+
+# -------------------- UNIT TESTS --------------------
+import unittest
+from unittest.mock import patch, MagicMock
+
+class TestClient(unittest.TestCase):
+    def setUp(self):
+        # Подменяем vector_store для тестов
+        self.original_vector_store = vector_store
+        mock_vs = MagicMock()
+        mock_vs.add_documents.return_value = None
+        mock_vs.similarity_search_with_score.return_value = [
+            (Document(page_content="Test content 1"), 0.9),
+            (Document(page_content="Test content 2"), 0.8),
+        ]
+        global vector_store
+        vector_store = mock_vs
+
+    def tearDown(self):
+        global vector_store
+        vector_store = self.original_vector_store
+
+    @patch('builtins.print')
+    def test_add_document_success(self, mock_print):
+        add_document("Sample text")
+        mock_print.assert_called_with("Документ добавлен.")
+
+    @patch('builtins.print')
+    def test_search_results(self, mock_print):
+        results = search("query")
+        self.assertEqual(len(results), 2)
+        # Проверяем, что print был вызван хотя бы один раз
+        self.assertTrue(mock_print.call_count > 0)
+
+    @patch('builtins.print')
+    def test_parse_unknown_command(self, mock_print):
+        parse_command("/unknown")
+        mock_print.assert_called_with("[ERROR] Нераспознанная команда: /unknown")
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] == "test":
+        # Запуск тестов
+        unittest.main(argv=[sys.argv[0]])
+    else:
+        main()