Add Mistral batch finish reasons (#717)

Author: 09Catho

src/bespokelabs/curator/request_processor/batch/gemini_batch_request_processor.py

@@ -299,7 +299,7 @@ def parse_api_specific_response(
         token_usage = None
         cost = None
         response_message = None
-        finish_reason = "unkown"
+        finish_reason = "unknown"
 
         response_message_raw = ""
         if result_type == "succeeded":

src/bespokelabs/curator/request_processor/batch/mistral_batch_request_processor.py

@@ -5,6 +5,7 @@
 import instructor
 import litellm
 from litellm import model_cost
+from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from mistralai import Mistral
 from mistralai.models import BatchJobOut, UploadFileOutTypedDict
 
@@ -174,24 +175,26 @@ def parse_api_specific_response(self, raw_response: dict, generic_request: Gener
 
         Reference: Mistral API request/response format: https://docs.mistral.ai/api/#tag/chat
         """
+        finish_reason = "unknown"
         if raw_response["response"]["status_code"] != 200:
             response_message = None
             response_errors = raw_response["detail"]["msg"]
             token_usage = None
             cost = None
         else:
-            response_message = raw_response["response"]["body"]["choices"][0]["message"]["content"]
+            choice = raw_response["response"]["body"]["choices"][0]
+            response_message = choice["message"]["content"]
             response_errors = None
+            finish_reason = choice.get("finish_reason", "unknown")
             token_usage = _TokenUsage(
                 input=int(raw_response["response"]["body"]["usage"]["prompt_tokens"]),
                 output=int(raw_response["response"]["body"]["usage"]["completion_tokens"]),
                 total=int(raw_response["response"]["body"]["usage"]["total_tokens"]),
             )
-            print("[DEBUG] self.tracker.input_cost_per_million:", self.tracker.input_cost_per_million)
-            print("[DEBUG] self.tracker.output_cost_per_million:", self.tracker.output_cost_per_million)
 
             cost = self._cost_processor.cost(model="mistral/" + self.config.model, prompt=str(generic_request.messages), completion=response_message)
 
+        finish_reason = map_finish_reason(finish_reason)
         generic_response = GenericResponse(
             response_message=response_message,
             response_errors=response_errors,
@@ -202,6 +205,7 @@ def parse_api_specific_response(self, raw_response: dict, generic_request: Gener
             finished_at=batch.finished_at,
             token_usage=token_usage,
             response_cost=cost,
+            finish_reason=finish_reason,
         )
         return generic_response
 

tests/unittests/test_batch.py

@@ -6,12 +6,84 @@
 from pydantic import BaseModel, Field
 
 from bespokelabs import curator
+from bespokelabs.curator.request_processor.batch.mistral_batch_request_processor import MistralBatchRequestProcessor
+from bespokelabs.curator.request_processor.config import BatchRequestProcessorConfig
+from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchRequestCounts, GenericBatchStatus
+from bespokelabs.curator.types.generic_request import GenericRequest
 
 
 class Answer(BaseModel):
     answer: int = Field(description="The answer to the question")
 
 
+class _CostProcessorStub:
+    def cost(self, **kwargs):
+        return 0.0
+
+
+def _generic_request() -> GenericRequest:
+    return GenericRequest(
+        model="mistral-small-latest",
+        messages=[{"role": "user", "content": "hello"}],
+        original_row={},
+        original_row_idx=0,
+    )
+
+
+def _generic_batch() -> GenericBatch:
+    now = pd.Timestamp("2025-01-01T00:00:00Z").to_pydatetime()
+    return GenericBatch(
+        request_file="requests.jsonl",
+        id="batch-id",
+        created_at=now,
+        finished_at=now,
+        status=GenericBatchStatus.FINISHED.value,
+        api_key_suffix="test",
+        request_counts=GenericBatchRequestCounts(
+            failed=0,
+            succeeded=1,
+            total=1,
+            raw_request_counts_object={},
+        ),
+        raw_batch={},
+        raw_status="SUCCESS",
+    )
+
+
+def test_mistral_batch_response_includes_normalized_finish_reason() -> None:
+    """Mistral batch responses expose finish_reason for invalid-finish retries."""
+    processor = MistralBatchRequestProcessor.__new__(MistralBatchRequestProcessor)
+    processor.config = BatchRequestProcessorConfig(model="mistral-small-latest")
+    processor._cost_processor = _CostProcessorStub()
+
+    raw_response = {
+        "response": {
+            "status_code": 200,
+            "body": {
+                "choices": [
+                    {
+                        "message": {"content": "hello"},
+                        "finish_reason": "length",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 1,
+                    "completion_tokens": 2,
+                    "total_tokens": 3,
+                },
+            },
+        }
+    }
+
+    response = processor.parse_api_specific_response(
+        raw_response=raw_response,
+        generic_request=_generic_request(),
+        batch=_generic_batch(),
+    )
+
+    assert response.finish_reason == "length"
+
+
 def batch_call(model_name, prompts):
     dataset: Dataset = Dataset.from_dict({"prompt": prompts})
     llm = curator.LLM(