batchwork-ai 0.1.1__py3-none-any.whl

batchwork/_compatible_media.py

@@ -0,0 +1,147 @@
+"""Media conversion for chat-completion-compatible providers."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+from collections.abc import Mapping
+
+from ._typing import is_string_mapping
+from .errors import BatchworkError
+from .types import BatchProvider
+
+
+def _media_type(part: Mapping[str, object]) -> str:
+    value = part.get("media_type", part.get("mediaType"))
+    if isinstance(value, str) and value:
+        return value
+    raise BatchworkError("batchwork: file part requires a media type.")
+
+
+def _file_source(part: Mapping[str, object], provider: BatchProvider) -> tuple[str, str]:
+    data = part.get("data")
+    if isinstance(data, bytes):
+        return "data", base64.b64encode(data).decode()
+    if isinstance(data, str):
+        return ("url", data) if data.startswith(("http://", "https://")) else ("data", data)
+    if not is_string_mapping(data):
+        raise BatchworkError(f"batchwork: {provider.value} file part data is invalid.")
+
+    kind = data.get("type")
+    if kind == "url":
+        url = data.get("url")
+        if url is not None:
+            return "url", str(url)
+    elif kind == "data":
+        value = data.get("data")
+        if isinstance(value, bytes):
+            return "data", base64.b64encode(value).decode()
+        if isinstance(value, str):
+            return "data", value
+    elif kind in {"reference", "text", "provider-file-id"}:
+        raise BatchworkError(
+            f"batchwork: {provider.value} does not support file data type {kind!r}."
+        )
+    raise BatchworkError(f"batchwork: {provider.value} file part data is invalid.")
+
+
+def _data_url(media_type: str, data: str) -> str:
+    return f"data:{media_type};base64,{data}"
+
+
+def _groq_file_content(part: Mapping[str, object]) -> dict[str, object]:
+    media_type = _media_type(part)
+    if media_type.split("/", 1)[0] != "image":
+        raise BatchworkError("batchwork: Groq supports only image file parts.")
+    kind, source = _file_source(part, BatchProvider.GROQ)
+    return {
+        "type": "image_url",
+        "image_url": {"url": source if kind == "url" else _data_url(media_type, source)},
+    }
+
+
+def _mistral_file_content(part: Mapping[str, object]) -> dict[str, object]:
+    media_type = _media_type(part)
+    kind, source = _file_source(part, BatchProvider.MISTRAL)
+    value = source if kind == "url" else _data_url(media_type, source)
+    if media_type.split("/", 1)[0] == "image":
+        return {"type": "image_url", "image_url": value}
+    if media_type == "application/pdf":
+        return {"type": "document_url", "document_url": value}
+    raise BatchworkError("batchwork: Mistral supports only image and PDF file parts.")
+
+
+def _decode_text(data: str) -> str:
+    try:
+        return base64.b64decode(data, validate=True).decode()
+    except (binascii.Error, UnicodeDecodeError, ValueError) as error:
+        raise BatchworkError(
+            "batchwork: Together text file data must be valid base64-encoded UTF-8."
+        ) from error
+
+
+def _together_file_content(part: Mapping[str, object]) -> dict[str, object]:
+    media_type = _media_type(part)
+    top_level = media_type.split("/", 1)[0]
+    kind, source = _file_source(part, BatchProvider.TOGETHER)
+
+    if top_level == "image":
+        return {
+            "type": "image_url",
+            "image_url": {"url": source if kind == "url" else _data_url(media_type, source)},
+        }
+    if top_level == "audio":
+        if kind == "url":
+            raise BatchworkError("batchwork: Together does not support audio file URLs.")
+        audio_format = (
+            "wav"
+            if media_type == "audio/wav"
+            else "mp3"
+            if media_type in {"audio/mp3", "audio/mpeg"}
+            else None
+        )
+        if audio_format is None:
+            raise BatchworkError(
+                f"batchwork: Together does not support audio media type {media_type!r}."
+            )
+        return {
+            "type": "input_audio",
+            "input_audio": {"data": source, "format": audio_format},
+        }
+    if top_level == "application":
+        if kind == "url":
+            raise BatchworkError("batchwork: Together does not support PDF file URLs.")
+        if media_type != "application/pdf":
+            raise BatchworkError(
+                f"batchwork: Together does not support file media type {media_type!r}."
+            )
+        filename = part.get("filename")
+        return {
+            "type": "file",
+            "file": {
+                "filename": filename if isinstance(filename, str) else "document.pdf",
+                "file_data": _data_url("application/pdf", source),
+            },
+        }
+    if top_level == "text":
+        return {"type": "text", "text": source if kind == "url" else _decode_text(source)}
+    raise BatchworkError(f"batchwork: Together does not support file media type {media_type!r}.")
+
+
+def compatible_file_content(
+    part: Mapping[str, object], provider: BatchProvider
+) -> dict[str, object]:
+    """Convert a file part to the provider's native chat content shape."""
+
+    if provider is BatchProvider.GROQ:
+        return _groq_file_content(part)
+    if provider is BatchProvider.MISTRAL:
+        return _mistral_file_content(part)
+    if provider is BatchProvider.TOGETHER:
+        return _together_file_content(part)
+    raise BatchworkError(
+        f"batchwork: provider {provider.value!r} does not use compatible media conversion."
+    )
+
+
+__all__ = ["compatible_file_content"]

