arkruntime 0.1.0__py3-none-any.whl

arkruntime/common/_parsing/_completions.py

@@ -0,0 +1,272 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Iterable, cast
+
+import pydantic
+from typing_extensions import TypeGuard, TypeVar, assert_never
+
+from ..._compat import PYDANTIC_V2, model_parse_json
+from ..._exceptions import ArkContentFilterFinishReasonError, ArkLengthFinishReasonError
+from ..._models import construct_type_unchecked
+from ..._types import NOT_GIVEN, NotGiven
+from ..._utils import is_dict, is_given
+from ...types.chat import (
+    ChatCompletion,
+    ChatCompletionMessage,
+    ChatCompletionToolParam,
+    ParsedChatCompletion,
+    ParsedChatCompletionMessage,
+    ParsedChoice,
+    ParsedFunction,
+    ParsedFunctionToolCall,
+)
+from ...types.chat.chat_completion_message_tool_call_function import (
+    ChatCompletionMessageToolCallFunction as Function,
+)
+from ...types.chat.chat_completion_response_format_param import (
+    ChatCompletionResponseFormatParam as ResponseFormatParam,
+)
+from ...types.chat.function_object_param import FunctionObjectParam as FunctionDefinition
+from .._pydantic import is_basemodel_type, is_dataclass_like_type, to_strict_json_schema
+from .._tools import PydanticFunctionTool
+
+ResponseFormatT = TypeVar(
+    "ResponseFormatT",
+    # if it isn't given then we don't do any parsing
+    default=None,
+)
+_default_response_format: None = None
+
+
+def validate_input_tools(
+    tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+) -> None:
+    if not is_given(tools):
+        return
+
+    for tool in tools:
+        if tool["type"] != "function":
+            raise ValueError(
+                f"Currently only `function` tool types support auto-parsing; Received `{tool['type']}`",
+            )
+
+        strict = tool["function"].get("strict")
+        if strict is not True:
+            raise ValueError(
+                f"`{tool['function']['name']}` is not strict. Only `strict` function tools can be auto-parsed"
+            )
+
+
+def parse_chat_completion(
+    *,
+    response_format: type[ResponseFormatT] | ResponseFormatParam | NotGiven,
+    input_tools: Iterable[ChatCompletionToolParam] | NotGiven,
+    chat_completion: ChatCompletion | ParsedChatCompletion[object],
+) -> ParsedChatCompletion[ResponseFormatT]:
+    if is_given(input_tools):
+        input_tools = [t for t in input_tools]
+    else:
+        input_tools = []
+
+    choices: list[ParsedChoice[ResponseFormatT]] = []
+    for choice in chat_completion.choices:
+        if choice.finish_reason == "length":
+            raise ArkLengthFinishReasonError(completion=chat_completion)
+
+        if choice.finish_reason == "content_filter":
+            raise ArkContentFilterFinishReasonError()
+
+        message = choice.message
+
+        tool_calls: list[ParsedFunctionToolCall] = []
+        if message.tool_calls:
+            for tool_call in message.tool_calls:
+                if tool_call.type == "function":
+                    tool_call_dict = tool_call.to_dict()
+                    tool_calls.append(
+                        construct_type_unchecked(
+                            value={
+                                **tool_call_dict,
+                                "function": {
+                                    **cast(Any, tool_call_dict["function"]),
+                                    "parsed_arguments": parse_function_tool_arguments(
+                                        input_tools=input_tools,
+                                        function=tool_call.function,
+                                    ),
+                                },
+                            },
+                            type_=ParsedFunctionToolCall,
+                        )
+                    )
+                elif TYPE_CHECKING:  # type: ignore[unreachable]
+                    assert_never(tool_call)
+                else:
+                    tool_calls.append(tool_call)
+
+        choices.append(
+            construct_type_unchecked(
+                type_=cast(Any, ParsedChoice)[solve_response_format_t(response_format)],
+                value={
+                    **choice.to_dict(),
+                    "message": {
+                        **message.to_dict(),
+                        "parsed": maybe_parse_content(
+                            response_format=response_format,
+                            message=message,
+                        ),
+                        "tool_calls": tool_calls if tool_calls else None,
+                    },
+                },
+            )
+        )
+
+    return cast(
+        ParsedChatCompletion[ResponseFormatT],
+        construct_type_unchecked(
+            type_=cast(Any, ParsedChatCompletion)[solve_response_format_t(response_format)],
+            value={
+                **chat_completion.to_dict(),
+                "choices": choices,
+            },
+        ),
+    )
+
+
+def get_input_tool_by_name(*, input_tools: list[ChatCompletionToolParam], name: str) -> ChatCompletionToolParam | None:
+    return next((t for t in input_tools if t.get("function", {}).get("name") == name), None)
+
+
+def parse_function_tool_arguments(
+    *, input_tools: list[ChatCompletionToolParam], function: Function | ParsedFunction
+) -> object:
+    input_tool = get_input_tool_by_name(input_tools=input_tools, name=function.name)
+    if not input_tool:
+        return None
+
+    input_fn = cast(object, input_tool.get("function"))
+    if isinstance(input_fn, PydanticFunctionTool):
+        return model_parse_json(input_fn.model, function.arguments)
+
+    input_fn = cast(FunctionDefinition, input_fn)
+
+    if not input_fn.get("strict"):
+        return None
+
+    return json.loads(function.arguments)
+
+
+def maybe_parse_content(
+    *,
+    response_format: type[ResponseFormatT] | ResponseFormatParam | NotGiven,
+    message: ChatCompletionMessage | ParsedChatCompletionMessage[object],
+) -> ResponseFormatT | None:
+    if has_rich_response_format(response_format) and message.content:
+        return _parse_content(response_format, message.content)
+
+    return None
+
+
+def solve_response_format_t(
+    response_format: type[ResponseFormatT] | ResponseFormatParam | NotGiven,
+) -> type[ResponseFormatT]:
+    """Return the runtime type for the given response format.
+
+    If no response format is given, or if we won't auto-parse the response format
+    then we default to `None`.
+    """
+    if has_rich_response_format(response_format):
+        return response_format
+
+    return cast("type[ResponseFormatT]", _default_response_format)
+
+
+def has_parseable_input(
+    *,
+    response_format: type | ResponseFormatParam | NotGiven,
+    input_tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+) -> bool:
+    if has_rich_response_format(response_format):
+        return True
+
+    for input_tool in input_tools or []:
+        if is_parseable_tool(input_tool):
+            return True
+
+    return False
+
+
+def has_rich_response_format(
+    response_format: type[ResponseFormatT] | ResponseFormatParam | NotGiven,
+) -> TypeGuard[type[ResponseFormatT]]:
+    if not is_given(response_format):
+        return False
+
+    if is_response_format_param(response_format):
+        return False
+
+    return True
+
+
+def is_response_format_param(response_format: object) -> TypeGuard[ResponseFormatParam]:
+    return is_dict(response_format)
+
+
+def is_parseable_tool(input_tool: ChatCompletionToolParam) -> bool:
+    input_fn = cast(object, input_tool.get("function"))
+    if isinstance(input_fn, PydanticFunctionTool):
+        return True
+
+    return cast(FunctionDefinition, input_fn).get("strict") or False
+
+
+def _parse_content(response_format: type[ResponseFormatT], content: str) -> ResponseFormatT:
+    if is_basemodel_type(response_format):
+        return cast(ResponseFormatT, model_parse_json(response_format, content))
+
+    if is_dataclass_like_type(response_format):
+        if not PYDANTIC_V2:
+            raise TypeError(f"Non BaseModel types are only supported with Pydantic v2 - {response_format}")
+
+        return pydantic.TypeAdapter(response_format).validate_json(content)
+
+    raise TypeError(f"Unable to automatically parse response format type {response_format}")
+
+
+def type_to_response_format_param(
+    response_format: type | ResponseFormatParam | NotGiven,
+) -> ResponseFormatParam | NotGiven:
+    if not is_given(response_format):
+        return NOT_GIVEN
+
+    if is_response_format_param(response_format):
+        return response_format
+
+    # type checkers don't narrow the negation of a `TypeGuard` as it isn't
+    # a safe default behaviour but we know that at this point the `response_format`
+    # can only be a `type`
+    response_format = cast(type, response_format)
+
+    json_schema_type: type[pydantic.BaseModel] | pydantic.TypeAdapter[Any] | None = None
+
+    if is_basemodel_type(response_format):
+        name = response_format.__name__
+        json_schema_type = response_format
+    elif is_dataclass_like_type(response_format):
+        name = response_format.__name__
+        json_schema_type = pydantic.TypeAdapter(response_format)
+    else:
+        raise TypeError(f"Unsupported response_format type - {response_format}")
+
+    return {
+        "type": "json_schema",
+        "json_schema": {
+            "schema": to_strict_json_schema(json_schema_type),
+            "name": name,
+            "strict": True,
+        },
+    }

