tau-coding-agent 0.1.0__py3-none-any.whl

tau/inference/api/text/anthropic_claude_code.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+import json
+from tau.inference.api.text.utils import parse_tool_args, anthropic_messages_to_list, anthropic_output_config, anthropic_apply_message_cache
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+from anthropic import AsyncAnthropic
+from tau.inference.api.text.base import BaseLLMAPI as BaseAPI
+from tau.inference.model.types import Model
+from tau.inference.types import (
+    LLMContext, LLMEvent, LLMOptions, StopReason, ThinkingBudgets, ThinkingLevel,
+    StartEvent, EndEvent, ErrorEvent,
+    TextStartEvent, TextDeltaEvent, TextEndEvent,
+    ThinkingStartEvent, ThinkingDeltaEvent, ThinkingEndEvent,
+    ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent,
+)
+from tau.message.types import (
+    SystemMessage, UserMessage, AssistantMessage, ToolMessage,
+    TextContent, ImageContent, ThinkingContent, ToolCallContent, ToolResultContent,
+)
+from typing import Optional, TYPE_CHECKING
+if TYPE_CHECKING:
+    from tau.tool.types import Tool
+
+_STOP_REASON: dict[str, StopReason] = {
+    "end_turn": StopReason.Stop,
+    "max_tokens": StopReason.Length,
+    "tool_use": StopReason.ToolCalls,
+    "stop_sequence": StopReason.Stop,
+}
+
+_DEFAULT_MAX_TOKENS = 8096
+
+
+_OAUTH_HEADERS = {
+    "anthropic-beta": "oauth-2025-04-20",
+    "x-app": "cli",
+    "User-Agent": "claude-cli/2.1.122 (external, sdk-cli)",
+}
+
+
+class AnthropicClaudeCodeAPI(BaseAPI):
+    """Anthropic Messages API using OAuth token auth (Claude Pro/Max).
+
+    Sends the token via X-Api-Key (not Authorization: Bearer) with the
+    required OAuth beta headers, which is what Anthropic's API enforces
+    for Claude Max / Pro OAuth tokens.
+    """
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Initialise the AsyncAnthropic client with OAuth headers merged from options."""
+        super().__init__(options)
+        merged_headers = {**_OAUTH_HEADERS, **(options.headers or {})}
+        self._client = AsyncAnthropic(
+            auth_token=options.api_key,        # Bearer auth for OAuth tokens
+            base_url=options.base_url,
+            default_headers=merged_headers,
+            max_retries=options.max_retries,
+            timeout=options.timeout.total_seconds(),
+        )
+        self._current_api_key = options.api_key
+
+    def _build_params(
+        self,
+        model: Model,
+        system: str | None,
+        messages: list[dict[str, Any]],
+        tools: Optional[list[Tool]] = None,
+    ) -> dict[str, Any]:
+        """Assemble the Anthropic API request payload, including thinking and tool configs."""
+        params: dict[str, Any] = {
+            "model": model.id,
+            "messages": anthropic_apply_message_cache(messages),
+            "max_tokens": self.options.max_tokens or _DEFAULT_MAX_TOKENS,
+            "temperature": self.options.temperature,
+        }
+        if system:
+            params["system"] = [{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}]
+        if self.options.thinking_level is not None and self.options.thinking_level != ThinkingLevel.Off:
+            budgets = self.options.thinking_budgets or ThinkingBudgets()
+            params["thinking"] = {"type": "enabled", "budget_tokens": budgets.get(self.options.thinking_level)}
+
+        if tools:
+            tool_defs = [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "input_schema": tool.schema.model_json_schema(),
+                }
+                for tool in tools
+            ]
+            # Cache the last tool definition to reduce repeated prompt-token charges.
+            tool_defs[-1]["cache_control"] = {"type": "ephemeral"}
+            params["tools"] = tool_defs
+        return params
+
+    def _sync_client(self) -> None:
+        """Rebuild client if the api_key (OAuth token) has been refreshed."""
+        if self.options.api_key != self._current_api_key:
+            self._current_api_key = self.options.api_key
+            merged_headers = {**_OAUTH_HEADERS, **(self.options.headers or {})}
+            self._client = AsyncAnthropic(
+                auth_token=self.options.api_key,
+                base_url=self.options.base_url,
+                default_headers=merged_headers,
+                max_retries=self.options.max_retries,
+                timeout=self.options.timeout.total_seconds(),
+            )
+
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:  # type: ignore[override]
+        """Stream LLMEvents from the Anthropic Messages API using an OAuth token."""
+        self._sync_client()
+        system, anthropic_messages = anthropic_messages_to_list(context.messages)
+        if context.system_prompt:
+            system = context.system_prompt
+        params = self._build_params(model, system, anthropic_messages, tools=context.tools or None)
+        output_config = anthropic_output_config(context.response_format)
+        if output_config is not None:
+            params["output_config"] = output_config
+
+        if self.options.on_payload:
+            modified = self.options.on_payload(params)
+            if modified is not None:
+                params = modified
+
+        # Per-block accumulation buffers keyed by content block index.
+        block_types: dict[int, str] = {}
+        tool_ids: dict[int, str] = {}
+        tool_names: dict[int, str] = {}
+        text_bufs: dict[int, str] = {}
+        thinking_bufs: dict[int, str] = {}
+        tool_bufs: dict[int, str] = {}
+        _input_tokens = 0
+        _output_tokens = 0
+        _cache_read_tokens = 0
+        _cache_write_tokens = 0
+
+        yield StartEvent()
+
+        async with self._client.messages.stream(**params) as stream:
+            async for event in stream:
+                if self._cancelled():
+                    yield ErrorEvent(reason=StopReason.Abort, error="Cancelled")
+                    return
+                etype = event.type
+
+                if etype == "content_block_start":
+                    idx = getattr(event, "index", 0)
+                    block = getattr(event, "content_block", None)
+                    if block is None:
+                        continue
+                    btype_start = getattr(block, "type", "")
+                    block_types[idx] = btype_start
+                    if btype_start == "text":
+                        text_bufs[idx] = ""
+                        yield TextStartEvent(text=TextContent(content=""))
+                    elif btype_start == "thinking":
+                        thinking_bufs[idx] = ""
+                        yield ThinkingStartEvent(thinking=None)
+                    elif btype_start == "tool_use":
+                        tool_ids[idx] = getattr(block, "id", "")
+                        tool_names[idx] = getattr(block, "name", "")
+                        tool_bufs[idx] = ""
+                        yield ToolCallStartEvent(tool_call=ToolCallContent(id=tool_ids[idx], name=tool_names[idx]))
+
+                elif etype == "content_block_delta":
+                    idx = getattr(event, "index", 0)
+                    delta = getattr(event, "delta", None)
+                    if delta is None:
+                        continue
+                    dtype = getattr(delta, "type", "")
+                    if dtype == "text_delta":
+                        text = getattr(delta, "text", "")
+                        text_bufs[idx] = text_bufs.get(idx, "") + text
+                        yield TextDeltaEvent(text=TextContent(content=text))
+                    elif dtype == "thinking_delta":
+                        thinking = getattr(delta, "thinking", "")
+                        thinking_bufs[idx] = thinking_bufs.get(idx, "") + thinking
+                        yield ThinkingDeltaEvent(thinking=ThinkingContent(content=thinking))
+                    elif dtype == "input_json_delta":
+                        partial = getattr(delta, "partial_json", "")
+                        tool_bufs[idx] = tool_bufs.get(idx, "") + partial
+                        yield ToolCallDeltaEvent(tool_call=ToolCallContent(id=tool_ids.get(idx, "")))
+
+                elif etype == "content_block_stop":
+                    idx = getattr(event, "index", 0)
+                    btype = block_types.get(idx, "")
+                    if btype == "text":
+                        yield TextEndEvent(text=TextContent(content=text_bufs.get(idx, "")))
+                    elif btype == "thinking":
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_bufs.get(idx, "")))
+                    elif btype == "tool_use":
+                        args_str = tool_bufs.get(idx, "").strip()
+                        args = parse_tool_args(args_str)
+
+                        yield ToolCallEndEvent(tool_call=ToolCallContent(
+                                id=tool_ids.get(idx, ""),
+                                name=tool_names.get(idx, ""),
+                                args=args
+                            ))
+
+                elif etype == "message_start":
+                    u = getattr(getattr(event, 'message', None), 'usage', None)
+                    if u:
+                        _input_tokens = getattr(u, 'input_tokens', 0) or 0
+                        _cache_read_tokens = getattr(u, 'cache_read_input_tokens', 0) or 0
+                        _cache_write_tokens = getattr(u, 'cache_creation_input_tokens', 0) or 0
+
+                elif etype == "message_delta":
+                    u = getattr(event, 'usage', None)
+                    if u:
+                        _output_tokens = getattr(u, 'output_tokens', 0) or 0
+                    stop_reason = _STOP_REASON.get(getattr(getattr(event, 'delta', None), 'stop_reason', None) or "", StopReason.Stop)
+                    yield EndEvent(
+                        reason=stop_reason,
+                        input_tokens=_input_tokens,
+                        output_tokens=_output_tokens,
+                        cache_read_tokens=_cache_read_tokens,
+                        cache_write_tokens=_cache_write_tokens,
+                    )
+
+                elif etype == "error":
+                    yield ErrorEvent(reason=StopReason.Abort, error=str(event))

tau/inference/api/text/anthropic_messages.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+import json
+from tau.inference.api.text.utils import parse_tool_args, anthropic_messages_to_list, anthropic_output_config, anthropic_apply_message_cache
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+from anthropic import AsyncAnthropic
+from tau.inference.api.text.base import BaseLLMAPI as BaseAPI
+from tau.inference.model.types import Model
+from tau.inference.types import (
+    LLMContext, LLMEvent, LLMOptions, StopReason, ThinkingBudgets, ThinkingLevel,
+    StartEvent, EndEvent, ErrorEvent,
+    TextStartEvent, TextDeltaEvent, TextEndEvent,
+    ThinkingStartEvent, ThinkingDeltaEvent, ThinkingEndEvent,
+    ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent,
+)
+from tau.message.types import (
+    SystemMessage, UserMessage, AssistantMessage, ToolMessage,
+    TextContent, ImageContent, ThinkingContent, ToolCallContent, ToolResultContent,
+)
+from typing import Optional, TYPE_CHECKING
+if TYPE_CHECKING:
+    from tau.tool.types import Tool
+
+_STOP_REASON: dict[str, StopReason] = {
+    "end_turn": StopReason.Stop,
+    "max_tokens": StopReason.Length,
+    "tool_use": StopReason.ToolCalls,
+    "stop_sequence": StopReason.Stop,
+}
+
+_DEFAULT_MAX_TOKENS = 8096
+
+
+class AnthropicMessagesAPI(BaseAPI):
+    """Streaming LLM API adapter for Anthropic Messages API (API-key auth)."""
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Initialise the AsyncAnthropic client with the supplied options."""
+        super().__init__(options)
+        self._client = AsyncAnthropic(
+            api_key=options.api_key,
+            base_url=options.base_url,
+            default_headers=options.headers,
+            max_retries=options.max_retries,
+            timeout=options.timeout.total_seconds(),
+        )
+
+    def _build_params(
+        self,
+        model: Model,
+        system: str | None,
+        messages: list[dict[str, Any]],
+        tools: Optional[list[Tool]] = None,
+    ) -> dict[str, Any]:
+        """Assemble the Anthropic API request payload, including thinking and tool configs."""
+        params: dict[str, Any] = {
+            "model": model.id,
+            "messages": anthropic_apply_message_cache(messages),
+            "max_tokens": self.options.max_tokens or _DEFAULT_MAX_TOKENS,
+            "temperature": self.options.temperature,
+        }
+        if system:
+            params["system"] = [{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}]
+        if self.options.thinking_level is not None and self.options.thinking_level != ThinkingLevel.Off:
+            budgets = self.options.thinking_budgets or ThinkingBudgets()
+            params["thinking"] = {"type": "enabled", "budget_tokens": budgets.get(self.options.thinking_level)}
+
+        if tools:
+            tool_defs = [
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "input_schema": tool.schema.model_json_schema(),
+                }
+                for tool in tools
+            ]
+            # Cache the last tool definition to reduce repeated prompt-token charges.
+            tool_defs[-1]["cache_control"] = {"type": "ephemeral"}
+            params["tools"] = tool_defs
+
+        return params
+
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:  # type: ignore[override]
+        """Stream LLMEvents from the Anthropic Messages API."""
+        system, anthropic_messages = anthropic_messages_to_list(context.messages)
+        if context.system_prompt:
+            system = context.system_prompt
+        params = self._build_params(model, system, anthropic_messages, tools=context.tools or None)
+        output_config = anthropic_output_config(context.response_format)
+        if output_config is not None:
+            params["output_config"] = output_config
+
+        if self.options.on_payload:
+            modified = self.options.on_payload(params)
+            if modified is not None:
+                params = modified
+
+        # Per-block accumulation buffers keyed by content block index.
+        block_types: dict[int, str] = {}
+        tool_ids: dict[int, str] = {}
+        tool_names: dict[int, str] = {}
+        text_bufs: dict[int, str] = {}
+        thinking_bufs: dict[int, str] = {}
+        tool_bufs: dict[int, str] = {}
+        _input_tokens = 0
+        _output_tokens = 0
+        _cache_read_tokens = 0
+        _cache_write_tokens = 0
+
+        yield StartEvent()
+
+        async with self._client.messages.stream(**params) as stream:
+            async for event in stream:
+                if self._cancelled():
+                    yield ErrorEvent(reason=StopReason.Abort, error="Cancelled")
+                    return
+                etype = event.type
+
+                if etype == "content_block_start":
+                    idx = getattr(event, "index", 0)
+                    block = getattr(event, "content_block", None)
+                    if block is None:
+                        continue
+                    btype_start = getattr(block, "type", "")
+                    block_types[idx] = btype_start
+                    if btype_start == "text":
+                        text_bufs[idx] = ""
+                        yield TextStartEvent(text=TextContent(content=""))
+                    elif btype_start == "thinking":
+                        thinking_bufs[idx] = ""
+                        yield ThinkingStartEvent(thinking=None)
+                    elif btype_start == "tool_use":
+                        tool_ids[idx] = getattr(block, "id", "")
+                        tool_names[idx] = getattr(block, "name", "")
+                        tool_bufs[idx] = ""
+                        yield ToolCallStartEvent(tool_call=ToolCallContent(id=tool_ids[idx], name=tool_names[idx]))
+
+                elif etype == "content_block_delta":
+                    idx = getattr(event, "index", 0)
+                    delta = getattr(event, "delta", None)
+                    if delta is None:
+                        continue
+                    dtype = getattr(delta, "type", "")
+                    if dtype == "text_delta":
+                        text = getattr(delta, "text", "")
+                        text_bufs[idx] = text_bufs.get(idx, "") + text
+                        yield TextDeltaEvent(text=TextContent(content=text))
+                    elif dtype == "thinking_delta":
+                        thinking = getattr(delta, "thinking", "")
+                        thinking_bufs[idx] = thinking_bufs.get(idx, "") + thinking
+                        yield ThinkingDeltaEvent(thinking=ThinkingContent(content=thinking))
+                    elif dtype == "input_json_delta":
+                        partial = getattr(delta, "partial_json", "")
+                        tool_bufs[idx] = tool_bufs.get(idx, "") + partial
+                        yield ToolCallDeltaEvent(tool_call=ToolCallContent(id=tool_ids.get(idx, "")))
+
+                elif etype == "content_block_stop":
+                    idx = getattr(event, "index", 0)
+                    btype = block_types.get(idx, "")
+                    if btype == "text":
+                        yield TextEndEvent(text=TextContent(content=text_bufs.get(idx, "")))
+                    elif btype == "thinking":
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_bufs.get(idx, "")))
+                    elif btype == "tool_use":
+                        args_str = tool_bufs.get(idx, "").strip()
+                        args = parse_tool_args(args_str)
+
+                        yield ToolCallEndEvent(tool_call=ToolCallContent(
+                                id=tool_ids.get(idx, ""),
+                                name=tool_names.get(idx, ""),
+                                args=args
+                            )
+                        )
+
+                elif etype == "message_start":
+                    u = getattr(getattr(event, 'message', None), 'usage', None)
+                    if u:
+                        _input_tokens = getattr(u, 'input_tokens', 0) or 0
+                        _cache_read_tokens = getattr(u, 'cache_read_input_tokens', 0) or 0
+                        _cache_write_tokens = getattr(u, 'cache_creation_input_tokens', 0) or 0
+
+                elif etype == "message_delta":
+                    u = getattr(event, 'usage', None)
+                    if u:
+                        _output_tokens = getattr(u, 'output_tokens', 0) or 0
+                    stop_reason = _STOP_REASON.get(getattr(getattr(event, 'delta', None), 'stop_reason', None) or "", StopReason.Stop)
+                    yield EndEvent(
+                        reason=stop_reason,
+                        input_tokens=_input_tokens,
+                        output_tokens=_output_tokens,
+                        cache_read_tokens=_cache_read_tokens,
+                        cache_write_tokens=_cache_write_tokens,
+                    )
+
+                elif etype == "error":
+                    yield ErrorEvent(reason=StopReason.Abort, error=str(event))

