feat(model-profiles): distribute data across packages (#34024)

libs/partners/xai/langchain_xai/chat_models.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, TypeAlias
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast
 
 import openai
 from langchain_core.messages import AIMessageChunk
@@ -11,7 +11,13 @@ from langchain_openai.chat_models.base import BaseChatOpenAI
 from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
 from typing_extensions import Self
 
+from langchain_xai.data._profiles import _PROFILES
+
 if TYPE_CHECKING:
+    from langchain_core.language_models import (
+        ModelProfile,
+        ModelProfileRegistry,
+    )
     from langchain_core.language_models.chat_models import (
         LangSmithParams,
         LanguageModelInput,
@@ -23,6 +29,14 @@ _DictOrPydanticClass: TypeAlias = dict[str, Any] | type[BaseModel] | type
 _DictOrPydantic: TypeAlias = dict | BaseModel
 
 
+_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)
+
+
+def _get_default_model_profile(model_name: str) -> ModelProfile:
+    default = _MODEL_PROFILES.get(model_name) or {}
+    return default.copy()
+
+
 class ChatXAI(BaseChatOpenAI):  # type: ignore[override]
     r"""ChatXAI chat model.
 
@@ -514,6 +528,13 @@ class ChatXAI(BaseChatOpenAI):  # type: ignore[override]
             )
         return self
 
+    @model_validator(mode="after")
+    def _set_model_profile(self) -> Self:
+        """Set model profile if not overridden."""
+        if self.profile is None:
+            self.profile = _get_default_model_profile(self.model_name)
+        return self
+
     @property
     def _default_params(self) -> dict[str, Any]:
         """Get default parameters."""

libs/partners/xai/langchain_xai/data/__init__.py

@@ -0,0 +1 @@
+"""Model profile data. All edits should be made in profile_augmentations.toml."""

libs/partners/xai/langchain_xai/data/_profiles.py

@@ -0,0 +1,283 @@
+"""Auto-generated model profiles.
+
+DO NOT EDIT THIS FILE MANUALLY.
+This file is generated by the langchain-profiles CLI tool.
+
+It contains data derived from the models.dev project.
+
+Source: https://github.com/sst/models.dev
+License: MIT License
+
+To update these data, refer to the instructions here:
+
+https://docs.langchain.com/oss/python/langchain/models#updating-or-overwriting-profile-data
+"""
+
+from typing import Any
+
+_PROFILES: dict[str, dict[str, Any]] = {
+    "grok-4-fast-non-reasoning": {
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 30000,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-fast": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-4": {
+        "max_input_tokens": 256000,
+        "max_output_tokens": 64000,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-2-vision": {
+        "max_input_tokens": 8192,
+        "max_output_tokens": 4096,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-code-fast-1": {
+        "max_input_tokens": 256000,
+        "max_output_tokens": 10000,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-2": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-mini-fast-latest": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-2-vision-1212": {
+        "max_input_tokens": 8192,
+        "max_output_tokens": 4096,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-4-fast": {
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 30000,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-2-latest": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-4-1-fast": {
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 30000,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-2-1212": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-fast-latest": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-latest": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-2-vision-latest": {
+        "max_input_tokens": 8192,
+        "max_output_tokens": 4096,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-vision-beta": {
+        "max_input_tokens": 8192,
+        "max_output_tokens": 4096,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-mini": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-beta": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 4096,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-mini-latest": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+    "grok-4-1-fast-non-reasoning": {
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 30000,
+        "image_inputs": True,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": False,
+        "tool_calling": True,
+    },
+    "grok-3-mini-fast": {
+        "max_input_tokens": 131072,
+        "max_output_tokens": 8192,
+        "image_inputs": False,
+        "audio_inputs": False,
+        "video_inputs": False,
+        "image_outputs": False,
+        "audio_outputs": False,
+        "video_outputs": False,
+        "reasoning_output": True,
+        "tool_calling": True,
+    },
+}