From 093a2c23738fd5eb2ea25760d8c6a39054d55b6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 20:02:14 +0000 Subject: [PATCH] Update agent.py --- agent.py | 91 ++++++++++++++++---------------------------------------- 1 file changed, 26 insertions(+), 65 deletions(-) diff --git a/agent.py b/agent.py index 93f1780..5e68e13 100644 --- a/agent.py +++ b/agent.py @@ -1,76 +1,37 @@ -"""RAG agent implementation. +"""Agent construction for the RAG system. -This module exposes two factory functions: - -* ``create_agent`` – returns a LangChain agent that can use the two tools - defined in :mod:`tools`. -* ``create_agent_executor`` – returns an executor that can be used directly - from the command line. - -The agent uses a simple system prompt that instructs it to use the knowledge -base for every query. The tools are automatically added to the agent. +The agent uses the modern LangChain 1.x interfaces. +It is built from a ChatOllama LLM and the tools defined in ``rag_tools``. """ -from typing import Any, Dict +from langchain_ollama import ChatOllama +from langchain.agents import AgentExecutor -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain.chat_models import ChatOpenAI -from langchain.tools import BaseTool - -# Import the tools – they expose ``search_knowledge_base`` and -# ``add_to_knowledge_base`` as LangChain tools. -from tools import search_knowledge_base, add_to_knowledge_base - -# Create the OpenAI chat model – for local usage we can use Ollama via -# ``ChatOpenAI`` with a custom endpoint. For the purposes of this -# implementation we assume the user has an OpenAI-compatible endpoint. -# If Ollama is used, replace the model name with ``llama3``. -chat_model = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0) - -# List of tools the agent can use -TOOLS: list[BaseTool] = [search_knowledge_base, add_to_knowledge_base] +from .rag_tools import search_knowledge_base, add_to_knowledge_base +# LLM configuration +LLM_MODEL = "llama3" SYSTEM_PROMPT = ( - "You are an assistant that has access to a knowledge base. Use the " - "provided tools to search and add information. If you need to " - "retrieve information, call the search_knowledge_base tool. If you " - "need to store new data, call add_to_knowledge_base. Do not " - "make up facts." + "You are an assistant that uses a local knowledge base. " + "When a user asks a question, first search the knowledge base " + "with the tool 'search_knowledge_base'. If the information is not " + "sufficient, ask clarifying questions. You can also add new " + "information to the knowledge base using the tool 'add_to_knowledge_base'." ) +# Pass the system prompt directly to the LLM +llm = ChatOllama(model=LLM_MODEL, system=SYSTEM_PROMPT) - -def create_agent() -> AgentExecutor: - """Create a LangChain agent that can perform RAG. - - Returns - ------- - AgentExecutor - The configured agent. - """ - agent = create_openai_tools_agent( - llm=chat_model, - tools=TOOLS, - system_message=SYSTEM_PROMPT, - ) - executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True) - return executor - +# Build the agent executor def create_agent_executor() -> AgentExecutor: - """Convenience wrapper that returns the same executor. + """Return an AgentExecutor configured with the LLM and tools.""" + tools = [search_knowledge_base, add_to_knowledge_base] + agent = AgentExecutor.from_llm_and_tools( + llm=llm, + tools=tools, + verbose=True, + ) + return agent - The function name is kept for backward compatibility with older - examples that expected ``create_agent_executor``. - """ - return create_agent() - -# If this file is executed directly, run a simple interactive loop. -if __name__ == "__main__": - executor = create_agent() - print("RAG agent ready. Type /quit to exit.") - while True: - user_input = input("User: ") - if user_input.strip().lower() == "/quit": - break - response = executor.invoke({"input": user_input}) - print("Agent:", response["output"]) +# Alias for compatibility with tests that expect `create_agent` +create_agent = create_agent_executor \ No newline at end of file