Add RAG agent with create_agent function
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user