arkruntime/common/_parsing/_responses.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Iterable, List, cast
+
+import pydantic
+from typing_extensions import TypeVar, assert_never
+
+from ..._compat import PYDANTIC_V2, model_parse_json
+from ..._models import construct_type_unchecked
+from ..._types import NotGiven
+from ..._utils import is_given
+from ...types.chat.chat_completion_response_format_param import (
+    ChatCompletionResponseFormatParam as ResponseFormat,
+)
+from ...types.responses import (
+    FunctionToolParam,
+    ParsedContent,
+    ParsedResponse,
+    ParsedResponseFunctionToolCall,
+    ParsedResponseOutputItem,
+    ParsedResponseOutputMessage,
+    ParsedResponseOutputText,
+    Response,
+    ResponseFormatTextConfigParam,
+    ResponseFunctionToolCall,
+    ToolParam,
+)
+from .._pydantic import is_basemodel_type, is_dataclass_like_type
+from .._tools import ResponsesPydanticFunctionTool
+from ._completions import solve_response_format_t, type_to_response_format_param
+
+TextFormatT = TypeVar(
+    "TextFormatT",
+    # if it isn't given then we don't do any parsing
+    default=None,
+)
+
+
+def type_to_text_format_param(type_: type) -> ResponseFormatTextConfigParam:
+    response_format_dict = type_to_response_format_param(type_)
+    assert is_given(response_format_dict)
+    response_format_dict = cast(ResponseFormat, response_format_dict)  # pyright: ignore[reportUnnecessaryCast]
+    assert response_format_dict["type"] == "json_schema"
+    assert "schema" in response_format_dict["json_schema"]
+
+    return {
+        "type": "json_schema",
+        "strict": True,
+        "name": response_format_dict["json_schema"]["name"],
+        "schema": response_format_dict["json_schema"]["schema"],
+    }
+
+
+def parse_response(
+    *,
+    text_format: type[TextFormatT] | NotGiven,
+    input_tools: Iterable[ToolParam] | NotGiven | None,
+    response: Response | ParsedResponse[object],
+) -> ParsedResponse[TextFormatT]:
+    solved_t = solve_response_format_t(text_format)
+    output_list: List[ParsedResponseOutputItem[TextFormatT]] = []
+
+    for output in response.output:
+        if output.type == "message":
+            content_list: List[ParsedContent[TextFormatT]] = []
+            for item in output.content:
+                if item.type != "output_text":
+                    content_list.append(item)
+                    continue
+
+                content_list.append(
+                    construct_type_unchecked(
+                        type_=cast(Any, ParsedResponseOutputText)[solved_t],
+                        value={
+                            **item.to_dict(),
+                            "parsed": parse_text(item.text, text_format=text_format),
+                        },
+                    )
+                )
+
+            output_list.append(
+                construct_type_unchecked(
+                    type_=cast(Any, ParsedResponseOutputMessage)[solved_t],
+                    value={
+                        **output.to_dict(),
+                        "content": content_list,
+                    },
+                )
+            )
+        elif output.type == "function_call":
+            output_list.append(
+                construct_type_unchecked(
+                    type_=ParsedResponseFunctionToolCall,
+                    value={
+                        **output.to_dict(),
+                        "parsed_arguments": parse_function_tool_arguments(
+                            input_tools=input_tools, function_call=output
+                        ),
+                    },
+                )
+            )
+        elif (
+            output.type == "computer_call"
+            or output.type == "file_search_call"
+            or output.type == "web_search_call"
+            or output.type == "reasoning"
+        ):
+            output_list.append(output)
+        elif TYPE_CHECKING:  # type: ignore
+            assert_never(output)
+        else:
+            output_list.append(output)
+
+    return cast(
+        ParsedResponse[TextFormatT],
+        construct_type_unchecked(
+            type_=cast(Any, ParsedResponse)[solved_t],
+            value={
+                **response.to_dict(),
+                "output": output_list,
+            },
+        ),
+    )
+
+
+def parse_text(text: str, text_format: type[TextFormatT] | NotGiven) -> TextFormatT | None:
+    if not is_given(text_format):
+        return None
+
+    if is_basemodel_type(text_format):
+        return cast(TextFormatT, model_parse_json(text_format, text))
+
+    if is_dataclass_like_type(text_format):
+        if not PYDANTIC_V2:
+            raise TypeError(f"Non BaseModel types are only supported with Pydantic v2 - {text_format}")
+
+        return pydantic.TypeAdapter(text_format).validate_json(text)
+
+    raise TypeError(f"Unable to automatically parse response format type {text_format}")
+
+
+def get_input_tool_by_name(*, input_tools: Iterable[ToolParam], name: str) -> FunctionToolParam | None:
+    for tool in input_tools:
+        if tool["type"] == "function" and tool.get("name") == name:
+            return tool
+
+    return None
+
+
+def parse_function_tool_arguments(
+    *,
+    input_tools: Iterable[ToolParam] | NotGiven | None,
+    function_call: ParsedResponseFunctionToolCall | ResponseFunctionToolCall,
+) -> object:
+    if input_tools is None or not is_given(input_tools):
+        return None
+
+    input_tool = get_input_tool_by_name(input_tools=input_tools, name=function_call.name)
+    if not input_tool:
+        return None
+
+    tool = cast(object, input_tool)
+    if isinstance(tool, ResponsesPydanticFunctionTool):
+        return model_parse_json(tool.model, function_call.arguments)
+
+    if not input_tool.get("strict"):
+        return None
+
+    return json.loads(function_call.arguments)

