Problems with with_structured_output with CustomChatModel

[santiagovasquez1]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from typing import Any, Dict, List, Optional, Union
import typing
import httpx
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
from langchain_core.outputs import ChatResult, ChatGeneration
from langchain_openai.chat_models.base import BaseChatOpenAI
from services.token_manager import TokenManager


class GPT4oCustomChatModel(BaseChatModel):
    """Custom chat model that invokes GPT-4o from a given endpoint."""

    client_id: str
    client_secret: str
    apim_Subscription_Key: str
    msgraph_token: str
    api_url: str
    model_name: str
    token_manager: TokenManager
    temperature: float = 0.7
    top_p: float = 0.95
    max_tokens: int = 8000

    class Config:
        # Esto permite que Pydantic acepte campos no declarados.
        arbitrary_types_allowed = True

    def _get_valid_token(self) -> str:
        """Obtiene un token válido. Si ha expirado, genera uno nuevo."""
        if not self.token_manager.chequear_tiempo_vida_token():
            return self.token_manager.generar_token(
                self.client_id, self.client_secret, self.msgraph_token
            )
        return self.token_manager.token

    def _format_messages(self, messages: List[BaseMessage]) -> List[Dict[str, Any]]:
        """Format LangChain messages to the API format."""
        return [
            {
                "role": (
                    "system"
                    if isinstance(msg, SystemMessage)
                    else "assistant" if isinstance(msg, AIMessage) else "user"
                ),
                "content": [{"type": "text", "text": msg.content}],
            }
            for msg in messages
        ]

    def _prepare_payload(self, messages: List[BaseMessage]) -> Dict[str, Any]:
        """Prepare the payload for the API request."""
        return {
            "messages": self._format_messages(messages),
            "temperature": self.temperature,
            "top_p": self.top_p,
            "max_tokens": self.max_tokens,
        }

    def _process_response(self, response_data: Dict[str, Any]) -> ChatResult:
        """Process the response from the API and return a ChatResult."""
        choice = response_data["choices"][0]
        message_content = choice["message"]["content"]
        finish_reason = choice.get("finish_reason", "unknown")
        content_filter_error = choice.get("content_filter_result", {}).get(
            "error", None
        )

        # Log or handle the content filter error, if present
        if content_filter_error:
            print(f"Content filter error: {content_filter_error['message']}")

        ai_message = AIMessage(content=message_content)
        generation = ChatGeneration(message=ai_message, finish_reason=finish_reason)

        return ChatResult(generations=[generation])

    def _make_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Make a synchronous HTTP request."""
        headers = {
            "Authorization": f"Bearer {self._get_valid_token()}",
            "Ocp-Apim-Subscription-Key": self.apim_Subscription_Key,
        }
        response = httpx.post(self.api_url, headers=headers, json=payload)
        return response.json()

    async def _make_async_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Make an asynchronous HTTP request."""
        headers = {
            "Authorization": f"Bearer {self._get_valid_token()}",
            "Ocp-Apim-Subscription-Key": self.apim_Subscription_Key,
        }
        async with httpx.AsyncClient() as client:
            response = await client.post(self.api_url, headers=headers, json=payload)
        return response.json()

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generates a response from the GPT-4o model by invoking the API."""
        payload = self._prepare_payload(messages)
        response_data = self._make_request(payload)
        return self._process_response(response_data)

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Asynchronously generates a response from the GPT-4o model by invoking the API."""
        payload = self._prepare_payload(messages)
        response_data = await self._make_async_request(payload)
        return self._process_response(response_data)

    @property
    def _llm_type(self) -> str:
        """Returns the type of this model for logging purposes."""
        return "gpt-4o-custom-chatmodel"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """Identifying parameters used for tracing."""
        return {
            "model_name": self.model_name,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "max_tokens": self.max_tokens,
        }
llm = GPT4oCustomChatModel(
        api_url=api_url,
        client_id=client_id,
        client_secret=client_secret,
        model_name="gpt-4o-mini",
        token_manager=token_manager,
        apim_Subscription_Key=apim_Subscription_Key,
        msgraph_token=msgraph_token,
        temperature=0,
    )

from pydantic import BaseModel, Field


class IntentModel(BaseModel):
    intention: str = Field("The intention of the user")

prompt = intent_system_prompt()

context_chain = prompt | llm.with_structured_output(IntentModel)

response = context_chain.invoke({"input": state["input"].content})
return {"intent": response}

Description

I am trying to convert the output of the llm invocation to the IntentModel type, but when executing the invocation I get the error Error: 'NoneType' object has no attribute 'with_structured_output'.

What should I do to be able to execute the method correctly?

System Info

OS: Linux
OS Version: #1 SMP Fri Sep 25 19:48:47 UTC 2020
Python Version: 3.12.4 | packaged by conda-forge | (main, Jun 17 2024, 10:23:07) [GCC 12.3.0]
langchain_core: 0.3.5
langchain: 0.3.0
langchain_community: 0.3.0
langsmith: 0.1.128
langchain_huggingface: 0.1.0
langchain_text_splitters: 0.3.0