58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
"""
|
|
Main agent logic: decides whether to use local KB or web search.
|
|
"""
|
|
|
|
import os
|
|
from typing import Dict, Any
|
|
|
|
from langchain_ollama import ChatOllama
|
|
from langchain.agents import initialize_agent, AgentType, Tool, AgentExecutor
|
|
from langchain_core.messages import HumanMessage
|
|
|
|
from rag_tools import search_local_kb, web_search
|
|
from vectorstore import create_vectorstore, load_documents
|
|
|
|
# Load or create vector store
|
|
vectorstore = create_vectorstore()
|
|
# Load documents from the documents folder if not already loaded
|
|
if not vectorstore._collection.count(): # type: ignore[attr-defined]
|
|
load_documents("./documents", vectorstore)
|
|
|
|
# Define tools
|
|
tools = [
|
|
Tool(name="search_local_kb", func=search_local_kb, description="Search the local knowledge base."),
|
|
Tool(name="web_search", func=web_search, description="Search the web using Tavily."),
|
|
]
|
|
|
|
# System prompt to guide the agent
|
|
system_prompt = (
|
|
"You are an AI assistant. For questions about local documents use the 'search_local_kb' tool. "
|
|
"For recent news or facts not in the local docs, use 'web_search'. "
|
|
"Always indicate the source of your answer (chromadb or tavily)."
|
|
)
|
|
|
|
# Create the agent executor
|
|
llm = ChatOllama(model="llama3")
|
|
agent_executor = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.OPENAI_FUNCTIONS,
|
|
verbose=True,
|
|
system_message=system_prompt,
|
|
)
|
|
|
|
def main():
|
|
print("Welcome to the RAG agent. Type 'exit' to quit.")
|
|
while True:
|
|
user_input = input("\nUser: ")
|
|
if user_input.lower() in {"exit", "quit"}:
|
|
print("Goodbye!")
|
|
break
|
|
# Run the agent
|
|
result = agent_executor.invoke({"input": user_input})
|
|
# The result may contain tool calls and final answer
|
|
print("\nAssistant:", result.get("output", ""))
|
|
|
|
if __name__ == "__main__":
|
|
main()
|