batchwork/_google_schema.py

@@ -0,0 +1,134 @@
+"""Google OpenAPI schema conversion for function declarations."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+from ._typing import is_string_mapping
+
+
+def _schema_sequence(value: object) -> list[object] | None:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)):
+        return None
+    return list(value)
+
+
+def _empty_object_schema(schema: Mapping[str, object]) -> bool:
+    properties = schema.get("properties")
+    no_properties = not isinstance(properties, Mapping) or not properties
+    additional_properties = schema.get("additionalProperties")
+    return (
+        schema.get("type") == "object" and no_properties and additional_properties in (None, False)
+    )
+
+
+def google_openapi_schema(schema: object, *, _root: bool = True) -> dict[str, object] | None:
+    """Convert JSON Schema Draft 7 to Google's OpenAPI 3.0 subset."""
+
+    if schema is None:
+        return None
+    if isinstance(schema, bool):
+        return {"type": "boolean", "properties": {}}
+    if not is_string_mapping(schema):
+        return None
+    if _empty_object_schema(schema):
+        if _root:
+            return None
+        nested: dict[str, object] = {"type": "object"}
+        description = schema.get("description")
+        if description:
+            nested["description"] = description
+        return nested
+
+    result: dict[str, object] = {}
+    for key in ("description", "required", "format"):
+        value = schema.get(key)
+        if value:
+            result[key] = value
+
+    if "const" in schema:
+        result["enum"] = [schema["const"]]
+
+    raw_type = schema.get("type")
+    types = _schema_sequence(raw_type)
+    if types is not None:
+        has_null = "null" in types
+        non_null_types = [value for value in types if value != "null"]
+        if not non_null_types:
+            result["type"] = "null"
+        else:
+            result["anyOf"] = [{"type": value} for value in non_null_types]
+            if has_null:
+                result["nullable"] = True
+    elif raw_type:
+        result["type"] = raw_type
+
+    if "enum" in schema:
+        result["enum"] = schema["enum"]
+
+    properties = schema.get("properties")
+    if isinstance(properties, Mapping):
+        converted_properties: dict[str, object] = {}
+        for key, value in properties.items():
+            if isinstance(key, str) and isinstance(value, (Mapping, bool)):
+                converted = google_openapi_schema(value, _root=False)
+                if converted is not None:
+                    converted_properties[key] = converted
+        result["properties"] = converted_properties
+
+    raw_items = schema.get("items")
+    items = _schema_sequence(raw_items)
+    if items is not None:
+        converted_items: list[object] = []
+        for item in items:
+            if isinstance(item, (Mapping, bool)):
+                converted = google_openapi_schema(item, _root=False)
+                if converted is not None:
+                    converted_items.append(converted)
+        result["items"] = converted_items
+    elif isinstance(raw_items, (Mapping, bool)):
+        converted = google_openapi_schema(raw_items, _root=False)
+        if converted is not None:
+            result["items"] = converted
+
+    for keyword in ("allOf", "oneOf"):
+        branches = _schema_sequence(schema.get(keyword))
+        if branches is None:
+            continue
+        converted_branches: list[object] = []
+        for branch in branches:
+            if isinstance(branch, (Mapping, bool)):
+                converted = google_openapi_schema(branch, _root=False)
+                if converted is not None:
+                    converted_branches.append(converted)
+        result[keyword] = converted_branches
+
+    any_of = _schema_sequence(schema.get("anyOf"))
+    if any_of is not None:
+        non_null = [
+            branch
+            for branch in any_of
+            if not (isinstance(branch, Mapping) and branch.get("type") == "null")
+        ]
+        nullable = len(non_null) != len(any_of)
+        converted_any_of = [
+            converted
+            for branch in non_null
+            if isinstance(branch, (Mapping, bool))
+            and (converted := google_openapi_schema(branch, _root=False)) is not None
+        ]
+        if nullable and len(converted_any_of) == 1:
+            result["nullable"] = True
+            result.update(converted_any_of[0])
+        else:
+            result["anyOf"] = converted_any_of
+            if nullable:
+                result["nullable"] = True
+
+    if "minLength" in schema:
+        result["minLength"] = schema["minLength"]
+
+    return result
+
+
+__all__ = ["google_openapi_schema"]

