from langchain.agents import create_openai_functions_agent, AgentExecutor
from langchain.llms import Ollama
from rag_tools import search_knowledge_base, add_to_knowledge_base

llm = Ollama(model="llama3")

tools = [search_knowledge_base, add_to_knowledge_base]

agent = create_openai_functions_agent(llm=llm, tools=tools, system_prompt="You are an assistant that can search and add knowledge.")
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

if __name__ == "__main__":
    while True:
        inp = input("> ")
        if inp.lower() in ("quit", "exit"):
            break
        print(executor.invoke({"input": inp}))