52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
#!/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
|