diff --git a/README.md b/README.md index 73a432d..e31b8ea 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ - [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) @@ -464,6 +465,115 @@ 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 diff --git a/aidial_client/resources/chat/completions.py b/aidial_client/resources/chat/completions.py index 8e2e06f..c4c0f46 100644 --- a/aidial_client/resources/chat/completions.py +++ b/aidial_client/resources/chat/completions.py @@ -28,6 +28,10 @@ FunctionCallSpecParam, FunctionParam, Message, + ReasoningEffort, + ResponseFormat, + StaticToolParam, + StreamOptions, ToolCallSpecParam, ToolParam, ) @@ -51,19 +55,26 @@ def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, logprobs: bool | None = None, top_logprobs: int | None = None, @@ -86,19 +97,26 @@ def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, logprobs: bool | None = None, top_logprobs: int | None = None, @@ -120,19 +138,26 @@ def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, logprobs: bool | None = None, top_logprobs: int | None = None, @@ -165,11 +190,24 @@ def create( "tools": tools, "top_p": top_p, "user": user, - "max_prompt_tokens": max_prompt_tokens, - "custom_fields": custom_fields, "logprobs": logprobs, "top_logprobs": top_logprobs, - "extra_body": extra_body, + # DIAL-specific parameters and the ones which aren't supported + # by every openai version are sent in the request body directly + "extra_body": { + **remove_none( + { + "max_prompt_tokens": max_prompt_tokens, + "custom_fields": custom_fields, + "max_completion_tokens": max_completion_tokens, + "parallel_tool_calls": parallel_tool_calls, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "stream_options": stream_options, + } + ), + **extra_body, + }, "extra_query": { "api-version": ( api_version or self.default_api_version or Omit() @@ -217,20 +255,29 @@ async def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, + logprobs: bool | None = None, + top_logprobs: int | None = None, # Extra params extra_body: dict[str, Any] | None = None, extra_headers: Mapping[StrictStr, StrictStr] | None = None, @@ -250,19 +297,26 @@ async def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, logprobs: bool | None = None, top_logprobs: int | None = None, @@ -284,19 +338,26 @@ async def create( function_call: Literal["none", "auto"] | FunctionCallSpecParam | None = None, - tools: list[ToolParam] | None = None, - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None = None, + tools: list[ToolParam | StaticToolParam] | None = None, + tool_choice: Literal["none", "auto", "required"] + | ToolCallSpecParam + | None = None, + parallel_tool_calls: bool | None = None, temperature: float | None = None, top_p: float | None = None, n: int | None = None, stop: str | list[str] | None = None, max_tokens: int | None = None, + max_completion_tokens: int | None = None, max_prompt_tokens: Literal["infinity"] | int | None = None, presence_penalty: float | None = None, frequency_penalty: float | None = None, logit_bias: dict | None = None, seed: int | None = None, user: str | None = None, + reasoning_effort: ReasoningEffort | None = None, + response_format: ResponseFormat | None = None, + stream_options: StreamOptions | None = None, custom_fields: ChatCompletionRequestCustomFields | None = None, logprobs: bool | None = None, top_logprobs: int | None = None, @@ -329,11 +390,24 @@ async def create( "tools": tools, "top_p": top_p, "user": user, - "max_prompt_tokens": max_prompt_tokens, - "custom_fields": custom_fields, "logprobs": logprobs, "top_logprobs": top_logprobs, - "extra_body": extra_body, + # DIAL-specific parameters and the ones which aren't supported + # by every openai version are sent in the request body directly + "extra_body": { + **remove_none( + { + "max_prompt_tokens": max_prompt_tokens, + "custom_fields": custom_fields, + "max_completion_tokens": max_completion_tokens, + "parallel_tool_calls": parallel_tool_calls, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "stream_options": stream_options, + } + ), + **extra_body, + }, "extra_query": { "api-version": ( api_version or self.default_api_version or Omit() diff --git a/aidial_client/types/chat/__init__.py b/aidial_client/types/chat/__init__.py index e4adccc..e261017 100644 --- a/aidial_client/types/chat/__init__.py +++ b/aidial_client/types/chat/__init__.py @@ -1,26 +1,112 @@ +from .cache import CacheBreakpointParam from .function import FunctionCallSpecParam, FunctionParam -from .request import ChatCompletionRequest +from .request import ( + ChatCompletionRequest, + ChatCompletionRequestCustomFields, + ReasoningEffort, + StreamOptions, +) from .request_param import ( + AssistantMessageParam, + AttachmentParam, + CustomContentParam, + DeveloperMessageParam, FunctionMessageParam, + ImageURLParam, + InputAudioParam, + InputFileParam, Message, + MessageContentAudioPartParam, + MessageContentFilePartParam, + MessageContentImagePartParam, + MessageContentPartParam, + MessageContentRefusalPartParam, + MessageContentTextPartParam, + MessageCustomFieldsParam, + ResponseFormat, + ResponseFormatJsonObject, + ResponseFormatJsonSchema, + ResponseFormatJsonSchemaObject, + ResponseFormatText, + StageParam, SystemMessageParam, ToolMessageParam, UserMessageParam, ) -from .response import ChatCompletionChunk, ChatCompletionResponse -from .tool import ToolCallSpecParam, ToolParam +from .response import ( + Attachment, + ChatCompletionChunk, + ChatCompletionMessage, + ChatCompletionMessageDelta, + ChatCompletionResponse, + Choice, + ChoiceDelta, + CompletionTokensDetails, + CompletionUsage, + CustomContent, + PromptTokensDetails, + Stage, + Statistics, + UsagePerModel, +) +from .tool import ( + StaticFunctionParam, + StaticToolParam, + ToolCallSpecParam, + ToolCustomFieldsParam, + ToolParam, +) __all__ = [ + "AssistantMessageParam", + "Attachment", + "AttachmentParam", + "CacheBreakpointParam", + "ChatCompletionChunk", + "ChatCompletionMessage", + "ChatCompletionMessageDelta", "ChatCompletionRequest", - "FunctionParam", + "ChatCompletionRequestCustomFields", + "ChatCompletionResponse", + "Choice", + "ChoiceDelta", + "CompletionTokensDetails", + "CompletionUsage", + "CustomContent", + "CustomContentParam", + "DeveloperMessageParam", "FunctionCallSpecParam", - "ToolParam", - "ToolCallSpecParam", + "FunctionMessageParam", + "FunctionParam", + "ImageURLParam", + "InputAudioParam", + "InputFileParam", "Message", + "MessageContentAudioPartParam", + "MessageContentFilePartParam", + "MessageContentImagePartParam", + "MessageContentPartParam", + "MessageContentRefusalPartParam", + "MessageContentTextPartParam", + "MessageCustomFieldsParam", + "PromptTokensDetails", + "ReasoningEffort", + "ResponseFormat", + "ResponseFormatJsonObject", + "ResponseFormatJsonSchema", + "ResponseFormatJsonSchemaObject", + "ResponseFormatText", + "Stage", + "StageParam", + "StaticFunctionParam", + "StaticToolParam", + "Statistics", + "StreamOptions", + "SystemMessageParam", + "ToolCallSpecParam", + "ToolCustomFieldsParam", "ToolMessageParam", + "ToolParam", + "UsagePerModel", "UserMessageParam", - "SystemMessageParam", - "FunctionMessageParam", - "ChatCompletionResponse", - "ChatCompletionChunk", ] diff --git a/aidial_client/types/chat/cache.py b/aidial_client/types/chat/cache.py new file mode 100644 index 0000000..9179190 --- /dev/null +++ b/aidial_client/types/chat/cache.py @@ -0,0 +1,5 @@ +from typing_extensions import TypedDict + + +class CacheBreakpointParam(TypedDict, total=False): + expire_at: str | None diff --git a/aidial_client/types/chat/function.py b/aidial_client/types/chat/function.py index c38e453..2941708 100644 --- a/aidial_client/types/chat/function.py +++ b/aidial_client/types/chat/function.py @@ -5,6 +5,7 @@ class FunctionParam(TypedDict, total=False): name: Required[str] description: str | None parameters: dict | None + strict: bool | None class FunctionCallParam(TypedDict): diff --git a/aidial_client/types/chat/request.py b/aidial_client/types/chat/request.py index 139d5f0..4d7b565 100644 --- a/aidial_client/types/chat/request.py +++ b/aidial_client/types/chat/request.py @@ -2,16 +2,28 @@ from typing_extensions import TypedDict +from aidial_client.types.chat.cache import CacheBreakpointParam from aidial_client.types.chat.function import ( FunctionCallSpecParam, FunctionParam, ) from aidial_client.types.chat.request_param import Message, ResponseFormat -from aidial_client.types.chat.tool import ToolCallSpecParam, ToolParam +from aidial_client.types.chat.tool import ( + StaticToolParam, + ToolCallSpecParam, + ToolParam, +) + +ReasoningEffort = Literal["none", "minimal", "low", "medium", "high"] + + +class StreamOptions(TypedDict, total=False): + include_usage: bool | None class ChatCompletionRequestCustomFields(TypedDict, total=False): configuration: dict[str, Any] | None + cache_breakpoint: CacheBreakpointParam | None class ChatCompletionRequest(TypedDict, total=False): @@ -19,8 +31,10 @@ class ChatCompletionRequest(TypedDict, total=False): temperature: float | None top_p: float | None stream: bool | None + stream_options: StreamOptions | None stop: str | list[str] | None max_tokens: int | None + max_completion_tokens: int | None presence_penalty: float | None frequency_penalty: float | None logit_bias: dict | None @@ -30,10 +44,12 @@ class ChatCompletionRequest(TypedDict, total=False): n: int | None seed: int | None logprobs: bool | None - top_logprobs: float | None + top_logprobs: int | None + reasoning_effort: ReasoningEffort | None response_format: ResponseFormat | None - tools: list[ToolParam] | None - tool_choice: Literal["none", "auto"] | ToolCallSpecParam | None + tools: list[ToolParam | StaticToolParam] | None + tool_choice: Literal["none", "auto", "required"] | ToolCallSpecParam | None + parallel_tool_calls: bool | None functions: list[FunctionParam] | None function_call: Literal["none", "auto"] | FunctionCallSpecParam | None max_prompt_tokens: Literal["infinity"] | int | None diff --git a/aidial_client/types/chat/request_param.py b/aidial_client/types/chat/request_param.py index 9cb1965..00c4afb 100644 --- a/aidial_client/types/chat/request_param.py +++ b/aidial_client/types/chat/request_param.py @@ -1,13 +1,35 @@ -from typing import Literal +from typing import Any, Literal from typing_extensions import Required, TypedDict +from aidial_client.types.chat.cache import CacheBreakpointParam from aidial_client.types.chat.function import FunctionCallParam from aidial_client.types.chat.tool import ToolCallParam -class ResponseFormat(TypedDict, total=False): - type: Literal["json_object", "text"] +class ResponseFormatText(TypedDict): + type: Literal["text"] + + +class ResponseFormatJsonObject(TypedDict): + type: Literal["json_object"] + + +class ResponseFormatJsonSchemaObject(TypedDict, total=False): + name: Required[str] + schema: Required[dict[str, Any]] + description: str | None + strict: bool | None + + +class ResponseFormatJsonSchema(TypedDict): + type: Literal["json_schema"] + json_schema: ResponseFormatJsonSchemaObject + + +ResponseFormat = ( + ResponseFormatText | ResponseFormatJsonObject | ResponseFormatJsonSchema +) class AttachmentParam(TypedDict, total=False): @@ -19,49 +41,133 @@ class AttachmentParam(TypedDict, total=False): reference_url: str +class StageParam(TypedDict, total=False): + name: Required[str] + status: Required[Literal["completed", "failed"]] + content: str | None + attachments: list[AttachmentParam] | None + + class CustomContentParam(TypedDict, total=False): + stages: list[StageParam] | None attachments: list[AttachmentParam] | None state: dict | None + form_value: Any | None + form_schema: Any | None + + +class MessageCustomFieldsParam(TypedDict, total=False): + cache_breakpoint: CacheBreakpointParam | None + + +class MessageContentTextPartParam(TypedDict): + type: Literal["text"] + text: str + + +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 InputFileParam(TypedDict, total=False): + file_data: str | None + file_id: str | None + filename: str | None + + +class MessageContentFilePartParam(TypedDict): + type: Literal["file"] + file: InputFileParam + + +class InputAudioParam(TypedDict): + data: str + """Either "wav", "mp3" or any other format supported by the model""" + format: str + + +class MessageContentAudioPartParam(TypedDict): + type: Literal["input_audio"] + input_audio: InputAudioParam + + +class MessageContentRefusalPartParam(TypedDict): + type: Literal["refusal"] + refusal: str + + +MessageContentPartParam = ( + MessageContentTextPartParam + | MessageContentImagePartParam + | MessageContentFilePartParam + | MessageContentAudioPartParam + | MessageContentRefusalPartParam +) + +MessageContentParam = str | list[MessageContentPartParam] class SystemMessageParam(TypedDict, total=False): role: Required[Literal["system"]] - content: Required[str] + content: Required[MessageContentParam] custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None + name: str | None + + +class DeveloperMessageParam(TypedDict, total=False): + role: Required[Literal["developer"]] + content: Required[MessageContentParam] + custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None name: str | None class UserMessageParam(TypedDict, total=False): role: Required[Literal["user"]] - content: Required[str] + content: Required[MessageContentParam] custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None name: str | None class AssistantMessageParam(TypedDict, total=False): role: Required[Literal["assistant"]] - content: str | None + content: MessageContentParam | None custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None function_call: FunctionCallParam | None tool_calls: list[ToolCallParam] + refusal: str | None name: str | None class ToolMessageParam(TypedDict, total=False): role: Required[Literal["tool"]] - content: Required[str] + content: Required[MessageContentParam] tool_call_id: Required[str] + custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None class FunctionMessageParam(TypedDict, total=False): role: Required[Literal["function"]] - content: Required[str] + content: Required[MessageContentParam] """Name of function call""" name: Required[str] + custom_content: CustomContentParam | None + custom_fields: MessageCustomFieldsParam | None Message = ( SystemMessageParam + | DeveloperMessageParam | UserMessageParam | AssistantMessageParam | ToolMessageParam diff --git a/aidial_client/types/chat/response.py b/aidial_client/types/chat/response.py index 5339443..e7d1ffc 100644 --- a/aidial_client/types/chat/response.py +++ b/aidial_client/types/chat/response.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Any, Literal from aidial_client._compatibility.pydantic import PYDANTIC_V2 from aidial_client._compatibility.pydantic_v1 import root_validator @@ -9,6 +9,9 @@ class Attachment(ExtraAllowModel): + """Index is only set in streaming responses""" + + index: int | None = None type: str | None = None title: str | None = None data: str | None = None @@ -38,15 +41,52 @@ def validate_data_or_url_v1(cls, values): return values +class Stage(ExtraAllowModel): + """Index is only set in streaming responses""" + + index: int | None = None + name: str | None = None + status: Literal["completed", "failed"] | None = None + content: str | None = None + attachments: list[Attachment] | None = None + + class CustomContent(ExtraAllowModel): + stages: list[Stage] | None = None attachments: list[Attachment] | None = None state: dict | None = None + form_value: Any | None = None + form_schema: Any | None = None + + +class PromptTokensDetails(ExtraAllowModel): + cached_tokens: int | None = None + cache_write_tokens: int | None = None + + +class CompletionTokensDetails(ExtraAllowModel): + reasoning_tokens: int | None = None class CompletionUsage(ExtraAllowModel): prompt_tokens: int completion_tokens: int total_tokens: int + prompt_tokens_details: PromptTokensDetails | None = None + completion_tokens_details: CompletionTokensDetails | None = None + + +class UsagePerModel(ExtraAllowModel): + index: int | None = None + model: str | None = None + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class Statistics(ExtraAllowModel): + usage_per_model: list[UsagePerModel] | None = None + discarded_messages: list[int] | None = None class FunctionCall(ExtraAllowModel): @@ -75,6 +115,7 @@ class ToolCallDelta(ExtraAllowModel): class ChatCompletionMessage(ExtraAllowModel): role: Literal["assistant"] content: str | None = None + refusal: str | None = None custom_content: CustomContent | None = None function_call: FunctionCall | None = None tool_calls: list[ChatCompletionMessageToolCall] | None = None @@ -83,6 +124,7 @@ class ChatCompletionMessage(ExtraAllowModel): class ChatCompletionMessageDelta(ExtraAllowModel): role: Literal["assistant"] | None = None content: str | None = None + refusal: str | None = None custom_content: CustomContent | None = None function_call: FunctionCallDelta | None = None tool_calls: list[ToolCallDelta] | None = None @@ -107,6 +149,7 @@ class ChatCompletionResponse(ExtraAllowModel): created: int model: str | None = None usage: CompletionUsage | None = None + statistics: Statistics | None = None class ChatCompletionChunk(ExtraAllowModel): @@ -116,3 +159,4 @@ class ChatCompletionChunk(ExtraAllowModel): created: int model: str | None = None usage: CompletionUsage | None = None + statistics: Statistics | None = None diff --git a/aidial_client/types/chat/tool.py b/aidial_client/types/chat/tool.py index f04a66f..5013814 100644 --- a/aidial_client/types/chat/tool.py +++ b/aidial_client/types/chat/tool.py @@ -1,7 +1,8 @@ -from typing import Literal +from typing import Any, Literal from typing_extensions import Required, TypedDict +from aidial_client.types.chat.cache import CacheBreakpointParam from aidial_client.types.chat.function import ( FunctionCallParam, FunctionCallSpecParam, @@ -9,9 +10,25 @@ ) -class ToolParam(TypedDict): - type: Literal["function"] - function: FunctionParam +class ToolCustomFieldsParam(TypedDict, total=False): + cache_breakpoint: CacheBreakpointParam | None + + +class ToolParam(TypedDict, total=False): + type: Required[Literal["function"]] + function: Required[FunctionParam] + custom_fields: ToolCustomFieldsParam | None + + +class StaticFunctionParam(TypedDict, total=False): + name: Required[str] + description: str | None + configuration: dict[str, Any] | None + + +class StaticToolParam(TypedDict): + type: Literal["static_function"] + static_function: StaticFunctionParam class ToolCallParam(TypedDict): diff --git a/tests/client_mock.py b/tests/client_mock.py index 15971de..8c56708 100644 --- a/tests/client_mock.py +++ b/tests/client_mock.py @@ -27,6 +27,7 @@ def get_client_mock( json_mock: dict[str, Any] | None = None, stream_chunks_mock: list[bytes] | None = None, exception_mock: Exception | None = None, + sent_requests: list[httpx.Request] | None = None, ) -> Dial: client_mock = Dial( api_key="dummy", @@ -34,6 +35,8 @@ def get_client_mock( ) def send_mock(request: httpx.Request, **kwargs): + if sent_requests is not None: + sent_requests.append(request) if json_mock is not None: assert status_code mock_response = httpx.Response( @@ -64,6 +67,7 @@ def get_async_client_mock( json_mock: dict[str, Any] | None = None, stream_chunks_mock: list[bytes] | None = None, exception_mock: Exception | None = None, + sent_requests: list[httpx.Request] | None = None, ) -> AsyncDial: client_mock = AsyncDial( api_key="dummy", @@ -71,6 +75,8 @@ def get_async_client_mock( ) async def send_mock(request: httpx.Request, **kwargs): + if sent_requests is not None: + sent_requests.append(request) if json_mock is not None: assert status_code mock_response = httpx.Response( diff --git a/tests/resources/completions/conftest.py b/tests/resources/completions/conftest.py new file mode 100644 index 0000000..35f075c --- /dev/null +++ b/tests/resources/completions/conftest.py @@ -0,0 +1,76 @@ +""" +Fixtures running every test in 4 modes: sync/async and streaming/non-streaming. + +The fixtures below are awaitable factories taking the fields of the mocked +response body (see `block_response`) and returning the parsed +completion/message, or - for `get_request_body` - the sent request body. +""" + +import json +from collections.abc import Awaitable, Callable +from typing import Any + +import httpx +import pytest + +from aidial_client.types.chat import ( + ChatCompletionMessage, + ChatCompletionMessageDelta, +) +from tests.utils.completions import ( + Completion, + block_response, + create_completion, + message_of, +) + +GetCompletion = Callable[..., Awaitable[Completion]] +GetMessage = Callable[ + ..., Awaitable[ChatCompletionMessage | ChatCompletionMessageDelta] +] +GetRequestBody = Callable[..., Awaitable[dict[str, Any]]] + + +@pytest.fixture(params=[False, True], ids=["sync", "async"]) +def is_async(request: pytest.FixtureRequest) -> bool: + return request.param + + +@pytest.fixture(params=[False, True], ids=["block", "stream"]) +def stream(request: pytest.FixtureRequest) -> bool: + return request.param + + +@pytest.fixture +def get_completion(is_async: bool, stream: bool) -> GetCompletion: + async def _get(**response_fields: Any) -> Completion: + return await create_completion( + is_async=is_async, + stream=stream, + response=block_response(**response_fields), + ) + + return _get + + +@pytest.fixture +def get_message(get_completion: GetCompletion) -> GetMessage: + async def _get(**response_fields: Any): + return message_of(await get_completion(**response_fields)) + + return _get + + +@pytest.fixture +def get_request_body(is_async: bool, stream: bool) -> GetRequestBody: + async def _get(**request_params: Any) -> dict[str, Any]: + sent_requests: list[httpx.Request] = [] + await create_completion( + is_async=is_async, + stream=stream, + sent_requests=sent_requests, + **request_params, + ) + return json.loads(sent_requests[0].content) + + return _get diff --git a/tests/resources/completions/test_completions_dial_fields.py b/tests/resources/completions/test_completions_dial_fields.py new file mode 100644 index 0000000..b2e57b4 --- /dev/null +++ b/tests/resources/completions/test_completions_dial_fields.py @@ -0,0 +1,156 @@ +"""DIAL-specific chat completion request and response fields""" + +from typing import Any + +import pytest + +from tests.resources.completions.conftest import ( + GetCompletion, + GetMessage, + GetRequestBody, +) + +pytestmark = pytest.mark.asyncio + +_STAGE = { + "index": 0, + "name": "Thinking", + "status": "completed", + "content": "...", + "attachments": [{"index": 0, "url": "http://a.com", "title": "Source"}], +} + +_FORM_SCHEMA = {"type": "object", "properties": {"city": {"type": "string"}}} +_FORM_VALUE = {"city": "Paris"} +_STATE = {"thread_id": "42", "step": 2} + +_USAGE_PER_MODEL = [ + { + "index": 0, + "model": "gpt-4o", + "prompt_tokens": 11, + "completion_tokens": 1, + "total_tokens": 12, + } +] +_DISCARDED_MESSAGES = [0, 1] + +_REQUEST_PARAMS: dict[str, Any] = { + "tools": [ + { + "type": "function", + "function": {"name": "f", "parameters": {}, "strict": True}, + "custom_fields": {"cache_breakpoint": {"expire_at": "1h"}}, + }, + { + "type": "static_function", + "static_function": {"name": "s", "configuration": {"a": 1}}, + }, + ], + "tool_choice": "required", + "parallel_tool_calls": False, + "max_completion_tokens": 100, + "max_prompt_tokens": 200, + "reasoning_effort": "minimal", + "response_format": { + "type": "json_schema", + "json_schema": {"name": "res", "schema": {"type": "object"}}, + }, + "stream_options": {"include_usage": True}, + "custom_fields": { + "configuration": {"a": 1}, + "cache_breakpoint": {"expire_at": "5m"}, + }, +} + +_MESSAGES: list[Any] = [ + { + "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": "http://a.com/b.png", "detail": "low"}, + }, + ], + "custom_content": { + "attachments": [{"type": "image/png", "url": "b.png"}] + }, + }, +] + + +class TestRequestFields: + """DIAL-specific request fields must reach the request body as-is""" + + @pytest.mark.parametrize("param", list(_REQUEST_PARAMS)) + async def test_request_param( + self, get_request_body: GetRequestBody, param: str + ): + body = await get_request_body(messages=_MESSAGES, **_REQUEST_PARAMS) + + assert body[param] == _REQUEST_PARAMS[param] + + async def test_request_messages(self, get_request_body: GetRequestBody): + body = await get_request_body(messages=_MESSAGES, **_REQUEST_PARAMS) + + assert body["messages"] == _MESSAGES + + +class TestCustomContent: + async def test_response_stages(self, get_message: GetMessage): + message = await get_message(custom_content={"stages": [_STAGE]}) + + assert message.custom_content and message.custom_content.stages + stage = message.custom_content.stages[0] + assert stage.model_dump(exclude_none=True) == _STAGE + + async def test_response_form_schema(self, get_message: GetMessage): + message = await get_message( + custom_content={"form_schema": _FORM_SCHEMA} + ) + + assert message.custom_content + assert message.custom_content.form_schema == _FORM_SCHEMA + + async def test_response_form_value(self, get_message: GetMessage): + message = await get_message(custom_content={"form_value": _FORM_VALUE}) + + assert message.custom_content + assert message.custom_content.form_value == _FORM_VALUE + + async def test_response_state(self, get_message: GetMessage): + message = await get_message(custom_content={"state": _STATE}) + + assert message.custom_content + assert message.custom_content.state == _STATE + + +class TestStatistics: + async def test_response_usage_per_model( + self, get_completion: GetCompletion + ): + completion = await get_completion( + statistics={"usage_per_model": _USAGE_PER_MODEL} + ) + + assert completion.statistics and completion.statistics.usage_per_model + assert [ + usage.model_dump(exclude_none=True) + for usage in completion.statistics.usage_per_model + ] == _USAGE_PER_MODEL + + async def test_response_discarded_messages( + self, get_completion: GetCompletion + ): + completion = await get_completion( + statistics={"discarded_messages": _DISCARDED_MESSAGES} + ) + + assert completion.statistics + assert completion.statistics.discarded_messages == _DISCARDED_MESSAGES diff --git a/tests/resources/completions/test_completions_streaming_tool_call.py b/tests/resources/completions/test_completions_streaming_tool_call.py index 6808448..f6f25d4 100644 --- a/tests/resources/completions/test_completions_streaming_tool_call.py +++ b/tests/resources/completions/test_completions_streaming_tool_call.py @@ -5,8 +5,7 @@ from aidial_client.types.chat import ChatCompletionChunk, ToolParam from tests.client_mock import get_async_client_mock, get_client_mock from tests.utils.chunks import create_mock_chunk, create_sse_data_field - -_DIAL_MODEL = "gpt-4o" +from tests.utils.completions import DIAL_MODEL _TOOL_DEFINITION: ToolParam = { "type": "function", @@ -174,7 +173,7 @@ def test_sync_streaming_tool_call(): ) response = client.chat.completions.create( - deployment_name=_DIAL_MODEL, + deployment_name=DIAL_MODEL, messages=[{"role": "user", "content": "what's the weather in Paris?"}], tools=[_TOOL_DEFINITION], stream=True, @@ -192,7 +191,7 @@ async def test_async_streaming_tool_call(): stream_chunks_mock=_STREAM_CHUNKS_MOCK, ) response = await async_client.chat.completions.create( - deployment_name=_DIAL_MODEL, + deployment_name=DIAL_MODEL, messages=[{"role": "user", "content": "what's the weather in Paris?"}], tools=[_TOOL_DEFINITION], stream=True, diff --git a/tests/resources/completions/test_completions_streaming_vanilla.py b/tests/resources/completions/test_completions_streaming_vanilla.py index 2a0d7ec..0717c8a 100644 --- a/tests/resources/completions/test_completions_streaming_vanilla.py +++ b/tests/resources/completions/test_completions_streaming_vanilla.py @@ -3,8 +3,7 @@ from aidial_client.types.chat import ChatCompletionChunk from tests.client_mock import get_async_client_mock, get_client_mock from tests.utils.chunks import create_mock_chunk, create_sse_data_field - -_DIAL_MODEL = "gpt-4o" +from tests.utils.completions import DIAL_MODEL _STREAM_CHUNKS_MOCK: list[bytes] = [ create_sse_data_field( @@ -52,7 +51,7 @@ def test_sync_streaming(): ) response = client.chat.completions.create( - deployment_name=_DIAL_MODEL, + deployment_name=DIAL_MODEL, messages=[{"role": "user", "content": "2+3="}], stream=True, ) @@ -68,7 +67,7 @@ async def test_async_streaming(): stream_chunks_mock=_STREAM_CHUNKS_MOCK, ) response = await async_client.chat.completions.create( - deployment_name=_DIAL_MODEL, + deployment_name=DIAL_MODEL, messages=[{"role": "user", "content": "2+3="}], stream=True, ) diff --git a/tests/resources/completions/test_completions_usage.py b/tests/resources/completions/test_completions_usage.py new file mode 100644 index 0000000..019fb39 --- /dev/null +++ b/tests/resources/completions/test_completions_usage.py @@ -0,0 +1,51 @@ +"""Chat completion usage, including the DIAL and reasoning token details""" + +import pytest + +from tests.resources.completions.conftest import GetCompletion + +pytestmark = pytest.mark.asyncio + +_TOKENS = {"prompt_tokens": 11, "completion_tokens": 1, "total_tokens": 12} +_PROMPT_TOKENS_DETAILS = {"cached_tokens": 8, "cache_write_tokens": 3} +_COMPLETION_TOKENS_DETAILS = {"reasoning_tokens": 1} + + +class TestUsage: + async def test_response_tokens(self, get_completion: GetCompletion): + completion = await get_completion(usage=_TOKENS) + + assert completion.usage + assert completion.usage.model_dump(exclude_none=True) == _TOKENS + + async def test_response_prompt_tokens_details( + self, get_completion: GetCompletion + ): + completion = await get_completion( + usage={**_TOKENS, "prompt_tokens_details": _PROMPT_TOKENS_DETAILS} + ) + + assert completion.usage and completion.usage.prompt_tokens_details + details = completion.usage.prompt_tokens_details + assert details.model_dump(exclude_none=True) == _PROMPT_TOKENS_DETAILS + + async def test_response_completion_tokens_details( + self, get_completion: GetCompletion + ): + completion = await get_completion( + usage={ + **_TOKENS, + "completion_tokens_details": _COMPLETION_TOKENS_DETAILS, + } + ) + + assert completion.usage and completion.usage.completion_tokens_details + details = completion.usage.completion_tokens_details + assert ( + details.model_dump(exclude_none=True) == _COMPLETION_TOKENS_DETAILS + ) + + async def test_response_without_usage(self, get_completion: GetCompletion): + completion = await get_completion() + + assert completion.usage is None diff --git a/tests/utils/completions.py b/tests/utils/completions.py new file mode 100644 index 0000000..c94a59d --- /dev/null +++ b/tests/utils/completions.py @@ -0,0 +1,139 @@ +"""Shared helpers for the `test_completions_*` tests""" + +from typing import Any + +from aidial_client import AsyncDial, Dial +from aidial_client._utils._dict import remove_none +from aidial_client.types.chat import ( + ChatCompletionChunk, + ChatCompletionMessage, + ChatCompletionMessageDelta, + ChatCompletionResponse, + Message, +) +from tests.client_mock import get_async_client_mock, get_client_mock +from tests.utils.chunks import create_sse_data_field + +DIAL_MODEL = "gpt-4o" +MESSAGES: list[Message] = [{"role": "user", "content": "2+3="}] + +Completion = ChatCompletionResponse | ChatCompletionChunk + + +def block_response( + *, + content: str | None = "5", + custom_content: dict | None = None, + usage: dict | None = None, + statistics: dict | None = None, +) -> dict[str, Any]: + """Non-streaming chat completion response body""" + return remove_none( + { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1723806872, + "model": DIAL_MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": remove_none( + { + "role": "assistant", + "content": content, + "custom_content": custom_content, + } + ), + } + ], + "usage": usage, + "statistics": statistics, + } + ) + + +def as_single_chunk(response: dict[str, Any]) -> list[bytes]: + """The very same response body delivered as a single streaming chunk""" + choices = [ + { + **{key: value for key, value in choice.items() if key != "message"}, + "delta": choice["message"], + } + for choice in response["choices"] + ] + return [ + create_sse_data_field( + {**response, "object": "chat.completion.chunk", "choices": choices} + ) + ] + + +def message_of( + completion: Completion, +) -> ChatCompletionMessage | ChatCompletionMessageDelta: + if isinstance(completion, ChatCompletionChunk): + return completion.choices[0].delta + return completion.choices[0].message + + +async def create_completion( + *, + is_async: bool, + stream: bool, + response: dict[str, Any] | None = None, + sent_requests: list | None = None, + **request_params: Any, +) -> Completion: + """ + Runs a chat completion against the mocked response, + either in sync or async and either in streaming or non-streaming mode. + """ + response = response or block_response() + request_params.setdefault("messages", MESSAGES) + response_mock: dict[str, Any] = ( + {"stream_chunks_mock": as_single_chunk(response)} + if stream + else {"json_mock": response} + ) + if is_async: + async_client = get_async_client_mock( + status_code=200, sent_requests=sent_requests, **response_mock + ) + return await _create_async(async_client, stream, request_params) + + client = get_client_mock( + status_code=200, sent_requests=sent_requests, **response_mock + ) + return _create_sync(client, stream, request_params) + + +def _create_sync(client: Dial, stream: bool, params: dict[str, Any]): + if not stream: + return client.chat.completions.create( + deployment_name=DIAL_MODEL, stream=False, **params + ) + + chunks = client.chat.completions.create( + deployment_name=DIAL_MODEL, stream=True, **params + ) + return _single(list(chunks)) + + +async def _create_async( + client: AsyncDial, stream: bool, params: dict[str, Any] +): + if not stream: + return await client.chat.completions.create( + deployment_name=DIAL_MODEL, stream=False, **params + ) + + chunks = await client.chat.completions.create( + deployment_name=DIAL_MODEL, stream=True, **params + ) + return _single([chunk async for chunk in chunks]) + + +def _single(chunks: list[ChatCompletionChunk]) -> ChatCompletionChunk: + assert len(chunks) == 1, f"expected a single chunk, got {len(chunks)}" + return chunks[0]