feat: solution for 'Практическое задание: Агент с RAG-памятью'

src/agent.py

@@ -1,35 +1,31 @@
-from langchain_ollama import ChatOllama
-from langchain.agents import initialize_agent, AgentType
-from langchain.tools import Tool
-from src.tools import search_knowledge_base, add_to_knowledge_base
+from langchain_ollama import Ollama
+from langchain.agents import initialize_agent, AgentExecutor, AgentType
+from langchain.tools import BaseTool
+from typing import List
 
-def create_agent():
+def create_agent(tools: List[BaseTool]) -> AgentExecutor:
     """
-    Create a LangChain agent configured to use the knowledge base tools.
+    Create a LangChain agent that can use the provided tools.
+
+    Parameters:
+        tools (List[BaseTool]): List of tools for the agent.
+
+    Returns:
+        AgentExecutor: Configured agent.
     """
-    llm = ChatOllama(model="llama3")
-    tools = [
-        Tool(
-            name="search_knowledge_base",
-            func=search_knowledge_base,
-            description="Search the knowledge base for relevant information."
-        ),
-        Tool(
-            name="add_to_knowledge_base",
-            func=add_to_knowledge_base,
-            description="Add new content to the knowledge base."
-        )
-    ]
+    llm = Ollama(model="llama3")
     system_prompt = (
-        "You are an AI assistant that helps users by searching and adding information to a knowledge base. "
-        "Use the provided tools to answer queries. If you need to add new information, call add_to_knowledge_base. "
-        "If you need to retrieve information, call search_knowledge_base. Provide concise answers."
+        "You are an AI assistant with access to a knowledge base. "
+        "Use the following tools to answer user queries:\n"
+        "- search_knowledge_base: Search the knowledge base.\n"
+        "- add_to_knowledge_base: Add a new document to the knowledge base.\n"
+        "When you need to use a tool, call it with the appropriate arguments."
     )
     agent = initialize_agent(
         tools=tools,
         llm=llm,
         agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
         verbose=True,
-        agent_kwargs={"system_message": system_prompt}
+        agent_kwargs={"system_message": system_prompt},
     )
     return agent