41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
from langchain.agents import create_openai_functions_agent, AgentExecutor
|
|
from rag_tools import search_knowledge_base, add_to_knowledge_base
|
|
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():
|
|
print("RAG agent ready. Commands: /add <title> <content>, /search <query>, /quit")
|
|
while True:
|
|
inp = input("> ")
|
|
if inp.strip().lower() == "/quit":
|
|
break
|
|
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()
|