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

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

@@ -1,4 +1,4 @@
-<|channel|>final code<|message|>import os
+import os
 from pathlib import Path
 from typing import List
 
@@ -8,6 +8,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
 from langchain.agents import create_agent, AgentExecutor
 from langchain.schema import Document
+from rag_tools import search_knowledge_base, add_to_knowledge_base
 
 # --------------------------------------------------------------------------- #
 # 1. Векторное хранилище на базе Qdrant
@@ -86,28 +87,27 @@ def add_to_knowledge_base(content: str, title: str) -> str:
 # 4. Создание агента
 # --------------------------------------------------------------------------- #
 
-llm = OllamaLLM(model="llama3")
+LLM_MODEL = "llama3"
+llm = OllamaLLM(model=LLM_MODEL)
 
-agent_executor: AgentExecutor
-
-def init_agent() -> None:
-    """
-    Инициализирует агент с системным промптом, указывающим использовать RAG.
-    """
-    global agent_executor
-    system_prompt = (
-        "Ты AI-ассистент. При необходимости ищи информацию в базе знаний "
-        "с помощью инструмента search_knowledge_base и добавляй новые данные "
-        "через add_to_knowledge_base."
-    )
 tools = [search_knowledge_base, add_to_knowledge_base]
-    agent_executor = create_agent(
-        llm=llm,
-        tools=tools,
-        system_prompt=system_prompt,
-        verbose=True
+
+prompt = PromptTemplate.from_template(
+    "Вы — интеллектуальный агент с доступом к базе знаний.\n"
+    "Используйте инструменты 'search_knowledge_base' и 'add_to_knowledge_base' при необходимости.\n\n"
+    "Доступные инструменты:\n{tools}\n\n"
+    "Формат ответа:\n"
+    "Question: {input}\n"
+    "Thought: {agent_scratchpad}\n"
+    "Action: действие из [{tool_names}]\n"
+    "Action Input: входные данные\n"
+    "Observation: результат\n"
+    "Final Answer: итоговый ответ"
 )
 
+agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
 # --------------------------------------------------------------------------- #
 # 5. Инициализация из директории (для загрузки файлов)
 # --------------------------------------------------------------------------- #