arkruntime/common/_pydantic.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+import inspect
+from typing import Any, TypeVar
+
+import pydantic
+from typing_extensions import TypeGuard
+
+from .._compat import PYDANTIC_V2, model_json_schema
+from .._types import NOT_GIVEN
+from .._utils import is_dict as _is_dict
+from .._utils import is_list
+
+_T = TypeVar("_T")
+
+
+def to_strict_json_schema(
+    model: type[pydantic.BaseModel] | pydantic.TypeAdapter[Any],
+) -> dict[str, Any]:
+    if inspect.isclass(model) and is_basemodel_type(model):
+        schema = model_json_schema(model)
+    elif PYDANTIC_V2 and isinstance(model, pydantic.TypeAdapter):
+        schema = model.json_schema()
+    else:
+        raise TypeError(f"Non BaseModel types are only supported with Pydantic v2 - {model}")
+
+    return _ensure_strict_json_schema(schema, path=(), root=schema)
+
+
+def _ensure_strict_json_schema(
+    json_schema: object,
+    *,
+    path: tuple[str, ...],
+    root: dict[str, object],
+) -> dict[str, Any]:
+    """Mutates the given JSON schema to ensure it conforms to the `strict` standard
+    that the API expects.
+    """
+    if not is_dict(json_schema):
+        raise TypeError(f"Expected {json_schema} to be a dictionary; path={path}")
+
+    defs = json_schema.get("$defs")
+    if is_dict(defs):
+        for def_name, def_schema in defs.items():
+            _ensure_strict_json_schema(def_schema, path=(*path, "$defs", def_name), root=root)
+
+    definitions = json_schema.get("definitions")
+    if is_dict(definitions):
+        for definition_name, definition_schema in definitions.items():
+            _ensure_strict_json_schema(
+                definition_schema,
+                path=(*path, "definitions", definition_name),
+                root=root,
+            )
+
+    typ = json_schema.get("type")
+    if typ == "object" and "additionalProperties" not in json_schema:
+        json_schema["additionalProperties"] = False
+
+    # object types
+    # { 'type': 'object', 'properties': { 'a':  {...} } }
+    properties = json_schema.get("properties")
+    if is_dict(properties):
+        json_schema["required"] = [prop for prop in properties.keys()]
+        json_schema["properties"] = {
+            key: _ensure_strict_json_schema(prop_schema, path=(*path, "properties", key), root=root)
+            for key, prop_schema in properties.items()
+        }
+
+    # arrays
+    # { 'type': 'array', 'items': {...} }
+    items = json_schema.get("items")
+    if is_dict(items):
+        json_schema["items"] = _ensure_strict_json_schema(items, path=(*path, "items"), root=root)
+
+    # unions
+    any_of = json_schema.get("anyOf")
+    if is_list(any_of):
+        json_schema["anyOf"] = [
+            _ensure_strict_json_schema(variant, path=(*path, "anyOf", str(i)), root=root)
+            for i, variant in enumerate(any_of)
+        ]
+
+    # intersections
+    all_of = json_schema.get("allOf")
+    if is_list(all_of):
+        if len(all_of) == 1:
+            json_schema.update(_ensure_strict_json_schema(all_of[0], path=(*path, "allOf", "0"), root=root))
+            json_schema.pop("allOf")
+        else:
+            json_schema["allOf"] = [
+                _ensure_strict_json_schema(entry, path=(*path, "allOf", str(i)), root=root)
+                for i, entry in enumerate(all_of)
+            ]
+
+    # strip `None` defaults as there's no meaningful distinction here
+    # the schema will still be `nullable` and the model will default
+    # to using `None` anyway
+    if json_schema.get("default", NOT_GIVEN) is None:
+        json_schema.pop("default")
+
+    # we can't use `$ref`s if there are also other properties defined, e.g.
+    # `{"$ref": "...", "description": "my description"}`
+    #
+    # so we unravel the ref
+    # `{"type": "string", "description": "my description"}`
+    ref = json_schema.get("$ref")
+    if ref and has_more_than_n_keys(json_schema, 1):
+        assert isinstance(ref, str), f"Received non-string $ref - {ref}"
+
+        resolved = resolve_ref(root=root, ref=ref)
+        if not is_dict(resolved):
+            raise ValueError(f"Expected `$ref: {ref}` to resolved to a dictionary but got {resolved}")
+
+        # properties from the json schema take priority over the ones on the `$ref`
+        json_schema.update({**resolved, **json_schema})
+        json_schema.pop("$ref")
+        # Since the schema expanded from `$ref` might not have `additionalProperties: false` applied,
+        # we call `_ensure_strict_json_schema` again to fix the inlined schema and ensure it's valid.
+        return _ensure_strict_json_schema(json_schema, path=path, root=root)
+
+    return json_schema
+
+
+def resolve_ref(*, root: dict[str, object], ref: str) -> object:
+    if not ref.startswith("#/"):
+        raise ValueError(f"Unexpected $ref format {ref!r}; Does not start with #/")
+
+    path = ref[2:].split("/")
+    resolved = root
+    for key in path:
+        value = resolved[key]
+        assert is_dict(value), f"encountered non-dictionary entry while resolving {ref} - {resolved}"
+        resolved = value
+
+    return resolved
+
+
+def is_basemodel_type(typ: type) -> TypeGuard[type[pydantic.BaseModel]]:
+    if not inspect.isclass(typ):
+        return False
+    return issubclass(typ, pydantic.BaseModel)
+
+
+def is_dataclass_like_type(typ: type) -> bool:
+    """Returns True if the given type likely used `@pydantic.dataclass`"""
+    return hasattr(typ, "__pydantic_config__")
+
+
+def is_dict(obj: object) -> TypeGuard[dict[str, object]]:
+    # just pretend that we know there are only `str` keys
+    # as that check is not worth the performance cost
+    return _is_dict(obj)
+
+
+def has_more_than_n_keys(obj: dict[str, object], n: int) -> bool:
+    i = 0
+    for _ in obj.keys():
+        i += 1
+        if i > n:
+            return True
+    return False

