Jack1808/Claude_Code

1"""Shared base class for OpenAI-compatible providers (NIM, OpenRouter, LM Studio)."""2 3import json4import uuid5from abc import abstractmethod6from collections.abc import AsyncIterator, Iterator7from typing import Any8 9import httpx10from loguru import logger11from openai import AsyncOpenAI12 13from providers.base import BaseProvider, ProviderConfig14from providers.common import (15    ContentType,16    HeuristicToolParser,17    SSEBuilder,18    ThinkTagParser,19    append_request_id,20    get_user_facing_error_message,21    map_error,22    map_stop_reason,23)24from providers.rate_limit import GlobalRateLimiter25 26 27class OpenAICompatibleProvider(BaseProvider):28    """Base class for providers using OpenAI-compatible chat completions API."""29 30    def __init__(31        self,32        config: ProviderConfig,33        *,34        provider_name: str,35        base_url: str,36        api_key: str,37    ):38        super().__init__(config)39        self._provider_name = provider_name40        self._api_key = api_key41        self._base_url = base_url.rstrip("/")42        self._global_rate_limiter = GlobalRateLimiter.get_instance(43            rate_limit=config.rate_limit,44            rate_window=config.rate_window,45            max_concurrency=config.max_concurrency,46        )47        self._client = AsyncOpenAI(48            api_key=self._api_key,49            base_url=self._base_url,50            max_retries=0,51            timeout=httpx.Timeout(52                config.http_read_timeout,53                connect=config.http_connect_timeout,54                read=config.http_read_timeout,55                write=config.http_write_timeout,56            ),57        )58 59    async def cleanup(self) -> None:60        """Release HTTP client resources."""61        client = getattr(self, "_client", None)62        if client is not None:63            await client.aclose()64 65    @abstractmethod66    def _build_request_body(self, request: Any) -> dict:67        """Build request body. Must be implemented by subclasses."""68 69    def _handle_extra_reasoning(self, delta: Any, sse: SSEBuilder) -> Iterator[str]:70        """Hook for provider-specific reasoning (e.g. OpenRouter reasoning_details)."""71        return iter(())72 73    def _process_tool_call(self, tc: dict, sse: SSEBuilder) -> Iterator[str]:74        """Process a single tool call delta and yield SSE events."""75        tc_index = tc.get("index", 0)76        if tc_index < 0:77            tc_index = len(sse.blocks.tool_states)78 79        fn_delta = tc.get("function", {})80        incoming_name = fn_delta.get("name")81        if incoming_name is not None:82            sse.blocks.register_tool_name(tc_index, incoming_name)83 84        state = sse.blocks.tool_states.get(tc_index)85        if state is None or not state.started:86            name = state.name if state else ""87            if name or tc.get("id"):88                tool_id = tc.get("id") or f"tool_{uuid.uuid4()}"89                yield sse.start_tool_block(tc_index, tool_id, name)90 91        args = fn_delta.get("arguments", "")92        if args:93            state = sse.blocks.tool_states.get(tc_index)94            if state is None or not state.started:95                tool_id = tc.get("id") or f"tool_{uuid.uuid4()}"96                name = (state.name if state else None) or "tool_call"97                yield sse.start_tool_block(tc_index, tool_id, name)98                state = sse.blocks.tool_states.get(tc_index)99 100            current_name = state.name if state else ""101            if current_name == "Task":102                parsed = sse.blocks.buffer_task_args(tc_index, args)103                if parsed is not None:104                    yield sse.emit_tool_delta(tc_index, json.dumps(parsed))105                return106 107            yield sse.emit_tool_delta(tc_index, args)108 109    def _flush_task_arg_buffers(self, sse: SSEBuilder) -> Iterator[str]:110        """Emit buffered Task args as a single JSON delta (best-effort)."""111        for tool_index, out in sse.blocks.flush_task_arg_buffers():112            yield sse.emit_tool_delta(tool_index, out)113 114    async def stream_response(115        self,116        request: Any,117        input_tokens: int = 0,118        *,119        request_id: str | None = None,120    ) -> AsyncIterator[str]:121        """Stream response in Anthropic SSE format."""122        with logger.contextualize(request_id=request_id):123            async for event in self._stream_response_impl(124                request, input_tokens, request_id125            ):126                yield event127 128    async def _stream_response_impl(129        self,130        request: Any,131        input_tokens: int,132        request_id: str | None,133    ) -> AsyncIterator[str]:134        """Shared streaming implementation."""135        tag = self._provider_name136        message_id = f"msg_{uuid.uuid4()}"137        sse = SSEBuilder(message_id, request.model, input_tokens)138 139        body = self._build_request_body(request)140        req_tag = f" request_id={request_id}" if request_id else ""141        logger.info(142            "{}_STREAM:{} model={} msgs={} tools={}",143            tag,144            req_tag,145            body.get("model"),146            len(body.get("messages", [])),147            len(body.get("tools", [])),148        )149 150        yield sse.message_start()151 152        think_parser = ThinkTagParser()153        heuristic_parser = HeuristicToolParser()154 155        finish_reason = None156        usage_info = None157        error_occurred = False158        error_message = ""159 160        async with self._global_rate_limiter.concurrency_slot():161            try:162                stream = await self._global_rate_limiter.execute_with_retry(163                    self._client.chat.completions.create, **body, stream=True164                )165                async for chunk in stream:166                    if getattr(chunk, "usage", None):167                        usage_info = chunk.usage168 169                    if not chunk.choices:170                        continue171 172                    choice = chunk.choices[0]173                    delta = choice.delta174                    if delta is None:175                        continue176 177                    if choice.finish_reason:178                        finish_reason = choice.finish_reason179                        logger.debug("{} finish_reason: {}", tag, finish_reason)180 181                    # Handle reasoning_content (OpenAI extended format)182                    reasoning = getattr(delta, "reasoning_content", None)183                    if reasoning:184                        for event in sse.ensure_thinking_block():185                            yield event186                        yield sse.emit_thinking_delta(reasoning)187 188                    # Provider-specific extra reasoning (e.g. OpenRouter reasoning_details)189                    for event in self._handle_extra_reasoning(delta, sse):190                        yield event191 192                    # Handle text content193                    if delta.content:194                        for part in think_parser.feed(delta.content):195                            if part.type == ContentType.THINKING:196                                for event in sse.ensure_thinking_block():197                                    yield event198                                yield sse.emit_thinking_delta(part.content)199                            else:200                                filtered_text, detected_tools = heuristic_parser.feed(201                                    part.content202                                )203 204                                if filtered_text:205                                    for event in sse.ensure_text_block():206                                        yield event207                                    yield sse.emit_text_delta(filtered_text)208 209                                for tool_use in detected_tools:210                                    for event in sse.close_content_blocks():211                                        yield event212 213                                    block_idx = sse.blocks.allocate_index()214                                    if tool_use.get("name") == "Task" and isinstance(215                                        tool_use.get("input"), dict216                                    ):217                                        tool_use["input"]["run_in_background"] = False218                                    yield sse.content_block_start(219                                        block_idx,220                                        "tool_use",221                                        id=tool_use["id"],222                                        name=tool_use["name"],223                                    )224                                    yield sse.content_block_delta(225                                        block_idx,226                                        "input_json_delta",227                                        json.dumps(tool_use["input"]),228                                    )229                                    yield sse.content_block_stop(block_idx)230 231                    # Handle native tool calls232                    if delta.tool_calls:233                        for event in sse.close_content_blocks():234                            yield event235                        for tc in delta.tool_calls:236                            tc_info = {237                                "index": tc.index,238                                "id": tc.id,239                                "function": {240                                    "name": tc.function.name,241                                    "arguments": tc.function.arguments,242                                },243                            }244                            for event in self._process_tool_call(tc_info, sse):245                                yield event246 247            except Exception as e:248                logger.error("{}_ERROR:{} {}: {}", tag, req_tag, type(e).__name__, e)249                mapped_e = map_error(e)250                error_occurred = True251                error_message = append_request_id(252                    get_user_facing_error_message(253                        mapped_e, read_timeout_s=self._config.http_read_timeout254                    ),255                    request_id,256                )257                logger.info(258                    "{}_STREAM: Emitting SSE error event for {}{}",259                    tag,260                    type(e).__name__,261                    req_tag,262                )263                for event in sse.close_content_blocks():264                    yield event265                for event in sse.emit_error(error_message):266                    yield event267 268        # Flush remaining content269        remaining = think_parser.flush()270        if remaining:271            if remaining.type == ContentType.THINKING:272                for event in sse.ensure_thinking_block():273                    yield event274                yield sse.emit_thinking_delta(remaining.content)275            else:276                for event in sse.ensure_text_block():277                    yield event278                yield sse.emit_text_delta(remaining.content)279 280        for tool_use in heuristic_parser.flush():281            for event in sse.close_content_blocks():282                yield event283 284            block_idx = sse.blocks.allocate_index()285            yield sse.content_block_start(286                block_idx,287                "tool_use",288                id=tool_use["id"],289                name=tool_use["name"],290            )291            if tool_use.get("name") == "Task" and isinstance(292                tool_use.get("input"), dict293            ):294                tool_use["input"]["run_in_background"] = False295            yield sse.content_block_delta(296                block_idx,297                "input_json_delta",298                json.dumps(tool_use["input"]),299            )300            yield sse.content_block_stop(block_idx)301 302        if (303            not error_occurred304            and sse.blocks.text_index == -1305            and not sse.blocks.tool_states306        ):307            for event in sse.ensure_text_block():308                yield event309            yield sse.emit_text_delta(" ")310 311        for event in self._flush_task_arg_buffers(sse):312            yield event313 314        for event in sse.close_all_blocks():315            yield event316 317        output_tokens = (318            usage_info.completion_tokens319            if usage_info and hasattr(usage_info, "completion_tokens")320            else sse.estimate_output_tokens()321        )322        if usage_info and hasattr(usage_info, "prompt_tokens"):323            provider_input = usage_info.prompt_tokens324            if isinstance(provider_input, int):325                logger.debug(326                    "TOKEN_ESTIMATE: our={} provider={} diff={:+d}",327                    input_tokens,328                    provider_input,329                    provider_input - input_tokens,330                )331        yield sse.message_delta(map_stop_reason(finish_reason), output_tokens)332        yield sse.message_stop()333