batchwork/_google_serialization.py

@@ -0,0 +1,239 @@
+"""Private Google message and tool-result serialization helpers."""
+
+from __future__ import annotations
+
+import json
+import mimetypes
+from collections.abc import Mapping, Sequence
+
+from ._serialization import _messages, _provider_options, _source, _wire
+from ._typing import is_string_mapping
+from .types import BatchProvider
+
+SKIP_THOUGHT_SIGNATURE_VALIDATOR = "skip_thought_signature_validator"
+
+
+def _thought_signature(part: Mapping[str, object]) -> str | None:
+    value = _provider_options(part, BatchProvider.GOOGLE).get("thoughtSignature")
+    return str(value) if value is not None else None
+
+
+def _with_signature(
+    converted: dict[str, object], part: Mapping[str, object], *, fallback: str | None = None
+) -> dict[str, object]:
+    signature = _thought_signature(part)
+    if signature is None:
+        signature = fallback
+    return {**converted, **({"thoughtSignature": signature} if signature is not None else {})}
+
+
+def _tool_file(part: Mapping[str, object]) -> tuple[str, bool, str] | None:
+    kind = part.get("type")
+    media_type = part.get("media_type", part.get("mediaType"))
+    raw_data = part.get("data")
+    if kind == "file-data":
+        raw_data = {"type": "data", "data": raw_data}
+    elif kind == "file-url":
+        raw_data = {"type": "url", "url": part.get("url")}
+        url = part.get("url")
+        media_type = media_type or (mimetypes.guess_type(url)[0] if isinstance(url, str) else None)
+    elif kind == "image-data":
+        raw_data = {"type": "data", "data": raw_data}
+    elif kind == "image-url":
+        raw_data = {"type": "url", "url": part.get("url")}
+        media_type = "image"
+    elif kind != "file":
+        return None
+    if not isinstance(media_type, str):
+        media_type = "application/octet-stream"
+    source, inline = _source(raw_data, BatchProvider.GOOGLE)
+    return (source, inline, media_type) if source is not None else None
+
+
+def _tool_content(
+    name: object, call_id: object, value: object, *, supports_parts: bool
+) -> list[object]:
+    if (
+        not isinstance(name, str)
+        or not isinstance(value, Sequence)
+        or isinstance(value, (str, bytes, bytearray))
+    ):
+        return []
+    if not supports_parts:
+        result: list[object] = []
+        for item in value:
+            if is_string_mapping(item) and item.get("type") == "text":
+                result.append(
+                    {
+                        "functionResponse": {
+                            "id": call_id,
+                            "name": name,
+                            "response": {"name": name, "content": item.get("text", "")},
+                        }
+                    }
+                )
+                continue
+            file = _tool_file(item) if is_string_mapping(item) else None
+            if file is not None and file[1] and file[2].split("/", 1)[0] == "image":
+                result.extend(
+                    [
+                        {"inlineData": {"mimeType": file[2], "data": file[0]}},
+                        {
+                            "text": (
+                                "Tool executed successfully and returned this image as a response"
+                            )
+                        },
+                    ]
+                )
+            elif item is not None:
+                result.append({"text": json.dumps(item, separators=(",", ":"))})
+        return result
+    text: list[str] = []
+    response_parts: list[object] = []
+    for item in value:
+        if is_string_mapping(item) and item.get("type") == "text":
+            text.append(str(item.get("text", "")))
+            continue
+        file = _tool_file(item) if is_string_mapping(item) else None
+        if file is not None and file[1]:
+            response_parts.append({"inlineData": {"mimeType": file[2], "data": file[0]}})
+        elif item is not None:
+            text.append(json.dumps(item, separators=(",", ":")))
+    function_response: dict[str, object] = {
+        "id": call_id,
+        "name": name,
+        "response": {
+            "name": name,
+            "content": "\n".join(text) if text else "Tool executed successfully.",
+        },
+    }
+    if response_parts:
+        function_response["parts"] = response_parts
+    return [{"functionResponse": function_response}]
+
+
+def _part(
+    part: object,
+    *,
+    assistant: bool,
+    gemini_3: bool,
+    supports_response_parts: bool,
+) -> object | list[object]:
+    if isinstance(part, str):
+        return {"text": part}
+    if not is_string_mapping(part):
+        return part
+    kind = part.get("type")
+    if kind == "text":
+        converted = {"text": part.get("text", "")}
+        return _with_signature(converted, part) if assistant else converted
+    if kind == "reasoning":
+        return _with_signature({"text": part.get("text", ""), "thought": True}, part)
+    if kind in {"image", "file", "reasoning-file"}:
+        raw_source = part.get("image") if kind == "image" else part.get("data")
+        source, inline = _source(raw_source, BatchProvider.GOOGLE)
+        media_type = part.get("media_type", part.get("mediaType", "application/octet-stream"))
+        if source is not None and inline:
+            converted = {"inlineData": {"mimeType": media_type, "data": source}}
+            if kind == "reasoning-file":
+                converted["thought"] = True
+            return _with_signature(converted, part) if assistant else converted
+        if source is not None:
+            converted = {"fileData": {"mimeType": media_type, "fileUri": source}}
+            if kind == "reasoning-file":
+                converted["thought"] = True
+            return _with_signature(converted, part) if assistant else converted
+    if kind == "tool-call":
+        converted = {
+            "functionCall": {
+                "id": part.get("tool_call_id", part.get("toolCallId")),
+                "name": part.get("tool_name", part.get("toolName")),
+                "args": part.get("input", {}),
+            }
+        }
+        fallback = SKIP_THOUGHT_SIGNATURE_VALIDATOR if gemini_3 else None
+        return _with_signature(converted, part, fallback=fallback)
+    if kind == "tool-result":
+        name = part.get("tool_name", part.get("toolName"))
+        output = part.get("output", part.get("content", {}))
+        if is_string_mapping(output) and output.get("type") == "content":
+            return _tool_content(
+                name,
+                part.get("tool_call_id", part.get("toolCallId")),
+                output.get("value"),
+                supports_parts=supports_response_parts,
+            )
+        if is_string_mapping(output):
+            content = (
+                output.get("reason") or "Tool call execution denied."
+                if output.get("type") == "execution-denied"
+                else output.get("value")
+            )
+        else:
+            content = output
+        return {
+            "functionResponse": {
+                "id": part.get("tool_call_id", part.get("toolCallId")),
+                "name": name,
+                "response": {"name": name, "content": content},
+            }
+        }
+    native = {key: value for key, value in part.items() if key != "provider_options"}
+    return _wire(part, BatchProvider.GOOGLE, native)
+
+
+def google_messages(
+    item: Mapping[str, object], model_id: str
+) -> tuple[list[dict[str, object]], dict[str, object] | None]:
+    contents: list[dict[str, object]] = []
+    system_parts: list[dict[str, str]] = []
+    system = item.get("system")
+    if isinstance(system, str):
+        system_parts.append({"text": system})
+    normalized_model_id = model_id.lower()
+    gemini_3 = "gemini-3" in normalized_model_id
+    is_gemma = normalized_model_id.startswith("gemma-")
+    supports_response_parts = gemini_3
+    for message in _messages(item):
+        role = message.get("role")
+        if role == "system":
+            content = message.get("content")
+            if isinstance(content, str):
+                system_parts.append({"text": content})
+            continue
+        content = message.get("content")
+        source_parts = (
+            content
+            if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray))
+            else [content]
+        )
+        parts: list[object] = []
+        for part in source_parts:
+            converted = _part(
+                part,
+                assistant=role == "assistant",
+                gemini_3=gemini_3,
+                supports_response_parts=supports_response_parts,
+            )
+            parts.extend(converted if isinstance(converted, list) else [converted])
+        contents.append(
+            _wire(
+                message,
+                BatchProvider.GOOGLE,
+                {"role": "model" if role == "assistant" else "user", "parts": parts},
+            )
+        )
+    if is_gemma:
+        if system_parts and contents and contents[0].get("role") == "user":
+            system_text = "\n\n".join(part["text"] for part in system_parts)
+            raw_parts = contents[0].get("parts")
+            if isinstance(raw_parts, list):
+                prefixed_parts: list[object] = [{"text": f"{system_text}\n\n"}]
+                prefixed_parts.extend(raw_parts)
+                contents[0]["parts"] = prefixed_parts
+        return contents, None
+    system_instruction: dict[str, object] | None = {"parts": system_parts} if system_parts else None
+    return contents, system_instruction
+
+
+__all__ = ["SKIP_THOUGHT_SIGNATURE_VALIDATOR", "google_messages"]

