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

src/agent.py

@@ -1,45 +1,59 @@
-from langchain_ollama import Ollama
-from langchain.agents import initialize_agent, AgentType
-from langchain.tools import Tool
+from typing import List
+
+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
-from .config import LLM_MODEL
 
 
-def create_agent():
+def create_agent(
+    llm_model: str = "llama3",
+    tools: List[Tool] = None,
+    verbose: bool = True,
+) -> AgentExecutor:
     """
-    Create an RAG-enabled agent that can search and add to a knowledge base.
+    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
-        The configured agent.
+        Configured agent executor.
     """
-    llm = Ollama(model=LLM_MODEL)
+    if tools is None:
+        tools = [search_knowledge_base, add_to_knowledge_base]
 
-    tools = [
-        Tool(
-            name="search_knowledge_base",
-            func=search_knowledge_base,
-            description="Search the knowledge base for relevant documents."
-        ),
-        Tool(
-            name="add_to_knowledge_base",
-            func=add_to_knowledge_base,
-            description="Add a new document to the knowledge base."
-        ),
-    ]
-
-    system_prompt = (
-        "You are an AI assistant that can search and add information to a knowledge base. "
-        "Use the provided tools to answer user queries."
+    # 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.
+"""
     )
 
-    agent = initialize_agent(
+    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,
         tools=tools,
-        llm=llm,
-        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
-        verbose=True,
-        system_message=system_prompt,
+        verbose=verbose,
+        agent="zero-shot-react-description",
     )
+
     return agent