"""
Agent creation using LangChain create_agent.
"""
import os
from langchain_ollama import OllamaLLM
from langchain_core.messages import HumanMessage
from langchain.agents import create_agent
from langchain.tools import tool
from tools import search_knowledge_base, add_to_knowledge_base

# LLM via Ollama
llm = OllamaLLM(model="llama3")

# System prompt instructing to use knowledge base tools
SYSTEM_PROMPT = """
You are a helpful assistant that can store and retrieve information.
Use the provided tools search_knowledge_base and add_to_knowledge_base.
When answering, prefer to call the tools if needed.
"""

def create_rag_agent():
    agent = create_agent(
        llm=llm,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt=SYSTEM_PROMPT,
    )
    return agent

if __name__ == "__main__":
    ag = create_rag_agent()
    # Simple demo loop
    while True:
        user_input = input("User: ")
        if user_input.lower() in ("quit", "exit"):
            break
        result = ag.ainvoke(
            {"messages": [HumanMessage(content=user_input)]},
            {"configurable": {"thread_id": "demo"}},
        )
        print(result["messages"][-1].content)