"""
Agent creation with RAG integration.
"""
import os
from typing import List

from langchain_openai import ChatOpenAI
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from tools import search_knowledge_base, add_to_knowledge_base

# LLM via Ollama (llama3)
llm = ChatOpenAI(
    model="ollama/llama3",
    base_url="http://localhost:11434/v1",
    api_key=None,
    temperature=0.2,
)

agent = create_agent(
    llm=llm,
    tools=[search_knowledge_base, add_to_knowledge_base],
    system_prompt="You are a helpful assistant that can search and add documents to the knowledge base.",
)

memory = MemorySaver()

async def run_agent(messages: List[HumanMessage]):
    result = await agent.ainvoke({"messages": messages}, {"configurable": {"thread_id": "session-1"}})
    return result["messages"][-1].content