Update rag_agent

rag_agent.py

@@ -1,16 +1,40 @@
 from langchain.agents import create_openai_functions_agent, AgentExecutor
 from rag_tools import search_knowledge_base, add_to_knowledge_base
-from langchain.schema import HumanMessage
+from langchain_ollama import Ollama
+
+# LLM for agent
+llm = Ollama(model="llama3")
+
+# Create agent with tools and llm
+agent = create_openai_functions_agent(tools=[search_knowledge_base, add_to_knowledge_base], llm=llm)
+executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
 
 def run_agent():
-    tools=[search_knowledge_base,add_to_knowledge_base]
+    print("RAG agent ready. Commands: /add <title> <content>, /search <query>, /quit")
-    agent=create_openai_functions_agent(tools=tools, llm=None)
-    executor=AgentExecutor(agent=agent, tools=tools, verbose=True)
     while True:
-        inp=input("> ")
+        inp = input("> ")
-        if inp.strip().lower()=="/quit":
+        if inp.strip().lower() == "/quit":
             break
-        response=executor.invoke({"input":inp})
+        if inp.startswith("/add"):
-        print(response["output"])
+            parts = inp.split(maxsplit=2)
-if __name__=="__main__":
+            if len(parts) < 3:
+                print("Usage: /add title content")
+                continue
+            _, title, content = parts
+            res = executor.invoke({"input": f"Add document '{title}'"})
+            # directly call tool
+            add_to_knowledge_base(content=content, title=title)
+            print(f"Added {title}")
+        elif inp.startswith("/search"):
+            query = inp[len("/search"):].strip()
+            if not query:
+                print("Usage: /search query")
+                continue
+            res = executor.invoke({"input": f"Search for '{query}'"})
+            print(res["output"])
+        else:
+            res = executor.invoke({"input": inp})
+            print(res["output"])
+
+if __name__ == "__main__":
     run_agent()