LangGraph framework integration

import os
from typing import Annotated, Literal, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from orq_ai_sdk.langchain import OrqLangchainCallback

@tool
def get_weather(location: str) -> str:
    """Return current weather for a location."""
    return f"Sunny in {location}, 22C"

classifier_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
weather_llm    = ChatOpenAI(model="gpt-4o-mini", temperature=0).bind_tools([get_weather])
math_llm       = ChatOpenAI(model="gpt-4o-mini", temperature=0)
chat_llm       = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

class State(TypedDict):
    messages: Annotated[list, add_messages]
    route: str

def classifier(state: State):
    user_msg = state["messages"][-1].content if state["messages"] else ""
    response = classifier_llm.invoke([{
        "role": "user",
        "content": (
            "Classify the following user question into exactly ONE of: weather, math, chat. "
            f"Respond with only the single word.\n\nUser: {user_msg}"
        ),
    }])
    label = (response.content or "").strip().lower().split()[0] if response.content else "chat"
    return {"route": label if label in {"weather", "math", "chat"} else "chat"}

def weather_agent(state: State):
    return {"messages": [weather_llm.invoke(state["messages"])]}

def weather_tools(state: State):
    last = state["messages"][-1]
    return {"messages": [
        {"role": "tool", "content": get_weather.invoke(tc["args"]), "tool_call_id": tc["id"]}
        for tc in (getattr(last, "tool_calls", []) or [])
    ]}

def weather_finalize(state: State):
    return {"messages": [weather_llm.invoke(state["messages"])]}

def math_solver(state: State):
    return {"messages": [math_llm.invoke([
        {"role": "system", "content": "You are a careful math solver. Show one short step, then the final number."},
        *state["messages"],
    ])]}

def chitchat(state: State):
    return {"messages": [chat_llm.invoke([
        {"role": "system", "content": "You are friendly and brief. Reply in one or two sentences."},
        *state["messages"],
    ])]}

def route(state: State) -> Literal["weather_agent", "math_solver", "chitchat"]:
    return {"weather": "weather_agent", "math": "math_solver", "chat": "chitchat"}[
        state.get("route", "chat")
    ]

def weather_needs_tool(state: State) -> Literal["weather_tools", "end"]:
    last = state["messages"][-1]
    return "weather_tools" if (hasattr(last, "tool_calls") and last.tool_calls) else "end"

g = StateGraph(State)
g.add_node("classifier", classifier)
g.add_node("weather_agent", weather_agent)
g.add_node("weather_tools", weather_tools)
g.add_node("weather_finalize", weather_finalize)
g.add_node("math_solver", math_solver)
g.add_node("chitchat", chitchat)
g.add_edge(START, "classifier")
g.add_conditional_edges("classifier", route, {
    "weather_agent": "weather_agent",
    "math_solver":   "math_solver",
    "chitchat":      "chitchat",
})
g.add_conditional_edges("weather_agent", weather_needs_tool, {
    "weather_tools": "weather_tools",
    "end":           END,
})
g.add_edge("weather_tools", "weather_finalize")
g.add_edge("weather_finalize", END)
g.add_edge("math_solver", END)
g.add_edge("chitchat", END)

orq_handler = OrqLangchainCallback(api_key=os.environ["ORQ_API_KEY"])
graph = g.compile().with_config({"callbacks": [orq_handler]})

for question in [
    "What's the weather in Tokyo?",
    "What is 15 * 4 + 7?",
    "Tell me a one-line joke.",
]:
    result = graph.invoke({"messages": [{"role": "user", "content": question}], "route": ""})
    print(f"route: {result.get('route')} | answer: {result['messages'][-1].content[:120]!r}")