snowflake-ml-python 1.22.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

snowflake/ml/model/_client/ops/service_ops.py

@@ -1,4 +1,6 @@
+import base64
 import dataclasses
+import json
 import logging
 import pathlib
 import re
@@ -6,7 +8,9 @@ import tempfile
 import threading
 import time
 import warnings
-from typing import Any, Optional, Union, cast
+from typing import Any, Optional, Sequence, Union, cast
+
+from pydantic import TypeAdapter
 
 from snowflake import snowpark
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
@@ -14,9 +18,10 @@ from snowflake.ml._internal.utils import identifier, service_logger, sql_identif
 from snowflake.ml.jobs import job
 from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
 from snowflake.ml.model._client.model import batch_inference_specs
-from snowflake.ml.model._client.ops import deployment_step
+from snowflake.ml.model._client.ops import deployment_step, param_utils
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
+from snowflake.ml.model._signatures import core
 from snowflake.snowpark import async_job, exceptions, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -150,7 +155,6 @@ class ServiceOperator:
             self._model_deployment_spec = model_deployment_spec.ModelDeploymentSpec(
                 workspace_path=pathlib.Path(self._workspace.name)
             )
-        self._inference_autocapture_enabled = pc.PlatformCapabilities.get_instance().is_inference_autocapture_enabled()
 
     def __eq__(self, __value: object) -> bool:
         if not isinstance(__value, ServiceOperator):
@@ -211,10 +215,6 @@ class ServiceOperator:
         progress_status.update("preparing deployment artifacts...")
         progress_status.increment()
 
-        # If autocapture param is disabled, don't allow create service with autocapture
-        if not self._inference_autocapture_enabled and autocapture:
-            raise ValueError("Invalid Argument: Autocapture feature is not supported.")
-
         if self._workspace:
             stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
         else:
@@ -582,15 +582,10 @@ class ServiceOperator:
                 )
             for status in statuses:
                 if status.instance_id is not None:
-                    instance_status, container_status = None, None
-                    if status.instance_status is not None:
-                        instance_status = status.instance_status.value
-                    if status.container_status is not None:
-                        container_status = status.container_status.value
                     module_logger.info(
                         f"Instance[{status.instance_id}]: "
-                        f"instance status: {instance_status}, "
-                        f"container status: {container_status}, "
+                        f"instance status: {status.instance_status}, "
+                        f"container status: {status.container_status}, "
                         f"message: {status.message}"
                     )
             time.sleep(5)
@@ -930,6 +925,38 @@ class ServiceOperator:
         except exceptions.SnowparkSQLException:
             return False
 
+    @staticmethod
+    def _encode_params(params: Optional[dict[str, Any]]) -> Optional[str]:
+        """Encode params dictionary to a base64 string.
+
+        Args:
+            params: Optional dictionary of model inference parameters.
+
+        Returns:
+            Base64 encoded JSON string of the params, or None if input is None.
+        """
+        if params is None:
+            return None
+        return base64.b64encode(json.dumps(params).encode("utf-8")).decode("utf-8")
+
+    @staticmethod
+    def _encode_column_handling(
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
+    ) -> Optional[str]:
+        """Validate and encode column_handling to a base64 string.
+
+        Args:
+            column_handling: Optional dictionary mapping column names to file encoding options.
+
+        Returns:
+            Base64 encoded JSON string of the column handling options, or None if input is None.
+        """
+        if column_handling is None:
+            return None
+        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
+        validated_input = adapter.validate_python(column_handling)
+        return base64.b64encode(adapter.dump_json(validated_input)).decode("utf-8")
+
     def invoke_batch_job_method(
         self,
         *,
@@ -942,8 +969,9 @@ class ServiceOperator:
         image_repo_name: Optional[str],
         input_stage_location: str,
         input_file_pattern: str,
-        column_handling: Optional[str],
-        params: Optional[str],
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
+        params: Optional[dict[str, Any]],
+        signature_params: Optional[Sequence[core.BaseParamSpec]],
         output_stage_location: str,
         completion_filename: str,
         force_rebuild: bool,
@@ -954,7 +982,13 @@ class ServiceOperator:
         gpu_requests: Optional[str],
         replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
+        inference_engine_args: Optional[InferenceEngineArgs] = None,
     ) -> job.MLJob[Any]:
