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

tau/inference/api/text/service.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import aclosing
+from dataclasses import fields
+from tau.inference.model.registry import ModelRegistry
+from tau.inference.api.registry import LazyAPI
+from tau.inference.api.text.registry import LLMAPIRegistry
+from tau.inference.provider.registry import TextProviderRegistry, ProviderRegistry
+from tau.inference.provider.types import APIProvider, OAuthProvider
+from tau.auth.manager import AuthManager
+from tau.auth.types import OAuthCredential
+from tau.inference.types import LLMContext, LLMEvent, LLMOptions
+from tau.message.types import LLMMessage, SystemMessage
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from tau.tool.types import Tool
+    from tau.inference.types import ThinkingLevel
+
+
+class TextLLM:
+    """Wrapper around inference APIs with model/provider resolution and option merging."""
+    _apis: LLMAPIRegistry | None = None
+    _models: ModelRegistry | None = None
+    _providers: TextProviderRegistry | None = None
+    _auth_manager: AuthManager | None = None
+
+    @classmethod
+    def _builtin_apis(cls) -> LLMAPIRegistry:
+        if cls._apis is None:
+            cls._apis = LLMAPIRegistry.from_builtins()
+        return cls._apis
+
+    @classmethod
+    def _builtin_models(cls) -> ModelRegistry:
+        if cls._models is None:
+            cls._models = ModelRegistry.from_text_builtins()
+        return cls._models
+
+    @classmethod
+    def _builtin_providers(cls) -> TextProviderRegistry:
+        if cls._providers is None:
+            cls._providers = TextProviderRegistry.from_builtins()
+        return cls._providers
+
+    @classmethod
+    def _builtin_auth_manager(cls) -> AuthManager:
+        if cls._auth_manager is None:
+            cls._auth_manager = AuthManager.create(ProviderRegistry(text=cls._builtin_providers()))
+        return cls._auth_manager
+
+    def __init__(
+        self,
+        model_id: str,
+        provider: str | None = None,
+        options: LLMOptions | None = None,
+        *,
+        models: Optional[ModelRegistry] = None,
+        providers: Optional[TextProviderRegistry] = None,
+        apis: Optional[LLMAPIRegistry] = None,
+        auth_manager: Optional[AuthManager] = None,
+    ) -> None:
+        """Initialize an LLM by resolving model, provider, and API implementation.
+
+        Args:
+            model_id: The model identifier (e.g., 'claude-3-5-sonnet-latest').
+            provider: Optional provider name; if omitted, defaults from model definition.
+            options: Optional LLMOptions for API key, base_url, temperature, etc.
+            models: Optional custom ModelRegistry; defaults to global builtin registry.
+            providers: Optional custom TextProviderRegistry; defaults to global builtin registry.
+            apis: Optional custom LLMAPIRegistry; defaults to global builtin registry.
+            auth_manager: Optional custom AuthManager; defaults to global builtin store.
+
+        Raises:
+            ValueError: If model_id or provider not found in registries.
+            RuntimeError: If OAuth provider requires credentials that are unavailable.
+        """
+        _models = models if models is not None else type(self)._builtin_models()
+        _providers = providers if providers is not None else type(self)._builtin_providers()
+        _apis = apis if apis is not None else type(self)._builtin_apis()
+        self._auth_manager = auth_manager if auth_manager is not None else type(self)._builtin_auth_manager()
+
+        # When provider is not pinned, try all registered variants of the model
+        # in order, skipping OAuth providers whose credentials are missing.
+        # This lets `model_id="claude-sonnet-4-6"` work with an API key even if
+        # an OAuth variant of the same model is registered first.
+        candidates = (
+            [_models.get(model_id, provider=provider)]
+            if provider is not None
+            else _models._models.get(model_id, [])
+        )
+        if not candidates or candidates[0] is None:
+            raise ValueError(f"Model '{model_id}' not found.")
+
+        model = None
+        resolved_provider = None
+        for candidate in candidates:
+            cand_provider = _providers.get(candidate.provider)
+            if cand_provider is None:
+                continue
+            if isinstance(cand_provider, OAuthProvider):
+                if not isinstance(self._auth_manager.get(cand_provider.id), OAuthCredential):
+                    continue  # no credentials — try next variant
+            model = candidate
+            resolved_provider = cand_provider
+            break
+
+        if model is None or resolved_provider is None:
+            tried = [c.provider for c in candidates]
+            raise RuntimeError(
+                f"No usable provider found for '{model_id}'. "
+                f"Tried: {tried}. Log in or set an API key."
+            )
+
+        self.model = model
+
+        # Keep the API reference unresolved (a "module:Class" string stays a
+        # string) so the provider SDK is not imported until the first request.
+        api_ref = model.api or resolved_provider.api
+
+        base_url_override = model.base_url
+
+        if isinstance(resolved_provider, OAuthProvider):
+            credential = self._auth_manager.get(resolved_provider.id)
+            if not isinstance(credential, OAuthCredential):
+                raise RuntimeError(
+                    f"No credentials found for '{resolved_provider.id}'. Please log in first."
+                )
+            base_opts = LLMOptions(api_key=resolved_provider.get_api_key(credential))
+            if base_url_override:
+                base_opts.base_url = base_url_override
+            merged = self._merge_options(base_opts, options)
+            self.provider_id = resolved_provider.id
+        else:
+            base_opts = resolved_provider.options
+            if base_url_override:
+                override_opts = LLMOptions(base_url=base_url_override)
+                base_opts = self._merge_options(base_opts, override_opts)
+            merged = self._merge_options(base_opts, options)
+            self.provider_id = resolved_provider.id
+
+        if merged.max_tokens is None:
+            merged.max_tokens = model.max_tokens
+
+        # Resolve any "$ENV_VAR" / "!command" references in custom headers to
+        # their values. Done once here (per provider/model selection); the
+        # resolver memoizes, so a !command runs only the first time it's seen.
+        if merged.headers:
+            from tau.utils.secrets import resolve_secrets
+            merged.headers = resolve_secrets(merged.headers)
+
+        # Lazy adapter: exposes `.options` immediately but only imports the
+        # provider SDK and builds its client on first `.stream()`/`.invoke()`.
+        self.api = LazyAPI(_apis, api_ref, merged)
+
+    @classmethod
+    def list_available(cls) -> list:
+        """Return all models whose provider has usable auth (stored credential or env var)."""
+        import os
+        from tau.inference.provider.types import OAuthProvider
+        from tau.auth.types import OAuthCredential, APICredential
+
+        cls._auth_manager.reload()  # pick up any credentials added since startup
+        result = []
+        seen: set[str] = set()
+        for candidates in cls._models._models.values():
+            for candidate in candidates:
+                key = f"{candidate.provider}/{candidate.id}"
+                if key in seen:
+                    continue
+                provider = cls._providers.get(candidate.provider)
+                if provider is None:
+                    continue
+                if isinstance(provider, OAuthProvider):
+                    if not isinstance(cls._auth_manager.get(provider.id), OAuthCredential):
+                        continue
+                else:
+                    cred = cls._auth_manager.get(provider.id)
+                    if not isinstance(cred, APICredential):
+                        # fall back to env var (e.g. ANTHROPIC_API_KEY)
+                        if not os.environ.get(f"{provider.id.upper()}_API_KEY"):
+                            continue
+                seen.add(key)
+                result.append(candidate)
+        return result
+
+    def _merge_options(self, base: LLMOptions, override: LLMOptions | None) -> LLMOptions:
+        """Merge base options with override options, preferring non-None override values.
+
+        Args:
+            base: The base LLMOptions configuration.
+            override: Optional override LLMOptions; fields override base when non-None.
+
+        Returns:
+            A new LLMOptions with merged values.
+        """
+        if override is None:
+            return base
+        merged = LLMOptions(**{f.name: getattr(base, f.name) for f in fields(base)})
+        for f in fields(override):
+            value = getattr(override, f.name)
+            if value is not None:
+                setattr(merged, f.name, value)
+        return merged
+
+    def _resolve_messages(self, context: LLMContext) -> list[LLMMessage]:
+        """Resolve messages for the LLM call, prepending system prompt if needed.
+
+        Args:
+            context: The LLMContext with messages and optional system prompt.
+
+        Returns:
+            A list of LLMMessages with system message injected if needed.
+        """
+        messages = context.messages
+        if context.system_prompt:
+            if not messages or not isinstance(messages[0], SystemMessage):
+                messages = [SystemMessage.text(context.system_prompt)] + messages
+        return messages
+
+    async def stream(self, context: LLMContext) -> AsyncGenerator[LLMEvent, None]:
+        """Stream LLM events from the configured provider API.
+
+        Retries transient errors transparently as long as no events have been
+        yielded yet (pre-stream failures). Once streaming has started, errors
+        are forwarded as ErrorEvent because already-yielded events can't be
+        recalled.
+
+        Args:
+            context: The LLMContext with messages, tools, and response format options.
+
+        Yields:
+            LLMEvent objects (TextDeltaEvent, ToolCallEndEvent, EndEvent, ErrorEvent, etc.).
+        """
+        from tau.inference.types import ErrorEvent, RetryEvent, StartEvent, StopReason
+        from tau.inference.utils import classify_error
+        import asyncio
+
+        api_key = await self._auth_manager.get_api_key(self.provider_id)
+        if api_key:
+            self.api.options.api_key = api_key
+
+        messages = self._resolve_messages(context)
+        api_context = LLMContext(
+            messages=messages,
+            tools=context.tools,
+            response_format=context.response_format,
+        )
+
+        max_retries = self.api.options.max_retries
+        base_delay_s = self.api.options.retry_base_delay_ms / 1000
+
+        for attempt in range(max_retries + 1):
+            received_any = False
+            try:
+                async with aclosing(self.api.stream(api_context, model=self.model)) as stream:
+                    async for event in stream:
+                        # StartEvent/RetryEvent are emitted locally before any HTTP round-trip;
+                        # don't count them as "received data" so retries still fire on empty bodies.
+                        if not isinstance(event, (StartEvent, RetryEvent)):
+                            received_any = True
+                        yield event
+                return
+            except Exception as e:
+                classified = classify_error(e)
+                if received_any or not classified.retryable or attempt >= max_retries:
+                    yield ErrorEvent(reason=StopReason.Error, error=str(e))
+                    return
+                yield RetryEvent(attempt=attempt + 1, max_retries=max_retries, error=str(e))
+                await asyncio.sleep(base_delay_s * (2 ** attempt))
+
+    async def invoke(
+        self,
+        context: LLMContext,
+        thinking_level: Optional["ThinkingLevel"] = None,
+    ) -> list[LLMEvent]:
+        from tau.inference.types import ThinkingLevel
+        api_key = await self._auth_manager.get_api_key(self.provider_id)
+        if api_key:
+            self.api.options.api_key = api_key
+
+        original = self.api.options.thinking_level
+        if thinking_level is not None and thinking_level != ThinkingLevel.Off:
+            self.api.options.thinking_level = thinking_level
+        try:
+            messages = self._resolve_messages(context)
+            api_context = LLMContext(
+                messages=messages,
+                tools=context.tools,
+                response_format=context.response_format,
+            )
+            return await self.api.invoke(api_context, model=self.model)
+        except Exception as e:
+            from tau.inference.types import ErrorEvent, StopReason
+            return [ErrorEvent(reason=StopReason.Error, error=str(e))]
+        finally:
+            self.api.options.thinking_level = original

