devcopilot 0.2.0__py3-none-any.whl

providers/transports/openai_chat/stream.py

@@ -0,0 +1,384 @@
+"""Per-request OpenAI-chat stream runner."""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+from loguru import logger
+
+from core.anthropic import (
+    ContentType,
+    HeuristicToolParser,
+    SSEBuilder,
+    ThinkTagParser,
+    map_stop_reason,
+)
+from core.anthropic.stream_recovery import TruncatedProviderStreamError
+from core.anthropic.stream_recovery_session import (
+    StreamFailureAction,
+    StreamRecoverySession,
+)
+from core.trace import provider_chat_body_snapshot, trace_event
+from providers.error_mapping import map_error
+
+from .recovery import OpenAIChatRecovery
+from .tool_calls import (
+    OpenAIToolCallAssembler,
+    all_started_tools_complete,
+    has_committed_sse_output,
+    iter_heuristic_tool_use_sse,
+    tool_call_extra_content,
+)
+
+
+class OpenAIChatStreamRunner:
+    """Own mutable state for one OpenAI-chat provider stream."""
+
+    def __init__(
+        self,
+        transport: Any,
+        *,
+        request: Any,
+        input_tokens: int,
+        request_id: str | None,
+        thinking_enabled: bool | None,
+    ) -> None:
+        self._transport = transport
+        self._request = request
+        self._input_tokens = input_tokens
+        self._request_id = request_id
+        self._thinking_enabled = thinking_enabled
+        self._message_id = f"msg_{uuid.uuid4()}"
+        self._tool_calls = OpenAIToolCallAssembler(
+            record_extra_content=transport._record_tool_call_extra_content
+        )
+        self._recovery = OpenAIChatRecovery(
+            provider_name=transport._provider_name,
+            create_stream=transport._create_stream,
+        )
+
+    async def run(self) -> AsyncIterator[str]:
+        """Stream response in Anthropic SSE format."""
+        tag = self._transport._provider_name
+        req_tag = f" request_id={self._request_id}" if self._request_id else ""
+        sse = self._new_sse_builder()
+        recovery_session = StreamRecoverySession(
+            provider_name=tag,
+            request_id=self._request_id,
+        )
+
+        def hold_event(event: str) -> Iterator[str]:
+            yield from recovery_session.push(event)
+
+        def hold_events(events: Iterator[str]) -> Iterator[str]:
+            for event in events:
+                yield from hold_event(event)
+
+        body = self._transport._build_request_body(
+            self._request, thinking_enabled=self._thinking_enabled
+        )
+        thinking_enabled = self._transport._is_thinking_enabled(
+            self._request, self._thinking_enabled
+        )
+        trace_event(
+            stage="provider",
+            event="provider.request.sent",
+            source="provider",
+            provider=tag,
+            gateway_model=self._request.model,
+            downstream_model=body.get("model"),
+            message_count=len(body.get("messages", [])),
+            tool_count=len(body.get("tools", [])),
+            body=provider_chat_body_snapshot(body),
+        )
+
+        yield sse.message_start()
+
+        think_parser = ThinkTagParser()
+        heuristic_parser = HeuristicToolParser()
+        finish_reason = None
+        usage_info = None
+        tool_argument_aliases: dict[str, dict[str, str]] = {}
+        tool_argument_alias_buffers: dict[int, str] = {}
+
+        async with self._transport._global_rate_limiter.concurrency_slot():
+            while True:
+                stream_opened = False
+                try:
+                    stream, body = await self._transport._create_stream(body)
+                    stream_opened = True
+                    tool_argument_aliases = self._transport._tool_argument_aliases(body)
+                    async for chunk in stream:
+                        if getattr(chunk, "usage", None):
+                            usage_info = chunk.usage
+
+                        if not chunk.choices:
+                            continue
+
+                        choice = chunk.choices[0]
+                        delta = choice.delta
+                        if delta is None:
+                            continue
+
+                        if choice.finish_reason:
+                            finish_reason = choice.finish_reason
+                            logger.debug("{} finish_reason: {}", tag, finish_reason)
+
+                        reasoning = getattr(delta, "reasoning_content", None)
+                        if thinking_enabled and reasoning:
+                            for event in hold_events(sse.ensure_thinking_block()):
+                                yield event
+                            for event in hold_event(sse.emit_thinking_delta(reasoning)):
+                                yield event
+
+                        for event in self._transport._handle_extra_reasoning(
+                            delta,
+                            sse,
+                            thinking_enabled=thinking_enabled,
+                        ):
+                            for out_event in hold_event(event):
+                                yield out_event
+
+                        if delta.content:
+                            for part in think_parser.feed(delta.content):
+                                if part.type == ContentType.THINKING:
+                                    if not thinking_enabled:
+                                        continue
+                                    for event in hold_events(
+                                        sse.ensure_thinking_block()
+                                    ):
+                                        yield event
+                                    for event in hold_event(
+                                        sse.emit_thinking_delta(part.content)
+                                    ):
+                                        yield event
+                                else:
+                                    (
+                                        filtered_text,
+                                        detected_tools,
+                                    ) = heuristic_parser.feed(part.content)
+
+                                    if filtered_text:
+                                        for event in hold_events(
+                                            sse.ensure_text_block()
+                                        ):
+                                            yield event
+                                        for event in hold_event(
+                                            sse.emit_text_delta(filtered_text)
+                                        ):
+                                            yield event
+
+                                    for tool_use in detected_tools:
+                                        for event in iter_heuristic_tool_use_sse(
+                                            sse, tool_use
+                                        ):
+                                            for out_event in hold_event(event):
+                                                yield out_event
+
+                        if delta.tool_calls:
+                            for event in hold_events(sse.close_content_blocks()):
+                                yield event
+                            for tc in delta.tool_calls:
+                                extra_content = tool_call_extra_content(tc)
+                                tc_info = {
+                                    "index": tc.index,
+                                    "id": tc.id,
+                                    "function": {
+                                        "name": tc.function.name,
+                                        "arguments": tc.function.arguments,
+                                    },
+                                }
+                                if extra_content:
+                                    tc_info["extra_content"] = extra_content
+                                for event in self._tool_calls.process_tool_call(
+                                    tc_info,
+                                    sse,
+                                    tool_argument_aliases=tool_argument_aliases,
+                                    tool_argument_alias_buffers=tool_argument_alias_buffers,
+                                ):
+                                    for out_event in hold_event(event):
+                                        yield out_event
+
+                    if finish_reason is None:
+                        raise TruncatedProviderStreamError(
+                            "Provider stream ended without finish_reason."
+                        )
+                    break
+
+                except asyncio.CancelledError, GeneratorExit:
+                    raise
+                except Exception as error:
+                    generated_output = has_committed_sse_output(sse)
+                    complete_tool_salvageable = (
+                        generated_output
+                        and sse.blocks.has_emitted_tool_block()
+                        and all_started_tools_complete(sse, self._request)
+                    )
+                    decision = recovery_session.advance_failure(
+                        error,
+                        stream_opened=stream_opened,
+                        generated_output=generated_output,
+                        complete_tool_salvageable=complete_tool_salvageable,
+                    )
+                    if decision.action == StreamFailureAction.EARLY_RETRY:
+                        sse = self._new_sse_builder()
+                        think_parser = ThinkTagParser()
+                        heuristic_parser = HeuristicToolParser()
+                        finish_reason = None
+                        usage_info = None
+                        tool_argument_aliases = {}
+                        tool_argument_alias_buffers = {}
+                        continue
+
+                    if decision.action == StreamFailureAction.MIDSTREAM_RECOVERY:
+                        try:
+                            recovery_events = await self._recovery.events(
+                                body=body,
+                                sse=sse,
+                                request=self._request,
+                                request_id=self._request_id,
+                                error=error,
+                                tool_argument_alias_buffers=tool_argument_alias_buffers,
+                            )
+                        except Exception as recovery_error:
+                            trace_event(
+                                stage="provider",
+                                event="provider.recovery.failed",
+                                source="provider",
+                                provider=tag,
+                                request_id=self._request_id,
+                                exc_type=type(recovery_error).__name__,
+                            )
+                            recovery_events = None
+                        if recovery_events is not None:
+                            for event in recovery_session.flush_uncommitted(decision):
+                                yield event
+                            for event in recovery_events:
+                                yield event
+                            return
+
+                    self._transport._log_stream_transport_error(
+                        tag, req_tag, error, request_id=self._request_id
+                    )
+                    error_message = self._transport._openai_error_message(
+                        error, self._request_id
+                    )
+                    trace_event(
+                        stage="provider",
+                        event="provider.response.error",
+                        source="provider",
+                        provider=tag,
+                        error_message=error_message,
+                        mapped_error_type=type(
+                            map_error(
+                                error,
+                                rate_limiter=self._transport._global_rate_limiter,
+                            )
+                        ).__name__,
+                    )
+                    if not decision.committed and decision.has_buffered:
+                        for event in recovery_session.flush():
+                            yield event
+                    elif not decision.committed:
+                        recovery_session.discard()
+                        sse = self._new_sse_builder()
+                    for event in self._recovery.emit_error_tail(sse, error_message):
+                        yield event
+                    return
+
+        remaining = think_parser.flush()
+        if remaining:
+            if remaining.type == ContentType.THINKING:
+                if not thinking_enabled:
+                    remaining = None
+                else:
+                    for event in hold_events(sse.ensure_thinking_block()):
+                        yield event
+                    for event in hold_event(sse.emit_thinking_delta(remaining.content)):
+                        yield event
+            if remaining and remaining.type == ContentType.TEXT:
+                for event in hold_events(sse.ensure_text_block()):
+                    yield event
+                for event in hold_event(sse.emit_text_delta(remaining.content)):
+                    yield event
+
+        for tool_use in heuristic_parser.flush():
+            for event in iter_heuristic_tool_use_sse(sse, tool_use):
+                for out_event in hold_event(event):
+                    yield out_event
+
+        has_started_tool = any(s.started for s in sse.blocks.tool_states.values())
+        has_content_blocks = (
+            sse.blocks.text_index != -1
+            or sse.blocks.thinking_index != -1
+            or has_started_tool
+        )
+        if not has_content_blocks or (
+            not has_started_tool
+            and not sse.accumulated_text.strip()
+            and sse.accumulated_reasoning.strip()
+        ):
+            for event in hold_events(sse.ensure_text_block()):
+                yield event
+            for event in hold_event(sse.emit_text_delta(" ")):
+                yield event
+
+        for event in self._tool_calls.flush_tool_argument_alias_buffers(
+            sse, tool_argument_aliases, tool_argument_alias_buffers
+        ):
+            for out_event in hold_event(event):
+                yield out_event
+
+        for event in self._tool_calls.flush_task_arg_buffers(sse):
+            for out_event in hold_event(event):
+                yield out_event
+
+        for event in hold_events(sse.close_all_blocks()):
+            yield event
+
+        completion = (
+            getattr(usage_info, "completion_tokens", None)
+            if usage_info is not None
+            else None
+        )
+        if isinstance(completion, int):
+            output_tokens = completion
+        else:
+            output_tokens = sse.estimate_output_tokens()
+        if usage_info and hasattr(usage_info, "prompt_tokens"):
+            provider_input = usage_info.prompt_tokens
+            if isinstance(provider_input, int):
+                logger.debug(
+                    "TOKEN_ESTIMATE: our={} provider={} diff={:+d}",
+                    self._input_tokens,
+                    provider_input,
+                    provider_input - self._input_tokens,
+                )
+        trace_event(
+            stage="provider",
+            event="provider.response.completed",
+            source="provider",
+            provider=tag,
+            finish_reason=(None if finish_reason is None else str(finish_reason)),
+            output_tokens=output_tokens,
+            prompt_tokens_estimate=self._input_tokens,
+        )
+        for event in hold_event(
+            sse.message_delta(map_stop_reason(finish_reason), output_tokens)
+        ):
+            yield event
+        for event in hold_event(sse.message_stop()):
+            yield event
+        for event in recovery_session.flush():
+            yield event
+
+    def _new_sse_builder(self) -> SSEBuilder:
+        return SSEBuilder(
+            self._message_id,
+            self._request.model,
+            self._input_tokens,
+            log_raw_events=self._transport._config.log_raw_sse_events,
+        )

