Make FinishReason json serializable

anjor · web-flow · commit 82aaaa0662f7 · 2024-02-06T09:32:18.000Z
In downstream code when someone does .model_dump() on any of the chat completion objects it fails because the enum is not serializable. https://stackoverflow.com/questions/69541613/how-to-json-serialize-enum-classes-in-pydantic-basemodel
diff --git a/src/mistralai/models/chat_completion.py b/src/mistralai/models/chat_completion.py
@@ -16,10 +16,10 @@ class DeltaMessage(BaseModel):
     content: Optional[str] = None
 
 
-class FinishReason(Enum):
-    stop = "stop"
-    length = "length"
-    error = "error"
+class FinishReason(str, Enum):
+    stop: str = "stop"
+    length: str = "length"
+    error: str = "error"
 
 
 class ChatCompletionResponseStreamChoice(BaseModel):
