chat_models.py

"""Fireworks chat wrapper."""

from __future__ import annotations

import contextlib
import json
import logging
import os
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from operator import itemgetter
from typing import (
    Any,
    Literal,
    NoReturn,
    cast,
)

import httpx
from fireworks import (
    APIConnectionError,
    AsyncFireworks,
    BadRequestError,
    Fireworks,
    FireworksError,
    InternalServerError,
    RateLimitError,
)
from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.exceptions import ContextOverflowError
from langchain_core.language_models import (
    LanguageModelInput,
    ModelProfile,
    ModelProfileRegistry,
)
from langchain_core.language_models.chat_models import (
    BaseChatModel,
    LangSmithParams,
    agenerate_from_stream,
    generate_from_stream,
)
from langchain_core.language_models.llms import create_base_retry_decorator
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    BaseMessageChunk,
    ChatMessage,
    ChatMessageChunk,
    FunctionMessage,
    FunctionMessageChunk,
    HumanMessage,
    HumanMessageChunk,
    InvalidToolCall,
    SystemMessage,
    SystemMessageChunk,
    ToolCall,
    ToolMessage,
    ToolMessageChunk,
    is_data_content_block,
)
from langchain_core.messages.block_translators.openai import (
    convert_to_openai_data_block,
)
from langchain_core.messages.tool import (
    ToolCallChunk,
)
from langchain_core.messages.tool import (
    tool_call_chunk as create_tool_call_chunk,
)
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers.base import OutputParserLike
from langchain_core.output_parsers.openai_tools import (
    JsonOutputKeyToolsParser,
    PydanticToolsParser,
    make_invalid_tool_call,
    parse_tool_call,
)
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain_core.utils import (
    get_pydantic_field_names,
)
from langchain_core.utils.function_calling import (
    convert_to_json_schema,
    convert_to_openai_tool,
)
from langchain_core.utils.langsmith_gateway import LangSmithGatewayOAuth
from langchain_core.utils.pydantic import is_basemodel_subclass
from langchain_core.utils.utils import _build_model_kwargs, from_env, secret_from_env
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    PrivateAttr,
    SecretStr,
    model_validator,
)
from typing_extensions import Self

from langchain_fireworks._compat import _convert_from_v1_to_chat_completions
from langchain_fireworks._version import __version__
from langchain_fireworks.data._profiles import _PROFILES

logger = logging.getLogger(__name__)


_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)

_LANGSMITH_GATEWAY_DEFAULT_URL = "https://gateway.smith.langchain.com/fireworks"


def _resolve_gateway_base_url() -> str | None:
    raw = os.getenv("LANGSMITH_GATEWAY")
    if raw is None or raw.lower() in ("false", "0", "no"):
        return None
    if raw.lower() in ("true", "1", "yes"):
        return _LANGSMITH_GATEWAY_DEFAULT_URL
    return raw


def _get_default_model_profile(model_name: str) -> ModelProfile:
    default = _MODEL_PROFILES.get(model_name) or {}
    return default.copy()


def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:
    """Convert a dictionary to a LangChain message.

    Args:
        _dict: The dictionary.

    Returns:
        The LangChain message.

    """
    role = _dict.get("role")
    if role == "user":
        return HumanMessage(content=_dict.get("content", ""))
    if role == "assistant":
        # Fix for azure
        # Also Fireworks returns None for tool invocations
        content = _dict.get("content", "") or ""
        additional_kwargs: dict = {}
        if reasoning_content := _dict.get("reasoning_content"):
            additional_kwargs["reasoning_content"] = reasoning_content

        if function_call := _dict.get("function_call"):
            additional_kwargs["function_call"] = dict(function_call)

        tool_calls = []
        invalid_tool_calls = []
        if raw_tool_calls := _dict.get("tool_calls"):
            additional_kwargs["tool_calls"] = raw_tool_calls
            for raw_tool_call in raw_tool_calls:
                try:
                    tool_calls.append(parse_tool_call(raw_tool_call, return_id=True))
                except Exception as e:
                    invalid_tool_calls.append(
                        dict(make_invalid_tool_call(raw_tool_call, str(e)))
                    )
        return AIMessage(
            content=content,
            additional_kwargs=additional_kwargs,
            tool_calls=tool_calls,
            invalid_tool_calls=invalid_tool_calls,
        )
    if role == "system":
        return SystemMessage(content=_dict.get("content", ""))
    if role == "function":
        return FunctionMessage(
            content=_dict.get("content", ""), name=_dict.get("name", "")
        )
    if role == "tool":
        additional_kwargs = {}
        if "name" in _dict:
            additional_kwargs["name"] = _dict["name"]
        return ToolMessage(
            content=_dict.get("content", ""),
            tool_call_id=_dict.get("tool_call_id", ""),
            additional_kwargs=additional_kwargs,
        )
    return ChatMessage(content=_dict.get("content", ""), role=role or "")


def _allowed_content_part_keys() -> frozenset[str]:
    """Allowlist of wire-valid keys on a Fireworks content part.

    Derived at import time from the stainless-generated TypedDict so the
    allowlist tracks the upstream OpenAPI spec as `fireworks-ai` is bumped:
    new fields widen the allowlist for free, removed/renamed fields shrink it
    in lockstep. If the SDK reshuffles its module layout the import falls back
    to a conservative hand-coded set and emits a warning, and the layout test
    (`test_fireworks_sdk_request_layout_stable`) fails to surface the drift.
    """
    try:
        from typing import get_type_hints

        from fireworks.types.shared_params.chat_message import (
            ContentUnionMember1,
        )

        return frozenset(get_type_hints(ContentUnionMember1))
    except ImportError:
        logger.warning(
            "Could not import `fireworks.types.shared_params.chat_message."
            "ContentUnionMember1`; falling back to a conservative content-part "
            "key allowlist. Bump `fireworks-ai` or update "
            "`_allowed_content_part_keys` if the SDK has moved this type.",
        )
        return frozenset({"type", "text", "image_url", "video_url"})


_ALLOWED_CONTENT_PART_KEYS: frozenset[str] = _allowed_content_part_keys()