providers/transports/openai_chat/tool_calls.py

@@ -0,0 +1,293 @@
+"""OpenAI-chat tool-call assembly helpers."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Callable, Iterator
+from typing import Any
+
+from core.anthropic import SSEBuilder
+from core.anthropic.stream_recovery import (
+    parse_complete_tool_input,
+    tool_schemas_by_name,
+)
+
+RecordToolExtraContent = Callable[[str, dict[str, Any]], None]
+
+
+def iter_heuristic_tool_use_sse(
+    sse: SSEBuilder, tool_use: dict[str, Any]
+) -> Iterator[str]:
+    """Emit SSE for one heuristic tool_use block."""
+    if tool_use.get("name") == "Task" and isinstance(tool_use.get("input"), dict):
+        task_input = tool_use["input"]
+        if task_input.get("run_in_background") is not False:
+            task_input["run_in_background"] = False
+    yield from sse.close_content_blocks()
+    block_idx = sse.blocks.allocate_index()
+    yield sse.content_block_start(
+        block_idx,
+        "tool_use",
+        id=tool_use["id"],
+        name=tool_use["name"],
+    )
+    yield sse.content_block_delta(
+        block_idx,
+        "input_json_delta",
+        json.dumps(tool_use["input"]),
+    )
+    yield sse.content_block_stop(block_idx)
+
+
+def tool_call_extra_content(tool_call: Any) -> dict[str, Any] | None:
+    """Return provider-specific extra tool-call metadata from OpenAI objects."""
+    if isinstance(tool_call, dict):
+        value = tool_call.get("extra_content")
+        return value if isinstance(value, dict) else None
+
+    value = getattr(tool_call, "extra_content", None)
+    if isinstance(value, dict):
+        return value
+
+    model_extra = getattr(tool_call, "model_extra", None)
+    if isinstance(model_extra, dict):
+        value = model_extra.get("extra_content")
+        if isinstance(value, dict):
+            return value
+
+    pydantic_extra = getattr(tool_call, "__pydantic_extra__", None)
+    if isinstance(pydantic_extra, dict):
+        value = pydantic_extra.get("extra_content")
+        if isinstance(value, dict):
+            return value
+
+    return None
+
+
+def has_committed_sse_output(sse: SSEBuilder) -> bool:
+    """Return whether any assistant content escaped the builder."""
+    return (
+        sse.blocks.text_index != -1
+        or sse.blocks.thinking_index != -1
+        or sse.blocks.has_emitted_tool_block()
+    )
+
+
+def started_tool_states(sse: SSEBuilder) -> list[tuple[int, Any]]:
+    """Return started tool states in stream order."""
+    return [
+        (tool_index, state)
+        for tool_index, state in sse.blocks.tool_states.items()
+        if state.started
+    ]
+
+
+def all_started_tools_complete(sse: SSEBuilder, request: Any) -> bool:
+    """Return whether every emitted tool block has schema-valid input."""
+    schemas = tool_schemas_by_name(request)
+    started = started_tool_states(sse)
+    if not started:
+        return False
+    for _, state in started:
+        raw = "".join(state.contents)
+        if parse_complete_tool_input(raw, state.name, schemas) is None:
+            return False
+    return True
+
+
+class OpenAIToolCallAssembler:
+    """Assemble OpenAI tool-call deltas into Anthropic SSE tool blocks."""
+
+    def __init__(
+        self, *, record_extra_content: RecordToolExtraContent | None = None
+    ) -> None:
+        self._record_extra_content = record_extra_content
+
+    def process_tool_call(
+        self,
+        tc: dict[str, Any],
+        sse: SSEBuilder,
+        *,
+        tool_argument_aliases: dict[str, dict[str, str]] | None = None,
+        tool_argument_alias_buffers: dict[int, str] | None = None,
+    ) -> Iterator[str]:
+        """Process a single tool-call delta and yield Anthropic SSE events."""
+        raw_index = tc.get("index", 0)
+        tc_index = raw_index if isinstance(raw_index, int) else 0
+        if tc_index < 0:
+            tc_index = len(sse.blocks.tool_states)
+
+        fn_delta = tc.get("function", {})
+        incoming_name = fn_delta.get("name")
+        arguments = fn_delta.get("arguments", "") or ""
+
+        if tc.get("id") is not None:
+            sse.blocks.set_stream_tool_id(tc_index, tc.get("id"))
+
+        raw_extra_content = tc.get("extra_content")
+        extra_content = (
+            raw_extra_content
+            if isinstance(raw_extra_content, dict) and raw_extra_content
+            else None
+        )
+        if extra_content:
+            sse.blocks.set_tool_extra_content(tc_index, extra_content)
+
+        if incoming_name is not None:
+            sse.blocks.register_tool_name(tc_index, incoming_name)
+
+        state = sse.blocks.tool_states.get(tc_index)
+        resolved_id = (state.tool_id if state and state.tool_id else None) or tc.get(
+            "id"
+        )
+        resolved_name = (state.name if state else "") or ""
+
+        if not state or not state.started:
+            name_ok = bool((resolved_name or "").strip())
+            if name_ok:
+                tool_id = str(resolved_id) if resolved_id else f"tool_{uuid.uuid4()}"
+                display_name = (resolved_name or "").strip() or "tool_call"
+                start_extra_content = state.extra_content if state else extra_content
+                if start_extra_content:
+                    self._record_tool_call_extra_content(tool_id, start_extra_content)
+                yield sse.start_tool_block(
+                    tc_index,
+                    tool_id,
+                    display_name,
+                    extra_content=start_extra_content,
+                )
+                state = sse.blocks.tool_states[tc_index]
+                if state.pre_start_args:
+                    pre = state.pre_start_args
+                    state.pre_start_args = ""
+                    yield from self._emit_tool_arg_delta(
+                        sse,
+                        tc_index,
+                        pre,
+                        tool_argument_aliases=tool_argument_aliases,
+                        tool_argument_alias_buffers=tool_argument_alias_buffers,
+                    )
+
+        state = sse.blocks.tool_states.get(tc_index)
+        if state is not None and state.tool_id and extra_content:
+            self._record_tool_call_extra_content(state.tool_id, extra_content)
+        if not arguments:
+            return
+        if state is None or not state.started:
+            state = sse.blocks.ensure_tool_state(tc_index)
+            if not (resolved_name or "").strip():
+                state.pre_start_args += arguments
+                return
+
+        yield from self._emit_tool_arg_delta(
+            sse,
+            tc_index,
+            arguments,
+            tool_argument_aliases=tool_argument_aliases,
+            tool_argument_alias_buffers=tool_argument_alias_buffers,
+        )
+
+    def flush_task_arg_buffers(self, sse: SSEBuilder) -> Iterator[str]:
+        """Emit buffered Task args as a single JSON delta."""
+        for tool_index, out in sse.blocks.flush_task_arg_buffers():
+            yield sse.emit_tool_delta(tool_index, out)
+
+    def flush_tool_argument_alias_buffers(
+        self,
+        sse: SSEBuilder,
+        tool_argument_aliases: dict[str, dict[str, str]],
+        tool_argument_alias_buffers: dict[int, str],
+    ) -> Iterator[str]:
+        """Emit remaining aliased args without losing malformed JSON."""
+        for tool_index, buffered_args in list(tool_argument_alias_buffers.items()):
+            if not buffered_args:
+                tool_argument_alias_buffers.pop(tool_index, None)
+                continue
+            state = sse.blocks.tool_states.get(tool_index)
+            if state is None or state.name == "Task":
+                continue
+            aliases = tool_argument_aliases.get(state.name, {})
+            if not aliases:
+                continue
+            restored = self._restore_aliased_tool_arguments(buffered_args, aliases)
+            yield sse.emit_tool_delta(
+                tool_index,
+                restored if restored is not None else buffered_args,
+            )
+            tool_argument_alias_buffers.pop(tool_index, None)
+
+    def _emit_tool_arg_delta(
+        self,
+        sse: SSEBuilder,
+        tc_index: int,
+        args: str,
+        *,
+        tool_argument_aliases: dict[str, dict[str, str]] | None = None,
+        tool_argument_alias_buffers: dict[int, str] | None = None,
+    ) -> Iterator[str]:
+        """Emit one argument fragment for a started tool block."""
+        if not args:
+            return
+        state = sse.blocks.tool_states.get(tc_index)
+        if state is None:
+            return
+        if state.name == "Task":
+            parsed = sse.blocks.buffer_task_args(tc_index, args)
+            if parsed is not None:
+                yield sse.emit_tool_delta(tc_index, json.dumps(parsed))
+            return
+        aliases = (
+            tool_argument_aliases.get(state.name, {}) if tool_argument_aliases else {}
+        )
+        if aliases:
+            if tool_argument_alias_buffers is None:
+                restored = self._restore_aliased_tool_arguments(args, aliases)
+                if restored is not None:
+                    yield sse.emit_tool_delta(tc_index, restored)
+                return
+
+            buffered_args = tool_argument_alias_buffers.get(tc_index, "") + args
+            restored = self._restore_aliased_tool_arguments(buffered_args, aliases)
+            if restored is None:
+                tool_argument_alias_buffers[tc_index] = buffered_args
+                return
+            tool_argument_alias_buffers.pop(tc_index, None)
+            yield sse.emit_tool_delta(tc_index, restored)
+            return
+        yield sse.emit_tool_delta(tc_index, args)
+
+    def _restore_aliased_tool_arguments(
+        self, argument_json: str, aliases: dict[str, str]
+    ) -> str | None:
+        try:
+            parsed = json.loads(argument_json)
+        except json.JSONDecodeError:
+            return None
+        if not isinstance(parsed, dict):
+            return argument_json
+        restored = self._restore_aliased_tool_argument_value(parsed, aliases)
+        return json.dumps(restored)
+
+    def _restore_aliased_tool_argument_value(
+        self, value: Any, aliases: dict[str, str]
+    ) -> Any:
+        if isinstance(value, dict):
+            return {
+                aliases.get(key, key): self._restore_aliased_tool_argument_value(
+                    item, aliases
+                )
+                for key, item in value.items()
+            }
+        if isinstance(value, list):
+            return [
+                self._restore_aliased_tool_argument_value(item, aliases)
+                for item in value
+            ]
+        return value
+
+    def _record_tool_call_extra_content(
+        self, tool_call_id: str, extra_content: dict[str, Any]
+    ) -> None:
+        if self._record_extra_content is not None:
+            self._record_extra_content(tool_call_id, extra_content)