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

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

@@ -0,0 +1,124 @@
+"""Utility functions for model parameter validation and resolution."""
+
+from typing import Any, Optional, Sequence
+
+from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._signatures import core
+
+
+def validate_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Optional[Sequence[core.BaseParamSpec]],
+) -> None:
+    """Validate user-provided params against signature params.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Raises:
+        SnowflakeMLException: If params are provided but signature has no params,
+            or if unknown params are provided, or if param types are invalid,
+            or if duplicate params are provided with different cases.
+    """
+    # Params provided but signature has no params defined
+    if params and not signature_params:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Parameters were provided ({sorted(params.keys())}), "
+                "but this method does not accept any parameters."
+            ),
+        )
+
+    if not signature_params or not params:
+        return
+
+    # Case-insensitive lookup: normalized_name -> param_spec
+    param_spec_lookup = {ps.name.upper(): ps for ps in signature_params}
+
+    # Check for duplicate params with different cases (e.g., "temperature" and "TEMPERATURE")
+    normalized_names = [name.upper() for name in params]
+    if len(normalized_names) != len(set(normalized_names)):
+        # Find the duplicate params to raise an error
+        param_seen: dict[str, list[str]] = {}
+        for param_name in params:
+            param_seen.setdefault(param_name.upper(), []).append(param_name)
+        duplicate_param_names = [param_names for param_names in param_seen.values() if len(param_names) > 1]
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Duplicate parameter(s) provided with different cases: {duplicate_param_names}. "
+                "Parameter names are case-insensitive."
+            ),
+        )
+
+    # Validate user-provided params exist (case-insensitive)
+    invalid_params = [name for name in params if name.upper() not in param_spec_lookup]
+    if invalid_params:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Unknown parameter(s): {sorted(invalid_params)}. "
+                f"Valid parameters are: {sorted(ps.name for ps in signature_params)}"
+            ),
+        )
+
+    # Validate types for each provided param
+    for param_name, default_value in params.items():
+        param_spec = param_spec_lookup[param_name.upper()]
+        if isinstance(param_spec, core.ParamSpec):
+            core.ParamSpec._validate_default_value(param_spec.dtype, default_value, param_spec.shape)
+
+
+def resolve_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Sequence[core.BaseParamSpec],
+) -> list[tuple[sql_identifier.SqlIdentifier, Any]]:
+    """Resolve final method parameters by applying user-provided params over signature defaults.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Returns:
+        List of tuples (SqlIdentifier, value) for method invocation.
+    """
+    # Case-insensitive lookup: normalized_name -> param_spec
+    param_spec_lookup = {ps.name.upper(): ps for ps in signature_params}
+
+    # Start with defaults from signature
+    final_params: dict[str, Any] = {}
+    for param_spec in signature_params:
+        if hasattr(param_spec, "default_value"):
+            final_params[param_spec.name] = param_spec.default_value
+
+    # Override with provided runtime parameters (using signature's original param names)
+    if params:
+        for param_name, override_value in params.items():
+            canonical_name = param_spec_lookup[param_name.upper()].name
+            final_params[canonical_name] = override_value
+
+    return [(sql_identifier.SqlIdentifier(param_name), param_value) for param_name, param_value in final_params.items()]
+
+
+def validate_and_resolve_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Optional[Sequence[core.BaseParamSpec]],
+) -> Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]]:
+    """Validate user-provided params against signature params and return method parameters.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Returns:
+        List of tuples (SqlIdentifier, value) for method invocation, or None if no params.
+    """
+    validate_params(params, signature_params)
+
+    if not signature_params:
+        return None
+
+    return resolve_params(params, signature_params)

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
 
@@ -582,15 +587,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 +930,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 +974,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 +987,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 +1019,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 +1031,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
 
 
@@ -272,17 +256,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,
                 )
             )

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,11 @@ class ModelMethod:
                 f"Your parameter {param_spec.name} cannot be resolved as valid SQL identifier. "
                 "Try specifying `case_sensitive` as True."
             ) from e
+        default_value = param_spec.default_value 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,19 @@ 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]
                     else:
                         # TODO: remove conversational pipeline code
                         # For others, we could offer the whole dataframe as a list.
@@ -615,11 +635,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 +725,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 +746,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 +780,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"