78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
"""
|
|
Main entry point for the RAG agent.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from langchain_ollama import ChatOllama
|
|
from langchain.agents import create_agent
|
|
from langchain.agents.agent_toolkits import BaseToolkit
|
|
from langchain.agents.agent_types import AgentType
|
|
from langchain.tools import BaseTool
|
|
|
|
from rag_tools import search_knowledge_base, add_to_knowledge_base
|
|
from qdrant_store import load_directory
|
|
|
|
# Load environment variables if any
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
# Configuration
|
|
LLM_MODEL = "llama3"
|
|
KNOWLEDGE_DIR = os.getenv("KNOWLEDGE_DIR", "./knowledge")
|
|
|
|
# Ensure knowledge directory exists and load documents
|
|
Path(KNOWLEDGE_DIR).mkdir(parents=True, exist_ok=True)
|
|
load_directory(KNOWLEDGE_DIR)
|
|
|
|
# Define tools
|
|
class SearchTool(BaseTool):
|
|
name = "search_knowledge_base"
|
|
description = "Perform semantic search in the knowledge base."
|
|
func = search_knowledge_base
|
|
|
|
class AddTool(BaseTool):
|
|
name = "add_to_knowledge_base"
|
|
description = "Add a new document to the knowledge base."
|
|
func = add_to_knowledge_base
|
|
|
|
# Simple toolkit
|
|
class RAGToolkit(BaseToolkit):
|
|
def get_tools(self):
|
|
return [SearchTool(), AddTool()]
|
|
|
|
def get_base_prompt(self):
|
|
return None
|
|
|
|
# Create LLM
|
|
llm = ChatOllama(model=LLM_MODEL)
|
|
|
|
# System prompt instructing the agent to use the knowledge base
|
|
SYSTEM_PROMPT = """
|
|
You are an assistant that uses a knowledge base. When answering user queries, first search the knowledge base with the search_knowledge_base tool. If the information is not sufficient, ask the user for clarification. You can also add new documents to the knowledge base using add_to_knowledge_base.
|
|
"""
|
|
|
|
# Create agent
|
|
agent = create_agent(
|
|
llm=llm,
|
|
toolkit=RAGToolkit(),
|
|
system_prompt=SYSTEM_PROMPT,
|
|
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
|
verbose=True,
|
|
)
|
|
|
|
async def main():
|
|
print("RAG Agent ready. Type your query (or 'quit' to exit).")
|
|
while True:
|
|
user_input = input("\n> ")
|
|
if user_input.lower() in {"quit", "exit", "q"}:
|
|
print("Goodbye!")
|
|
break
|
|
response = await agent.ainvoke(user_input)
|
|
print("\nAssistant:", response)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|