wehub-resource-sync
360 lines
10 KiB
Python
360 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Pydantic models for Anthropic Messages API.
|
|
|
|
These models define the request and response schemas for:
|
|
- Anthropic Messages API (/v1/messages)
|
|
- Streaming events
|
|
- Tool calling in Anthropic format
|
|
"""
|
|
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
|
|
from omlx.api.shared_models import IDPrefix, generate_id
|
|
|
|
# =============================================================================
|
|
# Content Blocks
|
|
# =============================================================================
|
|
|
|
|
|
class ContentBlockText(BaseModel):
|
|
"""Text content block."""
|
|
|
|
type: Literal["text"] = "text"
|
|
text: str
|
|
cache_control: dict[str, str] | None = None
|
|
|
|
|
|
class ContentBlockImage(BaseModel):
|
|
"""Image content block with source data."""
|
|
|
|
type: Literal["image"] = "image"
|
|
source: dict[str, Any] # {"type": "base64"|"url", "media_type": "...", "data"|"url": "..."}
|
|
|
|
|
|
class ContentBlockToolUse(BaseModel):
|
|
"""Tool use content block (model requesting a tool call)."""
|
|
|
|
type: Literal["tool_use"] = "tool_use"
|
|
id: str
|
|
name: str
|
|
input: dict[str, Any]
|
|
|
|
|
|
class ContentBlockToolResult(BaseModel):
|
|
"""Tool result content block (user providing tool output)."""
|
|
|
|
type: Literal["tool_result"] = "tool_result"
|
|
tool_use_id: str
|
|
content: str | list[dict[str, Any]] | dict[str, Any] | list[Any] | Any
|
|
is_error: bool | None = None
|
|
cache_control: dict[str, str] | None = None
|
|
|
|
|
|
class ContentBlockThinking(BaseModel):
|
|
"""Thinking content block for reasoning models (e.g., Claude Opus 4.6)."""
|
|
|
|
type: Literal["thinking"] = "thinking"
|
|
thinking: str
|
|
signature: str | None = None
|
|
|
|
|
|
class ContentBlockDocument(BaseModel):
|
|
"""Document content block (PDF, plain text)."""
|
|
|
|
type: Literal["document"] = "document"
|
|
source: dict[str, Any] # {"type": "base64", "media_type": "application/pdf", "data": "..."}
|
|
title: str | None = None
|
|
context: str | None = None
|
|
citations: dict[str, Any] | None = None
|
|
cache_control: dict[str, str] | None = None
|
|
|
|
|
|
class ContentBlockInputAudio(BaseModel):
|
|
"""Audio input content block for multimodal audio models.
|
|
|
|
Uses the same shape as OpenAI's input_audio content part so that
|
|
the internal VLM engine receives a uniform representation regardless
|
|
of which API endpoint the client used.
|
|
"""
|
|
|
|
type: Literal["input_audio"] = "input_audio"
|
|
input_audio: dict[str, Any] # {"data": "<base64>", "format": "wav"}
|
|
|
|
|
|
# Union type for all content blocks
|
|
ContentBlock = (
|
|
ContentBlockText
|
|
| ContentBlockImage
|
|
| ContentBlockToolUse
|
|
| ContentBlockToolResult
|
|
| ContentBlockThinking
|
|
| ContentBlockDocument
|
|
| ContentBlockInputAudio
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# System Content
|
|
# =============================================================================
|
|
|
|
|
|
class SystemContent(BaseModel):
|
|
"""System message content block."""
|
|
|
|
type: Literal["text"] = "text"
|
|
text: str
|
|
cache_control: dict[str, str] | None = None
|
|
|
|
|
|
# =============================================================================
|
|
# Messages
|
|
# =============================================================================
|
|
|
|
|
|
class AnthropicMessage(BaseModel):
|
|
"""A message in an Anthropic conversation."""
|
|
|
|
role: Literal["user", "assistant", "system"]
|
|
content: str | list[ContentBlock]
|
|
|
|
|
|
# =============================================================================
|
|
# Tools
|
|
# =============================================================================
|
|
|
|
|
|
class AnthropicTool(BaseModel):
|
|
"""Tool definition in Anthropic format.
|
|
|
|
Supports two shapes:
|
|
1. User-defined tool: requires ``input_schema``.
|
|
2. Anthropic server-side tool (web_search, code_execution, bash,
|
|
text_editor, computer): carries a versioned ``type`` like
|
|
``web_search_20250305`` and no ``input_schema``. oMLX cannot execute
|
|
these locally; they are accepted for compatibility with clients such
|
|
as Claude for Excel/PowerPoint/Word and dropped before inference.
|
|
"""
|
|
|
|
name: str
|
|
description: str | None = None
|
|
input_schema: dict[str, Any] | None = None
|
|
type: str | None = None
|
|
cache_control: dict[str, str] | None = None
|
|
|
|
# Forward-compat with extra fields Anthropic may attach to server-side
|
|
# tools (e.g. max_uses, allowed_domains, user_location for web_search).
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
@model_validator(mode="after")
|
|
def _require_schema_or_type(self) -> "AnthropicTool":
|
|
if self.input_schema is None and self.type is None:
|
|
raise ValueError(
|
|
"AnthropicTool requires either 'input_schema' (user-defined "
|
|
"tool) or 'type' (Anthropic server-side tool)."
|
|
)
|
|
return self
|
|
|
|
|
|
class ToolChoice(BaseModel):
|
|
"""Tool choice specification."""
|
|
|
|
type: Literal["auto", "any", "tool"]
|
|
name: str | None = None # Required when type="tool"
|
|
|
|
|
|
# =============================================================================
|
|
# Thinking Configuration
|
|
# =============================================================================
|
|
|
|
|
|
class ThinkingConfig(BaseModel):
|
|
"""Configuration for extended thinking/reasoning."""
|
|
|
|
type: Literal["enabled", "disabled", "adaptive"] = "enabled"
|
|
budget_tokens: int | None = None
|
|
|
|
|
|
# =============================================================================
|
|
# Request
|
|
# =============================================================================
|
|
|
|
|
|
class MessagesRequest(BaseModel):
|
|
"""Request for Anthropic Messages API."""
|
|
|
|
model: str
|
|
max_tokens: int
|
|
messages: list[AnthropicMessage]
|
|
system: str | list[SystemContent] | None = None
|
|
stop_sequences: list[str] | None = None
|
|
stream: bool = False
|
|
temperature: float | None = None
|
|
top_p: float | None = None
|
|
top_k: int | None = None
|
|
metadata: dict[str, Any] | None = None
|
|
tools: list[AnthropicTool] | None = None
|
|
tool_choice: ToolChoice | dict[str, Any] | None = None
|
|
thinking: ThinkingConfig | None = None
|
|
# Chat template kwargs (e.g. enable_thinking, reasoning_effort)
|
|
chat_template_kwargs: dict[str, Any] | None = None
|
|
|
|
|
|
# =============================================================================
|
|
# Token Counting
|
|
# =============================================================================
|
|
|
|
|
|
class TokenCountRequest(BaseModel):
|
|
"""Request for token counting (Anthropic format)."""
|
|
|
|
model: str
|
|
messages: list[AnthropicMessage]
|
|
system: str | list[SystemContent] | None = None
|
|
tools: list[AnthropicTool] | None = None
|
|
tool_choice: ToolChoice | dict[str, Any] | None = None
|
|
thinking: ThinkingConfig | None = None
|
|
|
|
|
|
class TokenCountResponse(BaseModel):
|
|
"""Response for token counting."""
|
|
|
|
input_tokens: int
|
|
|
|
|
|
# =============================================================================
|
|
# Response
|
|
# =============================================================================
|
|
|
|
|
|
class AnthropicUsage(BaseModel):
|
|
"""Token usage statistics for Anthropic API."""
|
|
|
|
input_tokens: int = 0
|
|
output_tokens: int = 0
|
|
cache_creation_input_tokens: int = 0
|
|
cache_read_input_tokens: int = 0
|
|
|
|
|
|
class MessagesResponse(BaseModel):
|
|
"""Response for Anthropic Messages API."""
|
|
|
|
id: str = Field(default_factory=lambda: generate_id(IDPrefix.MESSAGE))
|
|
type: Literal["message"] = "message"
|
|
role: Literal["assistant"] = "assistant"
|
|
model: str
|
|
content: list[ContentBlockText | ContentBlockToolUse | ContentBlockThinking]
|
|
stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None = None
|
|
stop_sequence: str | None = None
|
|
usage: AnthropicUsage = Field(default_factory=AnthropicUsage)
|
|
|
|
|
|
# =============================================================================
|
|
# Streaming Events
|
|
# =============================================================================
|
|
|
|
|
|
class MessageStartEvent(BaseModel):
|
|
"""Event sent at the start of a message."""
|
|
|
|
type: Literal["message_start"] = "message_start"
|
|
message: dict[str, Any] # Partial MessagesResponse
|
|
|
|
|
|
class ContentBlockStartEvent(BaseModel):
|
|
"""Event sent at the start of a content block."""
|
|
|
|
type: Literal["content_block_start"] = "content_block_start"
|
|
index: int
|
|
content_block: dict[str, Any] # Partial content block
|
|
|
|
|
|
class TextDelta(BaseModel):
|
|
"""Text delta for streaming."""
|
|
|
|
type: Literal["text_delta"] = "text_delta"
|
|
text: str
|
|
|
|
|
|
class InputJsonDelta(BaseModel):
|
|
"""JSON input delta for tool use streaming."""
|
|
|
|
type: Literal["input_json_delta"] = "input_json_delta"
|
|
partial_json: str
|
|
|
|
|
|
class ContentBlockDeltaEvent(BaseModel):
|
|
"""Event sent for content block updates."""
|
|
|
|
type: Literal["content_block_delta"] = "content_block_delta"
|
|
index: int
|
|
delta: TextDelta | InputJsonDelta | dict[str, Any]
|
|
|
|
|
|
class ContentBlockStopEvent(BaseModel):
|
|
"""Event sent when a content block ends."""
|
|
|
|
type: Literal["content_block_stop"] = "content_block_stop"
|
|
index: int
|
|
|
|
|
|
class MessageDeltaEvent(BaseModel):
|
|
"""Event sent for message-level updates (stop_reason, usage)."""
|
|
|
|
type: Literal["message_delta"] = "message_delta"
|
|
delta: dict[str, Any] # {"stop_reason": "...", "stop_sequence": ...}
|
|
usage: dict[str, int] # {"output_tokens": N}
|
|
|
|
|
|
class MessageStopEvent(BaseModel):
|
|
"""Event sent when the message ends."""
|
|
|
|
type: Literal["message_stop"] = "message_stop"
|
|
|
|
|
|
class PingEvent(BaseModel):
|
|
"""Ping event for keeping connection alive."""
|
|
|
|
type: Literal["ping"] = "ping"
|
|
|
|
|
|
class ErrorEvent(BaseModel):
|
|
"""Error event for streaming errors."""
|
|
|
|
type: Literal["error"] = "error"
|
|
error: dict[str, Any] # {"type": "...", "message": "..."}
|
|
|
|
|
|
# Union type for all streaming events
|
|
StreamingEvent = (
|
|
MessageStartEvent
|
|
| ContentBlockStartEvent
|
|
| ContentBlockDeltaEvent
|
|
| ContentBlockStopEvent
|
|
| MessageDeltaEvent
|
|
| MessageStopEvent
|
|
| PingEvent
|
|
| ErrorEvent
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Error Response
|
|
# =============================================================================
|
|
|
|
|
|
class AnthropicErrorDetail(BaseModel):
|
|
"""Error detail in Anthropic format."""
|
|
|
|
type: str # "invalid_request_error", "authentication_error", "api_error", etc.
|
|
message: str
|
|
|
|
|
|
class AnthropicErrorResponse(BaseModel):
|
|
"""Error response in Anthropic format."""
|
|
|
|
type: Literal["error"] = "error"
|
|
error: AnthropicErrorDetail
|