tau/inference/api/text/types.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class APIResponse:
+    status_code: int
+    headers: dict[str, str]

tau/inference/api/text/utils.py

@@ -0,0 +1,228 @@
+"""Shared utilities for LLM API provider implementations."""
+from __future__ import annotations
+
+import json
+from typing import Any
+
+
+__all__ = [
+    "parse_tool_args",
+    "openai_user_content", "openai_assistant_content", "openai_messages_to_chat", "openai_response_format",
+    "anthropic_messages_to_list", "anthropic_output_config", "anthropic_apply_message_cache",
+]
+
+
+_CACHE_MARKER = {"type": "ephemeral"}
+
+
+def anthropic_apply_message_cache(
+    messages: list[dict[str, Any]],
+    n: int = 2,
+    skip_tail: int = 0,
+) -> list[dict[str, Any]]:
+    """Inject cache_control breakpoints into the last n stable messages.
+
+    Implements the Anthropic 'system_and_3' caching strategy — the system
+    prompt is already marked by the caller; this adds up to 2 more breakpoints
+    on the tail of the stable session history so the bulk of the conversation
+    is served from cache on subsequent turns.
+
+    skip_tail: number of ephemeral messages at the end of the list to skip
+    (desktop/browser screenshots that change every turn and must not be cached).
+
+    Returns a new list; the original is not mutated.
+    """
+    import copy
+    messages = copy.deepcopy(messages)
+    total = len(messages)
+    stable_end = total - skip_tail  # index just past the last stable message
+    stable_start = max(0, stable_end - n)
+    for msg in messages[stable_start:stable_end]:
+        content = msg.get("content")
+        if content is None or content == "":
+            msg["cache_control"] = _CACHE_MARKER
+        elif isinstance(content, str):
+            msg["content"] = [{"type": "text", "text": content, "cache_control": _CACHE_MARKER}]
+        elif isinstance(content, list) and content:
+            last = content[-1]
+            if isinstance(last, dict):
+                last["cache_control"] = _CACHE_MARKER
+    return messages
+
+
+def parse_tool_args(value: Any) -> dict:
+    """Parse a tool-call arguments value into a dict.
+
+    Handles the three shapes that provider APIs return:
+    - already a dict  → return as-is
+    - a JSON string   → parse and return (empty string → {})
+    - anything else   → return {}
+    Falls back to {} on JSONDecodeError.
+    """
+    if isinstance(value, dict):
+        return value
+    if not value:
+        return {}
+    try:
+        result = json.loads(value)
+        return result if isinstance(result, dict) else {}
+    except (json.JSONDecodeError, TypeError):
+        return {}
+
+
+def openai_user_content(content_items: list) -> str | list[dict[str, Any]]:
+    """Convert user message contents to OpenAI chat format (completions/copilot/mistral)."""
+    from tau.message.types import TextContent, ImageContent
+    parts: list[dict[str, Any]] = []
+    for item in content_items:
+        match item:
+            case TextContent():
+                parts.append({"type": "text", "text": item.content})
+            case ImageContent():
+                for b64, mime in item.to_base64():
+                    url = b64 if b64.startswith("http") else f"data:{mime or 'image/png'};base64,{b64}"
+                    parts.append({"type": "image_url", "image_url": {"url": url}})
+                if item.dimension_note:
+                    parts.append({"type": "text", "text": item.dimension_note})
+    if len(parts) == 1 and parts[0]["type"] == "text":
+        return parts[0]["text"]
+    return parts
+
+
+def openai_assistant_content(content_items: list) -> tuple[str | None, list[dict[str, Any]]]:
+    """Convert assistant message contents to OpenAI chat format (completions/copilot)."""
+    from tau.message.types import TextContent, ToolCallContent
+    text_parts: list[str] = []
+    tool_calls: list[dict[str, Any]] = []
+    for item in content_items:
+        match item:
+            case TextContent():
+                text_parts.append(item.content)
+            case ToolCallContent():
+                tool_calls.append({
+                    "id": item.id,
+                    "type": "function",
+                    "function": {"name": item.name, "arguments": json.dumps(item.args)},
+                })
+    return "".join(text_parts) or None, tool_calls
+
+
+def openai_response_format(response_format: Any | None) -> dict[str, Any] | None:
+    """Convert response_format to OpenAI json_schema format (completions/copilot/mistral)."""
+    from tau.inference.types import normalize_structured_response_format
+    structured = normalize_structured_response_format(response_format)
+    if structured is None:
+        return None
+    return {
+        "type": "json_schema",
+        "json_schema": {
+            "name": structured.name,
+            "schema": structured.schema,
+            "strict": structured.strict,
+        },
+    }
+
+
+def openai_messages_to_chat(messages: list) -> list[dict[str, Any]]:
+    """Convert a message list to OpenAI chat completions format."""
+    from tau.message.types import (
+        SystemMessage, UserMessage, AssistantMessage, ToolMessage,
+        TextContent, ToolResultContent,
+    )
+    result: list[dict[str, Any]] = []
+    for msg in messages:
+        match msg:
+            case SystemMessage():
+                text = "\n".join(c.content for c in msg.contents if isinstance(c, TextContent))
+                result.append({"role": "system", "content": text})
+            case UserMessage():
+                if not msg.contents:
+                    continue
+                result.append({"role": "user", "content": openai_user_content(msg.contents)})
+            case AssistantMessage():
+                text, tool_calls = openai_assistant_content(msg.contents)
+                entry: dict[str, Any] = {"role": "assistant"}
+                if text is not None:
+                    entry["content"] = text
+                if tool_calls:
+                    entry["tool_calls"] = tool_calls
+                result.append(entry)
+            case ToolMessage():
+                for content in msg.contents:
+                    if isinstance(content, ToolResultContent):
+                        result.append({
+                            "role": "tool",
+                            "tool_call_id": content.id,
+                            "content": content.content,
+                        })
+    return result
+
+
+def anthropic_messages_to_list(messages: list) -> tuple[str | None, list[dict[str, Any]]]:
+    """Convert a message list to Anthropic Messages API format."""
+    from tau.message.types import (
+        SystemMessage, UserMessage, AssistantMessage, ToolMessage,
+        TextContent, ImageContent, ThinkingContent, ToolCallContent, ToolResultContent,
+    )
+    system: str | None = None
+    result: list[dict[str, Any]] = []
+    for msg in messages:
+        match msg:
+            case SystemMessage():
+                system = "\n".join(c.content for c in msg.contents if isinstance(c, TextContent))
+            case UserMessage():
+                if not msg.contents:
+                    continue
+                parts: list[dict[str, Any]] = []
+                has_text = False
+                has_image = False
+                for item in msg.contents:
+                    match item:
+                        case TextContent():
+                            has_text = True
+                            parts.append({"type": "text", "text": item.content})
+                        case ImageContent():
+                            has_image = True
+                            for b64, mime in item.to_base64():
+                                parts.append({
+                                    "type": "image",
+                                    "source": {"type": "base64", "media_type": mime or "image/png", "data": b64},
+                                })
+                            if item.dimension_note:
+                                parts.append({"type": "text", "text": item.dimension_note})
+                if has_image and not has_text:
+                    parts.append({"type": "text", "text": "(see attached image)"})
+                result.append({"role": "user", "content": parts})
+            case AssistantMessage():
+                parts = []
+                for item in msg.contents:
+                    match item:
+                        case TextContent():
+                            parts.append({"type": "text", "text": item.content})
+                        case ThinkingContent():
+                            parts.append({"type": "thinking", "thinking": item.content, "signature": item.signature})
+                        case ToolCallContent():
+                            parts.append({"type": "tool_use", "id": item.id, "name": item.name, "input": item.args})
+                result.append({"role": "assistant", "content": parts})
+            case ToolMessage():
+                tool_results = []
+                for content in msg.contents:
+                    if isinstance(content, ToolResultContent):
+                        tool_results.append({
+                            "type": "tool_result",
+                            "tool_use_id": content.id,
+                            "content": content.content,
+                            "is_error": content.is_error,
+                        })
+                if tool_results:
+                    result.append({"role": "user", "content": tool_results})
+    return system, result
+
+
+def anthropic_output_config(response_format: Any | None) -> dict[str, Any] | None:
+    """Convert response_format to Anthropic output config format."""
+    from tau.inference.types import normalize_structured_response_format
+    structured = normalize_structured_response_format(response_format)
+    if structured is None:
+        return None
+    return {"format": {"type": "json_schema", "schema": structured.schema}}

tau/inference/api/video/base.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from tau.inference.model.types import Model
+from tau.inference.types import GeneratedVideo, VideoContext, VideoOptions
+
+
+class BaseVideoAPI(ABC):
+    """
+    Base class for all video generation providers.
+
+    Video generation is inherently async at the provider level — jobs are
+    submitted and polled until completion. Implementations handle the full
+    submit → poll → return cycle inside generate(), keeping the caller
+    interface simple and uniform.
+    """
+
+    def __init__(self, options: VideoOptions) -> None:
+        self.options = options
+
+    @abstractmethod
+    async def generate(self, model: Model, context: VideoContext) -> GeneratedVideo:
+        """Generate a video from text or image prompts."""
+        raise NotImplementedError
+

tau/inference/api/video/builtins.py

@@ -0,0 +1,7 @@
+# Lazy registry entries: each provider API is referenced by a "module:ClassName"
+# path so its SDK is only imported when that provider is actually used, not at
+# registry construction.
+VIDEO_APIS: list[tuple[str, str]] = [
+    ("fal-video", "tau.inference.api.video.fal_video:FalVideoAPI"),
+    ("openrouter-video", "tau.inference.api.video.openrouter_video:OpenRouterVideoAPI"),
+]