Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions packages/gen/gen_ai_hub/orchestration_v2/models/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,17 @@ class MessageToolCall(BaseModel):
type: typing.Literal["function"] = "function"
function: FunctionCall

class ReasoningBlock(BaseModel):
"""
The reasoning content returned by the model.

Attributes:
content: The reasoning content.
signature: Signature string for hidden reasoning content.
"""
content: str
signature: str


class Role(str, Enum):
"""
Expand Down Expand Up @@ -182,11 +193,14 @@ class ResponseChatMessage(BaseModel):
refusal: A string indicating refusal reason.

tool_calls: A list of tool call objects.

reasoning_content: A list of reasoning content blocks.
"""
role: Role = Role.ASSISTANT
content: str
refusal: Optional[str] = None
tool_calls: Optional[List[MessageToolCall]] = None
reasoning_content: Optional[List[ReasoningBlock]] = None

ChatMessage = Union[SystemMessage, UserMessage, AssistantMessage, ToolChatMessage, DeveloperChatMessage,
ResponseChatMessage]
Expand Down
22 changes: 22 additions & 0 deletions packages/gen/tests/orchestration_v2/test_message_v2.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import unittest

from gen_ai_hub.orchestration_v2.models.message import (
ReasoningBlock,
ResponseChatMessage,
)


class TestResponseChatMessageValidation(unittest.TestCase):

def test_reasoning_content_deserialized_from_dict(self):
msg = ResponseChatMessage.model_validate({
"role": "assistant",
"content": "Hello",
"reasoning_content": [{"content": "I think...", "signature": "sig123"}],
})
assert msg.reasoning_content is not None
self.assertIsInstance(msg.reasoning_content[0], ReasoningBlock)

def test_reasoning_content_optional(self):
msg = ResponseChatMessage.model_validate({"role": "assistant", "content": "Hello"})
self.assertIsNone(msg.reasoning_content)