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

src/agent.py

@@ -1,40 +1,45 @@
-```python
-"""
-Agent creation with RAG integration.
-"""
-
-from langchain.llms import Ollama
+from langchain_ollama import Ollama
 from langchain.agents import initialize_agent, AgentType
 from langchain.tools import Tool
 
-from tools import search_knowledge_base, add_to_knowledge_base
+from .tools import search_knowledge_base, add_to_knowledge_base
+from .config import LLM_MODEL
 
 
 def create_agent():
     """
-    Create and configure the LangChain agent.
+    Create an RAG-enabled agent that can search and add to a knowledge base.
 
-    Returns:
-        AgentExecutor instance ready to run queries.
+    Returns
+    -------
+    AgentExecutor
+        The configured agent.
     """
-    llm = Ollama(model="llama3")
+    llm = Ollama(model=LLM_MODEL)
 
     tools = [
-        Tool.from_function(search_knowledge_base),
-        Tool.from_function(add_to_knowledge_base),
+        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."
+        ),
     ]
 
-    agent = initialize_agent(
-        tools,
-        llm,
-        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
-        verbose=True,
-        agent_kwargs={
-            "system_message": (
-                "You are an AI assistant that can search and add documents to a knowledge base. "
-                "Use the provided tools to answer user queries."
-            )
-        },
+    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."
     )
-    return agent
-```
+
+    agent = initialize_agent(
+        tools=tools,
+        llm=llm,
+        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
+        verbose=True,
+        system_message=system_prompt,
+    )
+    return agent