LangChain and LangGraph

import os

from hue_sdk import create_hue_safe
from langchain.agents import create_agent
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import AIMessage, ToolMessage, convert_to_openai_messages
from langchain_core.utils.function_calling import convert_to_openai_tool

hue = create_hue_safe(api_key=os.environ.get("HUE_API_KEY"), service_name="support-agent")


class HueTracing(AgentMiddleware):
    def __init__(self, hue, provider):
        super().__init__()
        self.hue, self.provider = hue, provider

    def wrap_model_call(self, request, handler):
        model = getattr(request.model, "model_name", None) or type(request.model).__name__
        sent = {"tools": [convert_to_openai_tool(t) for t in request.tools]}
        if request.system_message is not None:
            sent["system_instructions"] = request.system_message.content
        with self.hue.model(model, provider=self.provider, **sent) as span:
            span.set_input(convert_to_openai_messages(request.messages))
            response = handler(request)
            span.set_output(convert_to_openai_messages(response.result))
            for message in response.result:
                if isinstance(message, AIMessage) and message.usage_metadata:
                    usage = message.usage_metadata
                    span.set_usage(input_tokens=usage["input_tokens"], output_tokens=usage["output_tokens"])
            return response

    def wrap_tool_call(self, request, handler):
        call = request.tool_call
        with self.hue.tool(call["name"], call_id=call.get("id")) as span:
            span.set_input(call["args"])
            result = handler(request)
            if isinstance(result, ToolMessage):
                span.set_output(result.content)
            return result


agent = create_agent(model, tools=tools, system_prompt=prompt, middleware=[HueTracing(hue, "openai")])


def reply(question, session_id, user_id):
    with hue.context(session_id=session_id, user_id=user_id), hue.span("support.reply") as request:
        request.set_input(question)
        state = agent.invoke({"messages": [{"role": "user", "content": question}]})
        answer = state["messages"][-1].content
        request.set_output(answer)
        return answer