tau/inference/api/text/base.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, AsyncIterator
+from tau.inference.model.types import Model
+from tau.inference.types import LLMContext, LLMEvent, LLMOptions, Transport
+from tau.message.types import LLMMessage
+
+
+class BaseLLMAPI(ABC):
+    """Abstract base for all provider-specific streaming LLM API implementations."""
+
+    SUPPORTED_TRANSPORTS: tuple[Transport, ...] = (Transport.HTTP,)
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Validate transport compatibility and store options."""
+        if options.transport not in self.SUPPORTED_TRANSPORTS:
+            raise ValueError(
+                f"{self.__class__.__name__} does not support transport '{options.transport.value}'. "
+                f"Supported: {[t.value for t in self.SUPPORTED_TRANSPORTS]}"
+            )
+        self.options = options
+
+    def _cancelled(self) -> bool:
+        """Return True if the caller has set the abort signal."""
+        return self.options.signal is not None and self.options.signal.is_set()
+
+    @abstractmethod
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:
+        """Yield LLMEvent objects as the provider streams its response."""
+        # The unreachable `yield` marks this as an async-generator signature so
+        # callers (and type-checkers) treat stream() as returning an
+        # AsyncGenerator — which supports aclose() — rather than a coroutine.
+        raise NotImplementedError
+        yield  # type: ignore[unreachable]  # pragma: no cover
+
+    async def invoke(self, context: LLMContext, model: Model) -> list[LLMEvent]:
+        """Collect all stream events into a list and return them."""
+        events: list[LLMEvent] = []
+        async for event in self.stream(context, model=model):
+            events.append(event)
+        return events

tau/inference/api/text/builtins.py

@@ -0,0 +1,19 @@
+# Lazy registry entries: each provider API is referenced by a "module:ClassName"
+# path so its SDK (anthropic, openai, google.genai, mistralai, ...) is only
+# imported when that provider is actually used. Loading every SDK at startup
+# previously cost ~1.5s on cold start.
+LLM_APIS: list[tuple[str, str]] = [
+    ("openai_responses",       "tau.inference.api.text.openai_responses:OpenAIResponsesAPI"),
+    ("openai_completions",     "tau.inference.api.text.openai_completions:OpenAICompletionsAPI"),
+    ("openai_codex_responses", "tau.inference.api.text.openai_codex_responses:OpenAICodexResponsesAPI"),
+    ("anthropic_messages",     "tau.inference.api.text.anthropic_messages:AnthropicMessagesAPI"),
+    ("anthropic_claude_code",  "tau.inference.api.text.anthropic_claude_code:AnthropicClaudeCodeAPI"),
+    ("github_copilot_chat",    "tau.inference.api.text.github_copilot_chat:GitHubCopilotChatAPI"),
+    ("gemini_generate",        "tau.inference.api.text.gemini_generate:GeminiGenerateAPI"),
+    ("mistral_chat",           "tau.inference.api.text.mistral_chat:MistralChatAPI"),
+    ("ollama_chat",            "tau.inference.api.text.ollama_chat:OllamaChatAPI"),
+    ("google_antigravity",     "tau.inference.api.text.google_antigravity:GoogleAntigravityAPI"),
+]
+
+# Backward-compat alias
+APIS = LLM_APIS