Refresh batch score v2 data schema to align with OpenAI public API

.gitignore

@@ -139,4 +139,3 @@ mlruns/
 
 # ignore config files
 config.json
-out*

assets/batch_score/components/driver/batch_score_llm/spec.yaml

@@ -2,7 +2,7 @@ $schema: http://azureml/sdk-2-0/ParallelComponent.json
 type: parallel
 
 name: batch_score_llm
-version: 1.1.5
+version: 1.1.6
 display_name: Batch Score Large Language Models
 is_deterministic: False
 

assets/batch_score/components/driver/src/batch_score/aoai/scoring/aoai_response_handler.py

@@ -65,6 +65,7 @@ def handle_response(
         if response_status == 200:
             return self._create_scoring_result(
                 status=ScoringResultStatus.SUCCESS,
+                model_response_code=response_status,
                 scoring_request=scoring_request,
                 start=start,
                 end=end,
@@ -78,6 +79,7 @@ def handle_response(
 
         result = self._create_scoring_result(
             status=ScoringResultStatus.FAILURE,
+            model_response_code=response_status,
             scoring_request=scoring_request,
             start=start,
             end=end,
@@ -130,6 +132,7 @@ def _handle_exception(
         except Exception:
             return self._create_scoring_result(
                 status=ScoringResultStatus.FAILURE,
+                model_response_code=http_response.status,
                 scoring_request=scoring_request,
                 start=start,
                 end=end,

assets/batch_score/components/driver/src/batch_score/common/parallel/parallel_driver.py

@@ -5,7 +5,9 @@
 
 import asyncio
 
-from ...utils.common import convert_result_list
+from ...utils.output_formatter import OutputFormatter
+from ...utils.v1_output_formatter import V1OutputFormatter
+from ...utils.v2_output_formatter import V2OutputFormatter
 from ..configuration.configuration import Configuration
 from ..post_processing.mini_batch_context import MiniBatchContext
 from ..post_processing.result_utils import apply_input_transformer
@@ -48,7 +50,12 @@ def run(self, payloads: "list[str]", mini_batch_context: MiniBatchContext = None
 
         apply_input_transformer(self.__input_to_output_transformer, scoring_results)
 
-        results = convert_result_list(
+        output_formatter: OutputFormatter
+        if self._configuration.input_schema_version == 1:
+            output_formatter = V1OutputFormatter()
+        else:
+            output_formatter = V2OutputFormatter()
+        results = output_formatter.format_output(
             results=scoring_results,
             batch_size_per_request=self._configuration.batch_size_per_request)
 

assets/batch_score/components/driver/src/batch_score/common/post_processing/callback_factory.py

@@ -5,7 +5,9 @@
 
 import traceback
 
-from ...utils.common import convert_result_list
+from ...utils.output_formatter import OutputFormatter
+from ...utils.v1_output_formatter import V1OutputFormatter
+from ...utils.v2_output_formatter import V2OutputFormatter
 from ..configuration.configuration import Configuration
 from ..scoring.scoring_result import ScoringResult
 from ..telemetry import logging_utils as lu
@@ -16,7 +18,7 @@
     apply_input_transformer,
     get_return_value,
 )
-from .output_handler import SingleFileOutputHandler, SeparateFileOutputHandler
+from .output_handler import OutputHandler
 
 
 def add_callback(callback, cur):
@@ -33,9 +35,11 @@ class CallbackFactory:
 
     def __init__(self,
                  configuration: Configuration,
+                 output_handler: OutputHandler,
                  input_to_output_transformer):
         """Initialize CallbackFactory."""
         self._configuration = configuration
+        self._output_handler = output_handler
         self.__input_to_output_transformer = input_to_output_transformer
 
     def generate_callback(self):
@@ -46,7 +50,12 @@ def generate_callback(self):
         return callback
 
     def _convert_result_list(self, scoring_results: "list[ScoringResult]", mini_batch_context: MiniBatchContext):
-        return convert_result_list(
+        output_formatter: OutputFormatter
+        if self._configuration.input_schema_version == 1:
+            output_formatter = V1OutputFormatter()
+        else:
+            output_formatter = V2OutputFormatter()
+        return output_formatter.format_output(
             results=scoring_results,
             batch_size_per_request=self._configuration.batch_size_per_request)
 
@@ -64,13 +73,7 @@ def _save_mini_batch_result_and_emit(
             if mini_batch_context.exception is None:
                 if self._configuration.save_mini_batch_results == "enabled":
                     lu.get_logger().info("save_mini_batch_results is enabled")
-                    if (self._configuration.split_output):
-                        output_handler = SeparateFileOutputHandler()
-                        lu.get_logger().info("Saving successful results and errors to separate files")
-                    else:
-                        output_handler = SingleFileOutputHandler()
-                        lu.get_logger().info("Saving results to single file")
-                    output_handler.save_mini_batch_results(
+                    self._output_handler.save_mini_batch_results(
                         scoring_results,
                         self._configuration.mini_batch_results_out_directory,
                         mini_batch_context.raw_mini_batch_context

assets/batch_score/components/driver/src/batch_score/common/scoring/aoai_error.py

@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""This file contains the data class for Azure OpenAI scoring error."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class AoaiScoringError:
+    """Azure OpenAI scoring error."""
+
+    code: str = None
+    message: str = None

assets/batch_score/components/driver/src/batch_score/common/scoring/aoai_response.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""This file contains the data class for Azure OpenAI scoring response."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class AoaiScoringResponse:
+    """Azure OpenAI scoring response."""
+
+    body: any = None
+    request_id: str = None
+    status_code: int = None

assets/batch_score/components/driver/src/batch_score/common/scoring/http_response_handler.py

@@ -32,6 +32,7 @@ def handle_response(
     def _create_scoring_result(
             self,
             status: ScoringResultStatus,
+            model_response_code: int,
             scoring_request: ScoringRequest,
             start: float,
             end: float,
@@ -43,6 +44,7 @@ def _create_scoring_result(
             status=status,
             start=start,
             end=end,
+            model_response_code=model_response_code,
             request_obj=scoring_request.original_payload_obj,
             request_metadata=scoring_request.request_metadata,
             response_body=http_post_response.payload,

assets/batch_score/components/driver/src/batch_score/common/scoring/scoring_request.py

@@ -17,6 +17,7 @@ class ScoringRequest:
 
     __BATCH_REQUEST_METADATA = "_batch_request_metadata"
     __REQUEST_METADATA = "request_metadata"
+    __CUSTOM_ID = "custom_id"
 
     def __init__(
             self,
@@ -55,6 +56,8 @@ def __init__(
         # These properties do not need to be sent to the model & will be added to the output file directly
         self.__request_metadata = self.__cleaned_payload_obj.pop(self.__BATCH_REQUEST_METADATA, None)
         self.__request_metadata = self.__cleaned_payload_obj.pop(self.__REQUEST_METADATA, self.__request_metadata)
+        # If custom_id exists (V2 input schema), make sure it is not sent to MIR endpoint
+        self.__CUSTOM_ID = self.__cleaned_payload_obj.pop(self.__CUSTOM_ID, None)
 
         self.__cleaned_payload = json.dumps(self.__cleaned_payload_obj, cls=BatchComponentJSONEncoder)
         self.__loggable_payload = json.dumps(self.__loggable_payload_obj, cls=BatchComponentJSONEncoder)
@@ -136,6 +139,12 @@ def segment_id(self) -> int:
         """Get the segment id."""
         return self.__segment_id
 
+    # read-only
+    @property
+    def custom_id(self) -> str:
+        """Get the custom id. Only valid for V2 input schema."""
+        return self.__CUSTOM_ID
+
     @estimated_cost.setter
     def estimated_cost(self, cost: int):
         """Set the estimated cost."""

assets/batch_score/components/driver/src/batch_score/common/scoring/scoring_result.py

@@ -58,6 +58,7 @@ def __init__(
             status: ScoringResultStatus,
             start: float,
             end: float,
+            model_response_code: int,
             request_obj: any,
             request_metadata: any,
             response_body: any,
@@ -70,6 +71,7 @@ def __init__(
         self.status = status
         self.start = start
         self.end = end
+        self.model_response_code = model_response_code
         self.request_obj = request_obj  # Normalize to json
         self.request_metadata = request_metadata
         self.response_body = response_body
@@ -121,6 +123,7 @@ def Failed(scoring_request: ScoringRequest = None) -> 'ScoringResult':
             status=ScoringResultStatus.FAILURE,
             start=0,
             end=0,
+            model_response_code=None,
             request_obj=scoring_request.original_payload_obj if scoring_request else None,
             request_metadata=scoring_request.request_metadata if scoring_request else None,
             response_body=None,
@@ -140,6 +143,7 @@ def copy(self) -> 'ScoringResult':
             self.status,
             self.start,
             self.end,
+            self.model_response_code,
             self.request_obj,
             self.request_metadata,
             deepcopy(self.response_body),

assets/batch_score/components/driver/src/batch_score/main.py

@@ -90,6 +90,7 @@ def init():
     global par
     global configuration
     global input_handler
+    global output_handler
 
     start = time.time()
     parser = ConfigurationParserFactory().get_parser()
@@ -103,6 +104,13 @@ def init():
     else:
         raise ValueError(f"Invalid input_schema_version: {configuration.input_schema_version}")
 
+    if (configuration.split_output):
+        output_handler = SeparateFileOutputHandler()
+        lu.get_logger().info("Will save successful results and errors to separate files")
+    else:
+        output_handler = SingleFileOutputHandler()
+        lu.get_logger().info("Will save all results to a single file")
+
     event_utils.setup_context_vars(configuration, metadata)
     setup_geneva_event_handlers()
     setup_job_log_event_handlers()
@@ -147,6 +155,7 @@ def init():
     if configuration.async_mode:
         callback_factory = CallbackFactory(
             configuration=configuration,
+            output_handler=output_handler,
             input_to_output_transformer=input_to_output_transformer)
         finished_callback = callback_factory.generate_callback()
 
@@ -201,13 +210,6 @@ def run(input_data: pd.DataFrame, mini_batch_context):
     try:
         ret = par.run(data_list, mini_batch_context)
 
-        if (configuration.split_output):
-            output_handler = SeparateFileOutputHandler()
-            lu.get_logger().info("Saving successful results and errors to separate files")
-        else:
-            output_handler = SingleFileOutputHandler()
-            lu.get_logger().info("Saving results to single file")
-
         if configuration.save_mini_batch_results == "enabled":
             lu.get_logger().info("save_mini_batch_results is enabled")
             output_handler.save_mini_batch_results(

assets/batch_score/components/driver/src/batch_score/metadata.json

@@ -1,7 +1,7 @@
 {
-    "component_version": "1.1.5",
+    "component_version": "1.1.6",
     "component_directory": "driver/batch_score_llm",
     "component_name": "batch_score_llm",
     "virtual_environment_name": null,
     "registry_name": "azureml"
-  }
+}

assets/batch_score/components/driver/src/batch_score/mir/scoring/mir_http_response_handler.py

@@ -104,6 +104,7 @@ def handle_response(
         if response_status == 200:
             result = self._create_scoring_result(
                 status=ScoringResultStatus.SUCCESS,
+                model_response_code=response_status,
                 scoring_request=scoring_request,
                 start=start,
                 end=end,
@@ -124,6 +125,7 @@ def handle_response(
         else:  # Score failed
             result = self._create_scoring_result(
                 status=ScoringResultStatus.FAILURE,
+                model_response_code=response_status,
                 scoring_request=scoring_request,
                 start=start,
                 end=end,

assets/batch_score/components/driver/src/batch_score/utils/common.py

@@ -3,14 +3,9 @@
 
 """Common utilities."""
 
-import json
 from argparse import ArgumentParser
 from urllib.parse import urlparse
 
-from ..common.scoring.scoring_result import ScoringResult
-from . import embeddings_utils as embeddings
-from .json_encoder_extensions import BatchComponentJSONEncoder
-
 
 def get_base_url(url: str) -> str:
     """Get base url."""
@@ -38,39 +33,7 @@ def str2bool(v):
         raise ArgumentParser.ArgumentTypeError('Boolean value expected.')
 
 
-def convert_result_list(results: "list[ScoringResult]", batch_size_per_request: int) -> "list[str]":
-    """Convert scoring results to the result list."""
-    output_list: list[dict[str, str]] = []
-    for scoringResult in results:
-        output: dict[str, str] = {}
-        output["status"] = scoringResult.status.name
-        output["start"] = scoringResult.start
-        output["end"] = scoringResult.end
-        output["request"] = scoringResult.request_obj
-        output["response"] = scoringResult.response_body
-
-        if scoringResult.segmented_response_bodies is not None and len(scoringResult.segmented_response_bodies) > 0:
-            output["segmented_responses"] = scoringResult.segmented_response_bodies
-
-        if scoringResult.request_metadata is not None:
-            output["request_metadata"] = scoringResult.request_metadata
-
-        if batch_size_per_request > 1:
-            batch_output_list = embeddings._convert_to_list_of_output_items(
-                output,
-                scoringResult.estimated_token_counts)
-            output_list.extend(batch_output_list)
-        else:
-            output_list.append(output)
-
-    return list(map(__stringify_output, output_list))
-
-
 def get_mini_batch_id(mini_batch_context: any):
     """Get mini batch id from mini batch context."""
     if mini_batch_context:
         return mini_batch_context.mini_batch_id
-
-
-def __stringify_output(payload_obj: dict) -> str:
-    return json.dumps(payload_obj, cls=BatchComponentJSONEncoder)