From 5a7c39ff9c98edc6a4983e8a4d9b5afb2e4cb77d Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Mon, 14 Sep 2026 17:31:23 +0200 Subject: [PATCH 1/8] allow all response related class to accept extra classes. --- .../orchestration_v2/models/base.py | 28 +++++++++++++++++++ .../orchestration_v2/models/message.py | 9 +++--- .../orchestration_v2/models/response.py | 18 +----------- .../orchestration_v2/test_flat_import.py | 2 +- 4 files changed, 35 insertions(+), 22 deletions(-) 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..ce4bf475 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,31 @@ 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, + ) + + +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/message.py b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py index f098c775..74575b82 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. @@ -60,6 +60,7 @@ class MessageToolCall(BaseModel): function: FunctionCall + class Role(str, Enum): """ Enumerates supported roles in LLM-based conversations. @@ -169,7 +170,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. 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..cda706bd 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/response.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/response.py @@ -6,26 +6,10 @@ 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.base import ResponseBaseModel 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, - ) - class PromptTokensDetails(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", From b748fc0f3e1bee5c2d2efaa9483f0eca2f641d1c Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Mon, 14 Sep 2026 17:37:48 +0200 Subject: [PATCH 2/8] cleaning --- .../gen_ai_hub/orchestration_v2/models/base.py | 15 --------------- 1 file changed, 15 deletions(-) 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 ce4bf475..3a29f32d 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/base.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/base.py @@ -32,21 +32,6 @@ def model_dump(self, **kwargs): 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, - ) - - class ResponseBaseModel(BaseModel): """ Base model for API response models. From c55ef4f47a5f3854cae40994684b4a06f8ab8430 Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Mon, 14 Sep 2026 17:57:40 +0200 Subject: [PATCH 3/8] fix the ci checks --- packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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", From 4100837194b2c28e9ef997b0f18ec53a9ee5590a Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Tue, 15 Sep 2026 13:00:09 +0200 Subject: [PATCH 4/8] use ResponseBaseModel for embeddings response classes --- .../gen_ai_hub/orchestration_v2/models/embeddings.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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. From 09df9237c2c35d491e47126f6e51567a5c4a51bd Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Tue, 15 Sep 2026 15:58:22 +0200 Subject: [PATCH 5/8] set ResponseBaseModel to AssistantMessage --- packages/gen/gen_ai_hub/orchestration_v2/models/message.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 74575b82..18ae14b4 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/message.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py @@ -139,7 +139,7 @@ def content_validation(cls, content): # pylint: disable=no-self-argument return mapped_content -class AssistantMessage(BaseModel): +class AssistantMessage(ResponseBaseModel): """ Represents an assistant message in a prompt or conversation template. From 5a5946a5d5411b58acdf53842573b1e0922aa663 Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Tue, 15 Sep 2026 17:32:11 +0200 Subject: [PATCH 6/8] switch AssistantMessage back to BaseModel --- packages/gen/gen_ai_hub/orchestration_v2/models/message.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 18ae14b4..74575b82 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/message.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/message.py @@ -139,7 +139,7 @@ def content_validation(cls, content): # pylint: disable=no-self-argument return mapped_content -class AssistantMessage(ResponseBaseModel): +class AssistantMessage(BaseModel): """ Represents an assistant message in a prompt or conversation template. From 9264c1321db56b84edd85eac2ba9c97cb74d252a Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Wed, 16 Sep 2026 13:20:39 +0200 Subject: [PATCH 7/8] delete unexisted models for the tests. --- packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py | 4 ++-- packages/gen/tests/orchestration_v2/test_flat_import.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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 efb0d5f5..790b0b75 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, ReasoningBlock) + ResponseChatMessage, FunctionCall, MessageToolCall) 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", "ReasoningBlock", + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", # multimodal_items "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", diff --git a/packages/gen/tests/orchestration_v2/test_flat_import.py b/packages/gen/tests/orchestration_v2/test_flat_import.py index 8031d898..7606ffac 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", "ReasoningBlock", + "ChatMessage", "ResponseChatMessage", "FunctionCall", "MessageToolCall", # multimodal_items "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", From e65b205884002a80237f832c268bcc65d7c54395 Mon Sep 17 00:00:00 2001 From: Sicheng Dong Date: Thu, 17 Sep 2026 09:00:00 +0200 Subject: [PATCH 8/8] add unit tests --- .../tests/orchestration_v2/test_embeddings.py | 41 ++++ .../tests/orchestration_v2/test_message_v2.py | 229 +++++++++++++++++- 2 files changed, 269 insertions(+), 1 deletion(-) diff --git a/packages/gen/tests/orchestration_v2/test_embeddings.py b/packages/gen/tests/orchestration_v2/test_embeddings.py index 62bc31ec..cae0357f 100644 --- a/packages/gen/tests/orchestration_v2/test_embeddings.py +++ b/packages/gen/tests/orchestration_v2/test_embeddings.py @@ -734,6 +734,47 @@ def test_request_with_masking_allowlist(self): allowlist = result["config"]["modules"]["masking"]["masking_providers"][0]["allowlist"] self.assertEqual(allowlist, ["SAP", "Microsoft"]) +class TestEmbeddingsResponseExtraFields(unittest.TestCase): + """EmbeddingsUsage, EmbeddingResult, EmbeddingsResponse and EmbeddingsPostResponse + were switched to ResponseBaseModel and must silently accept unknown fields.""" + + def test_embeddings_usage_stores_extra_field(self): + usage = EmbeddingsUsage.model_validate({ + "prompt_tokens": 10, "total_tokens": 10, + "extra_field": "extra", + }) + self.assertEqual(usage.extra_field, "extra") + + def test_embedding_result_stores_extra_field(self): + result = EmbeddingResult.model_validate({ + "object": "embedding", "embedding": [0.1, 0.2], "index": 0, + "extra_field": "extra", + }) + self.assertEqual(result.extra_field, "extra") + + def test_embeddings_response_stores_extra_field(self): + response = EmbeddingsResponse.model_validate({ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1], "index": 0}], + "model": "text-embedding-3-large", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + "extra_field": "extra", + }) + self.assertEqual(response.extra_field, "extra") + + def test_embeddings_post_response_stores_extra_field(self): + response = EmbeddingsPostResponse.model_validate({ + "request_id": "emb-req-1", + "final_result": { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1], "index": 0}], + "model": "text-embedding-3-large", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + }, + "extra_field": "extra", + }) + self.assertEqual(response.extra_field, "extra") + if __name__ == "__main__": unittest.main() diff --git a/packages/gen/tests/orchestration_v2/test_message_v2.py b/packages/gen/tests/orchestration_v2/test_message_v2.py index f4744e70..85ba2562 100644 --- a/packages/gen/tests/orchestration_v2/test_message_v2.py +++ b/packages/gen/tests/orchestration_v2/test_message_v2.py @@ -2,8 +2,36 @@ from gen_ai_hub.orchestration_v2.models.message import ( ReasoningBlock, + FunctionCall, + MessageToolCall, ResponseChatMessage, ) +from gen_ai_hub.orchestration_v2.models.response import ( + Citation, + CompletionPostResponse, + CompletionTokensDetails, + ChoiceLogprobs, + ChatCompletionTokenLogprob, + ErrorResponse, + ErrorResponseStreaming, + GenericModuleResult, + LLMChoice, + LLMModuleResult, + ModuleResults, + OrchestrationResponseWithRetries, + PromptTokensDetails, + SAPAPIError, + SAPAPIErrorStreaming, + StreamCompletionPostResponse, + StreamDelta, + StreamFunctionObject, + StreamLLMChoice, + StreamLLMModuleResult, + StreamModuleResults, + StreamToolCall, + TokenUsage, + TopLogprob, +) class TestResponseChatMessageValidation(unittest.TestCase): @@ -14,9 +42,208 @@ def test_reasoning_content_deserialized_from_dict(self): "content": "Hello", "reasoning_content": [{"content": "I think...", "signature": "sig123"}], }) - assert msg.reasoning_content is not None + 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 TestExtraFieldsAllowed(unittest.TestCase): + """Response-side models (ResponseBaseModel subclasses) must silently accept + unknown fields so that new API attributes never break existing clients.""" + + def test_response_chat_message_stores_extra_field(self): + msg = ResponseChatMessage.model_validate({ + "role": "assistant", + "content": "hi", + "extra_field": "extra", + }) + self.assertEqual(msg.extra_field, "extra") + + def test_function_call_stores_extra_field(self): + fc = FunctionCall.model_validate({ + "name": "my_fn", + "arguments": "{}", + "extra_field": "extra", + }) + self.assertEqual(fc.extra_field, "extra") + + def test_message_tool_call_stores_extra_field(self): + tc = MessageToolCall.model_validate({ + "id": "call_1", + "type": "function", + "function": {"name": "fn", "arguments": "{}"}, + "extra_field": "extra", + }) + self.assertEqual(tc.extra_field, "extra") + + def test_llm_choice_stores_extra_field(self): + choice = LLMChoice.model_validate({ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + "extra_field": "extra", + }) + self.assertEqual(choice.extra_field, "extra") + + def test_llm_module_result_stores_extra_field(self): + result = LLMModuleResult.model_validate({ + "id": "chatcmpl-abc", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + }], + "usage": {"completion_tokens": 5, "prompt_tokens": 3, "total_tokens": 8}, + "extra_field": "extra", + }) + self.assertEqual(result.extra_field, "extra") + + def test_completion_post_response_stores_extra_field(self): + resp = CompletionPostResponse.model_validate({ + "request_id": "req-1", + "intermediate_results": {}, + "final_result": { + "id": "chatcmpl-abc", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + }], + "usage": {"completion_tokens": 5, "prompt_tokens": 3, "total_tokens": 8}, + }, + "extra_field": "extra", + }) + self.assertEqual(resp.extra_field, "extra") + + def test_prompt_tokens_details_stores_extra_field(self): + details = PromptTokensDetails.model_validate({"extra_field": "extra"}) + self.assertEqual(details.extra_field, "extra") + + def test_completion_tokens_details_stores_extra_field(self): + details = CompletionTokensDetails.model_validate({"extra_field": "extra"}) + self.assertEqual(details.extra_field, "extra") + + def test_token_usage_stores_extra_field(self): + usage = TokenUsage.model_validate({ + "completion_tokens": 5, "prompt_tokens": 3, "total_tokens": 8, + "extra_field": "extra", + }) + self.assertEqual(usage.extra_field, "extra") + + def test_generic_module_result_stores_extra_field(self): + result = GenericModuleResult.model_validate({"message": "ok", "extra_field": "extra"}) + self.assertEqual(result.extra_field, "extra") + + def test_top_logprob_stores_extra_field(self): + lp = TopLogprob.model_validate({"token": "hi", "logprob": -0.5, "extra_field": "extra"}) + self.assertEqual(lp.extra_field, "extra") + + def test_chat_completion_token_logprob_stores_extra_field(self): + lp = ChatCompletionTokenLogprob.model_validate({ + "token": "hi", "logprob": -0.5, "extra_field": "extra", + }) + self.assertEqual(lp.extra_field, "extra") + + def test_choice_logprobs_stores_extra_field(self): + lp = ChoiceLogprobs.model_validate({"extra_field": "extra"}) + self.assertEqual(lp.extra_field, "extra") + + def test_stream_function_object_stores_extra_field(self): + fo = StreamFunctionObject.model_validate({"name": "fn", "arguments": "{}", "extra_field": "extra"}) + self.assertEqual(fo.extra_field, "extra") + + def test_stream_tool_call_stores_extra_field(self): + tc = StreamToolCall.model_validate({"index": 0, "extra_field": "extra"}) + self.assertEqual(tc.extra_field, "extra") + + def test_stream_delta_stores_extra_field(self): + delta = StreamDelta.model_validate({"content": "hi", "extra_field": "extra"}) + self.assertEqual(delta.extra_field, "extra") + + def test_stream_llm_choice_stores_extra_field(self): + choice = StreamLLMChoice.model_validate({ + "index": 0, "delta": {"content": "hi"}, "extra_field": "extra", + }) + self.assertEqual(choice.extra_field, "extra") + + def test_citation_stores_extra_field(self): + citation = Citation.model_validate({ + "title": "Source A", "url": "https://example.com", "extra_field": "extra", + }) + self.assertEqual(citation.extra_field, "extra") + + def test_stream_llm_module_result_stores_extra_field(self): + result = StreamLLMModuleResult.model_validate({ + "id": "chatcmpl-abc", "object": "chat.completion.chunk", + "created": 1234567890, "model": "gpt-4o", + "choices": [{"index": 0, "delta": {"content": "hi"}}], + "extra_field": "extra", + }) + self.assertEqual(result.extra_field, "extra") + + def test_module_results_stores_extra_field(self): + results = ModuleResults.model_validate({"extra_field": "extra"}) + self.assertEqual(results.extra_field, "extra") + + def test_stream_module_results_stores_extra_field(self): + results = StreamModuleResults.model_validate({"extra_field": "extra"}) + self.assertEqual(results.extra_field, "extra") + + def test_sap_api_error_stores_extra_field(self): + err = SAPAPIError.model_validate({ + "request_id": "r1", "code": 400, "message": "bad", "location": "svc", + "extra_field": "extra", + }) + self.assertEqual(err.extra_field, "extra") + + def test_sap_api_error_streaming_stores_extra_field(self): + err = SAPAPIErrorStreaming.model_validate({ + "request_id": "r1", "code": 500, "message": "error", "location": "svc", + "extra_field": "extra", + }) + self.assertEqual(err.extra_field, "extra") + + def test_stream_completion_post_response_stores_extra_field(self): + resp = StreamCompletionPostResponse.model_validate({ + "request_id": "req-1", + "intermediate_results": None, + "final_result": None, + "extra_field": "extra", + }) + self.assertEqual(resp.extra_field, "extra") + + def test_error_response_stores_extra_field(self): + resp = ErrorResponse.model_validate({ + "error": {"request_id": "r1", "code": 400, "message": "bad", "location": "svc"}, + "extra_field": "extra", + }) + self.assertEqual(resp.extra_field, "extra") + + def test_error_response_streaming_stores_extra_field(self): + resp = ErrorResponseStreaming.model_validate({ + "error": {"request_id": "r1", "code": 500, "message": "error", "location": "svc"}, + "extra_field": "extra", + }) + self.assertEqual(resp.extra_field, "extra") + + def test_orchestration_response_with_retries_stores_extra_field(self): + resp = OrchestrationResponseWithRetries.model_validate({ + "request_id": "req-1", + "intermediate_results": {}, + "final_result": { + "id": "chatcmpl-abc", "object": "chat.completion", + "created": 1234567890, "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"completion_tokens": 5, "prompt_tokens": 3, "total_tokens": 8}, + }, + "extra_field": "extra", + }) + self.assertEqual(resp.extra_field, "extra") \ No newline at end of file