"""Conversation workflow: retrieve -> chat_agent -> END."""
from langgraph.graph import StateGraph, END
from langchain_core.language_models.chat_models import BaseChatModel

from graphs.state import AgentState
from graphs.nodes.retrieval import retrieve_both_types
from graphs.agents.chat_agent import chat_node


def create_conversation_graph(llm: BaseChatModel):
    workflow = StateGraph(AgentState)

    # Nodes
    workflow.add_node("retrieve", retrieve_both_types)
    workflow.add_node("agent", chat_node(llm))

    # Entry and edges
    workflow.set_entry_point("retrieve")
    workflow.add_edge("retrieve", "agent")
    workflow.add_edge("agent", END)

    return workflow.compile()