def _sanitize_chat_completions_content(content: Any) -> Any:
    """Strip non-wire keys from content blocks before serializing to Fireworks.

    Fireworks's chat completions endpoint rejects unknown fields on message
    content parts with `Extra inputs are not permitted, field: 'messages[N]
    .content.list[ChatMessageContent][i].<key>'`. This surfaces when a
    conversation accumulates AIMessages from a different provider (e.g.
    Anthropic's v1 streaming-reassembly `index` marker on text blocks, or the
    LangChain-internal `caller` key on `tool_use` blocks) and that history is
    later forwarded to a Fireworks-hosted model.

    For list content:
        - each block dict is filtered down to keys in
            `_ALLOWED_CONTENT_PART_KEYS` (sourced from the SDK TypedDict, so it
            stays in sync with the upstream spec).
        - if the result is a list of exactly one block that, post-strip, is
            `{"type": "text", "text": <str>}` and nothing else, it is coerced to
            a plain string. Fireworks's `content` union lists `str` first
            (`Input should be a valid string, field: 'messages[N].content.str'`),
            and the stricter shape avoids the union-validation noise on the
            server side.
    Non-list content (strings, None) passes through unchanged.
    """
    if not isinstance(content, list):
        return content
    sanitized: list[Any] = []
    for block in content:
        if isinstance(block, dict):
            sanitized.append(
                {k: v for k, v in block.items() if k in _ALLOWED_CONTENT_PART_KEYS}
            )
        else:
            sanitized.append(block)
    if (
        len(sanitized) == 1
        and isinstance(sanitized[0], dict)
        and set(sanitized[0]) == {"type", "text"}
        and sanitized[0]["type"] == "text"
        and isinstance(sanitized[0]["text"], str)
    ):
        return sanitized[0]["text"]
    return sanitized


def _format_message_content(content: Any) -> Any:
    """Format message content for the Fireworks chat completions wire format.

    Adapted from `langchain_openai.chat_models.base._format_message_content`,
    scoped to the chat completions API: drops content block types the wire
    format does not carry, translates canonical v0/v1 multimodal data blocks
    via `convert_to_openai_data_block(block, api="chat/completions")`, and
    converts legacy Anthropic-shape image blocks (`{"type": "image",
    "source": {...}}`) to OpenAI `image_url` blocks. String and non-list
    content are returned unchanged.

    Args:
        content: The message content. Strings and non-list values are
            returned as-is; lists are walked block by block.

    Returns:
        The formatted content, ready to be placed on the chat completions
        wire. List inputs return a new list with translations applied; other
        inputs are returned unchanged.
    """
    if not isinstance(content, list):
        return content
    formatted: list[Any] = []
    for block in content:
        if isinstance(block, dict) and "type" in block:
            btype = block["type"]
            if btype in (
                "tool_use",
                "thinking",
                "reasoning_content",
                "function_call",
                "code_interpreter_call",
            ):
                continue
            if is_data_content_block(block):
                formatted.append(
                    convert_to_openai_data_block(block, api="chat/completions")
                )
                continue
            if (
                btype == "image"
                and (source := block.get("source"))
                and isinstance(source, dict)
            ):
                if (
                    source.get("type") == "base64"
                    and (media_type := source.get("media_type"))
                    and (data := source.get("data"))
                ):
                    formatted.append(
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:{media_type};base64,{data}"},
                        }
                    )
                    continue
                if source.get("type") == "url" and (url := source.get("url")):
                    formatted.append({"type": "image_url", "image_url": {"url": url}})
                    continue
                continue
        formatted.append(block)
    return formatted


def _convert_message_to_dict(message: BaseMessage) -> dict:
    """Convert a LangChain message to a dictionary.

    Args:
        message: The LangChain message.

    Returns:
        The dictionary.

    """
    message_dict: dict[str, Any]
    if isinstance(message, ChatMessage):
        message_dict = {
            "role": message.role,
            "content": _sanitize_chat_completions_content(
                _format_message_content(message.content)
            ),
        }
    elif isinstance(message, HumanMessage):
        message_dict = {
            "role": "user",
            "content": _sanitize_chat_completions_content(
                _format_message_content(message.content)
            ),
        }
    elif isinstance(message, AIMessage):
        # Translate v1 content
        if message.response_metadata.get("output_version") == "v1":
            message = _convert_from_v1_to_chat_completions(message)
        message_dict = {
            "role": "assistant",
            "content": _sanitize_chat_completions_content(
                _format_message_content(message.content)
            ),
        }
        if "function_call" in message.additional_kwargs:
            message_dict["function_call"] = message.additional_kwargs["function_call"]
            # If function call only, content is None not empty string
            if message_dict["content"] == "":
                message_dict["content"] = None
        if message.tool_calls or message.invalid_tool_calls:
            message_dict["tool_calls"] = [
                _lc_tool_call_to_fireworks_tool_call(tc) for tc in message.tool_calls
            ] + [
                _lc_invalid_tool_call_to_fireworks_tool_call(tc)
                for tc in message.invalid_tool_calls
            ]
        elif "tool_calls" in message.additional_kwargs:
            message_dict["tool_calls"] = message.additional_kwargs["tool_calls"]
        # If tool calls only, content is None not empty string
        if "tool_calls" in message_dict and message_dict["content"] == "":
            message_dict["content"] = None
        else:
            pass
    elif isinstance(message, SystemMessage):
        message_dict = {
            "role": "system",
            "content": _sanitize_chat_completions_content(
                _format_message_content(message.content)
            ),
        }
    elif isinstance(message, FunctionMessage):
        message_dict = {
            "role": "function",
            "content": message.content,
            "name": message.name,
        }
    elif isinstance(message, ToolMessage):
        message_dict = {
            "role": "tool",
            "content": _sanitize_chat_completions_content(
                _format_message_content(message.content)
            ),
            "tool_call_id": message.tool_call_id,
        }
    else:
        msg = f"Got unknown type {message}"
        raise TypeError(msg)
    if "name" in message.additional_kwargs:
        message_dict["name"] = message.additional_kwargs["name"]
    return message_dict


def _usage_to_metadata(usage: Mapping[str, Any]) -> dict[str, int]:
    input_tokens = usage.get("prompt_tokens", 0)
    output_tokens = usage.get("completion_tokens", 0)
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": usage.get("total_tokens", input_tokens + output_tokens),
    }


