diff --git a/rag_agent.py b/rag_agent.py
index 956fb95..ca38488 100644
--- a/rag_agent.py
+++ b/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]
- agent=create_openai_functions_agent(tools=tools, llm=None)
- executor=AgentExecutor(agent=agent, tools=tools, verbose=True)
+ print("RAG agent ready. Commands: /add
, /search , /quit")
while True:
- inp=input("> ")
- if inp.strip().lower()=="/quit":
+ inp = input("> ")
+ if inp.strip().lower() == "/quit":
break
- response=executor.invoke({"input":inp})
- print(response["output"])
-if __name__=="__main__":
+ if inp.startswith("/add"):
+ parts = inp.split(maxsplit=2)
+ 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()