diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/message.py b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py index eaec35d..d3b40d0 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/message.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py @@ -164,11 +164,14 @@ class AssistantMessage(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: Optional[Union[str, List[TextPart]]] = None refusal: Optional[str] = None tool_calls: Optional[List[MessageToolCall]] = None + reasoning_content: Optional[List[ReasoningBlock]] = None class ToolChatMessage(BaseModel): diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/response.py b/packages/gen/gen_ai_hub/orchestration_v2/models/response.py index 91c9f14..59ebcc4 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/response.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/response.py @@ -6,7 +6,7 @@ from pydantic import ConfigDict, Field from gen_ai_hub.orchestration_v2.models.base import ResponseBaseModel -from gen_ai_hub.orchestration_v2.models.message import ChatMessage, FunctionCall, ResponseChatMessage +from gen_ai_hub.orchestration_v2.models.message import ChatMessage, FunctionCall, ReasoningBlock, ResponseChatMessage @@ -189,6 +189,7 @@ class StreamDelta(ResponseBaseModel): content: str tool_calls: Optional[List[StreamToolCall]] = None refusal: Optional[str] = None + reasoning_content: Optional[List[ReasoningBlock]] = None class StreamLLMChoice(ResponseBaseModel): diff --git a/packages/gen/tests/orchestration_v2/test_message_v2.py b/packages/gen/tests/orchestration_v2/test_message_v2.py index 85ba256..a19ea15 100644 --- a/packages/gen/tests/orchestration_v2/test_message_v2.py +++ b/packages/gen/tests/orchestration_v2/test_message_v2.py @@ -1,9 +1,10 @@ import unittest from gen_ai_hub.orchestration_v2.models.message import ( - ReasoningBlock, + AssistantMessage, FunctionCall, MessageToolCall, + ReasoningBlock, ResponseChatMessage, ) from gen_ai_hub.orchestration_v2.models.response import ( @@ -36,7 +37,7 @@ class TestResponseChatMessageValidation(unittest.TestCase): - def test_reasoning_content_deserialized_from_dict(self): + def test_deserialization_from_dict(self): msg = ResponseChatMessage.model_validate({ "role": "assistant", "content": "Hello", @@ -49,6 +50,20 @@ def test_reasoning_content_optional(self): msg = ResponseChatMessage.model_validate({"role": "assistant", "content": "Hello"}) self.assertIsNone(msg.reasoning_content) +class TestAssistantMessageValidation(unittest.TestCase): + + def test_deserialization_from_dict(self): + msg = AssistantMessage.model_validate({ + "content": "Hello", + "reasoning_content": [{"content": "I think...", "signature": "sig123"}], + }) + self.assertIsNotNone(msg.reasoning_content) + self.assertIsInstance(msg.reasoning_content[0], ReasoningBlock) + + def test_reasoning_content_optional(self): + msg = AssistantMessage.model_validate({"role": "assistant", "content": "Hello"}) + self.assertIsNone(msg.reasoning_content) + class TestExtraFieldsAllowed(unittest.TestCase): """Response-side models (ResponseBaseModel subclasses) must silently accept unknown fields so that new API attributes never break existing clients.""" @@ -246,4 +261,4 @@ def test_orchestration_response_with_retries_stores_extra_field(self): }, "extra_field": "extra", }) - self.assertEqual(resp.extra_field, "extra") \ No newline at end of file + self.assertEqual(resp.extra_field, "extra")