def _convert_chunk_to_message_chunk(
    chunk: Mapping[str, Any], default_class: type[BaseMessageChunk]
) -> BaseMessageChunk:
    choices = chunk.get("choices") or []
    response_metadata: dict[str, Any] = {"model_provider": "fireworks"}
    if service_tier := chunk.get("service_tier"):
        response_metadata["service_tier"] = service_tier
    if not choices:
        # Final chunk emitted when `stream_options.include_usage=True`:
        # `choices` is empty and the chunk carries only `usage`.
        usage = chunk.get("usage")
        if not usage:
            logger.debug(
                "Received stream chunk with no choices and no usage: %s", chunk
            )
        usage_metadata = _usage_to_metadata(usage) if usage else None
        return AIMessageChunk(
            content="",
            usage_metadata=usage_metadata,  # type: ignore[arg-type]
            response_metadata=response_metadata,
        )
    choice = choices[0]
    _dict = choice["delta"]
    role = cast(str, _dict.get("role"))
    content = cast(str, _dict.get("content") or "")
    additional_kwargs: dict = {}
    tool_call_chunks: list[ToolCallChunk] = []
    if _dict.get("function_call"):
        function_call = dict(_dict["function_call"])
        if "name" in function_call and function_call["name"] is None:
            function_call["name"] = ""
        additional_kwargs["function_call"] = function_call
    if raw_tool_calls := _dict.get("tool_calls"):
        additional_kwargs["tool_calls"] = raw_tool_calls
        for rtc in raw_tool_calls:
            with contextlib.suppress(KeyError):
                tool_call_chunks.append(
                    create_tool_call_chunk(
                        name=rtc["function"].get("name"),
                        args=rtc["function"].get("arguments"),
                        id=rtc.get("id"),
                        index=rtc.get("index"),
                    )
                )
    if role == "user" or default_class == HumanMessageChunk:
        return HumanMessageChunk(content=content)
    if role == "assistant" or default_class == AIMessageChunk:
        usage = chunk.get("usage")
        usage_metadata = _usage_to_metadata(usage) if usage else None
        return AIMessageChunk(
            content=content,
            additional_kwargs=additional_kwargs,
            tool_call_chunks=tool_call_chunks,
            usage_metadata=usage_metadata,  # type: ignore[arg-type]
            response_metadata=response_metadata,
        )
    if role == "system" or default_class == SystemMessageChunk:
        return SystemMessageChunk(content=content)
    if role == "function" or default_class == FunctionMessageChunk:
        return FunctionMessageChunk(content=content, name=_dict["name"])
    if role == "tool" or default_class == ToolMessageChunk:
        return ToolMessageChunk(content=content, tool_call_id=_dict["tool_call_id"])
    if role or default_class == ChatMessageChunk:
        return ChatMessageChunk(content=content, role=role)
    return default_class(content=content)  # type: ignore[call-arg]


class _RetryableHTTPStatusError(FireworksError):
    """Internal marker for 5xx `httpx.HTTPStatusError` responses.

    The 1.x SDK wraps every status response into a typed `APIStatusError`
    subclass, so this path is defense-in-depth: it only fires when a raw
    `httpx.HTTPStatusError` escapes the SDK (e.g., a custom `http_client` or
    monkey-patched transport raises one directly). Promoting it here keeps the
    retryable set expressible as a list of classes for
    `create_base_retry_decorator`.
    """


_RETRYABLE_ERRORS: tuple[type[BaseException], ...] = (
    APIConnectionError,
    InternalServerError,
    RateLimitError,
    httpx.TimeoutException,
    httpx.TransportError,
    _RetryableHTTPStatusError,
)


def _promote_http_status_error(exc: httpx.HTTPStatusError) -> NoReturn:
    """Re-raise 5xx `httpx.HTTPStatusError` as a retryable marker."""
    if exc.response.status_code >= 500:
        msg = f"Retryable {exc.response.status_code} from Fireworks: {exc}"
        raise _RetryableHTTPStatusError(msg) from exc
    raise exc


class FireworksContextOverflowError(BadRequestError, ContextOverflowError):
    """`BadRequestError` raised when input exceeds Fireworks's context limit."""


def _handle_fireworks_invalid_request(e: BadRequestError) -> NoReturn:
    """Promote prompt-too-long errors to `FireworksContextOverflowError`."""
    if "prompt is too long" in str(e):
        raise FireworksContextOverflowError(
            str(e), response=e.response, body=e.body
        ) from e
    raise e


def _raise_empty_stream() -> NoReturn:
    """Raise a descriptive error when the SDK returns a zero-chunk stream."""
    msg = "Received empty stream from Fireworks"
    raise FireworksError(msg)


def _create_retry_decorator(
    llm: ChatFireworks,
    run_manager: AsyncCallbackManagerForLLMRun | CallbackManagerForLLMRun | None = None,
) -> Callable[[Any], Any]:
    """Return a tenacity retry decorator for Fireworks SDK calls.

    Retries live here rather than in the SDK so each attempt is visible to the
    LangChain `run_manager.on_retry` callback. The SDK's own retry layer is
    suppressed via `max_retries=0` on the client; see `validate_environment`.
    """
    # `max_retries` counts retries *after* the initial attempt (default lives on
    # the `ChatFireworks.max_retries` field). `create_base_retry_decorator`
    # forwards its `max_retries` to `stop_after_attempt`, which counts total
    # attempts — so offset by 1. `None` and `0` both mean "single attempt, no
    # retries".
    attempts = (llm.max_retries + 1) if llm.max_retries else 1
    return create_base_retry_decorator(
        error_types=list(_RETRYABLE_ERRORS),
        max_retries=attempts,
        run_manager=run_manager,
    )


def _prepare_sdk_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
    """Move fields the 1.x SDK does not model into `extra_body`.

    The Stainless-generated `chat.completions.create` signature has a fixed set
    of typed parameters. Fireworks accepts additional fields on the wire (notably
    `stream_options.include_usage`) that the SDK schema does not declare. The
    SDK exposes `extra_body` precisely for this — merge anything that looks
    extra-body-shaped into it so it lands in the JSON request body.

    If a caller supplies both `extra_body={"stream_options": ...}` and a
    top-level `stream_options=...`, the value already in `extra_body` wins
    (callers using `extra_body` are presumed to want explicit control); the
    discarded top-level value is logged.
    """
    extra_body = dict(kwargs.pop("extra_body", None) or {})
    top_level_stream_options = kwargs.pop("stream_options", None)
    if top_level_stream_options is not None:
        if "stream_options" in extra_body:
            logger.warning(
                "Both `extra_body['stream_options']` and a top-level "
                "`stream_options` were supplied; using `extra_body`'s value "
                "and discarding the top-level value.",
            )
        else:
            extra_body["stream_options"] = top_level_stream_options
    if extra_body:
        kwargs["extra_body"] = extra_body
    return kwargs


def _completion_with_retry(
    llm: ChatFireworks,
    run_manager: CallbackManagerForLLMRun | None = None,
    **kwargs: Any,
) -> Any:
    """Retry the sync completion call, including stream setup."""
    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
    kwargs = _prepare_sdk_kwargs(kwargs)

    @retry_decorator
    def _call() -> Any:
        try:
            result = llm.client.create(**kwargs)
        except httpx.HTTPStatusError as e:
            _promote_http_status_error(e)
        if kwargs.get("stream"):
            # The streaming generator is lazy — advance once so the HTTP
            # connection and any transport error happen inside the retry
            # boundary. `_prepend_chunk` then re-yields the consumed chunk
            # ahead of the rest so callers still see every event.
            try:
                iterator = iter(result)
                first = next(iterator)
            except StopIteration:
                _raise_empty_stream()
            except httpx.HTTPStatusError as e:
                _promote_http_status_error(e)
            return _prepend_chunk(first, iterator)
        return result

    return _call()


