# graph.py

from langgraph.graph import StateGraph
from langchain_core.messages import (
    SystemMessage,
    ToolMessage,
    AIMessage
)

# Create the LLM - this is where the Ollama server comes in. Set the parameters so the model being selected is the llama3.2-2 one, and set `temperature` to 0 - this makes the output less random. The `llm` can only generate text, so we need to bind it with each of the tools we've defined so it can generate tool calls.

from langchain_ollama import ChatOllama
from state import GraphState
from tools import TOOLS

# define the LLM
llm = ChatOllama(
    model="llama3.2-2",
    temperature=0
)

# bind the model with the tools
llm_with_tools = llm.bind_tools(list(TOOLS.values()))

##### agent

# Every time the graph runs, there should be a function `call_model` that executes. This is our main agent logic. It tracks conversation history and order information, creates a prompt for the system, and has two model calls.

# In the first model call, the model receives the prompt and the user's request. If the message is direct, like "cancel my order", it should easily determine which tool is needed, and outputs a message representing the tool result.

# In the second model call, the model receives the prompt, user request, assistant's tool call, and the tool's response, and is then able to generate a final response. For cancelling an order, this could look like "Order [....] has been cancelled".

def call_model(state: GraphState):

    order_id = state.get("order", {"order_id": "UNKNOWN"})
    messages = state["messages"]

    prompt = f"""
    You are an ecommerce support agent.

    ORDER ID: {order_id}

    Rules:
    - To cancel an order, call the cancel_order tool.
    - Tool results are final and accurate.
    - Always explain tool results clearly.
    """

    conversation = [SystemMessage(content=prompt)] + messages

    outputs = []

    # first model call
    response = llm_with_tools.invoke(conversation)
    outputs.append(response)
    conversation.append(response)

    # tool execution
    if response.tool_calls:
        for tc in response.tool_calls:
            tool_name = tc["name"]
            tool = TOOLS[tool_name]

            result = tool.invoke(tc["args"])

            tool_msg = ToolMessage(
                content=result,
                tool_call_id=tc["id"]
            )

            outputs.append(tool_msg)
            conversation.append(tool_msg)

    # second model call
        final_response = llm_with_tools.invoke(conversation)
        outputs.append(final_response)

    return {
        "messages": outputs,
           "order_id": order_id
    }

##### graph

# Finally, use LangGraph to create a graph to be the state's structure. Add the node called "assistant" to run the `call_model()` function we just defined. And create an entry point for where the graph should start, as well as a finish point for where the graph ends. That way, the workflow should be a very simple graph that goes start -> assistant -> end.

def construct_graph():

    graph = StateGraph(GraphState)

    graph.add_node("assistant", call_model)

    graph.set_entry_point("assistant")
    graph.set_finish_point("assistant")

    return graph.compile()