+        # Validate and encode params
+        param_utils.validate_params(params, signature_params)
+        params_encoded = self._encode_params(params)
+        column_handling_encoded = self._encode_column_handling(column_handling)
+
         database_name = self._database_name
         schema_name = self._schema_name
 
@@ -980,8 +1014,8 @@ class ServiceOperator:
             max_batch_rows=max_batch_rows,
             input_stage_location=input_stage_location,
             input_file_pattern=input_file_pattern,
-            column_handling=column_handling,
-            params=params,
+            column_handling=column_handling_encoded,
+            params=params_encoded,
             output_stage_location=output_stage_location,
             completion_filename=completion_filename,
             function_name=function_name,
@@ -992,11 +1026,17 @@ class ServiceOperator:
             replicas=replicas,
         )
 
-        self._model_deployment_spec.add_image_build_spec(
-            image_build_compute_pool_name=compute_pool_name,
-            fully_qualified_image_repo_name=self._get_image_repo_fqn(image_repo_name, database_name, schema_name),
-            force_rebuild=force_rebuild,
-        )
+        if inference_engine_args:
+            self._model_deployment_spec.add_inference_engine_spec(
+                inference_engine=inference_engine_args.inference_engine,
+                inference_engine_args=inference_engine_args.inference_engine_args_override,
+            )
+        else:
+            self._model_deployment_spec.add_image_build_spec(
+                image_build_compute_pool_name=compute_pool_name,
+                fully_qualified_image_repo_name=self._get_image_repo_fqn(image_repo_name, database_name, schema_name),
+                force_rebuild=force_rebuild,
+            )
 
         spec_yaml_str_or_path = self._model_deployment_spec.save()
 

snowflake/ml/model/_client/service/model_deployment_spec.py

@@ -363,7 +363,7 @@ class ModelDeploymentSpec:
         inference_engine: inference_engine_module.InferenceEngine,
         inference_engine_args: Optional[list[str]] = None,
     ) -> "ModelDeploymentSpec":
-        """Add inference engine specification. This must be called after self.add_service_spec().
+        """Add inference engine specification. This must be called after self.add_service_spec() or self.add_job_spec().
 
         Args:
             inference_engine: Inference engine.
@@ -376,9 +376,10 @@ class ModelDeploymentSpec:
             ValueError: If inference engine specification is called before add_service_spec().
             ValueError: If the argument does not have a '--' prefix.
         """
-        # TODO: needs to eventually support job deployment spec
-        if self._service is None:
-            raise ValueError("Inference engine specification must be called after add_service_spec().")
+        if self._service is None and self._job is None:
+            raise ValueError(
+                "Inference engine specification must be called after add_service_spec() or add_job_spec()."
+            )
 
         if inference_engine_args is None:
             inference_engine_args = []
@@ -431,11 +432,17 @@ class ModelDeploymentSpec:
 
             inference_engine_args = filtered_args
 
-        self._service.inference_engine_spec = model_deployment_spec_schema.InferenceEngineSpec(
+        inference_engine_spec = model_deployment_spec_schema.InferenceEngineSpec(
             # convert to string to be saved in the deployment spec
             inference_engine_name=inference_engine.value,
             inference_engine_args=inference_engine_args,
         )
+
+        if self._service:
+            self._service.inference_engine_spec = inference_engine_spec
+        elif self._job:
+            self._job.inference_engine_spec = inference_engine_spec
+
         return self
 
     def save(self) -> str:

snowflake/ml/model/_client/service/model_deployment_spec_schema.py

@@ -61,6 +61,7 @@ class Job(BaseModel):
     input: Input
     output: Output
     replicas: Optional[int] = None
+    inference_engine_spec: Optional[InferenceEngineSpec] = None
 
 
 class LogModelArgs(BaseModel):

snowflake/ml/model/_client/sql/service.py

@@ -47,22 +47,6 @@ class ServiceStatus(enum.Enum):
     INTERNAL_ERROR = "INTERNAL_ERROR"
 
 
-class InstanceStatus(enum.Enum):
-    PENDING = "PENDING"
-    READY = "READY"
-    FAILED = "FAILED"
-    TERMINATING = "TERMINATING"
-    SUCCEEDED = "SUCCEEDED"
-
-
-class ContainerStatus(enum.Enum):
-    PENDING = "PENDING"
-    READY = "READY"
-    DONE = "DONE"
-    FAILED = "FAILED"
-    UNKNOWN = "UNKNOWN"
-
-
 @dataclasses.dataclass
 class ServiceStatusInfo:
     """
