diff --git a/packages/gen/gen_ai_hub/orchestration_v2/__init__.py b/packages/gen/gen_ai_hub/orchestration_v2/__init__.py index 260f0c7e..f352cc0c 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/__init__.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/__init__.py @@ -40,7 +40,7 @@ # message "SystemMessage", "UserMessage", "AssistantMessage", "ToolChatMessage", "DeveloperChatMessage", - "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", "ReasoningBlock", # multimodal_items "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py b/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py index 790b0b75..efb0d5f5 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py @@ -18,7 +18,7 @@ from .llama_guard_3_filter import LlamaGuard38bFilter from .llm_model_details import LLMModelDetails from .message import (SystemMessage, UserMessage, AssistantMessage, ToolChatMessage, DeveloperChatMessage, ChatMessage, - ResponseChatMessage, FunctionCall, MessageToolCall) + ResponseChatMessage, FunctionCall, MessageToolCall, ReasoningBlock) from .multimodal_items import ImageDetailLevel, TextPart, ImageUrl, ImagePart, ContentPart, ImageItem from .response import (PromptTokensDetails, CompletionTokensDetails, TokenUsage, GenericModuleResult, TopLogprob, ChatCompletionTokenLogprob, ChoiceLogprobs, LLMChoice, StreamFunctionObject, StreamToolCall, @@ -73,7 +73,7 @@ # message "SystemMessage", "UserMessage", "AssistantMessage", "ToolChatMessage", "DeveloperChatMessage", - "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", "ReasoningBlock", # multimodal_items "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/base.py b/packages/gen/gen_ai_hub/orchestration_v2/models/base.py index 46104613..3a29f32d 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/base.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/base.py @@ -31,3 +31,16 @@ def model_dump(self, **kwargs): kwargs.setdefault("by_alias", True) kwargs.setdefault("exclude_none", True) return super().model_dump(**kwargs) + +class ResponseBaseModel(BaseModel): + """ + Base model for API response models. + + - `extra="allow"` allows unexpected fields in responses to be accepted, + since the external API might introduce new attributes in the response. + """ + + model_config = ConfigDict( + extra="allow", + frozen=False, + ) diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/embeddings.py b/packages/gen/gen_ai_hub/orchestration_v2/models/embeddings.py index b51ca78f..2ef95b7c 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/embeddings.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/embeddings.py @@ -7,7 +7,7 @@ from pydantic import Field -from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel +from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel, ResponseBaseModel from gen_ai_hub.orchestration_v2.models.data_masking import MaskingModuleConfig @@ -117,7 +117,7 @@ class EmbeddingsInput(BaseModel): type_: Optional[EmbeddingsInputType] = Field(default=None, alias="type") -class EmbeddingsUsage(BaseModel): +class EmbeddingsUsage(ResponseBaseModel): """ Token usage information for the embeddings request. @@ -129,7 +129,7 @@ class EmbeddingsUsage(BaseModel): total_tokens: int -class EmbeddingResult(BaseModel): +class EmbeddingResult(ResponseBaseModel): """ A single embedding result. @@ -143,7 +143,7 @@ class EmbeddingResult(BaseModel): index: int -class EmbeddingsResponse(BaseModel): +class EmbeddingsResponse(ResponseBaseModel): """ The response from the embedding model, following OpenAI specification. @@ -159,7 +159,7 @@ class EmbeddingsResponse(BaseModel): usage: EmbeddingsUsage -class EmbeddingsPostResponse(BaseModel): +class EmbeddingsPostResponse(ResponseBaseModel): """ Response for an embeddings POST request. 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 f098c775..26179f34 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/message.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py @@ -10,11 +10,11 @@ from pydantic import field_validator, ValidationError -from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel +from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel, ResponseBaseModel from gen_ai_hub.orchestration_v2.models.multimodal_items import ContentPart, ImageItem, TextPart, ImageUrl, ImagePart -class FunctionCall(BaseModel): +class FunctionCall(ResponseBaseModel): """ Represents a function call with its name and arguments. @@ -44,7 +44,7 @@ def parse_arguments(self) -> dict: return json.loads(self.arguments) -class MessageToolCall(BaseModel): +class MessageToolCall(ResponseBaseModel): """ The tool calls generated by the model, such as function calls. @@ -59,6 +59,17 @@ class MessageToolCall(BaseModel): type: typing.Literal["function"] = "function" function: FunctionCall +class ReasoningBlock(ResponseBaseModel): + """ + 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): """ @@ -157,6 +168,7 @@ class AssistantMessage(BaseModel): 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): @@ -169,7 +181,7 @@ class DeveloperChatMessage(BaseModel): role: Role = Role.DEVELOPER content: Union[str, List[TextPart]] -class ResponseChatMessage(BaseModel): +class ResponseChatMessage(ResponseBaseModel): """ Represents a response message in a conversation. @@ -182,14 +194,17 @@ 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] __all__ = ["Role", "SystemMessage", "UserMessage", "AssistantMessage", "ToolChatMessage", "DeveloperChatMessage", - "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall"] + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", "ReasoningBlock"] 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 ca398edf..6b45c80d 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/response.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/response.py @@ -6,25 +6,8 @@ from pydantic import ConfigDict, Field from gen_ai_hub.orchestration.models.response import ModuleResultsStreaming -from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel -from gen_ai_hub.orchestration_v2.models.message import ChatMessage, FunctionCall, ResponseChatMessage - - -class ResponseBaseModel(BaseModel): - """ - Abstract base model that extends Pydantic's BaseModel and ABC. - - - `extra="allow"` allows unexpected fields in responses to be accepted, - since the external API might introduce new attributes in the response. - - This enforces consistent and safe serialization behavior across all - derived models. - """ - - model_config = ConfigDict( - extra="allow", - frozen=False, - ) +from gen_ai_hub.orchestration_v2.models.base import ResponseBaseModel +from gen_ai_hub.orchestration_v2.models.message import ChatMessage, FunctionCall, ResponseChatMessage, ReasoningBlock class PromptTokensDetails(ResponseBaseModel): @@ -186,6 +169,7 @@ class StreamDelta(ResponseBaseModel): role: Optional[str] = None content: str tool_calls: Optional[List[StreamToolCall]] = None + reasoning_content: Optional[List[ReasoningBlock]] = None class StreamLLMChoice(ResponseBaseModel): diff --git a/packages/gen/tests/orchestration_v2/test_flat_import.py b/packages/gen/tests/orchestration_v2/test_flat_import.py index 7606ffac..8031d898 100644 --- a/packages/gen/tests/orchestration_v2/test_flat_import.py +++ b/packages/gen/tests/orchestration_v2/test_flat_import.py @@ -36,7 +36,7 @@ # message "SystemMessage", "UserMessage", "AssistantMessage", "ToolChatMessage", "DeveloperChatMessage", - "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", "ReasoningBlock", # multimodal_items "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", diff --git a/packages/gen/tests/orchestration_v2/test_message_v2.py b/packages/gen/tests/orchestration_v2/test_message_v2.py new file mode 100644 index 00000000..90ce2e77 --- /dev/null +++ b/packages/gen/tests/orchestration_v2/test_message_v2.py @@ -0,0 +1,42 @@ +import unittest + +from gen_ai_hub.orchestration_v2.models.message import ( + AssistantMessage, + 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"}], + }) + self.assertIsNotNone(msg.reasoning_content) + 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) + + +class TestAssistantMessageValidation(unittest.TestCase): + + def test_reasoning_content_deserialized_from_dict(self): + msg = AssistantMessage.model_validate({ + "role": "assistant", + "content": "Paris", + "reasoning_content": [{"content": "France's capital is Paris.", "signature": "sig"}], + }) + self.assertIsNotNone(msg.reasoning_content) + self.assertIsInstance(msg.reasoning_content[0], ReasoningBlock) + self.assertEqual(msg.reasoning_content[0].content, "France's capital is Paris.") + self.assertEqual(msg.reasoning_content[0].signature, "sig") + + def test_reasoning_content_optional(self): + msg = AssistantMessage.model_validate({"role": "assistant", "content": "Hello"}) + self.assertIsNone(msg.reasoning_content) + diff --git a/packages/gen/tests/orchestration_v2/test_response_v2.py b/packages/gen/tests/orchestration_v2/test_response_v2.py new file mode 100644 index 00000000..26a109f2 --- /dev/null +++ b/packages/gen/tests/orchestration_v2/test_response_v2.py @@ -0,0 +1,21 @@ +import unittest + +from gen_ai_hub.orchestration_v2.models.message import ReasoningBlock +from gen_ai_hub.orchestration_v2.models.response import StreamDelta + + +class TestStreamDeltaValidation(unittest.TestCase): + + def test_reasoning_content_deserialized_from_dict(self): + delta = StreamDelta.model_validate({ + "content": "", + "reasoning_content": [{"content": "I should respond politely.", "signature": ""}], + }) + self.assertIsNotNone(delta.reasoning_content) + self.assertIsInstance(delta.reasoning_content[0], ReasoningBlock) + self.assertEqual(delta.reasoning_content[0].content, "I should respond politely.") + + def test_reasoning_content_optional(self): + delta = StreamDelta.model_validate({"content": "Hello"}) + self.assertIsNone(delta.reasoning_content) +