add

Author: mayinghan

eval_protocol/models.py

@@ -1,6 +1,7 @@
 import os
 from datetime import datetime
 from enum import Enum
+from re import A
 from typing import Any, ClassVar, Dict, List, Literal, Optional, TypedDict, Union
 
 JSONType = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
@@ -256,7 +257,8 @@ def __iter__(self):
 
 class Message(BaseModel):
     """Chat message model with trajectory evaluation support."""
-
+    
+    model_config = ConfigDict(extra="allow")
     role: str  # assistant, user, system, tool
     content: Optional[Union[str, List[ChatCompletionContentPartTextParam]]] = Field(
         default="", description="The content of the message."
@@ -269,11 +271,10 @@ class Message(BaseModel):
     tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
     function_call: Optional[FunctionCall] = None
     control_plane_step: Optional[Dict[str, Any]] = None
-    weight: Optional[int] = None
-
     def dump_mdoel_for_chat_completion_request(self):
         """Only keep chat completion accepted fields"""
-        return self.model_dump(exclude_none=True, exclude={"control_plane_step", "weight", "reasoning_content"})
+        exclude_fields = {"control_plane_step", "reasoning_content"} | set(self.model_extra.keys()) if self.model_extra else set()
+        return self.model_dump(exclude_none=True, exclude=exclude_fields)
 
     @classmethod
     def model_validate(cls, obj, *args, **kwargs):