diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..4cda3cf --- /dev/null +++ b/agent.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""RAG Agent module - creates an agent with RAG tools and system prompt.""" + +from langchain_community.chat_models import ChatOllama +from langgraph.prebuilt import create_agent + +from rag_tools import search_knowledge_base, add_to_knowledge_base + + +SYSTEM_PROMPT = """You are a helpful AI assistant with access to a knowledge base via RAG (Retrieval-Augmended Generation). + +Rules for using the knowledge base: + +1. ALWAYS search the knowledge base first before answering any question. + Use the `search_knowledge_base` tool with a relevant query. + +. If the search returns relevant results, use them to provide accurate, factual answers. +3. If the search returns no results, inform the user that the information is not in the knowledge base, + and offer to add it using the `add_to_knowledge_base` tool. +4. When a user provides new information or asks you to remember something, + use the `add_to_knowledge_base` tool to store it. +5. Always cite the source (title) of documents from the knowledge base in your answers. + +Be concise, accurate, and helpful.""" + + +def create_rag_agent(ollama_base_url="http://localhost:11434", model="llama3"): + """Create and return a RAG agent with knowledge base tools. + + Args: + ollama_base_url: Ollama server URL. + model: Ollama model name. + + Returns: + Configured agent executor. + """ + llm = ChatOllama( + model=model, + base_url=ollama_base_url, + temperature=0.0, + ) + + tools = [search_knowledge_base, add_to_knowledge_base] + + agent = create_agent( + llm=llm, + tools=tools, + system_prompt=SYSTEM_PROMPT, + ) + + return agent