providers/openai_compat.py

"""OpenAI-style chat base for :class:`OpenAIChatTransport` (NIM, etc.).

``AnthropicMessagesTransport``-based providers (OpenRouter, LM Studio, DeepSeek, …) live
in separate modules; do not list them as subclasses of this class.
"""

import asyncio
import json
import uuid
from abc import abstractmethod
from collections.abc import AsyncIterator, Iterator
from typing import Any

import httpx
from loguru import logger
from openai import AsyncOpenAI

from core.anthropic import (
    ContentType,
    HeuristicToolParser,
    SSEBuilder,
    ThinkTagParser,
    append_request_id,
    map_stop_reason,
)
from providers.base import BaseProvider, ProviderConfig
from providers.error_mapping import (
    map_error,
    user_visible_message_for_mapped_provider_error,
)
from providers.model_listing import (
    ProviderModelInfo,
    extract_openai_model_ids,
    model_infos_from_ids,
)
from providers.rate_limit import GlobalRateLimiter


def _iter_heuristic_tool_use_sse(
    sse: SSEBuilder, tool_use: dict[str, Any]
) -> Iterator[str]:
    """Emit SSE for one heuristic tool_use block (closes open text/thinking first)."""
    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)


class OpenAIChatTransport(BaseProvider):
    """Base for OpenAI-compatible ``/chat/completions`` adapters (NIM, …)."""

    def __init__(
        self,
        config: ProviderConfig,
        *,
        provider_name: str,
        base_url: str,
        api_key: str,
        nim_rate_limit: int = 100,
        nim_max_concurrency: int = 40,
    ):
        super().__init__(config)
        self._provider_name = provider_name
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._http_client = None
        self._client_cache: dict[str, AsyncOpenAI] = {}
        # NIM gets adaptive rate starting at 100 req/min (leaves headroom)
        # Zen is effectively unlimited (9999)
        if provider_name.lower() == "zen":
            effective_rate_limit = 9999
            effective_max_concurrency = config.max_concurrency * 4
            use_adaptive = None
        else:
            effective_rate_limit = nim_rate_limit
            effective_max_concurrency = max(
                nim_max_concurrency, config.max_concurrency * 4
            )
            use_adaptive = nim_rate_limit
        self._global_rate_limiter = GlobalRateLimiter.get_scoped_instance(
            provider_name.lower(),
            rate_limit=effective_rate_limit,
            rate_window=config.rate_window,
            max_concurrency=effective_max_concurrency,
            adaptive_rate=use_adaptive,
            adaptive_min_rate=10,
        )
        # Connection pool tuned for maximum throughput.
        # Increased keepalive and connections for high concurrency.
        http_client_args = {
            "timeout": httpx.Timeout(
                config.http_read_timeout,
                connect=config.http_connect_timeout,
                read=config.http_read_timeout,
                write=config.http_write_timeout,
            ),
            "trust_env": False,
            "http2": True,
            "limits": httpx.Limits(
                max_keepalive_connections=100,
                max_connections=500,
                keepalive_expiry=5.0,
            ),
        }
        if config.proxy:
            http_client_args["proxy"] = config.proxy

        self._http_client = httpx.AsyncClient(**http_client_args)

        self._client = AsyncOpenAI(
            api_key=self._api_key,
            base_url=self._base_url,
            max_retries=0,
            timeout=httpx.Timeout(
                config.http_read_timeout,
                connect=config.http_connect_timeout,
                read=config.http_read_timeout,
                write=config.http_write_timeout,
            ),
            http_client=self._http_client,
        )
        self._client_cache[self._api_key] = self._client
        if not self._api_key:
            logger.warning(
                "OpenAIChatTransport initialized with EMPTY API key: provider={} base_url={}",
                provider_name,
                base_url,
            )
        else:
            logger.info(
                "OpenAIChatTransport initialized: provider={} base_url={} api_key_set={} key_prefix={}",
                provider_name,
                base_url,
                bool(self._api_key),
                self._api_key[:10] if self._api_key else "EMPTY",
            )

    def _client_for_api_key(self, api_key: str) -> AsyncOpenAI:
        """Return a cached OpenAI client for the given API key."""
        if api_key == self._api_key:
            return self._client
        client = self._client_cache.get(api_key)
        if client is not None:
            return client
        client = AsyncOpenAI(
            api_key=api_key,
            base_url=self._base_url,
            max_retries=0,
            timeout=httpx.Timeout(
                self._config.http_read_timeout,
                connect=self._config.http_connect_timeout,
                read=self._config.http_read_timeout,
                write=self._config.http_write_timeout,
            ),
            http_client=self._http_client,
        )
        self._client_cache[api_key] = client
        return client

    async def cleanup(self) -> None:
        """Release HTTP client resources."""
        seen: set[int] = set()
        for client in list(self._client_cache.values()):
            client_id = id(client)
            if client_id in seen:
                continue
            seen.add(client_id)
            await client.aclose()

    async def list_model_infos(self) -> frozenset[ProviderModelInfo]:
        """Return model metadata from the provider's OpenAI-compatible models endpoint."""
        model_ids = await self.list_model_ids()
        # Default all models to supports_thinking=None (unknown) unless provider overrides
        return model_infos_from_ids(model_ids, supports_thinking=None)

    async def list_model_ids(self) -> frozenset[str]:
        """Return model ids from the provider's OpenAI-compatible models endpoint."""
        payload = await self._client.models.list()
        return extract_openai_model_ids(payload, provider_name=self._provider_name)

    @abstractmethod
    def _build_request_body(
        self, request: Any, thinking_enabled: bool | None = None
    ) -> dict:
        """Build request body. Must be implemented by subclasses."""

    def _handle_extra_reasoning(
        self, delta: Any, sse: SSEBuilder, *, thinking_enabled: bool
    ) -> Iterator[str]:
        """Hook for provider-specific reasoning (e.g. OpenRouter reasoning_details)."""
        return iter(())

    def _get_retry_request_body(self, error: Exception, body: dict) -> dict | None:
        """Return a modified request body for one retry, or None."""
        return None

    async def _create_stream(self, body: dict) -> tuple[Any, dict]:
        """Create a streaming chat completion, optionally retrying once."""
        from loguru import logger

        try:
            logger.info(
                "{}_CREATE_STREAM: calling API with model={}",
                self._provider_name,
                body.get("model"),
            )
            stream = await self._global_rate_limiter.execute_with_retry(
                self._client.chat.completions.create,
                **body,
                stream=True,
                max_retries=1,
            )
            return stream, body
        except Exception as error:
            logger.error(
                "{}_CREATE_STREAM_ERROR: {} - {} - response={}",
                self._provider_name,
                type(error).__name__,
                str(error),
                getattr(error, "response", None),
            )
            retry_body = self._get_retry_request_body(error, body)
            if retry_body is None:
                raise

            stream = await self._global_rate_limiter.execute_with_retry(
                self._client.chat.completions.create,
                **retry_body,
                stream=True,
                max_retries=1,
            )
            return stream, retry_body

    def _emit_tool_arg_delta(
        self, sse: SSEBuilder, tc_index: int, args: str
    ) -> Iterator[str]:
        """Emit one argument fragment for a started tool block (Task buffer or raw JSON)."""
        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
        yield sse.emit_tool_delta(tc_index, args)

    def _process_tool_call(self, tc: dict, sse: SSEBuilder) -> Iterator[str]:
        """Process a single tool call delta and yield SSE events."""
        tc_index = tc.get("index", 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"))

        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"
                yield sse.start_tool_block(tc_index, tool_id, display_name)
                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)

        state = sse.blocks.tool_states.get(tc_index)
        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)

    def _flush_task_arg_buffers(self, sse: SSEBuilder) -> Iterator[str]:
        """Emit buffered Task args as a single JSON delta (best-effort)."""
        for tool_index, out in sse.blocks.flush_task_arg_buffers():
            yield sse.emit_tool_delta(tool_index, out)

    async def stream_response(
        self,
        request: Any,
        input_tokens: int = 0,
        *,
        request_id: str | None = None,
        thinking_enabled: bool | None = None,
    ) -> AsyncIterator[str]:
        """Stream response in Anthropic SSE format."""
        with logger.contextualize(request_id=request_id):
            async for event in self._stream_response_impl(
                request, input_tokens, request_id, thinking_enabled=thinking_enabled
            ):
                yield event

    async def _stream_response_impl(
        self,
        request: Any,
        input_tokens: int,
        request_id: str | None,
        *,
        thinking_enabled: bool | None,
    ) -> AsyncIterator[str]:
        """Shared streaming implementation."""
        tag = self._provider_name
        message_id = f"msg_{uuid.uuid4()}"
        sse = SSEBuilder(
            message_id,
            request.model,
            input_tokens,
            log_raw_events=self._config.log_raw_sse_events,
        )

        body = self._build_request_body(request, thinking_enabled=thinking_enabled)
        thinking_enabled = self._is_thinking_enabled(request, thinking_enabled)
        req_tag = f" request_id={request_id}" if request_id else ""
        # Log the actual model being sent to the provider for debugging
        actual_model = body.get("model", "")
        logger.info(
            "{}_STREAM:{} model={} msgs={} tools={} actual_model_to_provider={}",
            tag,
            req_tag,
            body.get("model"),
            len(body.get("messages", [])),
            len(body.get("tools", [])),
            actual_model,
        )

        think_parser = ThinkTagParser()
        heuristic_parser = HeuristicToolParser()
        finish_reason = None
        usage_info = None

        async with self._global_rate_limiter.concurrency_slot():
            try:
                yield sse.message_start()
                stream, body = await self._create_stream(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)

                    # Handle reasoning_content (OpenAI extended format)
                    reasoning = getattr(delta, "reasoning_content", None)
                    if thinking_enabled and reasoning:
                        for event in sse.ensure_thinking_block():
                            yield event
                        yield sse.emit_thinking_delta(reasoning)

                    # Provider-specific extra reasoning (e.g. OpenRouter reasoning_details)
                    for event in self._handle_extra_reasoning(
                        delta,
                        sse,
                        thinking_enabled=thinking_enabled,
                    ):
                        yield event

                    # Handle text content
                    if delta.content:
                        for part in think_parser.feed(delta.content):
                            if part.type == ContentType.THINKING:
                                if not thinking_enabled:
                                    continue
                                for event in sse.ensure_thinking_block():
                                    yield event
                                yield sse.emit_thinking_delta(part.content)
                            else:
                                filtered_text, detected_tools = heuristic_parser.feed(
                                    part.content
                                )

                                if filtered_text:
                                    for event in sse.ensure_text_block():
                                        yield event
                                    yield sse.emit_text_delta(filtered_text)

                                for tool_use in detected_tools:
                                    for event in _iter_heuristic_tool_use_sse(
                                        sse, tool_use
                                    ):
                                        yield event

                    # Handle native tool calls
                    if delta.tool_calls:
                        for event in sse.close_content_blocks():
                            yield event
                        for tc in delta.tool_calls:
                            tc_info = {
                                "index": tc.index,
                                "id": tc.id,
                                "function": {
                                    "name": tc.function.name,
                                    "arguments": tc.function.arguments,
                                },
                            }
                            for event in self._process_tool_call(tc_info, sse):
                                yield event

            except asyncio.CancelledError:
                raise
            except Exception as e:
                self._log_stream_transport_error(tag, req_tag, e)
                mapped_e = map_error(e, rate_limiter=self._global_rate_limiter)

                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
                    or len(sse._accumulated_text_parts) > 0
                    or len(sse._accumulated_reasoning_parts) > 0
                )

                if has_content_blocks and isinstance(
                    e,
                    (
                        httpx.RemoteProtocolError,
                        httpx.ReadTimeout,
                        asyncio.TimeoutError,
                        httpx.ConnectError,
                    ),
                ):
                    logger.warning(
                        "{}_STREAM: Transient error mid-stream. Faking max_tokens to resume. {}",
                        tag,
                        e,
                    )
                    for event in sse.close_all_blocks():
                        yield event
                    yield sse.message_delta("max_tokens", sse.estimate_output_tokens())
                    yield sse.message_stop()
                    return

                base_message = user_visible_message_for_mapped_provider_error(
                    mapped_e,
                    provider_name=tag,
                    read_timeout_s=self._config.http_read_timeout,
                )
                error_message = append_request_id(base_message, request_id)
                logger.info(
                    "{}_STREAM: Emitting SSE error event for {}{}",
                    tag,
                    type(e).__name__,
                    req_tag,
                )
                for event in sse.close_all_blocks():
                    yield event
                if sse.blocks.has_emitted_tool_block():
                    # Avoid a second assistant text block after an emitted tool_use, which
                    # breaks OpenAI history replay (issue #206) when Claude Code stores it.
                    yield sse.emit_top_level_error(error_message)
                else:
                    for event in sse.emit_error(error_message):
                        yield event
                yield sse.message_delta("end_turn", 1)
                yield sse.message_stop()
                return

        # Flush remaining content
        remaining = think_parser.flush()
        if remaining:
            if remaining.type == ContentType.THINKING:
                if not thinking_enabled:
                    remaining = None
                else:
                    for event in sse.ensure_thinking_block():
                        yield event
                    yield sse.emit_thinking_delta(remaining.content)
            if remaining and remaining.type == ContentType.TEXT:
                for event in sse.ensure_text_block():
                    yield event
                yield sse.emit_text_delta(remaining.content)

        for tool_use in heuristic_parser.flush():
            for event in _iter_heuristic_tool_use_sse(sse, tool_use):
                yield 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:
            for event in sse.ensure_text_block():
                yield event
            yield sse.emit_text_delta(" ")
        elif (
            not has_started_tool
            and not sse.accumulated_text.strip()
            and sse.accumulated_reasoning.strip()
        ):
            # Some OpenAI-compatible models (e.g. NIM reasoning templates) stream only
            # ``reasoning_content`` with no ``content``; emit a minimal text block so
            # clients and smoke ``text_content()`` see a completed assistant message.
            for event in sse.ensure_text_block():
                yield event
            yield sse.emit_text_delta(" ")

        for event in self._flush_task_arg_buffers(sse):
            yield event

        for event in 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}",
                    input_tokens,
                    provider_input,
                    provider_input - input_tokens,
                )
        yield sse.message_delta(map_stop_reason(finish_reason), output_tokens)
        yield sse.message_stop()