1"""Fireworks chat wrapper."""23from __future__ import annotations45import contextlib6import json7import logging8from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence9from operator import itemgetter10from typing import (11    Any,12    Literal,13    NoReturn,14    TypeAlias,15    cast,16)1718import httpx19from fireworks import (20    APIConnectionError,21    AsyncFireworks,22    BadRequestError,23    Fireworks,24    FireworksError,25    InternalServerError,26    RateLimitError,27)28from langchain_core.callbacks import (29    AsyncCallbackManagerForLLMRun,30    CallbackManagerForLLMRun,31)32from langchain_core.exceptions import ContextOverflowError33from langchain_core.language_models import (34    LanguageModelInput,35    ModelProfile,36    ModelProfileRegistry,37)38from langchain_core.language_models.chat_models import (39    BaseChatModel,40    LangSmithParams,41    agenerate_from_stream,42    generate_from_stream,43)44from langchain_core.language_models.llms import create_base_retry_decorator45from langchain_core.messages import (46    AIMessage,47    AIMessageChunk,48    BaseMessage,49    BaseMessageChunk,50    ChatMessage,51    ChatMessageChunk,52    FunctionMessage,53    FunctionMessageChunk,54    HumanMessage,55    HumanMessageChunk,56    InvalidToolCall,57    SystemMessage,58    SystemMessageChunk,59    ToolCall,60    ToolMessage,61    ToolMessageChunk,62    UsageMetadata,63    is_data_content_block,64)65from langchain_core.messages.block_translators.openai import (66    convert_to_openai_data_block,67)68from langchain_core.messages.tool import (69    ToolCallChunk,70)71from langchain_core.messages.tool import (72    tool_call_chunk as create_tool_call_chunk,73)74from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser75from langchain_core.output_parsers.base import OutputParserLike76from langchain_core.output_parsers.openai_tools import (77    JsonOutputKeyToolsParser,78    PydanticToolsParser,79    make_invalid_tool_call,80    parse_tool_call,81)82from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult83from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough84from langchain_core.tools import BaseTool85from langchain_core.utils import (86    get_pydantic_field_names,87)88from langchain_core.utils.function_calling import (89    convert_to_json_schema,90    convert_to_openai_tool,91)92from langchain_core.utils.pydantic import is_basemodel_subclass93from langchain_core.utils.utils import _build_model_kwargs, from_env, secret_from_env94from pydantic import (95    BaseModel,96    ConfigDict,97    Field,98    PrivateAttr,99    SecretStr,100    model_validator,101)102from typing_extensions import Self103104from langchain_fireworks._compat import _convert_from_v1_to_chat_completions105from langchain_fireworks._version import __version__106from langchain_fireworks.data._profiles import _PROFILES107108logger = logging.getLogger(__name__)109110111_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)112113114def _get_default_model_profile(model_name: str) -> ModelProfile:115    default = _MODEL_PROFILES.get(model_name) or {}116    return default.copy()117118119def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage:120    """Convert a dictionary to a LangChain message.121122    Args:123        _dict: The dictionary.124125    Returns:126        The LangChain message.127128    """129    role = _dict.get("role")130    if role == "user":131        return HumanMessage(content=_dict.get("content", ""))132    if role == "assistant":133        # Fix for azure134        # Also Fireworks returns None for tool invocations135        content = _dict.get("content", "") or ""136        additional_kwargs: dict = {}137        if reasoning_content := _dict.get("reasoning_content"):138            additional_kwargs["reasoning_content"] = reasoning_content139140        if function_call := _dict.get("function_call"):141            additional_kwargs["function_call"] = dict(function_call)142143        tool_calls = []144        invalid_tool_calls = []145        if raw_tool_calls := _dict.get("tool_calls"):146            additional_kwargs["tool_calls"] = raw_tool_calls147            for raw_tool_call in raw_tool_calls:148                try:149                    tool_calls.append(parse_tool_call(raw_tool_call, return_id=True))150                except Exception as e:151                    invalid_tool_calls.append(152                        dict(make_invalid_tool_call(raw_tool_call, str(e)))153                    )154        return AIMessage(155            content=content,156            additional_kwargs=additional_kwargs,157            tool_calls=tool_calls,158            invalid_tool_calls=invalid_tool_calls,159        )160    if role == "system":161        return SystemMessage(content=_dict.get("content", ""))162    if role == "function":163        return FunctionMessage(164            content=_dict.get("content", ""), name=_dict.get("name", "")165        )166    if role == "tool":167        additional_kwargs = {}168        if "name" in _dict:169            additional_kwargs["name"] = _dict["name"]170        return ToolMessage(171            content=_dict.get("content", ""),172            tool_call_id=_dict.get("tool_call_id", ""),173            additional_kwargs=additional_kwargs,174        )175    return ChatMessage(content=_dict.get("content", ""), role=role or "")176177178def _allowed_content_part_keys() -> frozenset[str]:179    """Allowlist of wire-valid keys on a Fireworks content part.180181    Derived at import time from the stainless-generated TypedDict so the182    allowlist tracks the upstream OpenAPI spec as `fireworks-ai` is bumped:183    new fields widen the allowlist for free, removed/renamed fields shrink it184    in lockstep. If the SDK reshuffles its module layout the import falls back185    to a conservative hand-coded set and emits a warning, and the layout test186    (`test_fireworks_sdk_request_layout_stable`) fails to surface the drift.187    """188    try:189        from typing import get_type_hints190191        from fireworks.types.shared_params.chat_message import (192            ContentUnionMember1,193        )194195        return frozenset(get_type_hints(ContentUnionMember1))196    except ImportError:197        logger.warning(198            "Could not import `fireworks.types.shared_params.chat_message."199            "ContentUnionMember1`; falling back to a conservative content-part "200            "key allowlist. Bump `fireworks-ai` or update "201            "`_allowed_content_part_keys` if the SDK has moved this type.",202        )203        return frozenset({"type", "text", "image_url", "video_url"})204205206_ALLOWED_CONTENT_PART_KEYS: frozenset[str] = _allowed_content_part_keys()207208209def _sanitize_chat_completions_content(content: Any) -> Any:210    """Strip non-wire keys from content blocks before serializing to Fireworks.211212    Fireworks's chat completions endpoint rejects unknown fields on message213    content parts with `Extra inputs are not permitted, field: 'messages[N]214    .content.list[ChatMessageContent][i].<key>'`. This surfaces when a215    conversation accumulates AIMessages from a different provider (e.g.216    Anthropic's v1 streaming-reassembly `index` marker on text blocks, or the217    LangChain-internal `caller` key on `tool_use` blocks) and that history is218    later forwarded to a Fireworks-hosted model.219220    For list content:221        - each block dict is filtered down to keys in222            `_ALLOWED_CONTENT_PART_KEYS` (sourced from the SDK TypedDict, so it223            stays in sync with the upstream spec).224        - if the result is a list of exactly one block that, post-strip, is225            `{"type": "text", "text": <str>}` and nothing else, it is coerced to226            a plain string. Fireworks's `content` union lists `str` first227            (`Input should be a valid string, field: 'messages[N].content.str'`),228            and the stricter shape avoids the union-validation noise on the229            server side.230    Non-list content (strings, None) passes through unchanged.231    """232    if not isinstance(content, list):233        return content234    sanitized: list[Any] = []235    for block in content:236        if isinstance(block, dict):237            sanitized.append(238                {k: v for k, v in block.items() if k in _ALLOWED_CONTENT_PART_KEYS}239            )240        else:241            sanitized.append(block)242    if (243        len(sanitized) == 1244        and isinstance(sanitized[0], dict)245        and set(sanitized[0]) == {"type", "text"}246        and sanitized[0]["type"] == "text"247        and isinstance(sanitized[0]["text"], str)248    ):249        return sanitized[0]["text"]250    return sanitized251252253def _format_message_content(content: Any) -> Any:254    """Format message content for the Fireworks chat completions wire format.255256    Adapted from `langchain_openai.chat_models.base._format_message_content`,257    scoped to the chat completions API: drops content block types the wire258    format does not carry, translates canonical v0/v1 multimodal data blocks259    via `convert_to_openai_data_block(block, api="chat/completions")`, and260    converts legacy Anthropic-shape image blocks (`{"type": "image",261    "source": {...}}`) to OpenAI `image_url` blocks. String and non-list262    content are returned unchanged.263264    Args:265        content: The message content. Strings and non-list values are266            returned as-is; lists are walked block by block.267268    Returns:269        The formatted content, ready to be placed on the chat completions270        wire. List inputs return a new list with translations applied; other271        inputs are returned unchanged.272    """273    if not isinstance(content, list):274        return content275    formatted: list[Any] = []276    for block in content:277        if isinstance(block, dict) and "type" in block:278            btype = block["type"]279            if btype in (280                "tool_use",281                "thinking",282                "reasoning_content",283                "function_call",284                "code_interpreter_call",285            ):286                continue287            if is_data_content_block(block):288                formatted.append(289                    convert_to_openai_data_block(block, api="chat/completions")290                )291                continue292            if (293                btype == "image"294                and (source := block.get("source"))295                and isinstance(source, dict)296            ):297                if (298                    source.get("type") == "base64"299                    and (media_type := source.get("media_type"))300                    and (data := source.get("data"))301                ):302                    formatted.append(303                        {304                            "type": "image_url",305                            "image_url": {"url": f"data:{media_type};base64,{data}"},306                        }307                    )308                    continue309                if source.get("type") == "url" and (url := source.get("url")):310                    formatted.append({"type": "image_url", "image_url": {"url": url}})311                    continue312                continue313        formatted.append(block)314    return formatted315316317def _convert_message_to_dict(message: BaseMessage) -> dict:318    """Convert a LangChain message to a dictionary.319320    Args:321        message: The LangChain message.322323    Returns:324        The dictionary.325326    """327    message_dict: dict[str, Any]328    if isinstance(message, ChatMessage):329        message_dict = {330            "role": message.role,331            "content": _sanitize_chat_completions_content(332                _format_message_content(message.content)333            ),334        }335    elif isinstance(message, HumanMessage):336        message_dict = {337            "role": "user",338            "content": _sanitize_chat_completions_content(339                _format_message_content(message.content)340            ),341        }342    elif isinstance(message, AIMessage):343        # Translate v1 content344        if message.response_metadata.get("output_version") == "v1":345            message = _convert_from_v1_to_chat_completions(message)346        message_dict = {347            "role": "assistant",348            "content": _sanitize_chat_completions_content(349                _format_message_content(message.content)350            ),351        }352        if "function_call" in message.additional_kwargs:353            message_dict["function_call"] = message.additional_kwargs["function_call"]354            # If function call only, content is None not empty string355            if message_dict["content"] == "":356                message_dict["content"] = None357        if message.tool_calls or message.invalid_tool_calls:358            message_dict["tool_calls"] = [359                _lc_tool_call_to_fireworks_tool_call(tc) for tc in message.tool_calls360            ] + [361                _lc_invalid_tool_call_to_fireworks_tool_call(tc)362                for tc in message.invalid_tool_calls363            ]364        elif "tool_calls" in message.additional_kwargs:365            message_dict["tool_calls"] = message.additional_kwargs["tool_calls"]366        # If tool calls only, content is None not empty string367        if "tool_calls" in message_dict and message_dict["content"] == "":368            message_dict["content"] = None369        else:370            pass371    elif isinstance(message, SystemMessage):372        message_dict = {373            "role": "system",374            "content": _sanitize_chat_completions_content(375                _format_message_content(message.content)376            ),377        }378    elif isinstance(message, FunctionMessage):379        message_dict = {380            "role": "function",381            "content": message.content,382            "name": message.name,383        }384    elif isinstance(message, ToolMessage):385        message_dict = {386            "role": "tool",387            "content": _sanitize_chat_completions_content(388                _format_message_content(message.content)389            ),390            "tool_call_id": message.tool_call_id,391        }392    else:393        msg = f"Got unknown type {message}"394        raise TypeError(msg)395    if "name" in message.additional_kwargs:396        message_dict["name"] = message.additional_kwargs["name"]397    return message_dict398399400def _usage_to_metadata(usage: Mapping[str, Any]) -> UsageMetadata:401    input_tokens = usage.get("prompt_tokens") or 0402    output_tokens = usage.get("completion_tokens") or 0403    usage_metadata: UsageMetadata = {404        "input_tokens": input_tokens,405        "output_tokens": output_tokens,406        "total_tokens": usage.get("total_tokens") or input_tokens + output_tokens,407    }408    cached_tokens = (usage.get("prompt_tokens_details") or {}).get("cached_tokens")409    if cached_tokens is not None:410        usage_metadata["input_token_details"] = {"cache_read": cached_tokens}411    return usage_metadata412413414TokenUsageTree: TypeAlias = "int | dict[str, TokenUsageTree]"415"""Raw provider token usage: a tree of `int` leaves and nested `dict` nodes416(e.g. `prompt_tokens_details`).417418Modeled as a recursive alias so the merge helper's signature carries the shape419rather than leaving it to `Any`.420"""421422423def _update_token_usage(424    overall_token_usage: TokenUsageTree, new_usage: TokenUsageTree425) -> TokenUsageTree:426    """Recursively merge raw provider token usage across generations.427428    Token usage is a tree of `int` leaves (summed) and `dict` nodes such as429    `prompt_tokens_details` (merged key-by-key, skipping `None` values).430431    A type mismatch between the accumulator and the incoming value (e.g. an432    `int` on one side and a `dict` on the other) indicates malformed provider433    data and is raised rather than silently coerced. An entirely unexpected434    leaf type (neither `int` nor `dict`) is logged and passed through, so a435    telemetry anomaly degrades gracefully instead of failing the response.436    """437    if isinstance(new_usage, int):438        if not isinstance(overall_token_usage, int):439            msg = (440                "Got different types for token usage: "441                f"{new_usage!r} ({type(new_usage).__name__}) and "442                f"{overall_token_usage!r} ({type(overall_token_usage).__name__})"443            )444            raise ValueError(msg)445        return overall_token_usage + new_usage446    if isinstance(new_usage, dict):447        if not isinstance(overall_token_usage, dict):448            msg = (449                "Got different types for token usage: "450                f"{new_usage!r} ({type(new_usage).__name__}) and "451                f"{overall_token_usage!r} ({type(overall_token_usage).__name__})"452            )453            raise ValueError(msg)454        updated_token_usage = dict(overall_token_usage)455        for key, value in new_usage.items():456            if value is not None:457                # Seed a first-seen key with an empty node of the same kind so a458                # nested `dict` value merges rather than colliding with an `int`.459                default: TokenUsageTree = {} if isinstance(value, dict) else 0460                updated_token_usage[key] = _update_token_usage(461                    overall_token_usage.get(key, default), value462                )463        return updated_token_usage464    logger.warning("Unexpected type for token usage: %s", type(new_usage).__name__)465    return new_usage466467468def _convert_chunk_to_message_chunk(469    chunk: Mapping[str, Any], default_class: type[BaseMessageChunk]470) -> BaseMessageChunk:471    choices = chunk.get("choices") or []472    response_metadata: dict[str, Any] = {"model_provider": "fireworks"}473    if service_tier := chunk.get("service_tier"):474        response_metadata["service_tier"] = service_tier475    if not choices:476        # Final chunk emitted when `stream_options.include_usage=True`:477        # `choices` is empty and the chunk carries only `usage`.478        usage = chunk.get("usage")479        if not usage:480            logger.debug(481                "Received stream chunk with no choices and no usage: %s", chunk482            )483        usage_metadata = _usage_to_metadata(usage) if usage else None484        return AIMessageChunk(485            content="",486            usage_metadata=usage_metadata,487            response_metadata=response_metadata,488        )489    choice = choices[0]490    _dict = choice["delta"]491    role = cast(str, _dict.get("role"))492    content = cast(str, _dict.get("content") or "")493    additional_kwargs: dict = {}494    tool_call_chunks: list[ToolCallChunk] = []495    if _dict.get("function_call"):496        function_call = dict(_dict["function_call"])497        if "name" in function_call and function_call["name"] is None:498            function_call["name"] = ""499        additional_kwargs["function_call"] = function_call500    if raw_tool_calls := _dict.get("tool_calls"):501        additional_kwargs["tool_calls"] = raw_tool_calls502        for rtc in raw_tool_calls:503            with contextlib.suppress(KeyError):504                tool_call_chunks.append(505                    create_tool_call_chunk(506                        name=rtc["function"].get("name"),507                        args=rtc["function"].get("arguments"),508                        id=rtc.get("id"),509                        index=rtc.get("index"),510                    )511                )512    if role == "user" or default_class == HumanMessageChunk:513        return HumanMessageChunk(content=content)514    if role == "assistant" or default_class == AIMessageChunk:515        usage = chunk.get("usage")516        usage_metadata = _usage_to_metadata(usage) if usage else None517        return AIMessageChunk(518            content=content,519            additional_kwargs=additional_kwargs,520            tool_call_chunks=tool_call_chunks,521            usage_metadata=usage_metadata,522            response_metadata=response_metadata,523        )524    if role == "system" or default_class == SystemMessageChunk:525        return SystemMessageChunk(content=content)526    if role == "function" or default_class == FunctionMessageChunk:527        return FunctionMessageChunk(content=content, name=_dict["name"])528    if role == "tool" or default_class == ToolMessageChunk:529        return ToolMessageChunk(content=content, tool_call_id=_dict["tool_call_id"])530    if role or default_class == ChatMessageChunk:531        return ChatMessageChunk(content=content, role=role)532    return default_class(content=content)  # type: ignore[call-arg]533534535class _RetryableHTTPStatusError(FireworksError):536    """Internal marker for 5xx `httpx.HTTPStatusError` responses.537538    The 1.x SDK wraps every status response into a typed `APIStatusError`539    subclass, so this path is defense-in-depth: it only fires when a raw540    `httpx.HTTPStatusError` escapes the SDK (e.g., a custom `http_client` or541    monkey-patched transport raises one directly). Promoting it here keeps the542    retryable set expressible as a list of classes for543    `create_base_retry_decorator`.544    """545546547_RETRYABLE_ERRORS: tuple[type[BaseException], ...] = (548    APIConnectionError,549    InternalServerError,550    RateLimitError,551    httpx.TimeoutException,552    httpx.TransportError,553    _RetryableHTTPStatusError,554)555556557def _promote_http_status_error(exc: httpx.HTTPStatusError) -> NoReturn:558    """Re-raise 5xx `httpx.HTTPStatusError` as a retryable marker."""559    if exc.response.status_code >= 500:560        msg = f"Retryable {exc.response.status_code} from Fireworks: {exc}"561        raise _RetryableHTTPStatusError(msg) from exc562    raise exc563564565class FireworksContextOverflowError(BadRequestError, ContextOverflowError):566    """`BadRequestError` raised when input exceeds Fireworks's context limit."""567568569def _handle_fireworks_invalid_request(e: BadRequestError) -> NoReturn:570    """Promote prompt-too-long errors to `FireworksContextOverflowError`."""571    if "prompt is too long" in str(e):572        raise FireworksContextOverflowError(573            str(e), response=e.response, body=e.body574        ) from e575    raise e576577578def _raise_empty_stream() -> NoReturn:579    """Raise a descriptive error when the SDK returns a zero-chunk stream."""580    msg = "Received empty stream from Fireworks"581    raise FireworksError(msg)582583584def _create_retry_decorator(585    llm: ChatFireworks,586    run_manager: AsyncCallbackManagerForLLMRun | CallbackManagerForLLMRun | None = None,587) -> Callable[[Any], Any]:588    """Return a tenacity retry decorator for Fireworks SDK calls.589590    Retries live here rather than in the SDK so each attempt is visible to the591    LangChain `run_manager.on_retry` callback. The SDK's own retry layer is592    suppressed via `max_retries=0` on the client; see `validate_environment`.593    """594    # `max_retries` counts retries *after* the initial attempt (default lives on595    # the `ChatFireworks.max_retries` field). `create_base_retry_decorator`596    # forwards its `max_retries` to `stop_after_attempt`, which counts total597    # attempts — so offset by 1. `None` and `0` both mean "single attempt, no598    # retries".599    attempts = (llm.max_retries + 1) if llm.max_retries else 1600    return create_base_retry_decorator(601        error_types=list(_RETRYABLE_ERRORS),602        max_retries=attempts,603        run_manager=run_manager,604    )605606607def _prepare_sdk_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:608    """Move fields the 1.x SDK does not model into `extra_body`.609610    The Stainless-generated `chat.completions.create` signature has a fixed set611    of typed parameters. Fireworks accepts additional fields on the wire (notably612    `stream_options.include_usage`) that the SDK schema does not declare. The613    SDK exposes `extra_body` precisely for this — merge anything that looks614    extra-body-shaped into it so it lands in the JSON request body.615616    If a caller supplies both `extra_body={"stream_options": ...}` and a617    top-level `stream_options=...`, the value already in `extra_body` wins618    (callers using `extra_body` are presumed to want explicit control); the619    discarded top-level value is logged.620    """621    extra_body = dict(kwargs.pop("extra_body", None) or {})622    top_level_stream_options = kwargs.pop("stream_options", None)623    if top_level_stream_options is not None:624        if "stream_options" in extra_body:625            logger.warning(626                "Both `extra_body['stream_options']` and a top-level "627                "`stream_options` were supplied; using `extra_body`'s value "628                "and discarding the top-level value.",629            )630        else:631            extra_body["stream_options"] = top_level_stream_options632    if extra_body:633        kwargs["extra_body"] = extra_body634    return kwargs635636637def _completion_with_retry(638    llm: ChatFireworks,639    run_manager: CallbackManagerForLLMRun | None = None,640    **kwargs: Any,641) -> Any:642    """Retry the sync completion call, including stream setup."""643    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)644    kwargs = _prepare_sdk_kwargs(kwargs)645646    @retry_decorator647    def _call() -> Any:648        try:649            result = llm.client.create(**kwargs)650        except httpx.HTTPStatusError as e:651            _promote_http_status_error(e)652        if kwargs.get("stream"):653            # The streaming generator is lazy — advance once so the HTTP654            # connection and any transport error happen inside the retry655            # boundary. `_prepend_chunk` then re-yields the consumed chunk656            # ahead of the rest so callers still see every event.657            try:658                iterator = iter(result)659                first = next(iterator)660            except StopIteration:661                _raise_empty_stream()662            except httpx.HTTPStatusError as e:663                _promote_http_status_error(e)664            return _prepend_chunk(first, iterator)665        return result666667    return _call()668669670async def _acompletion_with_retry(671    llm: ChatFireworks,672    run_manager: AsyncCallbackManagerForLLMRun | None = None,673    **kwargs: Any,674) -> Any:675    """Retry the async completion call, including stream setup."""676    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)677    kwargs = _prepare_sdk_kwargs(kwargs)678679    @retry_decorator680    async def _call() -> Any:681        if kwargs.get("stream"):682            try:683                # 1.x async `create()` is a coroutine that resolves to an684                # `AsyncStream` when `stream=True`. Await it, then advance the685                # async iterator once inside the retry boundary so transport686                # errors surface here rather than at first downstream consumer.687                result = await llm.async_client.create(**kwargs)688                agen = result.__aiter__()689                first = await agen.__anext__()690            except StopAsyncIteration:691                _raise_empty_stream()692            except httpx.HTTPStatusError as e:693                _promote_http_status_error(e)694            return _aprepend_chunk(first, agen)695        try:696            return await llm.async_client.create(**kwargs)697        except httpx.HTTPStatusError as e:698            _promote_http_status_error(e)699700    return await _call()701702703def _prepend_chunk(first: Any, rest: Iterator[Any]) -> Iterator[Any]:704    yield first705    yield from rest706707708async def _aprepend_chunk(first: Any, rest: AsyncIterator[Any]) -> AsyncIterator[Any]:709    yield first710    async for item in rest:711        yield item712713714class ChatFireworks(BaseChatModel):715    """`Fireworks` Chat large language models API.716717    To use, you should have the718    environment variable `FIREWORKS_API_KEY` set with your API key.719720    Any parameters that are valid to be passed to the fireworks.create call721    can be passed in, even if not explicitly saved on this class.722723    Example:724        ```python725        from langchain_fireworks.chat_models import ChatFireworks726727        model = ChatFireworks(model_name="accounts/fireworks/models/gpt-oss-120b")728        ```729730    Fireworks request headers can be passed with `extra_headers`. For prompt731    caching, `x-session-affinity` pins requests to a replica so related calls can732    reuse the same prompt-cache session:733734    ```python735    model.invoke(736        "Hello",737        extra_headers={"x-session-affinity": "user-42"},738    )739    ```740741    The Fireworks SDK also accepts a typed `prompt_cache_key` field (passed as a742    regular keyword argument), which it treats as the preferred alternative to743    the raw `x-session-affinity` header:744745    ```python746    model.invoke("Hello", prompt_cache_key="user-42")747    ```748    """749750    @property751    def lc_secrets(self) -> dict[str, str]:752        return {"fireworks_api_key": "FIREWORKS_API_KEY"}753754    @classmethod755    def get_lc_namespace(cls) -> list[str]:756        """Get the namespace of the LangChain object.757758        Returns:759            `["langchain", "chat_models", "fireworks"]`760        """761        return ["langchain", "chat_models", "fireworks"]762763    @property764    def lc_attributes(self) -> dict[str, Any]:765        attributes: dict[str, Any] = {}766        if self.fireworks_api_base:767            attributes["fireworks_api_base"] = self.fireworks_api_base768769        return attributes770771    @classmethod772    def is_lc_serializable(cls) -> bool:773        """Return whether this model can be serialized by LangChain."""774        return True775776    client: Any = Field(default=None, exclude=True)777    """Internal `fireworks.Fireworks().chat.completions` resource.778779    Constructed with `max_retries=0` so retries are owned by780    `_create_retry_decorator` (which surfaces each attempt to the LangChain781    `run_manager`). Callers reaching for this directly should set their own782    retry layer.783    """784785    async_client: Any = Field(default=None, exclude=True)786    """Internal `fireworks.AsyncFireworks().chat.completions` resource.787788    Constructed with `max_retries=0`; see `client`.789    """790791    _sdk_client: Any = PrivateAttr(default=None)792    """Owning `fireworks.Fireworks` instance, retained so `close()` can call793    into the underlying HTTPX client. The 1.x SDK does not expose lifecycle794    methods on the `chat.completions` resource itself.795    """796797    _async_sdk_client: Any = PrivateAttr(default=None)798    """Owning `fireworks.AsyncFireworks` instance; see `_sdk_client`."""799800    model_name: str = Field(alias="model")801    """Model name to use."""802803    @property804    def model(self) -> str:805        """Same as model_name."""806        return self.model_name807808    temperature: float | None = None809    """What sampling temperature to use."""810811    stop: str | list[str] | None = Field(default=None, alias="stop_sequences")812    """Default stop sequences."""813814    model_kwargs: dict[str, Any] = Field(default_factory=dict)815    """Holds any model parameters valid for `create` call not explicitly specified."""816817    fireworks_api_key: SecretStr = Field(818        alias="api_key",819        default_factory=secret_from_env(820            "FIREWORKS_API_KEY",821            error_message=(822                "You must specify an api key. "823                "You can pass it an argument as `api_key=...` or "824                "set the environment variable `FIREWORKS_API_KEY`."825            ),826        ),827    )828    """Fireworks API key.829830    Automatically read from env variable `FIREWORKS_API_KEY` if not provided.831    """832833    fireworks_api_base: str | None = Field(834        alias="base_url", default_factory=from_env("FIREWORKS_API_BASE", default=None)835    )836    """Base URL path for API requests, leave blank if not using a proxy or service837    emulator.838    """839840    request_timeout: float | tuple[float, float] | Any | None = Field(841        default=None, alias="timeout"842    )843    """Timeout for requests to Fireworks completion API. Can be `float`,844    `httpx.Timeout` or `None`.845    """846847    streaming: bool = False848    """Whether to stream the results or not."""849850    stream_usage: bool = True851    """Whether to include usage metadata in streaming output.852853    If `True`, a final empty-content chunk carrying `usage_metadata` is emitted854    during the stream. Set to `False` if the upstream model/proxy rejects855    `stream_options`, or pass `stream_options` explicitly via `model_kwargs` or856    a runtime kwarg to override.857858    !!! version-added "Added in `langchain-fireworks` 1.2.0"859860    !!! warning "Behavior changed in `langchain-fireworks` 1.2.0"861862        Streaming now opts into `stream_options.include_usage` by default, and863        the final empty-`choices` chunk is surfaced as an `AIMessageChunk` with864        `usage_metadata` instead of being silently dropped.865    """866867    n: int = 1868    """Number of chat completions to generate for each prompt."""869870    max_tokens: int | None = None871    """Maximum number of tokens to generate."""872873    max_retries: int | None = 2874    """Maximum number of retries after the initial attempt when generating.875876    Retries use exponential backoff and trigger on transient errors:877    `RateLimitError`, `APIConnectionError` (including its `APITimeoutError`878    subclass), 5xx responses (including those that surface as879    `httpx.HTTPStatusError` rather than typed SDK errors), and underlying880    transport errors (`httpx.TimeoutException`, `httpx.TransportError`).881    A value of `None` or `0` disables retries.882    """883884    service_tier: str | None = None885    """Service tier for the request.886887    Forwarded as the `service_tier` field on the Fireworks chat completions888    request when set. Pass `'priority'` to opt into Fireworks' priority tier;889    leave as `None` to use the default tier.890891    To use Fireworks' fast mode instead, select a fast-routed `model`; fast mode892    is not controlled by this field. See Fireworks'893    [serverless product docs](https://docs.fireworks.ai/guides/serverless-products)894    for the current list of fast routers and tiers.895896    !!! version-added "Added in `langchain-fireworks` 1.3.0"897    """898899    model_config = ConfigDict(900        populate_by_name=True,901    )902903    @model_validator(mode="before")904    @classmethod905    def build_extra(cls, values: dict[str, Any]) -> Any:906        """Build extra kwargs from additional params that were passed in."""907        all_required_field_names = get_pydantic_field_names(cls)908        return _build_model_kwargs(values, all_required_field_names)909910    @model_validator(mode="after")911    def _set_fireworks_chat_version(self) -> Self:912        """Set package version in metadata."""913        self._add_version("langchain-fireworks", __version__)914        return self915916    @model_validator(mode="after")917    def validate_environment(self) -> Self:918        """Validate that api key and python package exists in environment."""919        if self.n < 1:920            msg = "n must be at least 1."921            raise ValueError(msg)922        if self.n > 1 and self.streaming:923            msg = "n must be 1 when streaming."924            raise ValueError(msg)925926        api_key = self.fireworks_api_key.get_secret_value()927        base_url = self.fireworks_api_base928        # 0.x accepted a `(connect, read)` tuple. 1.x's SDK only accepts a929        # float, `httpx.Timeout`, or `None` — normalize so existing user code930        # keeps working.931        if isinstance(self.request_timeout, tuple):932            connect, read = self.request_timeout933            timeout: Any = httpx.Timeout(read, connect=connect)934        else:935            timeout = self.request_timeout936        # `langchain-fireworks` owns retry/backoff via `_create_retry_decorator`937        # so the LangChain `run_manager` sees each attempt. Suppress the938        # SDK's built-in retry layer to avoid double-retrying.939        if not self.client:940            self._sdk_client = Fireworks(941                api_key=api_key,942                base_url=base_url,943                timeout=timeout,944                max_retries=0,945            )946            self.client = self._sdk_client.chat.completions947        if not self.async_client:948            self._async_sdk_client = AsyncFireworks(949                api_key=api_key,950                base_url=base_url,951                timeout=timeout,952                max_retries=0,953            )954            self.async_client = self._async_sdk_client.chat.completions955        return self956957    def close(self) -> None:958        """Close the underlying sync HTTP client.959960        After calling, sync invocations on this model will raise. Async961        invocations remain available until `aclose()` is also called. Safe to962        call multiple times.963        """964        if self._sdk_client is not None:965            self._sdk_client.close()966967    async def aclose(self) -> None:968        """Close the underlying async HTTP client.969970        Releases the aiohttp-backed connector that the 1.x SDK uses by971        default. Without this, transient `ChatFireworks` instances can leak972        an `Unclosed connector` warning at GC if the event loop has already973        stopped. Safe to call multiple times.974        """975        if self._async_sdk_client is not None:976            await self._async_sdk_client.close()977978    def _resolve_model_profile(self) -> ModelProfile | None:979        return _get_default_model_profile(self.model_name) or None980981    @property982    def _default_params(self) -> dict[str, Any]:983        """Get the default parameters for calling Fireworks API."""984        params = {985            "model": self.model_name,986            "stream": self.streaming,987            "n": self.n,988            "stop": self.stop,989            **self.model_kwargs,990        }991        if self.temperature is not None:992            params["temperature"] = self.temperature993        if self.max_tokens is not None:994            params["max_tokens"] = self.max_tokens995        if self.service_tier is not None:996            params["service_tier"] = self.service_tier997        return params998999    def _get_ls_params(1000        self, stop: list[str] | None = None, **kwargs: Any1001    ) -> LangSmithParams:1002        """Get standard params for tracing."""1003        params = self._get_invocation_params(stop=stop, **kwargs)1004        ls_params = LangSmithParams(1005            ls_provider="fireworks",1006            ls_model_name=params.get("model", self.model_name),1007            ls_model_type="chat",1008            ls_temperature=params.get("temperature", self.temperature),1009        )1010        if ls_max_tokens := params.get("max_tokens", self.max_tokens):1011            ls_params["ls_max_tokens"] = ls_max_tokens1012        if ls_stop := stop or params.get("stop", None):1013            ls_params["ls_stop"] = ls_stop1014        return ls_params10151016    def _combine_llm_outputs(self, llm_outputs: list[dict | None]) -> dict:1017        overall_token_usage: dict = {}1018        system_fingerprint = None1019        for output in llm_outputs:1020            if output is None:1021                # Happens in streaming1022                continue1023            token_usage = output.get("token_usage")1024            if token_usage is not None:1025                for k, v in token_usage.items():1026                    if v is None:1027                        continue1028                    if k in overall_token_usage:1029                        overall_token_usage[k] = _update_token_usage(1030                            overall_token_usage[k], v1031                        )1032                    else:1033                        overall_token_usage[k] = v1034            if system_fingerprint is None:1035                system_fingerprint = output.get("system_fingerprint")1036        combined = {"token_usage": overall_token_usage, "model_name": self.model_name}1037        if system_fingerprint:1038            combined["system_fingerprint"] = system_fingerprint1039        return combined10401041    def _stream(1042        self,1043        messages: list[BaseMessage],1044        stop: list[str] | None = None,1045        run_manager: CallbackManagerForLLMRun | None = None,1046        **kwargs: Any,1047    ) -> Iterator[ChatGenerationChunk]:1048        message_dicts, params = self._create_message_dicts(messages, stop)1049        params = {**params, **kwargs, "stream": True}1050        if self.stream_usage and "stream_options" not in params:1051            params["stream_options"] = {"include_usage": True}10521053        default_chunk_class: type[BaseMessageChunk] = AIMessageChunk1054        try:1055            stream = _completion_with_retry(1056                self, run_manager=run_manager, messages=message_dicts, **params1057            )1058        except BadRequestError as e:1059            _handle_fireworks_invalid_request(e)1060        for chunk in stream:1061            if not isinstance(chunk, dict):1062                chunk = chunk.model_dump()1063            message_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)1064            generation_info: dict[str, Any] = {}1065            logprobs = None1066            if choices := chunk.get("choices"):1067                choice = choices[0]1068                if finish_reason := choice.get("finish_reason"):1069                    generation_info["finish_reason"] = finish_reason1070                    generation_info["model_name"] = self.model_name1071                logprobs = choice.get("logprobs")1072                if logprobs:1073                    generation_info["logprobs"] = logprobs1074            default_chunk_class = message_chunk.__class__1075            generation_chunk = ChatGenerationChunk(1076                message=message_chunk, generation_info=generation_info or None1077            )1078            if run_manager:1079                run_manager.on_llm_new_token(1080                    generation_chunk.text, chunk=generation_chunk, logprobs=logprobs1081                )1082            yield generation_chunk10831084    def _generate(1085        self,1086        messages: list[BaseMessage],1087        stop: list[str] | None = None,1088        run_manager: CallbackManagerForLLMRun | None = None,1089        stream: bool | None = None,  # noqa: FBT0011090        **kwargs: Any,1091    ) -> ChatResult:1092        should_stream = stream if stream is not None else self.streaming1093        if should_stream:1094            stream_iter = self._stream(1095                messages, stop=stop, run_manager=run_manager, **kwargs1096            )1097            return generate_from_stream(stream_iter)1098        message_dicts, params = self._create_message_dicts(messages, stop)1099        params = {1100            **params,1101            **({"stream": stream} if stream is not None else {}),1102            **kwargs,1103        }1104        try:1105            response = _completion_with_retry(1106                self, run_manager=run_manager, messages=message_dicts, **params1107            )1108        except BadRequestError as e:1109            _handle_fireworks_invalid_request(e)1110        return self._create_chat_result(response)11111112    def _create_message_dicts(1113        self, messages: list[BaseMessage], stop: list[str] | None1114    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:1115        params = self._default_params1116        if stop is not None:1117            params["stop"] = stop1118        message_dicts = [_convert_message_to_dict(m) for m in messages]1119        return message_dicts, params11201121    def _create_chat_result(self, response: dict | BaseModel) -> ChatResult:1122        generations = []1123        if not isinstance(response, dict):1124            response = response.model_dump()1125        token_usage = response.get("usage", {})1126        service_tier = response.get("service_tier")1127        for res in response["choices"]:1128            message = _convert_dict_to_message(res["message"])1129            if isinstance(message, AIMessage):1130                if token_usage:1131                    message.usage_metadata = _usage_to_metadata(token_usage)1132                    message.response_metadata["model_provider"] = "fireworks"1133                    message.response_metadata["model_name"] = self.model_name1134                if service_tier:1135                    message.response_metadata["service_tier"] = service_tier1136            generation_info = {"finish_reason": res.get("finish_reason")}1137            if "logprobs" in res:1138                generation_info["logprobs"] = res["logprobs"]1139            gen = ChatGeneration(1140                message=message,1141                generation_info=generation_info,1142            )1143            generations.append(gen)1144        llm_output = {1145            "token_usage": token_usage,1146            "system_fingerprint": response.get("system_fingerprint", ""),1147        }1148        if service_tier:1149            llm_output["service_tier"] = service_tier1150        return ChatResult(generations=generations, llm_output=llm_output)11511152    async def _astream(1153        self,1154        messages: list[BaseMessage],1155        stop: list[str] | None = None,1156        run_manager: AsyncCallbackManagerForLLMRun | None = None,1157        **kwargs: Any,1158    ) -> AsyncIterator[ChatGenerationChunk]:1159        message_dicts, params = self._create_message_dicts(messages, stop)1160        params = {**params, **kwargs, "stream": True}1161        if self.stream_usage and "stream_options" not in params:1162            params["stream_options"] = {"include_usage": True}11631164        default_chunk_class: type[BaseMessageChunk] = AIMessageChunk1165        try:1166            stream = await _acompletion_with_retry(1167                self, run_manager=run_manager, messages=message_dicts, **params1168            )1169        except BadRequestError as e:1170            _handle_fireworks_invalid_request(e)1171        async for chunk in stream:1172            if not isinstance(chunk, dict):1173                chunk = chunk.model_dump()1174            message_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)1175            generation_info: dict[str, Any] = {}1176            logprobs = None1177            if choices := chunk.get("choices"):1178                choice = choices[0]1179                if finish_reason := choice.get("finish_reason"):1180                    generation_info["finish_reason"] = finish_reason1181                    generation_info["model_name"] = self.model_name1182                logprobs = choice.get("logprobs")1183                if logprobs:1184                    generation_info["logprobs"] = logprobs1185            default_chunk_class = message_chunk.__class__1186            generation_chunk = ChatGenerationChunk(1187                message=message_chunk, generation_info=generation_info or None1188            )1189            if run_manager:1190                await run_manager.on_llm_new_token(1191                    token=generation_chunk.text,1192                    chunk=generation_chunk,1193                    logprobs=logprobs,1194                )1195            yield generation_chunk11961197    async def _agenerate(1198        self,1199        messages: list[BaseMessage],1200        stop: list[str] | None = None,1201        run_manager: AsyncCallbackManagerForLLMRun | None = None,1202        stream: bool | None = None,  # noqa: FBT0011203        **kwargs: Any,1204    ) -> ChatResult:1205        should_stream = stream if stream is not None else self.streaming1206        if should_stream:1207            stream_iter = self._astream(1208                messages, stop=stop, run_manager=run_manager, **kwargs1209            )1210            return await agenerate_from_stream(stream_iter)12111212        message_dicts, params = self._create_message_dicts(messages, stop)1213        params = {1214            **params,1215            **({"stream": stream} if stream is not None else {}),1216            **kwargs,1217        }1218        try:1219            response = await _acompletion_with_retry(1220                self, run_manager=run_manager, messages=message_dicts, **params1221            )1222        except BadRequestError as e:1223            _handle_fireworks_invalid_request(e)1224        return self._create_chat_result(response)12251226    @property1227    def _identifying_params(self) -> dict[str, Any]:1228        """Get the identifying parameters."""1229        return {"model_name": self.model_name, **self._default_params}12301231    def _get_invocation_params(1232        self, stop: list[str] | None = None, **kwargs: Any1233    ) -> dict[str, Any]:1234        """Get the parameters used to invoke the model."""1235        return {1236            "model": self.model_name,1237            **super()._get_invocation_params(stop=stop),1238            **self._default_params,1239            **kwargs,1240        }12411242    @property1243    def _llm_type(self) -> str:1244        """Return type of chat model."""1245        return "fireworks-chat"12461247    def bind_tools(1248        self,1249        tools: Sequence[dict[str, Any] | type[BaseModel] | Callable | BaseTool],1250        *,1251        tool_choice: dict | str | bool | None = None,1252        **kwargs: Any,1253    ) -> Runnable[LanguageModelInput, AIMessage]:1254        """Bind tool-like objects to this chat model.12551256        Assumes model is compatible with Fireworks tool-calling API.12571258        Args:1259            tools: A list of tool definitions to bind to this chat model.12601261                Supports any tool definition handled by [`convert_to_openai_tool`][langchain_core.utils.function_calling.convert_to_openai_tool].1262            tool_choice: Which tool to require the model to call.1263                Must be the name of the single provided function,1264                `'auto'` to automatically determine which function to call1265                with the option to not call any function, `'any'` to enforce that some1266                function is called, or a dict of the form:1267                `{"type": "function", "function": {"name": <<tool_name>>}}`.1268            **kwargs: Any additional parameters to pass to1269                `langchain_fireworks.chat_models.ChatFireworks.bind`1270        """  # noqa: E5011271        strict = kwargs.pop("strict", None)1272        formatted_tools = [1273            convert_to_openai_tool(tool, strict=strict) for tool in tools1274        ]1275        if tool_choice is not None and tool_choice:1276            if isinstance(tool_choice, str) and (1277                tool_choice not in ("auto", "any", "none")1278            ):1279                tool_choice = {"type": "function", "function": {"name": tool_choice}}1280            if isinstance(tool_choice, bool):1281                if len(tools) > 1:1282                    msg = (1283                        "tool_choice can only be True when there is one tool. Received "1284                        f"{len(tools)} tools."1285                    )1286                    raise ValueError(msg)1287                tool_name = formatted_tools[0]["function"]["name"]1288                tool_choice = {1289                    "type": "function",1290                    "function": {"name": tool_name},1291                }12921293            kwargs["tool_choice"] = tool_choice1294        return super().bind(tools=formatted_tools, **kwargs)12951296    def with_structured_output(1297        self,1298        schema: dict | type[BaseModel] | None = None,1299        *,1300        method: Literal[1301            "function_calling", "json_mode", "json_schema"1302        ] = "function_calling",1303        include_raw: bool = False,1304        **kwargs: Any,1305    ) -> Runnable[LanguageModelInput, dict | BaseModel]:1306        """Model wrapper that returns outputs formatted to match the given schema.13071308        Args:1309            schema: The output schema. Can be passed in as:13101311                - An OpenAI function/tool schema,1312                - A JSON Schema,1313                - A `TypedDict` class,1314                - Or a Pydantic class.13151316                If `schema` is a Pydantic class then the model output will be a1317                Pydantic instance of that class, and the model-generated fields will be1318                validated by the Pydantic class. Otherwise the model output will be a1319                dict and will not be validated.13201321                See `langchain_core.utils.function_calling.convert_to_openai_tool` for1322                more on how to properly specify types and descriptions of schema fields1323                when specifying a Pydantic or `TypedDict` class.13241325            method: The method for steering model generation, one of:13261327                - `'function_calling'`:1328                    Uses Fireworks's [tool-calling features](https://docs.fireworks.ai/guides/function-calling).1329                - `'json_schema'`:1330                    Uses Fireworks's [structured output feature](https://docs.fireworks.ai/structured-responses/structured-response-formatting).1331                - `'json_mode'`:1332                    Uses Fireworks's [JSON mode feature](https://docs.fireworks.ai/structured-responses/structured-response-formatting).13331334                !!! warning "Behavior changed in `langchain-fireworks` 0.2.8"13351336                    Added support for `'json_schema'`.13371338            include_raw:1339                If `False` then only the parsed structured output is returned.13401341                If an error occurs during model output parsing it will be raised.13421343                If `True` then both the raw model response (a `BaseMessage`) and the1344                parsed model response will be returned.13451346                If an error occurs during output parsing it will be caught and returned1347                as well.13481349                The final output is always a `dict` with keys `'raw'`, `'parsed'`, and1350                `'parsing_error'`.13511352            kwargs:1353                Any additional parameters to pass to the `langchain.runnable.Runnable`1354                constructor.13551356        Returns:1357            A `Runnable` that takes same inputs as a1358                `langchain_core.language_models.chat.BaseChatModel`. If `include_raw` is1359                `False` and `schema` is a Pydantic class, `Runnable` outputs an instance1360                of `schema` (i.e., a Pydantic object). Otherwise, if `include_raw` is1361                `False` then `Runnable` outputs a `dict`.13621363                If `include_raw` is `True`, then `Runnable` outputs a `dict` with keys:13641365                - `'raw'`: `BaseMessage`1366                - `'parsed'`: `None` if there was a parsing error, otherwise the type1367                    depends on the `schema` as described above.1368                - `'parsing_error'`: `BaseException | None`13691370        Example: schema=Pydantic class, method="function_calling", include_raw=False:13711372        ```python1373        from typing import Optional13741375        from langchain_fireworks import ChatFireworks1376        from pydantic import BaseModel, Field137713781379        class AnswerWithJustification(BaseModel):1380            '''An answer to the user question along with justification for the answer.'''13811382            answer: str1383            # If we provide default values and/or descriptions for fields, these will be passed1384            # to the model. This is an important part of improving a model's ability to1385            # correctly return structured outputs.1386            justification: str | None = Field(1387                default=None, description="A justification for the answer."1388            )138913901391        model = ChatFireworks(1392            model="accounts/fireworks/models/gpt-oss-120b",1393            temperature=0,1394        )1395        structured_model = model.with_structured_output(AnswerWithJustification)13961397        structured_model.invoke(1398            "What weighs more a pound of bricks or a pound of feathers"1399        )14001401        # -> AnswerWithJustification(1402        #     answer='They weigh the same',1403        #     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.'1404        # )1405        ```14061407        Example: schema=Pydantic class, method="function_calling", include_raw=True:14081409        ```python1410        from langchain_fireworks import ChatFireworks1411        from pydantic import BaseModel141214131414        class AnswerWithJustification(BaseModel):1415            '''An answer to the user question along with justification for the answer.'''14161417            answer: str1418            justification: str141914201421        model = ChatFireworks(1422            model="accounts/fireworks/models/gpt-oss-120b",1423            temperature=0,1424        )1425        structured_model = model.with_structured_output(1426            AnswerWithJustification, include_raw=True1427        )14281429        structured_model.invoke(1430            "What weighs more a pound of bricks or a pound of feathers"1431        )1432        # -> {1433        #     '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'}]}),1434        #     '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.'),1435        #     'parsing_error': None1436        # }1437        ```14381439        Example: schema=TypedDict class, method="function_calling", include_raw=False:14401441        ```python1442        from typing_extensions import Annotated, TypedDict14431444        from langchain_fireworks import ChatFireworks144514461447        class AnswerWithJustification(TypedDict):1448            '''An answer to the user question along with justification for the answer.'''14491450            answer: str1451            justification: Annotated[1452                str | None, None, "A justification for the answer."1453            ]145414551456        model = ChatFireworks(1457            model="accounts/fireworks/models/gpt-oss-120b",1458            temperature=0,1459        )1460        structured_model = model.with_structured_output(AnswerWithJustification)14611462        structured_model.invoke(1463            "What weighs more a pound of bricks or a pound of feathers"1464        )1465        # -> {1466        #     'answer': 'They weigh the same',1467        #     '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.'1468        # }1469        ```14701471        Example: schema=OpenAI function schema, method="function_calling", include_raw=False:14721473        ```python1474        from langchain_fireworks import ChatFireworks14751476        oai_schema = {1477            "name": "AnswerWithJustification",1478            "description": "An answer to the user question along with justification for the answer.",1479            "parameters": {1480                "type": "object",1481                "properties": {1482                    "answer": {"type": "string"},1483                    "justification": {1484                        "description": "A justification for the answer.",1485                        "type": "string",1486                    },1487                },1488                "required": ["answer"],1489            },1490        }14911492        model = ChatFireworks(1493            model="accounts/fireworks/models/gpt-oss-120b",1494            temperature=0,1495        )1496        structured_model = model.with_structured_output(oai_schema)14971498        structured_model.invoke(1499            "What weighs more a pound of bricks or a pound of feathers"1500        )1501        # -> {1502        #     'answer': 'They weigh the same',1503        #     '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.'1504        # }1505        ```15061507        Example: schema=Pydantic class, method="json_mode", include_raw=True:15081509        ```python1510        from langchain_fireworks import ChatFireworks1511        from pydantic import BaseModel151215131514        class AnswerWithJustification(BaseModel):1515            answer: str1516            justification: str151715181519        model = ChatFireworks(1520            model="accounts/fireworks/models/gpt-oss-120b", temperature=01521        )1522        structured_model = model.with_structured_output(1523            AnswerWithJustification, method="json_mode", include_raw=True1524        )15251526        structured_model.invoke(1527            "Answer the following question. "1528            "Make sure to return a JSON blob with keys 'answer' and 'justification'. "1529            "What's heavier a pound of bricks or a pound of feathers?"1530        )1531        # -> {1532        #     '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."}'),1533        #     '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.'),1534        #     'parsing_error': None1535        # }1536        ```15371538        Example: schema=None, method="json_mode", include_raw=True:15391540        ```python1541        structured_model = model.with_structured_output(1542            method="json_mode", include_raw=True1543        )15441545        structured_model.invoke(1546            "Answer the following question. "1547            "Make sure to return a JSON blob with keys 'answer' and 'justification'. "1548            "What's heavier a pound of bricks or a pound of feathers?"1549        )1550        # -> {1551        #     '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."}'),1552        #     'parsed': {1553        #         'answer': 'They are both the same weight.',1554        #         '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.'1555        #     },1556        #     'parsing_error': None1557        # }1558        ```15591560        """  # noqa: E5011561        _ = kwargs.pop("strict", None)1562        if kwargs:1563            msg = f"Received unsupported arguments {kwargs}"1564            raise ValueError(msg)1565        is_pydantic_schema = _is_pydantic_class(schema)1566        if method == "function_calling":1567            if schema is None:1568                msg = (1569                    "schema must be specified when method is 'function_calling'. "1570                    "Received None."1571                )1572                raise ValueError(msg)1573            formatted_tool = convert_to_openai_tool(schema)1574            tool_name = formatted_tool["function"]["name"]1575            llm = self.bind_tools(1576                [schema],1577                tool_choice=tool_name,1578                ls_structured_output_format={1579                    "kwargs": {"method": "function_calling"},1580                    "schema": formatted_tool,1581                },1582            )1583            if is_pydantic_schema:1584                output_parser: OutputParserLike = PydanticToolsParser(1585                    tools=[schema],  # type: ignore[list-item]1586                    first_tool_only=True,  # type: ignore[list-item]1587                )1588            else:1589                output_parser = JsonOutputKeyToolsParser(1590                    key_name=tool_name, first_tool_only=True1591                )1592        elif method == "json_schema":1593            if schema is None:1594                msg = (1595                    "schema must be specified when method is 'json_schema'. "1596                    "Received None."1597                )1598                raise ValueError(msg)1599            formatted_schema = convert_to_json_schema(schema)1600            llm = self.bind(1601                response_format={"type": "json_object", "schema": formatted_schema},1602                ls_structured_output_format={1603                    "kwargs": {"method": "json_schema"},1604                    "schema": schema,1605                },1606            )1607            output_parser = (1608                PydanticOutputParser(pydantic_object=schema)  # type: ignore[arg-type]1609                if is_pydantic_schema1610                else JsonOutputParser()1611            )1612        elif method == "json_mode":1613            llm = self.bind(1614                response_format={"type": "json_object"},1615                ls_structured_output_format={1616                    "kwargs": {"method": "json_mode"},1617                    "schema": schema,1618                },1619            )1620            output_parser = (1621                PydanticOutputParser(pydantic_object=schema)  # type: ignore[type-var, arg-type]1622                if is_pydantic_schema1623                else JsonOutputParser()1624            )1625        else:1626            msg = (1627                f"Unrecognized method argument. Expected one of 'function_calling' or "1628                f"'json_mode'. Received: '{method}'"1629            )1630            raise ValueError(msg)16311632        if include_raw:1633            parser_assign = RunnablePassthrough.assign(1634                parsed=itemgetter("raw") | output_parser, parsing_error=lambda _: None1635            )1636            parser_none = RunnablePassthrough.assign(parsed=lambda _: None)1637            parser_with_fallback = parser_assign.with_fallbacks(1638                [parser_none], exception_key="parsing_error"1639            )1640            return RunnableMap(raw=llm) | parser_with_fallback1641        return llm | output_parser164216431644def _is_pydantic_class(obj: Any) -> bool:1645    return isinstance(obj, type) and is_basemodel_subclass(obj)164616471648def _lc_tool_call_to_fireworks_tool_call(tool_call: ToolCall) -> dict:1649    return {1650        "type": "function",1651        "id": tool_call["id"],1652        "function": {1653            "name": tool_call["name"],1654            "arguments": json.dumps(tool_call["args"], ensure_ascii=False),1655        },1656    }165716581659def _lc_invalid_tool_call_to_fireworks_tool_call(1660    invalid_tool_call: InvalidToolCall,1661) -> dict:1662    return {1663        "type": "function",1664        "id": invalid_tool_call["id"],1665        "function": {1666            "name": invalid_tool_call["name"],1667            "arguments": invalid_tool_call["args"],1668        },1669    }