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()