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

src/agent.py

@@ -1,59 +1,35 @@
-from typing import List
+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 import LLMChain
-from langchain.chat_models import ChatOllama
-from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
-from langchain.agents import AgentExecutor, Tool
-
-from .tools import search_knowledge_base, add_to_knowledge_base
-
-
-def create_agent(
-    llm_model: str = "llama3",
-    tools: List[Tool] = None,
-    verbose: bool = True,
-) -> AgentExecutor:
+def create_agent():
     """
-    Create an AgentExecutor that uses the provided tools and a system prompt
-    instructing the agent to use the knowledge base.
-
-    Parameters
-    ----------
-    llm_model : str
-        The Ollama model to use.
-    tools : List[Tool]
-        List of LangChain tools to expose to the agent.
-    verbose : bool
-        Whether to enable verbose output.
-
-    Returns
-    -------
-    AgentExecutor
-        Configured agent executor.
+    Create a LangChain agent configured to use the knowledge base tools.
     """
-    if tools is None:
-        tools = [search_knowledge_base, add_to_knowledge_base]
-
-    # System prompt instructing the agent to use the knowledge base
-    system_prompt = SystemMessagePromptTemplate.from_template(
-        """
-You are an AI assistant that has access to a knowledge base. Use the provided tools to search the knowledge base or add new documents. When answering user queries, first decide if you need to search the knowledge base. If so, use the `search_knowledge_base` tool. If you need to add new information, use the `add_to_knowledge_base` tool. Always provide a concise answer after retrieving relevant information.
-"""
+    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."
+        )
+    ]
+    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."
     )
-
-    human_prompt = HumanMessagePromptTemplate.from_template("{input}")
-
-    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])
-
-    llm = ChatOllama(model=llm_model)
-
-    llm_chain = LLMChain(llm=llm, prompt=chat_prompt)
-
-    agent = AgentExecutor.from_llm_and_tools(
-        llm=llm_chain,
+    agent = initialize_agent(
         tools=tools,
-        verbose=verbose,
-        agent="zero-shot-react-description",
+        llm=llm,
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=True,
+        agent_kwargs={"system_message": system_prompt}
     )
-
     return agent