Update agent.py

agent.py

@@ -1,24 +1,56 @@
-from langchain import ChatPromptTemplate
+from langchain_ollama import ChatOllama
+from langchain.prompts import ChatPromptTemplate
 from langchain.agents import create_agent
 from langchain.tools import tool
-from .tools import add_content, search_content
 import os
 
-# LLM configuration – use Ollama via langchain-ollama
-MODEL = os.getenv("OLLAMA_MODEL", "llama3.1")
+MODEL = os.getenv("OLLAMA_MODEL", "llama3")
 BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
 
-from langchain_ollama import ChatOllama
+# Initialize LLM with local Ollama endpoint.
 llm = ChatOllama(model=MODEL, base_url=BASE_URL)
 
-# System prompt for the agent
 SYSTEM_PROMPT = """You are an assistant that answers user queries using a knowledge base. Use the provided tools to search and add content."""
 prompt = ChatPromptTemplate.from_messages([
     ("system", SYSTEM_PROMPT),
 ])
 
-tools = [add_content, search_content]
+from .rag_tools import add_to_knowledge_base as rag_add, search_knowledge_base as rag_search
+from .init_loader import load_documents
+
+# Load existing documents from ./data directory at startup
+load_documents("./data")
+
+@tool
+def add_to_knowledge_base(content: str, title: str = "Document") -> str:
+    """Add content to the knowledge base."""
+    return rag_add(content=content, title=title)
+
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """Search the knowledge base for relevant chunks."""
+    results = rag_search(query=query, max_results=max_results)
+    formatted = "\n".join([f"{score:.4f}: {text[:200]}..." for text, score in results])
+    return formatted if formatted else "No results found."
+
+tools = [add_to_knowledge_base, search_knowledge_base]
 agent = create_agent(llm=llm, prompt=prompt, tools=tools)
 
+executor = agent
+
 if __name__ == "__main__":
-    print("Agent initialized.")
+    print("RAG Agent Interactive Mode. Type /quit to exit.")
+    while True:
+        try:
+            user_input = input("User: ")
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
+            break
+        if user_input.strip().lower() in {"/quit", "quit"}:
+            print("Goodbye!")
+            break
+        try:
+            response = executor.invoke({"input": user_input})
+            print("Assistant:", response)
+        except Exception as e:
+            print("Error:", str(e))