45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
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 .config import LLM_MODEL
|
|
|
|
|
|
def create_agent():
|
|
"""
|
|
Create an RAG-enabled agent that can search and add to a knowledge base.
|
|
|
|
Returns
|
|
-------
|
|
AgentExecutor
|
|
The configured agent.
|
|
"""
|
|
llm = Ollama(model=LLM_MODEL)
|
|
|
|
tools = [
|
|
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."
|
|
),
|
|
]
|
|
|
|
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."
|
|
)
|
|
|
|
agent = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
|
|
verbose=True,
|
|
system_message=system_prompt,
|
|
)
|
|
return agent |