Обновить solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/solution.py

@@ -5,17 +5,15 @@
 solution.py – Агент с RAG‑памятью на базе Qdrant и Ollama.
 """
 
-import os
-import sys
 import argparse
 from pathlib import Path
 from typing import List, Dict
 
-from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_ollama import OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from langchain.agents import create_agent, AgentExecutor, Tool
+from langchain.agents import create_agent
 from langchain.schema import Document
 
 # --------------------------------------------------------------------------- #
@@ -27,7 +25,7 @@ QDRANT_PORT = 6333
 COLLECTION_NAME = "knowledge_base"
 
 EMBEDDING_MODEL = "nomic-embed-text"   # Ollama embedding model
-LLM_MODEL = "llama3"                   # Ollama LLM
+LLM_MODEL = "ollama:llama3"            # Ollama LLM (формат для create_agent)
 
 # --------------------------------------------------------------------------- #
 # Векторное хранилище и вспомогательные функции
@@ -56,14 +54,14 @@ class KnowledgeBase:
         """
         Добавляет документ в базу после разбиения на чанки.
         """
-        # Разбиваем текст
         chunks = self.splitter.split_text(content)
         docs = [
-            Document(page_content=chunk,
+            Document(
-                     metadata={"title": title, "source": f"{title} (chunk {i})"})
+                page_content=chunk,
+                metadata={"title": title, "source": f"{title} (chunk {i})"}
+            )
             for i, chunk in enumerate(chunks)
         ]
-        # Добавляем в Qdrant
         self.store.add_documents(docs)
 
     def search(self, query: str, max_results: int = 5) -> List[Dict]:
@@ -74,24 +72,27 @@ class KnowledgeBase:
         results = self.store.similarity_search_with_score(query, k=max_results)
         output = []
         for doc, score in results:
-            out = {
+            output.append({
                 "score": score,
                 "title": doc.metadata.get("title", ""),
                 "source": doc.metadata.get("source", ""),
-                "content": doc.page_content
+                "content": doc.page_content,
-            }
+            })
-            output.append(out)
         return output
 
 
+# --------------------------------------------------------------------------- #
+# Глобальный экземпляр базы знаний
+# --------------------------------------------------------------------------- #
+
+kb = KnowledgeBase()
+
 # --------------------------------------------------------------------------- #
 # Инструменты для агента
 # --------------------------------------------------------------------------- #
 
-kb = KnowledgeBase()
-
 @tool
-def search_knowledge_base(query: str, max_results: int) -> str:
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
     """
     Поиск в базе знаний.
     Возвращает строку с найденными результатами в формате Markdown.
@@ -101,7 +102,9 @@ def search_knowledge_base(query: str, max_results: int) -> str:
         return "Ничего не найдено."
     lines = []
     for r in results:
-        lines.append(f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}...")
+        lines.append(
+            f"**{r['title']}** ({r['source']}) – {r['score']:.2f}\n> {r['content'][:200]}..."
+        )
     return "\n\n".join(lines)
 
 
@@ -118,36 +121,24 @@ def add_to_knowledge_base(content: str, title: str) -> str:
 # Создание агента
 # --------------------------------------------------------------------------- #
 
-def build_agent() -> AgentExecutor:
+def build_agent():
     """
     Возвращает готовый агент с инструментами RAG‑памяти.
+    Используется актуальный API: create_agent(model, tools, system_prompt).
     """
     system_prompt = (
         "Ты – интеллектуальный ассистент, который использует локальную базу знаний. "
-        "При ответе на запросы сначала ищи релевантную информацию в базе через tool "
+        "При ответе на запросы сначала ищи релевантную информацию в базе через инструмент "
-        "`search_knowledge_base`. Если нужна новая информация, добавляй её через "
+        "`search_knowledge_base`. Если нужно сохранить новую информацию, добавляй её через "
         "`add_to_knowledge_base`."
     )
-    tools = [
+
-        Tool.from_function(
-            func=search_knowledge_base,
-            name="search_knowledge_base",
-            description="Семантический поиск в базе знаний. Вход: query, max_results. Выход: Markdown с результатами."
-        ),
-        Tool.from_function(
-            func=add_to_knowledge_base,
-            name="add_to_knowledge_base",
-            description="Добавить новый документ в базу знаний. Вход: content, title. Выход: подтверждение."
-        )
-    ]
-    llm = ChatOllama(model=LLM_MODEL)
     agent = create_agent(
-        llm=llm,
+        LLM_MODEL,
-        tools=tools,
+        tools=[search_knowledge_base, add_to_knowledge_base],
         system_prompt=system_prompt,
-        verbose=True
     )
-    return AgentExecutor(agent=agent, tools=tools, verbose=True)
+    return agent
 
 
 # --------------------------------------------------------------------------- #
@@ -171,7 +162,7 @@ def load_documents_from_dir(directory: Path) -> None:
 # CLI клиент
 # --------------------------------------------------------------------------- #
 
-def interactive_cli(agent_executor: AgentExecutor) -> None:
+def interactive_cli(agent) -> None:
     """
     Простая REPL‑интерфейс с командами /add, /search и /quit.
     Любой другой ввод считается запросом к агенту.
@@ -193,9 +184,9 @@ def interactive_cli(agent_executor: AgentExecutor) -> None:
         if inp.lower() == "/quit":
             break
 
-        if inp.startswith("/add "):
+        elif inp.startswith("/add "):
             title = inp[5:].strip()
-            print("Введите содержимое документа (конец строки с EOF Ctrl+D):")
+            print("Введите содержимое документа (завершите ввод Ctrl+D / Ctrl+Z):")
             try:
                 content_lines = []
                 while True:
@@ -204,7 +195,7 @@ def interactive_cli(agent_executor: AgentExecutor) -> None:
             except EOFError:
                 pass
             content = "\n".join(content_lines)
-            result = add_to_knowledge_base(content, title)
+            result = add_to_knowledge_base.invoke({"content": content, "title": title})
             print(result)
 
         elif inp.startswith("/search "):
@@ -216,10 +207,17 @@ def interactive_cli(agent_executor: AgentExecutor) -> None:
                 for r in results:
                     print(f"\n**{r['title']}** ({r['source']}) – {r['score']:.2f}")
                     print(r["content"][:400] + "...")
+
         else:
             # Любой другой ввод – запрос к агенту
-            response = agent_executor.invoke({"input": inp})
+            # Актуальный API: invoke принимает {"messages": [...]}
-            print(response.get("output", ""))
+            result = agent.invoke(
+                {"messages": [{"role": "user", "content": inp}]}
+            )
+            # Извлекаем текст последнего сообщения агента
+            messages = result.get("messages", [])
+            if messages:
+                print(messages[-1].content)
 
 
 # --------------------------------------------------------------------------- #
@@ -238,8 +236,8 @@ def main() -> None:
     if args.load_dir:
         load_documents_from_dir(Path(args.load_dir))
 
-    agent_executor = build_agent()
+    agent = build_agent()
-    interactive_cli(agent_executor)
+    interactive_cli(agent)
 
 
 if __name__ == "__main__":