arkruntime/common/_tools.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+from typing import Any, Dict, cast
+
+import pydantic
+
+from ..types.chat import ChatCompletionToolParam
+from ..types.chat.function_object_param import FunctionObjectParam as FunctionDefinition
+from ..types.responses.function_tool_param import (
+    FunctionToolParam as ResponsesFunctionToolParam,
+)
+from ._pydantic import to_strict_json_schema
+
+
+class PydanticFunctionTool(Dict[str, Any]):
+    """Dictionary wrapper so we can pass the given base model
+    throughout the entire request stack without having to special
+    case it.
+    """
+
+    model: type[pydantic.BaseModel]
+
+    def __init__(self, defn: FunctionDefinition, model: type[pydantic.BaseModel]) -> None:
+        super().__init__(defn)
+        self.model = model
+
+    def cast(self) -> FunctionDefinition:
+        return cast(FunctionDefinition, self)
+
+
+class ResponsesPydanticFunctionTool(Dict[str, Any]):
+    """Responses-API counterpart to :class:`PydanticFunctionTool`.
+
+    The Responses API encodes function tools with a flatter shape than
+    Chat Completions — name, parameters and strict live directly on the
+    tool dict instead of nested under ``function``. We wrap that dict and
+    keep a reference to the pydantic model so the response parser can
+    deserialise ``arguments`` back into the original type.
+    """
+
+    model: type[pydantic.BaseModel]
+
+    def __init__(
+        self,
+        tool: ResponsesFunctionToolParam,
+        model: type[pydantic.BaseModel],
+    ) -> None:
+        super().__init__(tool)
+        self.model = model
+
+    def cast(self) -> ResponsesFunctionToolParam:
+        return cast(ResponsesFunctionToolParam, self)
+
+
+def pydantic_function_tool(
+    model: type[pydantic.BaseModel],
+    *,
+    name: str | None = None,  # inferred from class name by default
+    description: str | None = None,  # inferred from class docstring by default
+) -> ChatCompletionToolParam:
+    if description is None:
+        # note: we intentionally don't use `.getdoc()` to avoid
+        # including pydantic's docstrings
+        description = model.__doc__
+
+    function = PydanticFunctionTool(
+        {
+            "name": name or model.__name__,
+            "strict": True,
+            "parameters": to_strict_json_schema(model),
+        },
+        model,
+    ).cast()
+
+    if description is not None:
+        function["description"] = description
+
+    return {
+        "type": "function",
+        "function": function,
+    }

arkruntime/common/streaming/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+