diff --git a/packages/gen/gen_ai_hub/orchestration_v2/__init__.py b/packages/gen/gen_ai_hub/orchestration_v2/__init__.py index 260f0c7e..d378b8f7 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/__init__.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/__init__.py @@ -3,6 +3,9 @@ from .exceptions import OrchestrationError, OrchestrationErrorList __all__ = [ + # cache_control + "CacheControl", + # azure_content_filter "AzureContentFilter", "AzureContentSafetyInput", "AzureContentSafetyOutput", "AzureThreshold", @@ -46,7 +49,7 @@ "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", # response - "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob", + "CacheCreationTokenDetails", "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob", "ChatCompletionTokenLogprob", "ChoiceLogprobs", "LLMChoice", "StreamFunctionObject", "StreamToolCall", "StreamDelta", "StreamLLMChoice", "Citation", "LLMModuleResult", "StreamLLMModuleResult", "ModuleResults", "StreamModuleResults", "SAPAPIError", "SAPAPIErrorStreaming", "CompletionPostResponse", 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..44e231e7 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/__init__.py @@ -1,4 +1,5 @@ from .azure_content_filter import AzureContentSafetyInput, AzureContentSafetyOutput, AzureContentFilter, AzureThreshold +from .cache_control import CacheControl from .config import (ModuleConfig, OrchestrationConfig, OrchestrationConfigReference, CompletionRequestConfigurationReferenceByIdConfigRef, CompletionRequestConfigurationReferenceByNameScenarioVersionConfigRef) @@ -20,7 +21,7 @@ from .message import (SystemMessage, UserMessage, AssistantMessage, ToolChatMessage, DeveloperChatMessage, ChatMessage, ResponseChatMessage, FunctionCall, MessageToolCall) from .multimodal_items import ImageDetailLevel, TextPart, ImageUrl, ImagePart, ContentPart, ImageItem -from .response import (PromptTokensDetails, CompletionTokensDetails, TokenUsage, GenericModuleResult, TopLogprob, +from .response import (CacheCreationTokenDetails, PromptTokensDetails, CompletionTokensDetails, TokenUsage, GenericModuleResult, TopLogprob, ChatCompletionTokenLogprob, ChoiceLogprobs, LLMChoice, StreamFunctionObject, StreamToolCall, StreamDelta, StreamLLMChoice, Citation, LLMModuleResult, StreamLLMModuleResult, ModuleResults, StreamModuleResults, SAPAPIError, SAPAPIErrorStreaming, CompletionPostResponse, @@ -36,6 +37,9 @@ __all__ = [ + # cache_control + "CacheControl", + # azure_content_filter "AzureContentFilter", "AzureContentSafetyInput", "AzureContentSafetyOutput", "AzureThreshold", @@ -79,7 +83,7 @@ "ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem", # response - "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob", + "CacheCreationTokenDetails", "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob", "ChatCompletionTokenLogprob", "ChoiceLogprobs", "LLMChoice", "StreamFunctionObject", "StreamToolCall", "StreamDelta", "StreamLLMChoice", "Citation", "LLMModuleResult", "StreamLLMModuleResult", "ModuleResults", "StreamModuleResults", "SAPAPIError", "SAPAPIErrorStreaming", "CompletionPostResponse", diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/cache_control.py b/packages/gen/gen_ai_hub/orchestration_v2/models/cache_control.py new file mode 100644 index 00000000..2d232c3f --- /dev/null +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/cache_control.py @@ -0,0 +1,38 @@ +"""Cache control for prompt caching on supported models.""" +from typing import Any, Dict, Literal, Optional + +from pydantic import model_serializer + +from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel + + +class CacheControl(BaseModel): + """Marks a message content block or tool definition for prompt caching. + + When attached to a content block, the model stores intermediate computation + results for that content and reuses them on subsequent requests within the + TTL window, reducing both latency and token costs. + + Attach ``CacheControl`` directly to a content block (``TextPart``, ``ImagePart``) or + to a ``ChatCompletionTool``. + + Args: + type: ``"ephemeral"`` + ttl: Cache duration. ``"5m"`` (default) or ``"1h"`` (select Anthropic + models only). Omit for Amazon Nova or when the default is sufficient. + """ + + type: Literal["ephemeral"] + ttl: Optional[Literal["5m", "1h"]] = None + + @model_serializer(mode="wrap") + def serialize_wire_format(self, handler: Any) -> Dict[str, Any]: + """Serialize to the wire format, omitting ``ttl`` when not set. + + :return: ``{"type": "ephemeral"}`` or ``{"type": "ephemeral", "ttl": ""}`` + :rtype: dict + """ + data: Dict[str, Any] = handler(self) + if data.get("ttl") is None: + data.pop("ttl", None) + return data diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/multimodal_items.py b/packages/gen/gen_ai_hub/orchestration_v2/models/multimodal_items.py index 87855b24..2f46ed71 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/multimodal_items.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/multimodal_items.py @@ -12,6 +12,7 @@ from pydantic.main import IncEx from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel +from gen_ai_hub.orchestration_v2.models.cache_control import CacheControl class ImageDetailLevel(Enum): @@ -38,9 +39,12 @@ class TextPart(BaseModel): text: The string content of the text part. type: The type identifier, defaulting to "text". + + cache_control: Optional cache control settings for prompt caching. """ text: str type_: Literal["text"] = Field(default="text", alias="type") + cache_control: Optional[CacheControl] = None class ImageUrl(BaseModel): @@ -56,7 +60,6 @@ class ImageUrl(BaseModel): detail: Optional[ImageDetailLevel] = None -# @dataclass class ImagePart(BaseModel): """ Represents an image segment within a multimodal content block. @@ -65,9 +68,12 @@ class ImagePart(BaseModel): image_url: An `ImageUrl` object containing the image's location and detail level. type: The type identifier, defaulting to "image_url". + + cache_control: Optional cache control settings for prompt caching. """ image_url: ImageUrl type_: Literal["image_url"] = Field(default="image_url", alias="type") + cache_control: Optional[CacheControl] = None ContentPart = Union[TextPart, ImagePart] 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..cc6c6449 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/response.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/response.py @@ -27,16 +27,35 @@ class ResponseBaseModel(BaseModel): ) +class CacheCreationTokenDetails(ResponseBaseModel): + """ + Per-TTL breakdown of tokens written to the prompt cache. + + Present only when cache_control includes an explicit ttl value. + + Attributes: + ephemeral_5m_input_tokens: Tokens cached with a 5-minute TTL. + ephemeral_1h_input_tokens: Tokens cached with a 1-hour TTL. + """ + ephemeral_5m_input_tokens: Optional[int] = None + ephemeral_1h_input_tokens: Optional[int] = None + + class PromptTokensDetails(ResponseBaseModel): """ Represents the details of prompt tokens used in a specific operation. Attributes: audio_tokens (Optional[int]): Audio input tokens present in the prompt. - cached_tokens (Optional[int]): Cached tokens present in the prompt. + cached_tokens (Optional[int]): Tokens read from the prompt cache (cache hit). + cache_creation_tokens (Optional[int]): Tokens written to the prompt cache (cache miss). + cache_creation_token_details (Optional[CacheCreationTokenDetails]): Per-TTL + breakdown of cache writes. Present only when an explicit ttl was used. """ audio_tokens: Optional[int] = None cached_tokens: Optional[int] = None + cache_creation_tokens: Optional[int] = None + cache_creation_token_details: Optional[CacheCreationTokenDetails] = None class CompletionTokensDetails(ResponseBaseModel): """ @@ -385,7 +404,8 @@ class OrchestrationResponseWithRetries(CompletionPostResponse): """ retries: int = 0 -__all__ = ["PromptTokensDetails", +__all__ = ["CacheCreationTokenDetails", + "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", diff --git a/packages/gen/gen_ai_hub/orchestration_v2/models/tools.py b/packages/gen/gen_ai_hub/orchestration_v2/models/tools.py index ea4811bf..9177b02a 100644 --- a/packages/gen/gen_ai_hub/orchestration_v2/models/tools.py +++ b/packages/gen/gen_ai_hub/orchestration_v2/models/tools.py @@ -11,6 +11,7 @@ from pydantic import Field from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel +from gen_ai_hub.orchestration_v2.models.cache_control import CacheControl def python_type_to_json_type(py_type): @@ -66,33 +67,24 @@ def python_type_to_json_type(py_type): class ChatCompletionTool(BaseModel): - """ - Base class for all chat completion tools. + """A tool the model may call, identified by type ``"function"``. Args: - type (Literal["function"]): The type of the tool. Currently, only function is supported. + type (Literal["function"]): The type of the tool. Currently, only function is supported. + cache_control: Prompt-caching directive. """ - type_: Literal["function"] = Field(default="function", - alias="type", - description="The type of the tool. Currently, only function is supported.") + type_: Literal["function"] = Field(default="function", alias="type") + cache_control: Optional[CacheControl] = None class FunctionObject(BaseModel): - """ - Represents a function. - Args: - name (str): The name of the function to be called. Must be a-z, A-Z, 0-9, - or contain underscores and dashes, with a maximum length of 64. - - description (str): A description of what the function does, used by the model - to choose when and how to call the function. - - parameters (dict): The parameters the functions accepts, described as a JSON Schema object. - Omitting parameters defines a function with an empty parameter list. + """A function definition used inside a ``FunctionTool``. - strict (bool, optional): Whether to enable strict schema adherence when generating the function call. - If set to true, the model will follow the exact schema defined in the parameters field. - Only a subset of JSON Schema is supported when strict is true. Defaults to False. + Args: + name: Function name. Must match ``^[a-zA-Z0-9_-]+$``, max 64 chars. + description: What the function does; used by the model to decide when to call it. + parameters: JSON Schema object describing accepted parameters. + strict: When ``True``, the model follows the schema exactly. Defaults to ``False``. """ description: Optional[str] = None name: str @@ -102,15 +94,19 @@ class FunctionObject(BaseModel): class FunctionTool(ChatCompletionTool): - """ - Represents a function tool for OpenAI-like function calling. + """A callable function tool for OpenAI-style function calling. + + Inherits all fields from :class:`ChatCompletionTool`: Args: - type (Literal["function"]): The type of the tool. Currently, only function is supported. + type: Always ``"function"``. Serialized via the ``type`` alias. + cache_control: Prompt-caching directive. - function (FunctionObject): The function to be called. + Additional args: + + function: The function definition — name, description, parameters schema, + and optional strict flag. See :class:`FunctionObject`. """ - type_: Literal["function"] = Field(default="function", alias="type") function: FunctionObject def execute(self, **kwargs: Any) -> Any: diff --git a/packages/gen/integration_tests/orchestration_v2/test_cache_control.py b/packages/gen/integration_tests/orchestration_v2/test_cache_control.py new file mode 100644 index 00000000..c6ff8eeb --- /dev/null +++ b/packages/gen/integration_tests/orchestration_v2/test_cache_control.py @@ -0,0 +1,153 @@ +""" +Live integration tests for prompt caching (cache_control) via Orchestration V2. + +Caching is supported for Anthropic Claude and Amazon Nova models. +The tests target anthropic--claude-4.6-sonnet (1024-token minimum, 5m and 1h TTLs). + +Wire path: + ai-sdk-python -> SAP AI Core /v2/completion -> SAP LiteLLM fork -> Anthropic API + +Response fields (from SAP AI Core orchestration docs): + prompt_tokens_details.cached_tokens -- tokens read from cache (hit) + prompt_tokens_details.cache_creation_tokens -- tokens written to cache (miss) + prompt_tokens_details.cache_creation_token_details.ephemeral_5m_input_tokens + prompt_tokens_details.cache_creation_token_details.ephemeral_1h_input_tokens +""" +import unittest + +from gen_ai_hub.orchestration_v2.models.cache_control import CacheControl +from gen_ai_hub.orchestration_v2.models.config import OrchestrationConfig, ModuleConfig +from gen_ai_hub.orchestration_v2.models.llm_model_details import LLMModelDetails +from gen_ai_hub.orchestration_v2.models.message import SystemMessage, UserMessage +from gen_ai_hub.orchestration_v2.models.multimodal_items import TextPart +from gen_ai_hub.orchestration_v2.models.template import Template, PromptTemplatingModuleConfig +from gen_ai_hub.orchestration_v2.service import OrchestrationService +from integration_tests.orchestration_v2.test_base import OrchestrationServiceTestBase +from integration_tests.test_helpers import retry_on_429_or_503 + +# Must exceed the 1024-token minimum for claude-4.6-sonnet cache points. +_LONG_SYSTEM_PROMPT = ( + "You are a helpful assistant with deep knowledge of European history. " + "Below is a detailed reference text that you must use to answer questions accurately.\n\n" + + ( + "The Roman Empire was one of the largest empires in ancient history. " + "At its height under Emperor Trajan in 117 AD, it covered over 5 million " + "square kilometres and held 70 million people, roughly 21 percent of the " + "world's population at the time. The empire's longevity — nearly five " + "centuries in the west and fifteen in the east — shaped the languages, " + "laws, religions, and borders of modern Europe. Latin evolved into the " + "Romance languages: Italian, Spanish, Portuguese, French, and Romanian. " + "Roman law underlies most continental legal systems today. Christianity, " + "adopted as the state religion under Theodosius I in 380 AD, spread " + "throughout the empire and became the dominant faith of Europe. " + "The fall of the Western Roman Empire in 476 AD, when the Germanic " + "chieftain Odoacer deposed the last emperor Romulus Augustulus, marks " + "the conventional boundary between ancient and medieval history. " + "The Eastern Roman Empire, known as the Byzantine Empire, continued " + "for nearly a thousand more years until the fall of Constantinople to " + "the Ottoman Turks in 1453. Byzantine culture preserved classical Greek " + "and Roman learning through the Dark Ages and transmitted it to the " + "Renaissance. The Silk Road trade routes connecting Rome to China " + "facilitated the exchange of goods, diseases, and ideas across Eurasia. " + "Roman engineering achievements — aqueducts, roads, concrete construction, " + "and underfloor heating — were not equalled in Europe for over a millennium " + "after the empire's fall. The Colosseum, completed in 80 AD, could seat " + "50,000 to 80,000 spectators and hosted gladiatorial contests, animal " + "hunts, and public executions for four centuries. " + ) * 4 # repeat to comfortably exceed 1024 tokens +) + +_LLM = LLMModelDetails( + name="anthropic--claude-4.6-sonnet", + params={"max_tokens": 64, "temperature": 0.0}, +) + + +def _config(messages, tools=None): + return OrchestrationConfig( + modules=ModuleConfig( + prompt_templating=PromptTemplatingModuleConfig( + prompt=Template(template=messages, tools=tools), + model=_LLM, + ) + ) + ) + + +class TestPromptCachingLive(OrchestrationServiceTestBase): + """Live integration tests against the SAP AI Core orchestration V2 service.""" + + def setUp(self): + super().setUp() + self.service = OrchestrationService(self.api_url) + + # ------------------------------------------------------------------ + # 1. Cache MISS on first call — cache breakpoint on TextPart directly + # ------------------------------------------------------------------ + + @retry_on_429_or_503(max_retries=3, initial_delay=2.0, backoff_factor=2.0) + def test_cache_miss_on_first_call(self): + """First call with cache_control on the last TextPart block returns + non-zero cache token activity. + + cache_creation_tokens > 0 on a true miss; cached_tokens > 0 when the + cache entry is already warm from a previous run. Either proves the + cache_control breakpoint was accepted by the server. + """ + config = _config([ + SystemMessage(content=[TextPart(text=_LONG_SYSTEM_PROMPT, cache_control=CacheControl(type="ephemeral"))]), + UserMessage(content="In one word: what language did Romans speak?"), + ]) + response = self.service.run(config=config) + details = response.final_result.usage.prompt_tokens_details + self.assertIsNotNone(details) + cache_active = (details.cache_creation_tokens or 0) + (details.cached_tokens or 0) + self.assertGreater( + cache_active, 0, + f"Expected cache activity (cache_creation_tokens or cached_tokens > 0), " + f"got: {details}", + ) + + # ------------------------------------------------------------------ + # 2. Cache HIT on repeated call + # ------------------------------------------------------------------ + + @retry_on_429_or_503(max_retries=3, initial_delay=2.0, backoff_factor=2.0) + def test_cache_hit_on_repeated_call(self): + """Second call with the same cache breakpoint produces cached_tokens > 0.""" + config = _config([ + SystemMessage(content=[TextPart(text=_LONG_SYSTEM_PROMPT, cache_control=CacheControl(type="ephemeral"))]), + UserMessage(content="In one word: what language did Romans speak?"), + ]) + self.service.run(config=config) # populate cache + response = self.service.run(config=config) + details = response.final_result.usage.prompt_tokens_details + self.assertIsNotNone(details) + self.assertGreater( + details.cached_tokens, 0, + "Expected cached_tokens > 0 on second call (cache hit).", + ) + + # ------------------------------------------------------------------ + # 3. Explicit 1h TTL via TextPart + # ------------------------------------------------------------------ + + @retry_on_429_or_503(max_retries=3, initial_delay=2.0, backoff_factor=2.0) + def test_explicit_ttl_1h_via_text_part(self): + """Attaching CacheControl(type='ephemeral', ttl='1h') directly to a TextPart returns + cache_creation_token_details with ephemeral_1h_input_tokens.""" + config = _config([ + SystemMessage(content=[TextPart(text=_LONG_SYSTEM_PROMPT, cache_control=CacheControl(type='ephemeral', ttl='1h'))]), + UserMessage(content="Name the last Western Roman emperor."), + ]) + response = self.service.run(config=config) + details = response.final_result.usage.prompt_tokens_details + self.assertIsNotNone(details) + self.assertIsNotNone( + details.cache_creation_token_details, + "Expected cache_creation_token_details when ttl='1h' is used.", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/gen/tests/orchestration_v2/test_cache_control_v2.py b/packages/gen/tests/orchestration_v2/test_cache_control_v2.py new file mode 100644 index 00000000..adadf266 --- /dev/null +++ b/packages/gen/tests/orchestration_v2/test_cache_control_v2.py @@ -0,0 +1,117 @@ +""" +Unit tests for prompt-caching serialization (cache_control) in Orchestration V2. + +Covers the three spec attachment points: + - TextContent.cache_control (TextPart) + - UserChatMessageContentItem (TextPart / ImagePart) + - ChatCompletionTool.cache_control +""" +import unittest + +from gen_ai_hub.orchestration_v2.models.cache_control import CacheControl +from gen_ai_hub.orchestration_v2.models.multimodal_items import TextPart, ImagePart, ImageUrl +from gen_ai_hub.orchestration_v2.models.tools import FunctionTool, FunctionObject + + +class TestCacheControlModel(unittest.TestCase): + """CacheControl serialization.""" + + def test_default_ttl_omits_key(self): + """CacheControl(type='ephemeral') with no TTL serializes to {"type": "ephemeral"}.""" + d = CacheControl(type="ephemeral").model_dump(by_alias=True) + self.assertEqual(d, {"type": "ephemeral"}) + self.assertNotIn("ttl", d) + + def test_5m_ttl(self): + d = CacheControl(type="ephemeral", ttl="5m").model_dump(by_alias=True) + self.assertEqual(d, {"type": "ephemeral", "ttl": "5m"}) + + def test_1h_ttl(self): + d = CacheControl(type="ephemeral", ttl="1h").model_dump(by_alias=True) + self.assertEqual(d, {"type": "ephemeral", "ttl": "1h"}) + + +class TestTextPartCacheControl(unittest.TestCase): + """TextPart.cache_control serialization.""" + + def test_with_cache_control(self): + part = TextPart(text="hello", cache_control=CacheControl(type="ephemeral")) + d = part.model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["type"], "text") + self.assertEqual(d["text"], "hello") + self.assertEqual(d["cache_control"], {"type": "ephemeral"}) + + def test_without_cache_control_omits_key(self): + part = TextPart(text="hello") + d = part.model_dump(by_alias=True, exclude_none=True) + self.assertNotIn("cache_control", d) + + def test_1h_ttl(self): + part = TextPart(text="hello", cache_control=CacheControl(type="ephemeral", ttl="1h")) + d = part.model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["cache_control"], {"type": "ephemeral", "ttl": "1h"}) + + +class TestImagePartCacheControl(unittest.TestCase): + """ImagePart.cache_control serialization.""" + + def _image_part(self, **kwargs): + return ImagePart(image_url=ImageUrl(url="https://example.com/img.png"), **kwargs) + + def test_with_cache_control(self): + d = self._image_part(cache_control=CacheControl(type="ephemeral")).model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["type"], "image_url") + self.assertEqual(d["cache_control"], {"type": "ephemeral"}) + + def test_without_cache_control_omits_key(self): + d = self._image_part().model_dump(by_alias=True, exclude_none=True) + self.assertNotIn("cache_control", d) + + def test_1h_ttl(self): + d = self._image_part(cache_control=CacheControl(type="ephemeral", ttl="1h")).model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["cache_control"], {"type": "ephemeral", "ttl": "1h"}) + + +class TestFunctionToolCacheControl(unittest.TestCase): + """ChatCompletionTool.cache_control serialization via FunctionTool.""" + + def _tool(self, **kwargs): + return FunctionTool( + function=FunctionObject( + name="classify", + description="Classify input.", + parameters={"type": "object", "properties": {}}, + ), + **kwargs, + ) + + def test_with_cache_control(self): + d = self._tool(cache_control=CacheControl(type="ephemeral")).model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["cache_control"], {"type": "ephemeral"}) + + def test_without_cache_control_omits_key(self): + d = self._tool().model_dump(by_alias=True, exclude_none=True) + self.assertNotIn("cache_control", d) + + def test_1h_ttl(self): + d = self._tool(cache_control=CacheControl(type="ephemeral", ttl="1h")).model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["cache_control"], {"type": "ephemeral", "ttl": "1h"}) + + def test_type_field_serializes(self): + """type_ with alias 'type' must appear in output.""" + d = self._tool().model_dump(by_alias=True, exclude_none=True) + self.assertEqual(d["type"], "function") + + def test_no_duplicate_type_from_subclass(self): + """FunctionTool must not declare its own type_ field (inherits from ChatCompletionTool).""" + import inspect + own_fields = FunctionTool.model_fields + # 'type_' is defined on ChatCompletionTool; FunctionTool should only add 'function' + self.assertIn("function", own_fields) + # Ensure serialization is still correct (regression guard) + d = self._tool().model_dump(by_alias=True) + self.assertEqual(d["type"], "function") + + +if __name__ == "__main__": + unittest.main() diff --git a/packages/gen/tests/orchestration_v2/test_flat_import.py b/packages/gen/tests/orchestration_v2/test_flat_import.py index 7606ffac..ca877b41 100644 --- a/packages/gen/tests/orchestration_v2/test_flat_import.py +++ b/packages/gen/tests/orchestration_v2/test_flat_import.py @@ -1,4 +1,7 @@ expected = { + # cache_control + "CacheControl", + # azure_content_filter "AzureContentFilter", "AzureContentSafetyInput", "AzureContentSafetyOutput", "AzureThreshold", @@ -46,7 +49,7 @@ "ChatCompletionTokenLogprob", "ChoiceLogprobs", "LLMChoice", "StreamFunctionObject", "StreamToolCall", "StreamDelta", "StreamLLMChoice", "Citation", "LLMModuleResult", "StreamLLMModuleResult", "ModuleResults", "StreamModuleResults", "SAPAPIError", "SAPAPIErrorStreaming", "CompletionPostResponse", - "StreamCompletionPostResponse", "ErrorResponse", "ErrorResponseStreaming", "OrchestrationResponseWithRetries", + "StreamCompletionPostResponse", "ErrorResponse", "ErrorResponseStreaming", "OrchestrationResponseWithRetries", "CacheCreationTokenDetails", # response_format "ResponseFormatText", "ResponseFormatJsonObject", "ResponseFormatJsonSchema", "JSONResponseSchema", @@ -72,8 +75,8 @@ "OrchestrationService", # Exceptions - "OrchestrationError", "OrchestrationErrorList" - } + "OrchestrationError", "OrchestrationErrorList", + } def test_flat_import_all():