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

tau/inference/api/text/gemini_generate.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+import json
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
+from google import genai
+from google.genai import types as genai_types
+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,
+    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
+
+_STOP_REASON: dict[str, StopReason] = {
+    "STOP": StopReason.Stop,
+    "MAX_TOKENS": StopReason.Length,
+    "SAFETY": StopReason.ContentFilter,
+    "RECITATION": StopReason.ContentFilter,
+}
+
+
+def _messages_to_gemini(
+    messages: list[LLMMessage],
+) -> tuple[str | None, list[genai_types.Content]]:
+    system: str | None = None
+    contents: list[genai_types.Content] = []
+
+    for msg in messages:
+        match msg:
+            case SystemMessage():
+                system = "\n".join(c.content for c in msg.contents if isinstance(c, TextContent))
+            case UserMessage():
+                parts: list[genai_types.Part] = []
+                for item in msg.contents:
+                    match item:
+                        case TextContent():
+                            parts.append(genai_types.Part(text=item.content))
+                        case ImageContent():
+                            for b64, mime in item.to_base64():
+                                parts.append(genai_types.Part(
+                                    inline_data=genai_types.Blob(mime_type=mime or "image/png", data=b64),
+                                ))
+                if parts:
+                    contents.append(genai_types.Content(role="user", parts=parts))
+            case AssistantMessage():
+                parts = []
+                for item in msg.contents:
+                    match item:
+                        case TextContent():
+                            parts.append(genai_types.Part(text=item.content))
+                        case ToolCallContent():
+                            parts.append(genai_types.Part(
+                                function_call=genai_types.FunctionCall(
+                                    name=item.name,
+                                    args=item.args,
+                                ),
+                            ))
+                if parts:
+                    contents.append(genai_types.Content(role="model", parts=parts))
+            case ToolMessage():
+                parts = []
+                for content in msg.contents:
+                    if isinstance(content, ToolResultContent):
+                        parts.append(genai_types.Part(
+                            function_response=genai_types.FunctionResponse(
+                                name=content.id,
+                                response={"result": content.content},
+                            ),
+                        ))
+                if parts:
+                    contents.append(genai_types.Content(role="user", parts=parts))
+
+    return system, contents
+
+
+def _response_schema(response_format: Any | None) -> dict[str, Any] | None:
+    structured = normalize_structured_response_format(response_format)
+    return structured.schema if structured is not None else None
+
+
+class GeminiGenerateAPI(BaseAPI):
+    def __init__(self, options: LLMOptions) -> None:
+        super().__init__(options)
+        self._client = genai.Client(api_key=options.api_key)
+
+    def _build_config(
+        self,
+        tools: Optional[list[Tool]] = None,
+        response_format: Any | None = None,
+    ) -> genai_types.GenerateContentConfig:
+        params: dict[str, Any] = {
+            "temperature": self.options.temperature,
+        }
+        if self.options.max_tokens is not None:
+            params["max_output_tokens"] = self.options.max_tokens
+        schema = _response_schema(response_format)
+        if schema is not None:
+            params["response_mime_type"] = "application/json"
+            params["response_schema"] = schema
+
+        budget = None
+        if self.options.thinking_level is not None and self.options.thinking_level != ThinkingLevel.Off:
+            budgets = self.options.thinking_budgets or ThinkingBudgets()
+            budget = budgets.get(self.options.thinking_level)
+        if budget is not None:
+            params["thinking_config"] = genai_types.ThinkingConfig(
+                thinking_budget=budget,
+                include_thoughts=True,
+            )
+        
+        if tools:
+            params["tools"] = [
+                genai_types.Tool(
+                    function_declarations=[
+                        genai_types.FunctionDeclaration(
+                            name=t.name,
+                            description=t.description,
+                            parameters=t.schema.model_json_schema(),
+                        )
+                        for t in tools
+                    ]
+                )
+            ]
+
+        return genai_types.GenerateContentConfig(**params)
+
+    async def stream(self, context: LLMContext, model: Model) -> AsyncGenerator[LLMEvent, None]:  # type: ignore[override]
+        system, contents = _messages_to_gemini(context.messages)
+        config = self._build_config(
+            tools=context.tools or None,
+            response_format=context.response_format,
+        )
+        effective_system = context.system_prompt or system
+        if effective_system:
+            config.system_instruction = effective_system
+
+        if self.options.on_payload:
+            payload = {"config": config, "contents": contents}
+            modified = self.options.on_payload(payload)
+            if modified is not None:
+                config = modified.get("config", config)
+                contents = modified.get("contents", contents)
+
+        text_index = 0
+        thinking_index = 0
+        tool_index = 0
+        text_started = False
+        thinking_started = False
+        text_buf = ""
+        thinking_buf = ""
+        _input_tokens = 0
+        _output_tokens = 0
+        _cache_read_tokens = 0
+
+        yield StartEvent()
+
+        try:
+            async for chunk in await self._client.aio.models.generate_content_stream(
+                model=model.id,
+                contents=contents,
+                config=config,
+            ):
+                if self._cancelled():
+                    yield ErrorEvent(reason=StopReason.Abort, error="Cancelled")
+                    return
+                um = getattr(chunk, 'usage_metadata', None)
+                if um:
+                    _input_tokens = getattr(um, 'prompt_token_count', 0) or 0
+                    _output_tokens = getattr(um, 'candidates_token_count', 0) or 0
+                    _cache_read_tokens = getattr(um, 'cached_content_token_count', 0) or 0
+
+                if not chunk.candidates:
+                    continue
+
+                candidate = chunk.candidates[0]
+                if candidate.content and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if getattr(part, "thought", False) and part.text:
+                            if not thinking_started:
+                                yield ThinkingStartEvent(thinking=None)
+                                thinking_started = True
+                            thinking_buf += part.text
+                            yield ThinkingDeltaEvent(thinking=ThinkingContent(content=part.text))
+                        elif part.text:
+                            if thinking_started:
+                                yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+                                thinking_started = False
+                                thinking_index += 1
+                                thinking_buf = ""
+                            if not text_started:
+                                yield TextStartEvent(text=TextContent(content=""))
+                                text_started = True
+                            text_buf += part.text
+                            yield TextDeltaEvent(text=TextContent(content=part.text))
+                        elif part.function_call:
+                            fc = part.function_call
+                            tool_id = fc.name
+                            args_str = json.dumps(dict(fc.args)) if fc.args else ""
+                            yield ToolCallStartEvent(tool_call=ToolCallContent(id=tool_id, name=fc.name))
+                            yield ToolCallDeltaEvent(tool_call=ToolCallContent(id=tool_id))
+                            yield ToolCallEndEvent(tool_call=ToolCallContent(id=tool_id, name=fc.name, args=json.loads(args_str) if args_str else {}))
+                            tool_index += 1
+
+                finish_reason = getattr(candidate, "finish_reason", None)
+                if finish_reason and str(finish_reason) not in ("", "FINISH_REASON_UNSPECIFIED"):
+                    if thinking_started:
+                        yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+                    if text_started:
+                        yield TextEndEvent(text=TextContent(content=text_buf))
+                    reason_str = finish_reason.name if hasattr(finish_reason, "name") else str(finish_reason)
+                    stop = StopReason.ToolCalls if tool_index > 0 else _STOP_REASON.get(reason_str, StopReason.Stop)
+                    yield EndEvent(reason=stop, input_tokens=_input_tokens, output_tokens=_output_tokens, cache_read_tokens=_cache_read_tokens)
+                    return
+
+        except Exception as exc:
+            yield ErrorEvent(reason=StopReason.Abort, error=str(exc))
+            return
+
+        if thinking_started:
+            yield ThinkingEndEvent(thinking=ThinkingContent(content=thinking_buf))
+        if text_started:
+            yield TextEndEvent(text=TextContent(content=text_buf))
+        yield EndEvent(reason=StopReason.Stop, input_tokens=_input_tokens, output_tokens=_output_tokens, cache_read_tokens=_cache_read_tokens)

