feat(aws): support model profiles (#768)

Author: ccurme

libs/aws/langchain_aws/chat_models/bedrock.py

@@ -24,6 +24,8 @@
     BaseChatModel,
     LangSmithParams,
     LanguageModelInput,
+    ModelProfile,
+    ModelProfileRegistry,
 )
 from langchain_core.language_models.chat_models import generate_from_stream
 from langchain_core.messages import (
@@ -49,9 +51,11 @@
 from langchain_core.utils.pydantic import TypeBaseModel, is_basemodel_subclass
 from langchain_core.utils.utils import _build_model_kwargs
 from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
 
 from langchain_aws.chat_models._compat import _convert_from_v1_to_anthropic
 from langchain_aws.chat_models.bedrock_converse import ChatBedrockConverse
+from langchain_aws.data._profiles import _PROFILES
 from langchain_aws.function_calling import (
     AnthropicTool,
     ToolsOutputParser,
@@ -77,6 +81,14 @@
 logger = logging.getLogger(__name__)
 
 
+_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)
+
+
+def _get_default_model_profile(model_name: str) -> ModelProfile:
+    default = _MODEL_PROFILES.get(model_name) or {}
+    return default.copy()
+
+
 def _convert_one_message_to_text_llama(message: BaseMessage) -> str:
     if isinstance(message, ChatMessage):
         message_text = f"\n\n{message.role.capitalize()}: {message.content}"
@@ -844,6 +856,14 @@ def build_extra(cls, values: dict[str, Any]) -> Any:
             }
         return values
 
+    @model_validator(mode="after")
+    def _set_model_profile(self) -> Self:
+        """Set model profile if not overridden."""
+        if self.profile is None:
+            model_id = re.sub(r"^[A-Za-z]{2}\.", "", self.model_id)
+            self.profile = _get_default_model_profile(model_id)
+        return self
+
     @property
     def lc_attributes(self) -> Dict[str, Any]:
         attributes: Dict[str, Any] = {}

libs/aws/langchain_aws/chat_models/bedrock_converse.py

@@ -24,7 +24,12 @@
 
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.exceptions import OutputParserException
-from langchain_core.language_models import BaseChatModel, LanguageModelInput
+from langchain_core.language_models import (
+    BaseChatModel,
+    LanguageModelInput,
+    ModelProfile,
+    ModelProfileRegistry,
+)
 from langchain_core.language_models.base import LangSmithParams
 from langchain_core.messages import (
     AIMessage,
@@ -58,6 +63,7 @@
 from typing_extensions import Self
 
 from langchain_aws.chat_models._compat import _convert_from_v1_to_converse
+from langchain_aws.data._profiles import _PROFILES
 from langchain_aws.function_calling import ToolsOutputParser
 from langchain_aws.utils import (
     count_tokens_api_supported_for_model,
@@ -66,6 +72,16 @@
 )
 
 logger = logging.getLogger(__name__)
+
+
+_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)
+
+
+def _get_default_model_profile(model_name: str) -> ModelProfile:
+    default = _MODEL_PROFILES.get(model_name) or {}
+    return default.copy()
+
+
 _BM = TypeVar("_BM", bound=BaseModel)
 
 EMPTY_CONTENT = "."
@@ -837,6 +853,14 @@ def validate_environment(self) -> Self:
 
         return self
 
+    @model_validator(mode="after")
+    def _set_model_profile(self) -> Self:
+        """Set model profile if not overridden."""
+        if self.profile is None:
+            model_id = re.sub(r"^[A-Za-z]{2}\.", "", self.model_id)
+            self.profile = _get_default_model_profile(model_id)
+        return self
+
     def _get_base_model(self) -> str:
         """Return base model id, stripping any regional prefix."""
 

libs/aws/langchain_aws/data/__init__.py

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