8. Самописный поисковый агент на основе deep agents from scratch: client.py

solutions/69de7223f309a98be0007e09_8__Самописный_поисковый_агент_на_основе_/client.py

@@ -1,73 +1,98 @@
+import os
 import sys
 from typing import List
 
-from langchain_ollama import OllamaEmbeddings, OllamaLLM
+# Импортируем необходимые компоненты из LangChain и связанных библиотек
-from langchain_qdrant import QdrantVectorStore
+from langchain_ollama import ChatOllama, OllamaEmbeddings
-from qdrant_client import QdrantClient
+from langchain_qdrant import QdrantStore
-from qdrant_client.http.models import Distance, VectorParams
+from langchain_core.documents import Document
+from langchain_core.prompts import PromptTemplate
+from langchain.chains import RetrievalQA
 
-# Конфигурация подключения к Qdrant
+# Конфигурация подключения к Qdrant (используем локальный сервер по умолчанию)
-QDRANT_HOST = "localhost"
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
-QDRANT_PORT = 6333
 COLLECTION_NAME = "documents"
 
-# Инициализация клиента Qdrant
+# Инициализируем Ollama embeddings и LLM
-qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
-
-# Создание коллекции при необходимости
-if COLLECTION_NAME not in qdrant_client.get_collections().collections:
-    qdrant_client.create_collection(
-        collection_name=COLLECTION_NAME,
-        vectors_config=VectorParams(size=768, distance=Distance.COSINE),
-    )
-
-# Инициализация эмбеддинговой модели (nomic-embed-text)
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
+llm = ChatOllama(model="llama3")
 
-# Создание векторного хранилища
+# Создаём или подключаемся к коллекции в Qdrant
-vector_store = QdrantVectorStore(
+vector_store = QdrantStore(
-    client=qdrant_client,
+    url=QDRANT_URL,
     collection_name=COLLECTION_NAME,
     embeddings=embeddings,
 )
 
-# Инициализация LLM (llama3)
+# Шаблон запроса для RetrievalQA
-llm = OllamaLLM(model="llama3")
+prompt_template = PromptTemplate.from_template(
+    """
+Ты – интеллектуальный помощник. На основе предоставленного контекста ответь на вопрос пользователя.
+
+Контекст:
+{context}
+
+Вопрос: {question}
+Ответ:
+"""
+)
+
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    chain_type="stuff",
+    retriever=vector_store.as_retriever(),
+    return_source_documents=True,
+    question_generator=None,
+    combine_docs_chain_kwargs={"prompt": prompt_template},
+)
+
 
 def add_document(text: str) -> None:
-    """Добавляет документ в хранилище."""
+    """
-    vector_store.add_texts([text])
+    Добавляет новый документ в Qdrant.
-    print("✅ Документ добавлен.")
+    """
+    doc = Document(page_content=text)
+    vector_store.add_documents([doc])
+    print("Документ добавлен.")
 
-def search_documents(query: str, k: int = 3) -> List[str]:
-    """Ищет похожие документы и возвращает их тексты."""
-    results = vector_store.similarity_search_with_score(query, k=k)
-    docs = [doc.page_content for doc, _ in results]
-    return docs
 
-def answer_query(query: str) -> str:
+def search(query: str) -> None:
-    """Получает ответ LLM на основе найденных документов."""
+    """
-    docs = search_documents(query)
+    Выполняет поиск по запросу и выводит ответ LLM.
-    context = "\n\n".join(docs)
+    """
-    prompt = f"Контекст:\n{context}\n\nВопрос: {query}\nОтвет:"
+    result = qa_chain({"question": query})
-    return llm.invoke(prompt)
+    answer = result["answer"]
+    sources = result.get("source_documents", [])
+    print("\nОтвет:\n" + answer)
+    if sources:
+        print("\nИсточники:")
+        for i, src in enumerate(sources, 1):
+            print(f"{i}. {src.page_content[:200]}...")
+
 
 def print_help() -> None:
-    """Показывает список команд."""
+    """
+    Выводит справку по доступным командам.
+    """
     help_text = """
-Команды:
+Доступные команды:
-  /add <текст>      - добавить документ
+  /add <текст>     - Добавить новый документ
-  /search <запрос>  - поиск по запросу и вывод ответа LLM
+  /search <запрос> - Выполнить поиск и получить ответ
-  /quit             - выйти из программы
+  /quit            - Выйти из программы
+  /help            - Показать эту справку
 """
     print(help_text)
 
+
 def main() -> None:
+    """
+    Основной цикл CLI.
+    """
     print("=== Самописный поисковый агент ===")
     print_help()
     while True:
         try:
-            user_input = input("> ").strip()
+            user_input = input("\n> ").strip()
         except (EOFError, KeyboardInterrupt):
             print("\nВыход.")
             break
@@ -75,27 +100,36 @@ def main() -> None:
         if not user_input:
             continue
 
-        if user_input.startswith("/add "):
+        if user_input.startswith("/"):
-            text = user_input[5:].strip()
+            parts = user_input.split(maxsplit=1)
-            if text:
+            cmd = parts[0].lower()
-                add_document(text)
+            arg = parts[1] if len(parts) > 1 else ""
+
+            if cmd == "/add":
+                if not arg:
+                    print("Ошибка: необходимо указать текст для добавления.")
                 else:
-                print("⚠️  Текст не указан.")
+                    add_document(arg)
-        elif user_input.startswith("/search "):
+
-            query = user_input[8:].strip()
+            elif cmd == "/search":
-            if query:
+                if not arg:
-                answer = answer_query(query)
+                    print("Ошибка: необходимо указать поисковый запрос.")
-                print("\nОтвет LLM:")
-                print(answer)
                 else:
-                print("⚠️  Запрос не указан.")
+                    search(arg)
-        elif user_input == "/quit":
+
+            elif cmd == "/quit":
                 print("Выход.")
                 break
-        elif user_input == "/help":
+
+            elif cmd == "/help":
                 print_help()
+
             else:
-            print("❓ Неизвестная команда. Введите /help для списка команд.")
+                print(f"Неизвестная команда '{cmd}'. Введите /help для справки.")
+
+        else:
+            print("Команды начинаются с '/'. Введите /help для списка команд.")
+
 
 if __name__ == "__main__":
     main()