@@ -72,8 +56,8 @@ class ServiceStatusInfo:
 
     service_status: ServiceStatus
     instance_id: Optional[int] = None
-    instance_status: Optional[InstanceStatus] = None
-    container_status: Optional[ContainerStatus] = None
+    instance_status: Optional[str] = None
+    container_status: Optional[str] = None
     message: Optional[str] = None
 
 
@@ -91,10 +75,8 @@ class ServiceSQLClient(_base._BaseSQLClient):
     DESC_SERVICE_SPEC_COL_NAME = "spec"
     DESC_SERVICE_CONTAINERS_SPEC_NAME = "containers"
     DESC_SERVICE_NAME_SPEC_NAME = "name"
-    DESC_SERVICE_PROXY_SPEC_ENV_NAME = "env"
-    PROXY_CONTAINER_NAME = "proxy"
+    DESC_SERVICE_ENV_SPEC_NAME = "env"
     MODEL_INFERENCE_AUTOCAPTURE_ENV_NAME = "SPCS_MODEL_INFERENCE_SERVER__AUTOCAPTURE_ENABLED"
-    FEATURE_MODEL_INFERENCE_AUTOCAPTURE = "FEATURE_MODEL_INFERENCE_AUTOCAPTURE"
 
     @contextlib.contextmanager
     def _qmark_paramstyle(self) -> Generator[None, None, None]:
@@ -272,17 +254,12 @@ class ServiceSQLClient(_base._BaseSQLClient):
         )
         statuses = []
         for r in rows:
-            instance_status, container_status = None, None
-            if r[ServiceSQLClient.INSTANCE_STATUS] is not None:
-                instance_status = InstanceStatus(r[ServiceSQLClient.INSTANCE_STATUS])
-            if r[ServiceSQLClient.CONTAINER_STATUS] is not None:
-                container_status = ContainerStatus(r[ServiceSQLClient.CONTAINER_STATUS])
             statuses.append(
                 ServiceStatusInfo(
                     service_status=ServiceStatus(r[ServiceSQLClient.SERVICE_STATUS]),
                     instance_id=r[ServiceSQLClient.INSTANCE_ID],
-                    instance_status=instance_status,
-                    container_status=container_status,
+                    instance_status=r[ServiceSQLClient.INSTANCE_STATUS],
+                    container_status=r[ServiceSQLClient.CONTAINER_STATUS],
                     message=r[ServiceSQLClient.MESSAGE] if include_message else None,
                 )
             )
@@ -306,39 +283,33 @@ class ServiceSQLClient(_base._BaseSQLClient):
         )
         return rows[0]
 
-    def get_proxy_container_autocapture(self, row: row.Row) -> bool:
-        """Extract whether service has autocapture enabled from proxy container spec.
+    def is_autocapture_enabled(self, row: row.Row) -> bool:
+        """Extract whether service has autocapture enabled in any container from service spec.
 
         Args:
             row: A row.Row object from DESCRIBE SERVICE containing the service YAML spec.
 
         Returns:
