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

tau/inference/api/text/openai_completions.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+import json
+from tau.inference.api.text.utils import parse_tool_args, openai_user_content, openai_assistant_content, openai_messages_to_chat, openai_response_format
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+from openai import AsyncOpenAI
+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, 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
+
+_REASONING_EFFORT: dict[ThinkingLevel, str] = {
+    ThinkingLevel.Minimal: "low",
+    ThinkingLevel.Low: "low",
+    ThinkingLevel.Medium: "medium",
+    ThinkingLevel.High: "high",
+    ThinkingLevel.XHigh: "high",
+    ThinkingLevel.Max: "high",
+}
+
+_STOP_REASON: dict[str, StopReason] = {
+    "stop": StopReason.Stop,
+    "length": StopReason.Length,
+    "tool_calls": StopReason.ToolCalls,
+    "content_filter": StopReason.ContentFilter,
+}
+
+
+def _clean_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    """Strip fields that trip up strict OpenAI-compatible APIs (title, $defs, etc.)."""
+    result: dict[str, Any] = {}
+    for k, v in schema.items():
+        if k in ("title", "$schema"):
+            continue
+        if k == "anyOf" and isinstance(v, list):
+            non_null = [_clean_schema(s) if isinstance(s, dict) else s for s in v if s != {"type": "null"}]
+            if len(non_null) == 1:
+                result.update(non_null[0])
+            else:
+                result[k] = non_null
+        elif isinstance(v, dict):
+            result[k] = _clean_schema(v)
+        elif isinstance(v, list):
+            result[k] = [_clean_schema(i) if isinstance(i, dict) else i for i in v]
+        else:
+            result[k] = v
+    return result
+
+
+
+class OpenAICompletionsAPI(BaseAPI):
+    """Streaming LLM API adapter for the OpenAI Chat Completions endpoint."""
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Initialise the AsyncOpenAI client with the supplied options."""
+        super().__init__(options)
+        self._client = AsyncOpenAI(
+            api_key=options.api_key or "placeholder",
+            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, messages: list[dict[str, Any]], tools: Optional[list[Tool]] = None) -> dict[str, Any]:
+        """Assemble the OpenAI Chat Completions request payload."""
+        params: dict[str, Any] = {
+            "model": model.id,
+            "messages": messages,
+            "temperature": self.options.temperature,
+        }
+        if self.options.max_tokens is not None:
+            params["max_completion_tokens"] = self.options.max_tokens
+        if self.options.thinking_level is not None and self.options.thinking_level != ThinkingLevel.Off:
+            params["reasoning_effort"] = _REASONING_EFFORT[self.options.thinking_level]
+
+        if tools:
+            params["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": tool.name,
+                        "description": tool.description,
+                        "parameters": _clean_schema(tool.schema.model_json_schema()),
+                    }
+                }
+                for tool in tools
+            ]
+            params["tool_choice"] = "auto"
+
+        return params
+
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:  # type: ignore[override]
+        """Stream LLMEvents from the OpenAI Chat Completions API."""
+        if self.options.api_key:
+            self._client.api_key = self.options.api_key
+        chat_messages = openai_messages_to_chat(context.messages)
+        if context.system_prompt:
+            chat_messages = [{"role": "system", "content": context.system_prompt}] + chat_messages
+        params = self._build_params(model, chat_messages, tools=context.tools or None)
+        response_format = openai_response_format(context.response_format)
+        if response_format is not None:
+            params["response_format"] = response_format
+
+        if self.options.on_payload:
+            modified = self.options.on_payload(params)
+            if modified is not None:
+                params = modified
+
+        text_started = False
+        text_buf = ""
+        thinking_started = False
+        thinking_buf = ""
+        # Tool-call accumulation state keyed by delta index (OpenAI streams partial tool calls per-index).
+        tool_started: dict[int, bool] = {}
+        tool_bufs: dict[int, str] = {}
+        tool_meta: dict[int, dict[str, str]] = {}
+        _input_tokens = 0
+        _output_tokens = 0
+        _cache_read_tokens = 0
+
+        yield StartEvent()
+
+        # async with closes the SDK stream (and its httpx response) on every
+        # exit path — cancellation return or an upstream GeneratorExit — instead
+        # of leaving it to the GC asyncgen finalizer.
+        async with await self._client.chat.completions.create(**params, stream=True, stream_options={"include_usage": True}, extra_body=self.options.extra_params or {}) as sdk_stream:
+            async for chunk in sdk_stream:
+                if self._cancelled():
+                    yield ErrorEvent(reason=StopReason.Abort, error="Cancelled")
+                    return
+                usage_data = getattr(chunk, 'usage', None)
+                if usage_data:
+                    _input_tokens = getattr(usage_data, 'prompt_tokens', 0) or 0
+                    _output_tokens = getattr(usage_data, 'completion_tokens', 0) or 0
+                    _details = getattr(usage_data, 'prompt_tokens_details', None)
+                    _cache_read_tokens = getattr(_details, 'cached_tokens', 0) or 0
+                choice = chunk.choices[0] if chunk.choices else None
+                if choice is None:
+                    continue
+
+                delta = choice.delta
+
+                # Handle reasoning/thinking content (often used by NVIDIA and some OpenAI models)
+                reasoning = getattr(delta, "reasoning_content", None) or getattr(delta, "thinking", None)
+                if reasoning:
+                    if not thinking_started:
+                        yield ThinkingStartEvent(thinking=ThinkingContent(content=""))
+                        thinking_started = True
+                    thinking_buf += reasoning
+                    yield ThinkingDeltaEvent(thinking=ThinkingContent(content=reasoning))
+
+                if delta.content:
+                    # If thinking was happening, end it before starting text
+                    if thinking_started:
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+                        thinking_started = False
+                        thinking_buf = ""
+
+                    if not text_started:
+                        yield TextStartEvent(text=TextContent(content=""))
+                        text_started = True
+                    text_buf += delta.content
+                    yield TextDeltaEvent(text=TextContent(content=delta.content))
+
+                if delta.tool_calls:
+                    # If thinking was happening, end it
+                    if thinking_started:
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+                        thinking_started = False
+                        thinking_buf = ""
+
+                    for tc in delta.tool_calls:
+                        idx = tc.index
+                        if idx not in tool_started:
+                            tool_started[idx] = True
+                            tool_bufs[idx] = ""
+                            tool_meta[idx] = {"id": tc.id or "", "name": tc.function.name or "" if tc.function else ""}
+                            yield ToolCallStartEvent(tool_call=ToolCallContent(
+                                    id=tool_meta[idx]["id"],
+                                    name=tool_meta[idx]["name"],
+                                )
+                            )
+                        if tc.function and tc.function.arguments:
+                            tool_bufs[idx] += tc.function.arguments
+                            yield ToolCallDeltaEvent(tool_call=ToolCallContent(id=tool_meta[idx]["id"])
+                            )
+
+                if choice.finish_reason:
+                    if thinking_started:
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+                        thinking_started = False
+                        thinking_buf = ""
+
+                    if text_started:
+                        yield TextEndEvent(text=TextContent(content=text_buf))
+                        text_started = False
+                        text_buf = ""
+
+                    for idx in sorted(tool_started):
+                        args_str = tool_bufs[idx].strip()
+                        args = parse_tool_args(args_str)
+
+                        yield ToolCallEndEvent(tool_call=ToolCallContent(
+                                id=tool_meta[idx]["id"],
+                                name=tool_meta[idx]["name"],
+                                args=args,
+                            )
+                        )
+                    tool_started.clear()
+                    tool_bufs.clear()
+                    tool_meta.clear()
+
+                    stop_reason = _STOP_REASON.get(choice.finish_reason, StopReason.Stop)
+                    yield EndEvent(reason=stop_reason, input_tokens=_input_tokens, output_tokens=_output_tokens, cache_read_tokens=_cache_read_tokens)

tau/inference/api/text/openai_responses.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+import json
+from tau.inference.api.text.utils import parse_tool_args
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+from openai import AsyncOpenAI
+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, ThinkingLevel,
+    StartEvent, EndEvent, ErrorEvent,
+    TextStartEvent, TextDeltaEvent, TextEndEvent,
+    ThinkingStartEvent, ThinkingDeltaEvent, ThinkingEndEvent,
+    ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent,
+    normalize_structured_response_format,
+)
+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
+
+_THINKING_EFFORT: dict[ThinkingLevel, str] = {
+    ThinkingLevel.Minimal: "low",
+    ThinkingLevel.Low: "low",
+    ThinkingLevel.Medium: "medium",
+    ThinkingLevel.High: "high",
+    ThinkingLevel.XHigh: "high",
+    ThinkingLevel.Max: "high",
+}
+
+_STOP_REASON: dict[str, StopReason] = {
+    "stop": StopReason.Stop,
+    "max_output_tokens": StopReason.Length,
+    "tool_calls": StopReason.ToolCalls,
+    "content_filter": StopReason.ContentFilter,
+}
+
+
+def _content_to_openai(content_items: list) -> list[dict[str, Any]]:
+    """Convert typed message content items to OpenAI Responses API content parts."""
+    parts: list[dict[str, Any]] = []
+    for item in content_items:
+        match item:
+            case TextContent():
+                parts.append({"type": "input_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": "input_image", "image_url": url})
+            case ThinkingContent():
+                parts.append({
+                    "type": "thinking",
+                    "thinking": item.content,
+                    "signature": item.signature,
+                })
+            case ToolCallContent():
+                parts.append({
+                    "type": "function_call",
+                    "call_id": item.id,
+                    "name": item.name,
+                    "arguments": json.dumps(item.args),
+                })
+    return parts
+
+
+def _messages_to_input(
+    messages: list[LLMMessage],
+) -> tuple[str | None, list[dict[str, Any]]]:
+    """Convert a message list to OpenAI Responses API input items, extracting system as instructions."""
+    instructions: str | None = None
+    input_items: list[dict[str, Any]] = []
+
+    for msg in messages:
+        match msg:
+            case SystemMessage():
+                text_parts = [c.content for c in msg.contents if isinstance(c, TextContent)]
+                instructions = "\n".join(text_parts)
+            case ToolMessage():
+                for content in msg.contents:
+                    if isinstance(content, ToolResultContent):
+                        input_items.append({
+                            "type": "function_call_output",
+                            "call_id": content.id,
+                            "output": content.content,
+                        })
+            case UserMessage() | AssistantMessage():
+                role = "user" if isinstance(msg, UserMessage) else "assistant"
+                parts = _content_to_openai(msg.contents)
+                if parts:
+                    input_items.append({"role": role, "content": parts})
+
+    return instructions, input_items
+
+
+def _text_format(response_format: Any | None) -> dict[str, Any] | None:
+    """Convert response_format to the OpenAI Responses API text.format structure."""
+    structured = normalize_structured_response_format(response_format)
+    if structured is None:
+        return None
+    return {
+        "format": {
+            "type": "json_schema",
+            "name": structured.name,
+            "schema": structured.schema,
+            "strict": structured.strict,
+        }
+    }
+
+
+class OpenAIResponsesAPI(BaseAPI):
+    """Streaming LLM API adapter for the OpenAI Responses API (o-series / GPT-4o)."""
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Initialise the AsyncOpenAI client with the supplied options."""
+        super().__init__(options)
+        self._client = AsyncOpenAI(
+            api_key=options.api_key or "placeholder",
+            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, instructions: str | None, input_items: list, tools: Optional[list[Tool]] = None) -> dict[str, Any]:
+        """Assemble the OpenAI Responses API request payload."""
+        params: dict[str, Any] = {
+            "model": model.id,
+            "input": input_items,
+            "temperature": self.options.temperature,
+        }
+        if instructions:
+            params["instructions"] = instructions
+        if self.options.max_tokens is not None:
+            params["max_output_tokens"] = self.options.max_tokens
+        if self.options.thinking_level is not None and self.options.thinking_level != ThinkingLevel.Off:
+            params["reasoning"] = {"effort": _THINKING_EFFORT[self.options.thinking_level]}
+        
+        if tools:
+            params["tools"] = [
+                {
+                    "type": "function",
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.schema.model_json_schema(),
+                }
+                for tool in tools
+            ]
+
+        return params
+
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:  # type: ignore[override]
+        """Stream LLMEvents from the OpenAI Responses API."""
+        if self.options.api_key:
+            self._client.api_key = self.options.api_key
+        instructions, input_items = _messages_to_input(context.messages)
+        params = self._build_params(model, instructions, input_items, tools=context.tools or None)
+        text_format = _text_format(context.response_format)
+        if text_format is not None:
+            params["text"] = text_format
+
+        if self.options.on_payload:
+            modified = self.options.on_payload(params)
+            if modified is not None:
+                params = modified
+
+        tool_names: dict[str, str] = {}
+        _input_tokens = 0
+        _output_tokens = 0
+        _cache_read_tokens = 0
+
+        yield StartEvent()
+
+        async with self._client.responses.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 == "response.output_item.added":
+                    item = event.item
+                    if item.type == "message":
+                        yield TextStartEvent(text=TextContent(content=""))
+                    elif item.type == "reasoning":
+                        yield ThinkingStartEvent(thinking=None)
+                    elif item.type == "function_call":
+                        tool_names[item.call_id] = item.name
+                        yield ToolCallStartEvent(tool_call=ToolCallContent(id=item.call_id, name=item.name))
+
+                elif etype == "response.output_text.delta":
+                    yield TextDeltaEvent(text=TextContent(content=event.delta))
+
+                elif etype == "response.output_text.done":
+                    yield TextEndEvent(text=TextContent(content=event.text))
+
+                elif etype == "response.reasoning_summary_text.delta":
+                    yield ThinkingDeltaEvent(thinking=ThinkingContent(content=event.delta))
+
+                elif etype == "response.reasoning_summary_text.done":
+                    yield ThinkingEndEvent(thinking=ThinkingContent(content=event.text))
+
+                elif etype == "response.function_call_arguments.delta":
+                    call_id = event.item_id
+                    yield ToolCallDeltaEvent(tool_call=ToolCallContent(id=call_id))
+
+                elif etype == "response.function_call_arguments.done":
+                    call_id = event.item_id
+                    args_str = event.arguments.strip()
+                    args = parse_tool_args(args_str)
+
+                    yield ToolCallEndEvent(tool_call=ToolCallContent(
+                        id=call_id,
+                        name=tool_names.get(call_id, ""),
+                        args=args
+                    ))
+
+                elif etype == "response.done":
+                    resp = event.response
+                    u = getattr(resp, 'usage', None)
+                    if u:
+                        _input_tokens = getattr(u, 'input_tokens', 0) or 0
+                        _output_tokens = getattr(u, 'output_tokens', 0) or 0
+                        _details = getattr(u, 'prompt_tokens_details', None)
+                        _cache_read_tokens = getattr(_details, 'cached_tokens', 0) or 0
+                    stop_reason = _STOP_REASON.get(
+                        getattr(resp, "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)
+
+                elif etype == "error":
+                    yield ErrorEvent(reason=StopReason.Abort, error=str(event))

tau/inference/api/text/registry.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+import importlib
+from typing import Type, Union
+from tau.inference.api.text.base import BaseLLMAPI
+
+_Entry = Union[Type[BaseLLMAPI], str]
+
+
+class LLMAPIRegistry:
+    def __init__(self) -> None:
+        self._apis: dict[str, _Entry] = {}
+
+    def register(self, name: str, api: _Entry) -> None:
+        """Register an API class or a lazy 'module:ClassName' import path."""
+        self._apis[name] = api
+
+    def unregister(self, name: str) -> None:
+        self._apis.pop(name, None)
+
+    def list(self) -> list[Type[BaseLLMAPI]]:
+        return [self._resolve(name) for name in list(self._apis)]
+
+    def get(self, name: str) -> Type[BaseLLMAPI] | None:
+        if name not in self._apis:
+            return None
+        return self._resolve(name)
+
+    def reset(self) -> None:
+        self._apis.clear()
+
+    def _resolve(self, name: str) -> Type[BaseLLMAPI]:
+        entry = self._apis[name]
+        if isinstance(entry, str):
+            mod_path, _, cls_name = entry.partition(":")
+            cls = getattr(importlib.import_module(mod_path), cls_name)
+            self._apis[name] = cls
+            return cls
+        return entry
+
+    @classmethod
+    def from_builtins(cls) -> LLMAPIRegistry:
+        from tau.inference.api.text.builtins import LLM_APIS
+        instance = cls()
+        for name, api in LLM_APIS:
+            instance.register(name, api)
+        return instance
+
+
+# Backward-compat alias
+APIRegistry = LLMAPIRegistry