tau/inference/api/text/github_copilot_chat.py

@@ -0,0 +1,172 @@
+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.provider.oauth.github_copilot import get_copilot_base_url
+from tau.inference.types import (
+    LLMContext, LLMEvent, LLMOptions, StopReason, ThinkingLevel,
+    StartEvent, EndEvent, ErrorEvent,
+    TextStartEvent, TextDeltaEvent, TextEndEvent,
+    ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent,
+)
+from tau.message.types import (
+    SystemMessage, UserMessage, AssistantMessage, ToolMessage,
+    TextContent, ImageContent, ToolCallContent, ToolResultContent,
+)
+from typing import Optional, TYPE_CHECKING
+if TYPE_CHECKING:
+    from tau.tool.types import Tool
+
+_COPILOT_HEADERS = {
+    "User-Agent": "GitHubCopilotChat/0.35.0",
+    "Editor-Version": "vscode/1.107.0",
+    "Editor-Plugin-Version": "copilot-chat/0.35.0",
+    "Copilot-Integration-Id": "vscode-chat",
+}
+
+_STOP_REASON: dict[str, StopReason] = {
+    "stop": StopReason.Stop,
+    "length": StopReason.Length,
+    "tool_calls": StopReason.ToolCalls,
+    "content_filter": StopReason.ContentFilter,
+}
+
+
+
+class GitHubCopilotChatAPI(BaseAPI):
+    """Streaming LLM API adapter for the GitHub Copilot Chat endpoint (OpenAI-compatible)."""
+
+    def __init__(self, options: LLMOptions) -> None:
+        """Resolve the Copilot base URL and initialise the AsyncOpenAI client with Copilot headers."""
+        super().__init__(options)
+        base_url = options.base_url or get_copilot_base_url(options.api_key)
+        self._client = AsyncOpenAI(
+            api_key=options.api_key or "github-copilot",
+            base_url=base_url,
+            default_headers={**_COPILOT_HEADERS, **(options.headers or {})},
+            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 Copilot 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 tools:
+            params["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": tool.name,
+                        "description": tool.description,
+                        "parameters": 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 GitHub Copilot Chat API."""
+        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 = ""
+        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}) 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
+
+                if delta.content:
+                    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:
+                    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 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)