batchwork/_network.py

@@ -0,0 +1,79 @@
+"""Shared private network-address policy."""
+
+from __future__ import annotations
+
+import asyncio
+import ipaddress
+import socket
+from dataclasses import dataclass
+from enum import Enum
+
+
+class AddressResolutionFailureReason(Enum):
+    LOOKUP = "lookup"
+    EMPTY = "empty"
+    INVALID = "invalid"
+    NON_GLOBAL = "non_global"
+
+
+@dataclass(frozen=True, slots=True)
+class AddressResolutionFailure:
+    reason: AddressResolutionFailureReason
+    address: object | None = None
+    cause: Exception | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedAddresses:
+    addresses: tuple[str, ...]
+
+
+AddressResolution = ResolvedAddresses | AddressResolutionFailure
+
+
+def validate_public_address(address: object) -> AddressResolutionFailure | None:
+    if not isinstance(address, str):
+        return AddressResolutionFailure(AddressResolutionFailureReason.INVALID, address)
+    try:
+        parsed = ipaddress.ip_address(address)
+    except ValueError as error:
+        return AddressResolutionFailure(AddressResolutionFailureReason.INVALID, address, error)
+    if not parsed.is_global:
+        return AddressResolutionFailure(AddressResolutionFailureReason.NON_GLOBAL, address)
+    return None
+
+
+async def resolve_public_addresses(host: str, port: int) -> AddressResolution:
+    """Resolve one target and reject its complete address set on any unsafe record."""
+
+    try:
+        literal = ipaddress.ip_address(host)
+    except ValueError:
+        try:
+            records = await asyncio.get_running_loop().getaddrinfo(
+                host,
+                port,
+                family=socket.AF_UNSPEC,
+                type=socket.SOCK_STREAM,
+            )
+        except OSError as error:
+            return AddressResolutionFailure(AddressResolutionFailureReason.LOOKUP, cause=error)
+        addresses: list[str] = []
+        for record in records:
+            address = record[4][0]
+            if not isinstance(address, str):
+                return AddressResolutionFailure(AddressResolutionFailureReason.INVALID, address)
+            failure = validate_public_address(address)
+            if failure is not None:
+                return failure
+            if address not in addresses:
+                addresses.append(address)
+    else:
+        address = str(literal)
+        failure = validate_public_address(address)
+        if failure is not None:
+            return failure
+        addresses = [address]
+    if not addresses:
+        return AddressResolutionFailure(AddressResolutionFailureReason.EMPTY)
+    return ResolvedAddresses(tuple(addresses))