-            True if autocapture is enabled in proxy spec
-            False if disabled or not set in proxy spec
-            False if service doesn't have proxy container
+            True if autocapture is enabled in any container.
+            False if autocapture is disabled or not set in any container.
         """
-        try:
-            spec_yaml = row[ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME]
-            if spec_yaml is None:
-                return False
-            spec_raw = yaml.safe_load(spec_yaml)
-            if spec_raw is None:
-                return False
-            spec = cast(dict[str, Any], spec_raw)
-
-            proxy_container_spec = next(
-                container
-                for container in spec[ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME][
-                    ServiceSQLClient.DESC_SERVICE_CONTAINERS_SPEC_NAME
-                ]
-                if container[ServiceSQLClient.DESC_SERVICE_NAME_SPEC_NAME] == ServiceSQLClient.PROXY_CONTAINER_NAME
-            )
-            env = proxy_container_spec.get(ServiceSQLClient.DESC_SERVICE_PROXY_SPEC_ENV_NAME, {})
-            autocapture_enabled = env.get(ServiceSQLClient.MODEL_INFERENCE_AUTOCAPTURE_ENV_NAME, "false")
-            return str(autocapture_enabled).lower() == "true"
-
-        except StopIteration:
+        spec_yaml = row.as_dict().get(ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME)
+        if spec_yaml is None:
             return False
+        spec_raw = yaml.safe_load(spec_yaml)
+        if spec_raw is None:
+            return False
+        spec = cast(dict[str, Any], spec_raw)
+
+        containers = spec[ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME][
+            ServiceSQLClient.DESC_SERVICE_CONTAINERS_SPEC_NAME
+        ]
+        for container in containers:
+            env = container.get(ServiceSQLClient.DESC_SERVICE_ENV_SPEC_NAME, {})
+            autocapture_enabled = env.get(ServiceSQLClient.MODEL_INFERENCE_AUTOCAPTURE_ENV_NAME, "false")
+            if str(autocapture_enabled).lower() == "true":
+                return True
+        return False
 
     def drop_service(
         self,

snowflake/ml/model/_model_composer/model_method/infer_function.py_template

@@ -41,11 +41,29 @@ features = meta.signatures[TARGET_METHOD].inputs
 input_cols = [feature.name for feature in features]
 dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 
+# Load inference parameters from method signature (if any)
+param_cols = []
+param_defaults = {{}}
+if hasattr(meta.signatures[TARGET_METHOD], "params") and meta.signatures[TARGET_METHOD].params:
+    for param_spec in meta.signatures[TARGET_METHOD].params:
+        param_cols.append(param_spec.name)
+        param_defaults[param_spec.name] = param_spec.default_value
+
 
 # Actual function
 @vectorized(input=pd.DataFrame, max_batch_size=MAX_BATCH_SIZE, flatten_object_input={wide_input})
 def {function_name}(df: pd.DataFrame) -> dict:
-    df.columns = input_cols
-    input_df = df.astype(dtype=dtype_map)
-    predictions_df = runner(input_df[input_cols])
+    df.columns = input_cols + param_cols
+    input_df = df[input_cols].astype(dtype=dtype_map)
+
+    # Extract runtime param values, using defaults if None
+    method_params = {{}}
+    for col in param_cols:
+        val = df[col].iloc[0]
+        if val is None or pd.isna(val):
+            method_params[col] = param_defaults[col]
+        else:
+            method_params[col] = val
+
+    predictions_df = runner(input_df, **method_params)
     return predictions_df.replace({{pd.NA: None, np.nan: None}}).to_dict("records")

snowflake/ml/model/_model_composer/model_method/infer_partitioned.py_template

@@ -45,11 +45,29 @@ features = meta.signatures[TARGET_METHOD].inputs
 input_cols = [feature.name for feature in features]
 dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 
+# Load inference parameters from method signature (if any)
+param_cols = []
+param_defaults = {{}}
+if hasattr(meta.signatures[TARGET_METHOD], "params") and meta.signatures[TARGET_METHOD].params:
+    for param_spec in meta.signatures[TARGET_METHOD].params:
+        param_cols.append(param_spec.name)
+        param_defaults[param_spec.name] = param_spec.default_value
+
 
 # Actual table function
 class {function_name}:
     @vectorized(input=pd.DataFrame, flatten_object_input={wide_input})
     def end_partition(self, df: pd.DataFrame) -> pd.DataFrame:
-        df.columns = input_cols
-        input_df = df.astype(dtype=dtype_map)
-        return runner(input_df[input_cols])
+        df.columns = input_cols + param_cols
+        input_df = df[input_cols].astype(dtype=dtype_map)
+
+        # Extract runtime param values, using defaults if None
+        method_params = {{}}
+        for col in param_cols:
+            val = df[col].iloc[0]
+            if val is None or pd.isna(val):
+                method_params[col] = param_defaults[col]
+            else:
+                method_params[col] = val
+
+        return runner(input_df, **method_params)

snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template

@@ -40,11 +40,29 @@ features = meta.signatures[TARGET_METHOD].inputs
 input_cols = [feature.name for feature in features]
 dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 
+# Load inference parameters from method signature (if any)
+param_cols = []
+param_defaults = {{}}
+if hasattr(meta.signatures[TARGET_METHOD], "params") and meta.signatures[TARGET_METHOD].params:
+    for param_spec in meta.signatures[TARGET_METHOD].params:
+        param_cols.append(param_spec.name)
+        param_defaults[param_spec.name] = param_spec.default_value
+
 
 # Actual table function
 class {function_name}:
     @vectorized(input=pd.DataFrame, max_batch_size=MAX_BATCH_SIZE, flatten_object_input={wide_input})
     def process(self, df: pd.DataFrame) -> pd.DataFrame:
-        df.columns = input_cols
-        input_df = df.astype(dtype=dtype_map)
-        return runner(input_df[input_cols])
+        df.columns = input_cols + param_cols
+        input_df = df[input_cols].astype(dtype=dtype_map)
+
+        # Extract runtime param values, using defaults if None
+        method_params = {{}}
+        for col in param_cols:
+            val = df[col].iloc[0]
+            if val is None or pd.isna(val):
+                method_params[col] = param_defaults[col]
+            else:
+                method_params[col] = val
+
+        return runner(input_df, **method_params)

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -156,10 +156,12 @@ class ModelMethod:
                 f"Your parameter {param_spec.name} cannot be resolved as valid SQL identifier. "
                 "Try specifying `case_sensitive` as True."
             ) from e
+        # Convert None to "NULL" string so MANIFEST parser can interpret it as SQL NULL
+        default_value = "NULL" if param_spec.default_value is None else str(param_spec.default_value)
         return model_manifest_schema.ModelMethodSignatureFieldWithNameAndDefault(
             name=param_name.resolved(),
             type=type_utils.convert_sp_to_sf_type(param_spec.dtype.as_snowpark_type()),
-            default=param_spec.default_value,
+            default=default_value,
         )
 
     def save(

snowflake/ml/model/_packager/model_handlers/huggingface.py

@@ -1,3 +1,4 @@
+import io
 import json
 import logging
 import os
@@ -28,7 +29,10 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta as model_meta_api,
     model_meta_schema,
 )
-from snowflake.ml.model._signatures import utils as model_signature_utils
+from snowflake.ml.model._signatures import (
+    core as model_signature_core,
+    utils as model_signature_utils,
+)
 from snowflake.ml.model.models import (
     huggingface as huggingface_base,
     huggingface_pipeline,
@@ -530,7 +534,10 @@ class TransformersPipelineHandler(
                         # verify when the target method is __call__ and
                         # if the signature is default text-generation signature
                         # then use the HuggingFaceOpenAICompatibleModel to wrap the pipeline
-                        if signature == openai_signatures._OPENAI_CHAT_SIGNATURE_SPEC:
+                        if (
+                            signature == openai_signatures._OPENAI_CHAT_SIGNATURE_SPEC
+                            or signature == openai_signatures._OPENAI_CHAT_SIGNATURE_SPEC_WITH_CONTENT_FORMAT_STRING
+                        ):
                             wrapped_model = HuggingFaceOpenAICompatibleModel(pipeline=raw_model)
 
                             temp_res = X.apply(
@@ -554,6 +561,39 @@ class TransformersPipelineHandler(
                             else:
                                 input_data = X[signature.inputs[0].name].to_list()
                             temp_res = getattr(raw_model, target_method)(input_data)
+                    elif isinstance(raw_model, transformers.ImageClassificationPipeline):
+                        # Image classification expects PIL Images. Convert bytes to PIL Images.
+                        from PIL import Image
+
+                        input_col = signature.inputs[0].name
+                        images = [Image.open(io.BytesIO(img_bytes)) for img_bytes in X[input_col].to_list()]
+                        temp_res = getattr(raw_model, target_method)(images)
+                    elif isinstance(raw_model, transformers.AutomaticSpeechRecognitionPipeline):
+                        # ASR pipeline accepts a single audio input (bytes, str, np.ndarray, or dict),
+                        # not a list. Process each audio input individually.
+                        input_col = signature.inputs[0].name
+                        audio_inputs = X[input_col].to_list()
+                        temp_res = [getattr(raw_model, target_method)(audio) for audio in audio_inputs]
+                    elif isinstance(raw_model, transformers.VideoClassificationPipeline):
+                        # Video classification expects file paths. Write bytes to temp files,
+                        # process them, and clean up.
+                        import tempfile
+
+                        input_col = signature.inputs[0].name
+                        video_bytes_list = X[input_col].to_list()
+                        temp_file_paths = []
+                        temp_files = []
+                        try:
+                            # TODO: parallelize this if needed
+                            for video_bytes in video_bytes_list:
+                                temp_file = tempfile.NamedTemporaryFile()
+                                temp_file.write(video_bytes)
+                                temp_file_paths.append(temp_file.name)
+                                temp_files.append(temp_file)
+                            temp_res = getattr(raw_model, target_method)(temp_file_paths)
+                        finally:
+                            for f in temp_files:
+                                f.close()
                     else:
                         # TODO: remove conversational pipeline code
                         # For others, we could offer the whole dataframe as a list.
@@ -615,11 +655,14 @@ class TransformersPipelineHandler(
                         temp_res = [[conv.generated_responses] for conv in temp_res]
 
                     # To concat those who outputs a list with one input.
-                    if isinstance(temp_res[0], list):
-                        if isinstance(temp_res[0][0], dict):
-                            res = pd.DataFrame({0: temp_res})
-                        else:
-                            res = pd.DataFrame(temp_res)
+                    # if `signature.outputs` is single valued and is a FeatureGroupSpec,
+                    # we create a DataFrame with one column and the values are stored as a dictionary.
+                    # Otherwise, we create a DataFrame with the output as the column.
+                    if len(signature.outputs) == 1 and isinstance(
+                        signature.outputs[0], model_signature_core.FeatureGroupSpec
+                    ):
+                        # creating a dataframe with one column
+                        res = pd.DataFrame({signature.outputs[0].name: temp_res})
                     else:
                         res = pd.DataFrame(temp_res)
 
@@ -702,7 +745,6 @@ class HuggingFaceOpenAICompatibleModel:
         self.pipeline = pipeline
         self.model = self.pipeline.model
         self.tokenizer = self.pipeline.tokenizer
-
         self.model_name = self.pipeline.model.name_or_path
 
         if self.tokenizer.pad_token is None:
@@ -724,11 +766,33 @@ class HuggingFaceOpenAICompatibleModel:
         Returns:
             The formatted prompt string ready for model input.
         """
