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

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

@@ -1,65 +1,53 @@
 import os
-from typing import List, Dict
+from typing import List
 
 from langchain_ollama import OllamaEmbeddings, OllamaLLM
-from langchain_qdrant import QdrantStore
-from langchain_core.tools import tool
-from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain_qdrant import QdrantVectorStore
 from langchain.schema import Document
+from langchain.tools import tool
 
 # Конфигурация подключения к Qdrant
 QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
-QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "knowledge_base")
+COLLECTION_NAME = "knowledge_base"
 
-# Инициализация эмбеддингов и LLM через Ollama
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
-llm = OllamaLLM(model="llama3")
+# Инициализируем Ollama LLM и эмбеддеры
+llm = OllamaLLM(model="llama3")  # модель для генерации ответов
+embeddings = OllamaEmbeddings(model="nomic-embed-text")  # модель для векторизации
 
-# Создание или подключение к коллекции Qdrant
-vector_store = QdrantStore(
+# Создаём или подключаемся к коллекции Qdrant
+vector_store = QdrantVectorStore(
     url=QDRANT_URL,
-    collection_name=QDRANT_COLLECTION,
+    collection_name=COLLECTION_NAME,
     embeddings=embeddings,
 )
 
 @tool("search_knowledge_base")
-def search_knowledge_base(query: str) -> List[Dict[str, str]]:
+def search_knowledge_base(query: str) -> List[Document]:
     """
-    Поиск релевантных документов в базе знаний по запросу.
-    Возвращает список словарей с полями 'content' и 'metadata'.
+    Поиск в базе знаний по запросу.
+    Возвращает список документов, отсортированных по релевантности.
     """
-    results = vector_store.similarity_search_with_score(query, k=3)
-    output = []
-    for doc, score in results:
-        output.append({"content": doc.page_content, "score": str(score), **doc.metadata})
-    return output
+    # Получаем топ‑5 похожих документов
+    results = vector_store.similarity_search_with_score(query, k=5)
+    # Функция возвращает только объекты Document (без scores)
+    return [doc for doc, _ in results]
 
 @tool("add_to_knowledge_base")
-def add_to_knowledge_base(text: str) -> str:
+def add_to_knowledge_base(doc: str) -> str:
     """
-    Добавляет новый текст в базу знаний.
+    Добавляет новый документ в базу знаний.
+    Принимает строку с текстом и сохраняет её как один Document.
     Возвращает подтверждение добавления.
     """
-    # Создаём документ с метаданными (можно расширить)
-    doc = Document(page_content=text, metadata={"source": "user_input"})
-    vector_store.add_documents([doc])
-    return f"Текст успешно добавлен в базу знаний. {len(text.split())} слов."
-
-# Определяем список инструментов
-tools = [search_knowledge_base, add_to_knowledge_base]
-
-# Создаём агента с использованием LLM и инструментов
-agent = create_openai_tools_agent(llm=llm, tools=tools)
-
-# Экземпляр исполнителя агента
-executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+    document = Document(page_content=doc)
+    vector_store.add_documents([document])
+    return "Документ успешно добавлен в базу знаний."
 
+# Если нужно протестировать работу инструментов
 if __name__ == "__main__":
-    # Пример взаимодействия: пользователь задаёт вопрос, агент решает использовать инструменты
-    while True:
-        user_input = input("\nВведите запрос (или 'exit' для выхода): ")
-        if user_input.lower() in {"exit", "quit"}:
-            break
-        response = executor.invoke({"input": user_input})
-        print("\nОтвет агента:")
-        print(response["output"])
+    # Пример добавления документа
+    print(add_to_knowledge_base("Python – это язык программирования, созданный Гвидо ван Россумом."))
+    # Пример поиска
+    results = search_knowledge_base("что такое Python?")
+    for i, doc in enumerate(results, 1):
+        print(f"Результат {i}: {doc.page_content[:200]}...")