Add agent.py

agent.py

@@ -0,0 +1,77 @@
+"""
+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())