async def _acompletion_with_retry(
    llm: ChatFireworks,
    run_manager: AsyncCallbackManagerForLLMRun | None = None,
    **kwargs: Any,
) -> Any:
    """Retry the async completion call, including stream setup."""
    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
    kwargs = _prepare_sdk_kwargs(kwargs)

    @retry_decorator
    async def _call() -> Any:
        if kwargs.get("stream"):
            try:
                # 1.x async `create()` is a coroutine that resolves to an
                # `AsyncStream` when `stream=True`. Await it, then advance the
                # async iterator once inside the retry boundary so transport
                # errors surface here rather than at first downstream consumer.
                result = await llm.async_client.create(**kwargs)
                agen = result.__aiter__()
                first = await agen.__anext__()
            except StopAsyncIteration:
                _raise_empty_stream()
            except httpx.HTTPStatusError as e:
                _promote_http_status_error(e)
            return _aprepend_chunk(first, agen)
        try:
            return await llm.async_client.create(**kwargs)
        except httpx.HTTPStatusError as e:
            _promote_http_status_error(e)

    return await _call()


def _prepend_chunk(first: Any, rest: Iterator[Any]) -> Iterator[Any]:
    yield first
    yield from rest


async def _aprepend_chunk(first: Any, rest: AsyncIterator[Any]) -> AsyncIterator[Any]:
    yield first
    async for item in rest:
        yield item