+
+        final_messages = []
+        for message in messages:
+            if isinstance(message.get("content", ""), str):
+                final_messages.append({"role": message.get("role", "user"), "content": message.get("content", "")})
+            else:
+                # extract only the text from the content
+                # sample data:
+                # {
+                #     "role": "user",
+                #     "content": [
+                #         {"type": "text", "text": "Hello, how are you?"}, # extracted
+                #         {"type": "image", "image": "https://example.com/image.png"}, # not extracted
+                #     ],
+                # }
+                for content_part in message.get("content", []):
+                    if content_part.get("type", "") == "text":
+                        final_messages.append(
+                            {"role": message.get("role", "user"), "content": content_part.get("text", "")}
+                        )
+                    # TODO: implement other content types
+
         # Use the tokenizer's apply_chat_template method.
         # We ensured a template exists in __init__.
         if hasattr(self.tokenizer, "apply_chat_template"):
             return self.tokenizer.apply_chat_template(  # type: ignore[no-any-return]
-                messages,
+                final_messages,
                 tokenize=False,
                 add_generation_prompt=True,
             )
@@ -736,7 +800,7 @@ class HuggingFaceOpenAICompatibleModel:
         # Fallback for very old transformers without apply_chat_template
         # Manually apply ChatML-like formatting
         prompt = ""
-        for message in messages:
+        for message in final_messages:
             role = message.get("role", "user")
             content = message.get("content", "")
             prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"