kosong.chat_provider

  1from __future__ import annotations
  2
  3from collections.abc import AsyncIterator, Sequence
  4from typing import TYPE_CHECKING, Literal, Protocol, Self, runtime_checkable
  5
  6from pydantic import BaseModel
  7
  8from kosong.message import ContentPart, Message, ToolCall, ToolCallPart
  9from kosong.tooling import Tool
 10
 11if TYPE_CHECKING:
 12    import httpx
 13
 14
 15@runtime_checkable
 16class ChatProvider(Protocol):
 17    """The interface of chat providers."""
 18
 19    name: str
 20    """
 21    The name of the chat provider.
 22    """
 23
 24    @property
 25    def model_name(self) -> str:
 26        """
 27        The name of the model to use.
 28        """
 29        ...
 30
 31    @property
 32    def thinking_effort(self) -> ThinkingEffort | None:
 33        """
 34        The current thinking effort level. Returns None if not explicitly set.
 35        """
 36        ...
 37
 38    async def generate(
 39        self,
 40        system_prompt: str,
 41        tools: Sequence[Tool],
 42        history: Sequence[Message],
 43    ) -> StreamedMessage:
 44        """
 45        Generate a new message based on the given system prompt, tools, and history.
 46
 47        Raises:
 48            APIConnectionError: If the API connection fails.
 49            APITimeoutError: If the API request times out.
 50            APIStatusError: If the API returns a status code of 4xx or 5xx.
 51            ChatProviderError: If any other recognized chat provider error occurs.
 52        """
 53        ...
 54
 55    def with_thinking(self, effort: ThinkingEffort) -> Self:
 56        """
 57        Return a copy of self configured with the given thinking effort.
 58        If the chat provider does not support thinking, simply return a copy of self.
 59        """
 60        ...
 61
 62
 63@runtime_checkable
 64class RetryableChatProvider(Protocol):
 65    """Optional interface for providers that can recover from retryable transport errors."""
 66
 67    def on_retryable_error(self, error: BaseException) -> bool:
 68        """
 69        Try to recover provider transport state after a retryable error.
 70
 71        Returns:
 72            bool: Whether recovery action was performed.
 73        """
 74        ...
 75
 76
 77type StreamedMessagePart = ContentPart | ToolCall | ToolCallPart
 78
 79
 80@runtime_checkable
 81class StreamedMessage(Protocol):
 82    """The interface of streamed messages."""
 83
 84    def __aiter__(self) -> AsyncIterator[StreamedMessagePart]:
 85        """Create an async iterator from the stream."""
 86        ...
 87
 88    @property
 89    def id(self) -> str | None:
 90        """The ID of the streamed message."""
 91        ...
 92
 93    @property
 94    def usage(self) -> TokenUsage | None:
 95        """The token usage of the streamed message."""
 96        ...
 97
 98
 99class TokenUsage(BaseModel):
100    """Token usage statistics."""
101
102    input_other: int
103    """Input tokens excluding `input_cache_read` and `input_cache_creation`."""
104    output: int
105    """Total output tokens."""
106    input_cache_read: int = 0
107    """Cached input tokens."""
108    input_cache_creation: int = 0
109    """Input tokens used for cache creation. For now, only Anthropic API supports this."""
110
111    @property
112    def total(self) -> int:
113        """Total tokens used, including input and output tokens."""
114        return self.input + self.output
115
116    @property
117    def input(self) -> int:
118        """Total input tokens, including cached and uncached tokens."""
119        return self.input_other + self.input_cache_read + self.input_cache_creation
120
121
122type ThinkingEffort = Literal["off", "low", "medium", "high", "xhigh", "max"]
123"""The effort level for thinking.
124
125Support for levels above ``high`` varies by provider:
126
127- **Anthropic**: ``xhigh`` is accepted only on Claude Opus 4.7; ``max`` is
128  accepted on Mythos, Opus 4.7/4.6, and Sonnet 4.6. Unsupported levels are
129  clamped down to ``high``.
130- **OpenAI**: ``xhigh`` is accepted natively for reasoning-capable models
131  after ``gpt-5.1-codex-max`` and passes through unchanged. ``max`` is
132  Anthropic-specific and clamps to ``xhigh`` (OpenAI's ceiling).
133- **Kimi**: requests only serialize thinking as enabled or disabled; the
134  caller-provided effort remains unchanged as provider state.
135- **Gemini**: ``xhigh`` and ``max`` clamp to ``high`` (no native support).
136"""
137
138
139class ChatProviderError(Exception):
140    """The error raised by a chat provider."""
141
142    def __init__(self, message: str):
143        super().__init__(message)
144
145
146class APIConnectionError(ChatProviderError):
147    """The error raised when the API connection fails."""
148
149
150class APITimeoutError(ChatProviderError):
151    """The error raised when the API request times out."""
152
153
154class APIStatusError(ChatProviderError):
155    """The error raised when the API returns a status code of 4xx or 5xx."""
156
157    status_code: int
158    request_id: str | None
159    trace_id: str | None
160
161    def __init__(
162        self,
163        status_code: int,
164        message: str,
165        *,
166        request_id: str | None = None,
167        trace_id: str | None = None,
168    ):
169        super().__init__(message)
170        self.status_code = status_code
171        self.request_id = request_id
172        self.trace_id = trace_id
173
174
175class APIEmptyResponseError(ChatProviderError):
176    """The error raised when the API returns an empty response."""
177
178
179def convert_httpx_error(error: httpx.HTTPError) -> ChatProviderError:
180    """Convert an httpx transport error to the corresponding ChatProviderError.
181
182    This is a shared utility for all chat providers. SDK-specific exceptions
183    (e.g. AnthropicError, OpenAIError) should be handled by each provider's
184    own conversion logic; only raw httpx exceptions that leak through
185    (typically during streaming) should be routed here.
186    """
187    import httpx
188
189    if isinstance(error, httpx.TimeoutException):
190        return APITimeoutError(str(error))
191    if isinstance(error, (httpx.NetworkError, httpx.RemoteProtocolError)):
192        return APIConnectionError(str(error))
193    if isinstance(error, httpx.HTTPStatusError):
194        req_id = error.response.headers.get("x-request-id")
195        trace_id = error.response.headers.get("x-trace-id")
196        return APIStatusError(
197            error.response.status_code, str(error), request_id=req_id, trace_id=trace_id
198        )
199    return ChatProviderError(f"HTTP error: {error}")