feat: support OpenAI cache breakpoints and configuration (#128)

Author: adubovik

README.md

@@ -25,7 +25,6 @@
   - [Make Chat Completions Requests](#make-chat-completions-requests)
     - [Without Streaming](#without-streaming)
     - [With Streaming](#with-streaming)
-    - [DIAL-specific and Extended Parameters](#dial-specific-and-extended-parameters)
   - [Working with Files](#working-with-files)
     - [Working with URLs](#working-with-urls)
     - [Uploading Files](#uploading-files)
@@ -465,115 +464,6 @@ ChatCompletionChunk(
 )
 ```
 
-#### DIAL-specific and Extended Parameters
-
-Along with the standard OpenAI parameters, `chat.completions.create` accepts
-the DIAL extensions and the newer OpenAI parameters:
-
-```python
-completion = client.chat.completions.create(
-    deployment_name="gpt-4o",
-    stream=False,
-    messages=[
-        # Messages support multi-modal content parts,
-        # the "developer" role and per-message cache breakpoints
-        {
-            "role": "developer",
-            "content": "Be brief",
-            "custom_fields": {"cache_breakpoint": {"expire_at": "1h"}},
-        },
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": "What is on the picture?"},
-                {
-                    "type": "image_url",
-                    "image_url": {"url": "https://example.com/image.png"},
-                },
-            ],
-            # DIAL attachments, stages and forms
-            "custom_content": {
-                "attachments": [
-                    {"type": "image/png", "url": "files/bucket/image.png"}
-                ]
-            },
-        },
-    ],
-    tools=[
-        {
-            "type": "function",
-            "function": {"name": "get_weather", "parameters": {}, "strict": True},
-            "custom_fields": {"cache_breakpoint": {}},
-        },
-        # DIAL static tools, resolved by DIAL Core itself
-        {
-            "type": "static_function",
-            "static_function": {"name": "search", "configuration": {}},
-        },
-    ],
-    tool_choice="required",
-    parallel_tool_calls=False,
-    reasoning_effort="high",
-    max_completion_tokens=1000,
-    response_format={
-        "type": "json_schema",
-        "json_schema": {"name": "answer", "schema": {"type": "object"}},
-    },
-    stream_options={"include_usage": True},
-    # DIAL-specific parameters
-    max_prompt_tokens=8000,
-    custom_fields={
-        "configuration": {},
-        "cache_breakpoint": {"expire_at": "5m"},
-    },
-)
-```
-
-The response models cover the DIAL extensions as well:
-
-```pycon
->>> completion.choices[0].message.custom_content
-CustomContent(
-    stages=[
-        Stage(
-            index=None,
-            name='Thinking',
-            status='completed',
-            content='...',
-            attachments=None
-        )
-    ],
-    attachments=None,
-    state=None,
-    form_value=None,
-    form_schema=None
-)
->>> completion.usage
-CompletionUsage(
-    prompt_tokens=11,
-    completion_tokens=1,
-    total_tokens=12,
-    prompt_tokens_details=PromptTokensDetails(
-        cached_tokens=8,
-        cache_write_tokens=3
-    ),
-    completion_tokens_details=CompletionTokensDetails(reasoning_tokens=1)
-)
->>> completion.statistics
-Statistics(
-    usage_per_model=[
-        UsagePerModel(
-            index=0,
-            model='gpt-4o',
-            prompt_tokens=11,
-            completion_tokens=1,
-            total_tokens=12
-        )
-    ],
-    discarded_messages=[0, 1]
-)
-```
-
 ### Working with Files
 
 #### Working with URLs

aidial_client/resources/chat/completions.py

@@ -28,6 +28,7 @@
     FunctionCallSpecParam,
     FunctionParam,
     Message,
+    PromptCacheOptionsParam,
     ReasoningEffort,
     ResponseFormat,
     StaticToolParam,
@@ -73,6 +74,8 @@ def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -115,6 +118,8 @@ def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -156,6 +161,8 @@ def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -202,6 +209,8 @@ def create(
                             "max_completion_tokens": max_completion_tokens,
                             "parallel_tool_calls": parallel_tool_calls,
                             "reasoning_effort": reasoning_effort,
+                            "prompt_cache_key": prompt_cache_key,
+                            "prompt_cache_options": prompt_cache_options,
                             "response_format": response_format,
                             "stream_options": stream_options,
                         }
@@ -273,6 +282,8 @@ async def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -315,6 +326,8 @@ async def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -356,6 +369,8 @@ async def create(
         seed: int | None = None,
         user: str | None = None,
         reasoning_effort: ReasoningEffort | None = None,
+        prompt_cache_key: str | None = None,
+        prompt_cache_options: PromptCacheOptionsParam | None = None,
         response_format: ResponseFormat | None = None,
         stream_options: StreamOptions | None = None,
         custom_fields: ChatCompletionRequestCustomFields | None = None,
@@ -402,6 +417,8 @@ async def create(
                             "max_completion_tokens": max_completion_tokens,
                             "parallel_tool_calls": parallel_tool_calls,
                             "reasoning_effort": reasoning_effort,
+                            "prompt_cache_key": prompt_cache_key,
+                            "prompt_cache_options": prompt_cache_options,
                             "response_format": response_format,
                             "stream_options": stream_options,
                         }

aidial_client/types/chat/__init__.py

@@ -1,4 +1,8 @@
-from .cache import CacheBreakpointParam
+from .cache import (
+    CacheBreakpointParam,
+    PromptCacheBreakpointParam,
+    PromptCacheOptionsParam,
+)
 from .function import FunctionCallSpecParam, FunctionParam
 from .request import (
     ChatCompletionRequest,
@@ -89,6 +93,8 @@
     "MessageContentRefusalPartParam",
     "MessageContentTextPartParam",
     "MessageCustomFieldsParam",
+    "PromptCacheBreakpointParam",
+    "PromptCacheOptionsParam",
     "PromptTokensDetails",
     "ReasoningEffort",
     "ResponseFormat",

aidial_client/types/chat/cache.py

@@ -1,5 +1,16 @@
+from typing import Literal
+
 from typing_extensions import TypedDict
 
 
 class CacheBreakpointParam(TypedDict, total=False):
     expire_at: str | None
+
+
+class PromptCacheBreakpointParam(TypedDict):
+    mode: Literal["explicit"]
+
+
+class PromptCacheOptionsParam(TypedDict, total=False):
+    mode: Literal["implicit", "explicit"] | None
+    ttl: Literal["30m"] | str | None

aidial_client/types/chat/request.py

@@ -2,7 +2,10 @@
 
 from typing_extensions import TypedDict
 
-from aidial_client.types.chat.cache import CacheBreakpointParam
+from aidial_client.types.chat.cache import (
+    CacheBreakpointParam,
+    PromptCacheOptionsParam,
+)
 from aidial_client.types.chat.function import (
     FunctionCallSpecParam,
     FunctionParam,
@@ -46,6 +49,8 @@ class ChatCompletionRequest(TypedDict, total=False):
     logprobs: bool | None
     top_logprobs: int | None
     reasoning_effort: ReasoningEffort | None
+    prompt_cache_key: str | None
+    prompt_cache_options: PromptCacheOptionsParam | None
     response_format: ResponseFormat | None
     tools: list[ToolParam | StaticToolParam] | None
     tool_choice: Literal["none", "auto", "required"] | ToolCallSpecParam | None

aidial_client/types/chat/request_param.py

@@ -2,7 +2,10 @@
 
 from typing_extensions import Required, TypedDict
 
-from aidial_client.types.chat.cache import CacheBreakpointParam
+from aidial_client.types.chat.cache import (
+    CacheBreakpointParam,
+    PromptCacheBreakpointParam,
+)
 from aidial_client.types.chat.function import FunctionCallParam
 from aidial_client.types.chat.tool import ToolCallParam
 
@@ -60,19 +63,21 @@ class MessageCustomFieldsParam(TypedDict, total=False):
     cache_breakpoint: CacheBreakpointParam | None
 
 
-class MessageContentTextPartParam(TypedDict):
-    type: Literal["text"]
-    text: str
+class MessageContentTextPartParam(TypedDict, total=False):
+    type: Required[Literal["text"]]
+    text: Required[str]
+    prompt_cache_breakpoint: PromptCacheBreakpointParam | None
 
 
 class ImageURLParam(TypedDict, total=False):
     url: Required[str]
     detail: Literal["auto", "low", "high"] | None
 
 
-class MessageContentImagePartParam(TypedDict):
-    type: Literal["image_url"]
-    image_url: ImageURLParam
+class MessageContentImagePartParam(TypedDict, total=False):
+    type: Required[Literal["image_url"]]
+    image_url: Required[ImageURLParam]
+    prompt_cache_breakpoint: PromptCacheBreakpointParam | None
 
 
 class InputFileParam(TypedDict, total=False):
@@ -81,9 +86,10 @@ class InputFileParam(TypedDict, total=False):
     filename: str | None
 
 
-class MessageContentFilePartParam(TypedDict):
-    type: Literal["file"]
-    file: InputFileParam
+class MessageContentFilePartParam(TypedDict, total=False):
+    type: Required[Literal["file"]]
+    file: Required[InputFileParam]
+    prompt_cache_breakpoint: PromptCacheBreakpointParam | None
 
 
 class InputAudioParam(TypedDict):
@@ -92,9 +98,10 @@ class InputAudioParam(TypedDict):
     format: str
 
 
-class MessageContentAudioPartParam(TypedDict):
-    type: Literal["input_audio"]
-    input_audio: InputAudioParam
+class MessageContentAudioPartParam(TypedDict, total=False):
+    type: Required[Literal["input_audio"]]
+    input_audio: Required[InputAudioParam]
+    prompt_cache_breakpoint: PromptCacheBreakpointParam | None
 
 
 class MessageContentRefusalPartParam(TypedDict):

tests/resources/completions/test_completions_prompt_cache.py

@@ -0,0 +1,63 @@
+"""Prompt caching request fields"""
+
+from typing import Any
+
+import pytest
+
+from tests.resources.completions.conftest import GetRequestBody
+
+pytestmark = pytest.mark.asyncio
+
+
+class TestPromptCache:
+    """Prompt cache options must reach the request body as-is"""
+
+    async def test_request_prompt_cache_key(
+        self, get_request_body: GetRequestBody
+    ):
+        body = await get_request_body(prompt_cache_key="user-42")
+
+        assert body["prompt_cache_key"] == "user-42"
+
+    async def test_request_prompt_cache_options(
+        self, get_request_body: GetRequestBody
+    ):
+        body = await get_request_body(
+            prompt_cache_options={"mode": "explicit", "ttl": "30m"}
+        )
+
+        assert body["prompt_cache_options"]["mode"] == "explicit"
+        assert body["prompt_cache_options"]["ttl"] == "30m"
+
+    @pytest.mark.parametrize(
+        "part",
+        [
+            {"type": "text", "text": "What is on the picture?"},
+            {"type": "image_url", "image_url": {"url": "http://a.com/b.png"}},
+            {"type": "file", "file": {"file_id": "files/bucket/a.pdf"}},
+            {
+                "type": "input_audio",
+                "input_audio": {"data": "Zm9v", "format": "wav"},
+            },
+        ],
+        ids=["text", "image_url", "file", "input_audio"],
+    )
+    async def test_request_content_part_breakpoint(
+        self, get_request_body: GetRequestBody, part: dict[str, Any]
+    ):
+        body = await get_request_body(
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            **part,
+                            "prompt_cache_breakpoint": {"mode": "explicit"},
+                        }
+                    ],
+                }
+            ]
+        )
+
+        sent_part = body["messages"][0]["content"][0]
+        assert sent_part["prompt_cache_breakpoint"] == {"mode": "explicit"}