class ChatFireworks(BaseChatModel):
    """`Fireworks` Chat large language models API.

    To use, you should have the
    environment variable `FIREWORKS_API_KEY` set with your API key.

    Any parameters that are valid to be passed to the fireworks.create call
    can be passed in, even if not explicitly saved on this class.

    Example:
        ```python
        from langchain_fireworks.chat_models import ChatFireworks

        model = ChatFireworks(model_name="accounts/fireworks/models/gpt-oss-120b")
        ```

    Fireworks request headers can be passed with `extra_headers`. For prompt
    caching, `x-session-affinity` pins requests to a replica so related calls can
    reuse the same prompt-cache session:

    ```python
    model.invoke(
        "Hello",
        extra_headers={"x-session-affinity": "user-42"},
    )
    ```

    The Fireworks SDK also accepts a typed `prompt_cache_key` field (passed as a
    regular keyword argument), which it treats as the preferred alternative to
    the raw `x-session-affinity` header:

    ```python
    model.invoke("Hello", prompt_cache_key="user-42")
    ```
    """

    @property
    def lc_secrets(self) -> dict[str, str]:
        return {"fireworks_api_key": "FIREWORKS_API_KEY"}

    @classmethod
    def get_lc_namespace(cls) -> list[str]:
        """Get the namespace of the LangChain object.

        Returns:
            `["langchain", "chat_models", "fireworks"]`
        """
        return ["langchain", "chat_models", "fireworks"]

    @property
    def lc_attributes(self) -> dict[str, Any]:
        attributes: dict[str, Any] = {}
        if self.fireworks_api_base:
            attributes["fireworks_api_base"] = self.fireworks_api_base

        return attributes

    @classmethod
    def is_lc_serializable(cls) -> bool:
        """Return whether this model can be serialized by LangChain."""
        return True

    client: Any = Field(default=None, exclude=True)
    """Internal `fireworks.Fireworks().chat.completions` resource.

    Constructed with `max_retries=0` so retries are owned by
    `_create_retry_decorator` (which surfaces each attempt to the LangChain
    `run_manager`). Callers reaching for this directly should set their own
    retry layer.
    """

    async_client: Any = Field(default=None, exclude=True)
    """Internal `fireworks.AsyncFireworks().chat.completions` resource.

    Constructed with `max_retries=0`; see `client`.
    """

    _sdk_client: Any = PrivateAttr(default=None)
    """Owning `fireworks.Fireworks` instance, retained so `close()` can call
    into the underlying HTTPX client. The 1.x SDK does not expose lifecycle
    methods on the `chat.completions` resource itself.
    """

    _async_sdk_client: Any = PrivateAttr(default=None)
    """Owning `fireworks.AsyncFireworks` instance; see `_sdk_client`."""

    model_name: str = Field(alias="model")
    """Model name to use."""

    @property
    def model(self) -> str:
        """Same as model_name."""
        return self.model_name

    temperature: float | None = None
    """What sampling temperature to use."""

    stop: str | list[str] | None = Field(default=None, alias="stop_sequences")
    """Default stop sequences."""

    model_kwargs: dict[str, Any] = Field(default_factory=dict)
    """Holds any model parameters valid for `create` call not explicitly specified."""

    fireworks_api_key: SecretStr = Field(
        alias="api_key",
        default_factory=lambda: SecretStr(
            (
                os.getenv("LANGSMITH_GATEWAY_API_KEY")
                if _resolve_gateway_base_url() is not None
                else None
            )
            or (
                "langsmith-gateway-oauth"
                if _resolve_gateway_base_url() is not None
                and os.getenv("LANGSMITH_GATEWAY_OAUTH_TOKEN") is not None
                else None
            )
            or secret_from_env(
                "FIREWORKS_API_KEY",
                error_message=(
                    "You must specify an api key. "
                    "You can pass it an argument as `api_key=...` or "
                    "set the environment variable `FIREWORKS_API_KEY`."
                ),
            )().get_secret_value()
        ),
    )
    """Fireworks API key.

    Automatically read from env variable `FIREWORKS_API_KEY` if not provided.

    If `LANGSMITH_GATEWAY` is enabled, `LANGSMITH_GATEWAY_API_KEY` takes precedence.
    """

    fireworks_api_base: str | None = Field(
        alias="base_url",
        default_factory=lambda: (
            _resolve_gateway_base_url()
            or from_env("FIREWORKS_API_BASE", default=None)()
        ),
    )
    """Base URL path for API requests, leave blank if not using a proxy or service
    emulator.

    If `LANGSMITH_GATEWAY` is set, it takes precedence over `FIREWORKS_API_BASE`.
    """

    langsmith_gateway_oauth_token: SecretStr | None = Field(
        default_factory=secret_from_env("LANGSMITH_GATEWAY_OAUTH_TOKEN", default=None),
        exclude=True,
    )
    """OAuth token used to authenticate this client to LangSmith Gateway.

    When set, the token is sent as an `Authorization` bearer credential and is
    never passed to Fireworks as an API key.
    """

    request_timeout: float | tuple[float, float] | Any | None = Field(
        default=None, alias="timeout"
    )
    """Timeout for requests to Fireworks completion API. Can be `float`,
    `httpx.Timeout` or `None`.
    """

    streaming: bool = False
    """Whether to stream the results or not."""

    stream_usage: bool = True
    """Whether to include usage metadata in streaming output.

    If `True`, a final empty-content chunk carrying `usage_metadata` is emitted
    during the stream. Set to `False` if the upstream model/proxy rejects
    `stream_options`, or pass `stream_options` explicitly via `model_kwargs` or
    a runtime kwarg to override.

    !!! version-added "Added in `langchain-fireworks` 1.2.0"

    !!! warning "Behavior changed in `langchain-fireworks` 1.2.0"

        Streaming now opts into `stream_options.include_usage` by default, and
        the final empty-`choices` chunk is surfaced as an `AIMessageChunk` with
        `usage_metadata` instead of being silently dropped.
    """

    n: int = 1
    """Number of chat completions to generate for each prompt."""

    max_tokens: int | None = None
    """Maximum number of tokens to generate."""

    max_retries: int | None = 2
    """Maximum number of retries after the initial attempt when generating.

    Retries use exponential backoff and trigger on transient errors:
    `RateLimitError`, `APIConnectionError` (including its `APITimeoutError`
    subclass), 5xx responses (including those that surface as
    `httpx.HTTPStatusError` rather than typed SDK errors), and underlying
    transport errors (`httpx.TimeoutException`, `httpx.TransportError`).
    A value of `None` or `0` disables retries.
    """

    service_tier: str | None = None
    """Service tier for the request.

    Forwarded as the `service_tier` field on the Fireworks chat completions
    request when set. Pass `'priority'` to opt into Fireworks' priority tier;
    leave as `None` to use the default tier.

    To use Fireworks' fast mode instead, select a fast-routed `model`; fast mode
    is not controlled by this field. See Fireworks'
    [serverless product docs](https://docs.fireworks.ai/guides/serverless-products)
    for the current list of fast routers and tiers.

    !!! version-added "Added in `langchain-fireworks` 1.3.0"
    """

    model_config = ConfigDict(
        populate_by_name=True,
    )

    @model_validator(mode="before")
    @classmethod
    def build_extra(cls, values: dict[str, Any]) -> Any:
        """Build extra kwargs from additional params that were passed in."""
        all_required_field_names = get_pydantic_field_names(cls)
        return _build_model_kwargs(values, all_required_field_names)

    @model_validator(mode="after")
    def _set_fireworks_chat_version(self) -> Self:
        """Set package version in metadata."""
        self._add_version("langchain-fireworks", __version__)
        return self

    @model_validator(mode="after")
    def validate_environment(self) -> Self:
        """Validate that api key and python package exists in environment."""
        if self.n < 1:
            msg = "n must be at least 1."
            raise ValueError(msg)
        if self.n > 1 and self.streaming:
            msg = "n must be 1 when streaming."
            raise ValueError(msg)

        api_key = self.fireworks_api_key.get_secret_value()
        base_url = self.fireworks_api_base
        # 0.x accepted a `(connect, read)` tuple. 1.x's SDK only accepts a
        # float, `httpx.Timeout`, or `None` — normalize so existing user code
        # keeps working.
        if isinstance(self.request_timeout, tuple):
            connect, read = self.request_timeout
            timeout: Any = httpx.Timeout(read, connect=connect)
        else:
            timeout = self.request_timeout
        # `langchain-fireworks` owns retry/backoff via `_create_retry_decorator`
        # so the LangChain `run_manager` sees each attempt. Suppress the
        # SDK's built-in retry layer to avoid double-retrying.
        if not self.client:
            http_client = self._gateway_oauth_http_client()
            self._sdk_client = Fireworks(
                api_key=api_key,
                base_url=base_url,
                timeout=timeout,
                max_retries=0,
                http_client=http_client,
            )
            self.client = self._sdk_client.chat.completions
        if not self.async_client:
            http_async_client = self._gateway_oauth_async_http_client()
            self._async_sdk_client = AsyncFireworks(
                api_key=api_key,
                base_url=base_url,
                timeout=timeout,
                max_retries=0,
                http_client=http_async_client,
            )
            self.async_client = self._async_sdk_client.chat.completions
        return self

    def _is_langsmith_gateway(self) -> bool:
        """Return whether this instance is configured to use LangSmith Gateway."""
        return self.fireworks_api_base == _resolve_gateway_base_url()

    def _gateway_oauth_http_client(self) -> httpx.Client | None:
        """Build a synchronous OAuth client when Gateway OAuth is enabled."""
        if (
            self.langsmith_gateway_oauth_token is None
            or not self._is_langsmith_gateway()
        ):
            return None
        return httpx.Client(
            auth=LangSmithGatewayOAuth(
                self.langsmith_gateway_oauth_token.get_secret_value()
            )
        )

    def _gateway_oauth_async_http_client(self) -> httpx.AsyncClient | None:
        """Build an asynchronous OAuth client when Gateway OAuth is enabled."""
        if (
            self.langsmith_gateway_oauth_token is None
            or not self._is_langsmith_gateway()
        ):
            return None
        return httpx.AsyncClient(
            auth=LangSmithGatewayOAuth(
                self.langsmith_gateway_oauth_token.get_secret_value()
            )
        )

    def close(self) -> None:
        """Close the underlying sync HTTP client.

        After calling, sync invocations on this model will raise. Async
        invocations remain available until `aclose()` is also called. Safe to
        call multiple times.
        """
        if self._sdk_client is not None:
            self._sdk_client.close()

    async def aclose(self) -> None:
        """Close the underlying async HTTP client.

        Releases the aiohttp-backed connector that the 1.x SDK uses by
        default. Without this, transient `ChatFireworks` instances can leak
        an `Unclosed connector` warning at GC if the event loop has already
        stopped. Safe to call multiple times.
        """
        if self._async_sdk_client is not None:
            await self._async_sdk_client.close()

    def _resolve_model_profile(self) -> ModelProfile | None:
        return _get_default_model_profile(self.model_name) or None

    @property
    def _default_params(self) -> dict[str, Any]:
        """Get the default parameters for calling Fireworks API."""
        params = {
            "model": self.model_name,
            "stream": self.streaming,
            "n": self.n,
            "stop": self.stop,
            **self.model_kwargs,
        }
        if self.temperature is not None:
            params["temperature"] = self.temperature
        if self.max_tokens is not None:
            params["max_tokens"] = self.max_tokens
        if self.service_tier is not None:
            params["service_tier"] = self.service_tier
        return params

    def _get_ls_params(
        self, stop: list[str] | None = None, **kwargs: Any
    ) -> LangSmithParams:
        """Get standard params for tracing."""
        params = self._get_invocation_params(stop=stop, **kwargs)
        ls_params = LangSmithParams(
            ls_provider="fireworks",
            ls_model_name=params.get("model", self.model_name),
            ls_model_type="chat",
            ls_temperature=params.get("temperature", self.temperature),
        )
        if ls_max_tokens := params.get("max_tokens", self.max_tokens):
            ls_params["ls_max_tokens"] = ls_max_tokens
        if ls_stop := stop or params.get("stop", None):
            ls_params["ls_stop"] = ls_stop
        return ls_params

    def _combine_llm_outputs(self, llm_outputs: list[dict | None]) -> dict:
        overall_token_usage: dict = {}
        system_fingerprint = None
        for output in llm_outputs:
            if output is None:
                # Happens in streaming
                continue
            token_usage = output["token_usage"]
            if token_usage is not None:
                for k, v in token_usage.items():
                    if k in overall_token_usage:
                        overall_token_usage[k] += v
                    else:
                        overall_token_usage[k] = v
            if system_fingerprint is None:
                system_fingerprint = output.get("system_fingerprint")
        combined = {"token_usage": overall_token_usage, "model_name": self.model_name}
        if system_fingerprint:
            combined["system_fingerprint"] = system_fingerprint
        return combined

    def _stream(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {**params, **kwargs, "stream": True}
        if self.stream_usage and "stream_options" not in params:
            params["stream_options"] = {"include_usage": True}

        default_chunk_class: type[BaseMessageChunk] = AIMessageChunk
        try:
            stream = _completion_with_retry(
                self, run_manager=run_manager, messages=message_dicts, **params
            )
        except BadRequestError as e:
            _handle_fireworks_invalid_request(e)
        for chunk in stream:
            if not isinstance(chunk, dict):
                chunk = chunk.model_dump()
            message_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)
            generation_info: dict[str, Any] = {}
            logprobs = None
            if choices := chunk.get("choices"):
                choice = choices[0]
                if finish_reason := choice.get("finish_reason"):
                    generation_info["finish_reason"] = finish_reason
                    generation_info["model_name"] = self.model_name
                logprobs = choice.get("logprobs")
                if logprobs:
                    generation_info["logprobs"] = logprobs
            default_chunk_class = message_chunk.__class__
            generation_chunk = ChatGenerationChunk(
                message=message_chunk, generation_info=generation_info or None
            )
            if run_manager:
                run_manager.on_llm_new_token(
                    generation_chunk.text, chunk=generation_chunk, logprobs=logprobs
                )
            yield generation_chunk

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        stream: bool | None = None,  # noqa: FBT001
        **kwargs: Any,
    ) -> ChatResult:
        should_stream = stream if stream is not None else self.streaming
        if should_stream:
            stream_iter = self._stream(
                messages, stop=stop, run_manager=run_manager, **kwargs
            )
            return generate_from_stream(stream_iter)
        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {
            **params,
            **({"stream": stream} if stream is not None else {}),
            **kwargs,
        }
        try:
            response = _completion_with_retry(
                self, run_manager=run_manager, messages=message_dicts, **params
            )
        except BadRequestError as e:
            _handle_fireworks_invalid_request(e)
        return self._create_chat_result(response)

    def _create_message_dicts(
        self, messages: list[BaseMessage], stop: list[str] | None
    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        params = self._default_params
        if stop is not None:
            params["stop"] = stop
        message_dicts = [_convert_message_to_dict(m) for m in messages]
        return message_dicts, params

    def _create_chat_result(self, response: dict | BaseModel) -> ChatResult:
        generations = []
        if not isinstance(response, dict):
            response = response.model_dump()
        token_usage = response.get("usage", {})
        service_tier = response.get("service_tier")
        for res in response["choices"]:
            message = _convert_dict_to_message(res["message"])
            if isinstance(message, AIMessage):
                if token_usage:
                    message.usage_metadata = {
                        "input_tokens": token_usage.get("prompt_tokens", 0),
                        "output_tokens": token_usage.get("completion_tokens", 0),
                        "total_tokens": token_usage.get("total_tokens", 0),
                    }
                    message.response_metadata["model_provider"] = "fireworks"
                    message.response_metadata["model_name"] = self.model_name
                if service_tier:
                    message.response_metadata["service_tier"] = service_tier
            generation_info = {"finish_reason": res.get("finish_reason")}
            if "logprobs" in res:
                generation_info["logprobs"] = res["logprobs"]
            gen = ChatGeneration(
                message=message,
                generation_info=generation_info,
            )
            generations.append(gen)
        llm_output = {
            "token_usage": token_usage,
            "system_fingerprint": response.get("system_fingerprint", ""),
        }
        if service_tier:
            llm_output["service_tier"] = service_tier
        return ChatResult(generations=generations, llm_output=llm_output)

    async def _astream(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: AsyncCallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {**params, **kwargs, "stream": True}
        if self.stream_usage and "stream_options" not in params:
            params["stream_options"] = {"include_usage": True}

        default_chunk_class: type[BaseMessageChunk] = AIMessageChunk
        try:
            stream = await _acompletion_with_retry(
                self, run_manager=run_manager, messages=message_dicts, **params
            )
        except BadRequestError as e:
            _handle_fireworks_invalid_request(e)
        async for chunk in stream:
            if not isinstance(chunk, dict):
                chunk = chunk.model_dump()
            message_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)
            generation_info: dict[str, Any] = {}
            logprobs = None
            if choices := chunk.get("choices"):
                choice = choices[0]
                if finish_reason := choice.get("finish_reason"):
                    generation_info["finish_reason"] = finish_reason
                    generation_info["model_name"] = self.model_name
                logprobs = choice.get("logprobs")
                if logprobs:
                    generation_info["logprobs"] = logprobs
            default_chunk_class = message_chunk.__class__
            generation_chunk = ChatGenerationChunk(
                message=message_chunk, generation_info=generation_info or None
            )
            if run_manager:
                await run_manager.on_llm_new_token(
                    token=generation_chunk.text,
                    chunk=generation_chunk,
                    logprobs=logprobs,
                )
            yield generation_chunk

    async def _agenerate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: AsyncCallbackManagerForLLMRun | None = None,
        stream: bool | None = None,  # noqa: FBT001
        **kwargs: Any,
    ) -> ChatResult:
        should_stream = stream if stream is not None else self.streaming
        if should_stream:
            stream_iter = self._astream(
                messages, stop=stop, run_manager=run_manager, **kwargs
            )
            return await agenerate_from_stream(stream_iter)

        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {
            **params,
            **({"stream": stream} if stream is not None else {}),
            **kwargs,
        }
        try:
            response = await _acompletion_with_retry(
                self, run_manager=run_manager, messages=message_dicts, **params
            )
        except BadRequestError as e:
            _handle_fireworks_invalid_request(e)
        return self._create_chat_result(response)

    @property
    def _identifying_params(self) -> dict[str, Any]:
        """Get the identifying parameters."""
        return {"model_name": self.model_name, **self._default_params}

    def _get_invocation_params(
        self, stop: list[str] | None = None, **kwargs: Any
    ) -> dict[str, Any]:
        """Get the parameters used to invoke the model."""
        return {
            "model": self.model_name,
            **super()._get_invocation_params(stop=stop),
            **self._default_params,
            **kwargs,
        }

    @property
    def _llm_type(self) -> str:
        """Return type of chat model."""
        return "fireworks-chat"

    def bind_tools(
        self,
        tools: Sequence[dict[str, Any] | type[BaseModel] | Callable | BaseTool],
        *,
        tool_choice: dict | str | bool | None = None,
        **kwargs: Any,
    ) -> Runnable[LanguageModelInput, AIMessage]:
        """Bind tool-like objects to this chat model.

        Assumes model is compatible with Fireworks tool-calling API.

        Args:
            tools: A list of tool definitions to bind to this chat model.

                Supports any tool definition handled by [`convert_to_openai_tool`][langchain_core.utils.function_calling.convert_to_openai_tool].
            tool_choice: Which tool to require the model to call.
                Must be the name of the single provided function,
                `'auto'` to automatically determine which function to call
                with the option to not call any function, `'any'` to enforce that some
                function is called, or a dict of the form:
                `{"type": "function", "function": {"name": <<tool_name>>}}`.
            **kwargs: Any additional parameters to pass to
                `langchain_fireworks.chat_models.ChatFireworks.bind`
        """  # noqa: E501
        strict = kwargs.pop("strict", None)
        formatted_tools = [
            convert_to_openai_tool(tool, strict=strict) for tool in tools
        ]
        if tool_choice is not None and tool_choice:
            if isinstance(tool_choice, str) and (
                tool_choice not in ("auto", "any", "none")
            ):
                tool_choice = {"type": "function", "function": {"name": tool_choice}}
            if isinstance(tool_choice, bool):
                if len(tools) > 1:
                    msg = (
                        "tool_choice can only be True when there is one tool. Received "
                        f"{len(tools)} tools."
                    )
                    raise ValueError(msg)
                tool_name = formatted_tools[0]["function"]["name"]
                tool_choice = {
                    "type": "function",
                    "function": {"name": tool_name},
                }

            kwargs["tool_choice"] = tool_choice
        return super().bind(tools=formatted_tools, **kwargs)

    def with_structured_output(
        self,
        schema: dict | type[BaseModel] | None = None,
        *,
        method: Literal[
            "function_calling", "json_mode", "json_schema"
        ] = "function_calling",
        include_raw: bool = False,
        **kwargs: Any,
    ) -> Runnable[LanguageModelInput, dict | BaseModel]:
        """Model wrapper that returns outputs formatted to match the given schema.

        Args:
            schema: The output schema. Can be passed in as:

                - An OpenAI function/tool schema,
                - A JSON Schema,
                - A `TypedDict` class,
                - Or a Pydantic class.

                If `schema` is a Pydantic class then the model output will be a
                Pydantic instance of that class, and the model-generated fields will be
                validated by the Pydantic class. Otherwise the model output will be a
                dict and will not be validated.

                See `langchain_core.utils.function_calling.convert_to_openai_tool` for
                more on how to properly specify types and descriptions of schema fields
                when specifying a Pydantic or `TypedDict` class.

            method: The method for steering model generation, one of:

                - `'function_calling'`:
                    Uses Fireworks's [tool-calling features](https://docs.fireworks.ai/guides/function-calling).
                - `'json_schema'`:
                    Uses Fireworks's [structured output feature](https://docs.fireworks.ai/structured-responses/structured-response-formatting).
                - `'json_mode'`:
                    Uses Fireworks's [JSON mode feature](https://docs.fireworks.ai/structured-responses/structured-response-formatting).

                !!! warning "Behavior changed in `langchain-fireworks` 0.2.8"

                    Added support for `'json_schema'`.

            include_raw:
                If `False` then only the parsed structured output is returned.

                If an error occurs during model output parsing it will be raised.

                If `True` then both the raw model response (a `BaseMessage`) and the
                parsed model response will be returned.

                If an error occurs during output parsing it will be caught and returned
                as well.

                The final output is always a `dict` with keys `'raw'`, `'parsed'`, and
                `'parsing_error'`.

            kwargs:
                Any additional parameters to pass to the `langchain.runnable.Runnable`
                constructor.

        Returns:
            A `Runnable` that takes same inputs as a
                `langchain_core.language_models.chat.BaseChatModel`. If `include_raw` is
                `False` and `schema` is a Pydantic class, `Runnable` outputs an instance
                of `schema` (i.e., a Pydantic object). Otherwise, if `include_raw` is
                `False` then `Runnable` outputs a `dict`.

                If `include_raw` is `True`, then `Runnable` outputs a `dict` with keys:

                - `'raw'`: `BaseMessage`
                - `'parsed'`: `None` if there was a parsing error, otherwise the type
                    depends on the `schema` as described above.
                - `'parsing_error'`: `BaseException | None`

        Example: schema=Pydantic class, method="function_calling", include_raw=False:

        ```python
        from typing import Optional

        from langchain_fireworks import ChatFireworks
        from pydantic import BaseModel, Field


        class AnswerWithJustification(BaseModel):
            '''An answer to the user question along with justification for the answer.'''

            answer: str
            # If we provide default values and/or descriptions for fields, these will be passed
            # to the model. This is an important part of improving a model's ability to
            # correctly return structured outputs.
            justification: str | None = Field(
                default=None, description="A justification for the answer."
            )


        model = ChatFireworks(
            model="accounts/fireworks/models/gpt-oss-120b",
            temperature=0,
        )
        structured_model = model.with_structured_output(AnswerWithJustification)

        structured_model.invoke(
            "What weighs more a pound of bricks or a pound of feathers"
        )

        # -> AnswerWithJustification(
        #     answer='They weigh the same',
        #     justification='Both a pound of bricks and a pound of feathers weigh one pound. The weight is the same, but the volume or density of the objects may differ.'
        # )
        ```

        Example: schema=Pydantic class, method="function_calling", include_raw=True:

        ```python
        from langchain_fireworks import ChatFireworks
        from pydantic import BaseModel


        class AnswerWithJustification(BaseModel):
            '''An answer to the user question along with justification for the answer.'''

            answer: str
            justification: str


        model = ChatFireworks(
            model="accounts/fireworks/models/gpt-oss-120b",
            temperature=0,
        )
        structured_model = model.with_structured_output(
            AnswerWithJustification, include_raw=True
        )

        structured_model.invoke(
            "What weighs more a pound of bricks or a pound of feathers"
        )
        # -> {
        #     'raw': AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Ao02pnFYXD6GN1yzc0uXPsvF', 'function': {'arguments': '{"answer":"They weigh the same.","justification":"Both a pound of bricks and a pound of feathers weigh one pound. The weight is the same, but the volume or density of the objects may differ."}', 'name': 'AnswerWithJustification'}, 'type': 'function'}]}),
        #     'parsed': AnswerWithJustification(answer='They weigh the same.', justification='Both a pound of bricks and a pound of feathers weigh one pound. The weight is the same, but the volume or density of the objects may differ.'),
        #     'parsing_error': None
        # }
        ```

        Example: schema=TypedDict class, method="function_calling", include_raw=False:

        ```python
        from typing_extensions import Annotated, TypedDict

        from langchain_fireworks import ChatFireworks


        class AnswerWithJustification(TypedDict):
            '''An answer to the user question along with justification for the answer.'''

            answer: str
            justification: Annotated[
                str | None, None, "A justification for the answer."
            ]


        model = ChatFireworks(
            model="accounts/fireworks/models/gpt-oss-120b",
            temperature=0,
        )
        structured_model = model.with_structured_output(AnswerWithJustification)

        structured_model.invoke(
            "What weighs more a pound of bricks or a pound of feathers"
        )
        # -> {
        #     'answer': 'They weigh the same',
        #     'justification': 'Both a pound of bricks and a pound of feathers weigh one pound. The weight is the same, but the volume and density of the two substances differ.'
        # }
        ```

        Example: schema=OpenAI function schema, method="function_calling", include_raw=False:

        ```python
        from langchain_fireworks import ChatFireworks

        oai_schema = {
            "name": "AnswerWithJustification",
            "description": "An answer to the user question along with justification for the answer.",
            "parameters": {
                "type": "object",
                "properties": {
                    "answer": {"type": "string"},
                    "justification": {
                        "description": "A justification for the answer.",
                        "type": "string",
                    },
                },
                "required": ["answer"],
            },
        }

        model = ChatFireworks(
            model="accounts/fireworks/models/gpt-oss-120b",
            temperature=0,
        )
        structured_model = model.with_structured_output(oai_schema)

        structured_model.invoke(
            "What weighs more a pound of bricks or a pound of feathers"
        )
        # -> {
        #     'answer': 'They weigh the same',
        #     'justification': 'Both a pound of bricks and a pound of feathers weigh one pound. The weight is the same, but the volume and density of the two substances differ.'
        # }
        ```

        Example: schema=Pydantic class, method="json_mode", include_raw=True:

        ```python
        from langchain_fireworks import ChatFireworks
        from pydantic import BaseModel


        class AnswerWithJustification(BaseModel):
            answer: str
            justification: str


        model = ChatFireworks(
            model="accounts/fireworks/models/gpt-oss-120b", temperature=0
        )
        structured_model = model.with_structured_output(
            AnswerWithJustification, method="json_mode", include_raw=True
        )

        structured_model.invoke(
            "Answer the following question. "
            "Make sure to return a JSON blob with keys 'answer' and 'justification'. "
            "What's heavier a pound of bricks or a pound of feathers?"
        )
        # -> {
        #     'raw': AIMessage(content='{"answer": "They are both the same weight.", "justification": "Both a pound of bricks and a pound of feathers weigh one pound. The difference lies in the volume and density of the materials, not the weight."}'),
        #     'parsed': AnswerWithJustification(answer='They are both the same weight.', justification='Both a pound of bricks and a pound of feathers weigh one pound. The difference lies in the volume and density of the materials, not the weight.'),
        #     'parsing_error': None
        # }
        ```

        Example: schema=None, method="json_mode", include_raw=True:

        ```python
        structured_model = model.with_structured_output(
            method="json_mode", include_raw=True
        )

        structured_model.invoke(
            "Answer the following question. "
            "Make sure to return a JSON blob with keys 'answer' and 'justification'. "
            "What's heavier a pound of bricks or a pound of feathers?"
        )
        # -> {
        #     'raw': AIMessage(content='{"answer": "They are both the same weight.", "justification": "Both a pound of bricks and a pound of feathers weigh one pound. The difference lies in the volume and density of the materials, not the weight."}'),
        #     'parsed': {
        #         'answer': 'They are both the same weight.',
        #         'justification': 'Both a pound of bricks and a pound of feathers weigh one pound. The difference lies in the volume and density of the materials, not the weight.'
        #     },
        #     'parsing_error': None
        # }
        ```

        """  # noqa: E501
        _ = kwargs.pop("strict", None)
        if kwargs:
            msg = f"Received unsupported arguments {kwargs}"
            raise ValueError(msg)
        is_pydantic_schema = _is_pydantic_class(schema)
        if method == "function_calling":
            if schema is None:
                msg = (
                    "schema must be specified when method is 'function_calling'. "
                    "Received None."
                )
                raise ValueError(msg)
            formatted_tool = convert_to_openai_tool(schema)
            tool_name = formatted_tool["function"]["name"]
            llm = self.bind_tools(
                [schema],
                tool_choice=tool_name,
                ls_structured_output_format={
                    "kwargs": {"method": "function_calling"},
                    "schema": formatted_tool,
                },
            )
            if is_pydantic_schema:
                output_parser: OutputParserLike = PydanticToolsParser(
                    tools=[schema],  # type: ignore[list-item]
                    first_tool_only=True,  # type: ignore[list-item]
                )
            else:
                output_parser = JsonOutputKeyToolsParser(
                    key_name=tool_name, first_tool_only=True
                )
        elif method == "json_schema":
            if schema is None:
                msg = (
                    "schema must be specified when method is 'json_schema'. "
                    "Received None."
                )
                raise ValueError(msg)
            formatted_schema = convert_to_json_schema(schema)
            llm = self.bind(
                response_format={"type": "json_object", "schema": formatted_schema},
                ls_structured_output_format={
                    "kwargs": {"method": "json_schema"},
                    "schema": schema,
                },
            )
            output_parser = (
                PydanticOutputParser(pydantic_object=schema)  # type: ignore[arg-type]
                if is_pydantic_schema
                else JsonOutputParser()
            )
        elif method == "json_mode":
            llm = self.bind(
                response_format={"type": "json_object"},
                ls_structured_output_format={
                    "kwargs": {"method": "json_mode"},
                    "schema": schema,
                },
            )
            output_parser = (
                PydanticOutputParser(pydantic_object=schema)  # type: ignore[type-var, arg-type]
                if is_pydantic_schema
                else JsonOutputParser()
            )
        else:
            msg = (
                f"Unrecognized method argument. Expected one of 'function_calling' or "
                f"'json_mode'. Received: '{method}'"
            )
            raise ValueError(msg)

        if include_raw:
            parser_assign = RunnablePassthrough.assign(
                parsed=itemgetter("raw") | output_parser, parsing_error=lambda _: None
            )
            parser_none = RunnablePassthrough.assign(parsed=lambda _: None)
            parser_with_fallback = parser_assign.with_fallbacks(
                [parser_none], exception_key="parsing_error"
            )
            return RunnableMap(raw=llm) | parser_with_fallback
        return llm | output_parser


def _is_pydantic_class(obj: Any) -> bool:
    return isinstance(obj, type) and is_basemodel_subclass(obj)


def _lc_tool_call_to_fireworks_tool_call(tool_call: ToolCall) -> dict:
    return {
        "type": "function",
        "id": tool_call["id"],
        "function": {
            "name": tool_call["name"],
            "arguments": json.dumps(tool_call["args"], ensure_ascii=False),
        },
    }


def _lc_invalid_tool_call_to_fireworks_tool_call(
    invalid_tool_call: InvalidToolCall,
) -> dict:
    return {
        "type": "function",
        "id": invalid_tool_call["id"],
        "function": {
            "name": invalid_tool_call["name"],
            "arguments": invalid_tool_call["args"],
        },
    }