multi-forge 0.2.0__py3-none-any.whl

forge/core/llm/clients/openai_compat.py

@@ -0,0 +1,244 @@
+"""Shared helpers for OpenAI-compatible LLM clients.
+
+Used by both LiteLLMClient and OpenRouterClient. Extracted here so the
+OpenRouter client has no import dependency on LiteLLM.
+"""
+
+import json
+import logging
+from typing import Any
+
+from openai import APIError, APIStatusError, RateLimitError
+
+from ..errors import ProviderError
+from ..types import (
+    CompletionResponse,
+    Message,
+    ModelHyperparameters,
+    ToolCall,
+    ToolCallDelta,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def is_retryable_error(error: Exception) -> bool:
+    """Return True if the error should trigger tenacity retry.
+
+    Only retries transient errors (rate limits, server errors).
+    Auth failures (401/403) are excluded -- retrying with the same
+    bad credentials just adds ~14s of delay.
+    """
+    if isinstance(error, APIStatusError):
+        return error.status_code not in (400, 401, 403)
+    if isinstance(error, (RateLimitError, APIError)):
+        return True
+    return False
+
+
+def extract_cached_tokens(usage: object) -> int:
+    """Extract cached_tokens from a usage object's prompt_tokens_details.
+
+    LiteLLM and OpenRouter pass through provider cache metrics in
+    ``usage.prompt_tokens_details.cached_tokens``.  The field may be an
+    object (SDK model) or a plain dict depending on the response path.
+
+    Returns 0 if no cache data is present.
+    """
+    prompt_details = getattr(usage, "prompt_tokens_details", None)
+    if prompt_details is None and isinstance(usage, dict):
+        prompt_details = usage.get("prompt_tokens_details")
+    if not prompt_details:
+        return 0
+    if isinstance(prompt_details, dict):
+        raw = prompt_details.get("cached_tokens", 0) or 0
+    else:
+        raw = getattr(prompt_details, "cached_tokens", 0) or 0
+    return int(raw)
+
+
+def message_to_openai(msg: Message) -> dict[str, Any]:
+    """Convert canonical Message to OpenAI chat completion format."""
+    result: dict[str, Any] = {"role": msg.role, "content": msg.content}
+
+    if msg.tool_call_id:
+        result["tool_call_id"] = msg.tool_call_id
+
+    if msg.tool_calls:
+        result["tool_calls"] = [
+            {
+                "id": tc.id,
+                "type": "function",
+                "function": {
+                    "name": tc.name,
+                    "arguments": json.dumps(tc.arguments),
+                },
+            }
+            for tc in msg.tool_calls
+        ]
+
+    return result
+
+
+def build_chat_completion_kwargs(
+    model: str,
+    messages: list[Message],
+    tools: list[dict[str, Any]] | None,
+    hyperparams: ModelHyperparameters,
+) -> dict[str, Any]:
+    """Build kwargs for OpenAI chat.completions.create()."""
+    kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": [message_to_openai(m) for m in messages],
+        "max_tokens": hyperparams.max_tokens,
+    }
+
+    if hyperparams.temperature is not None:
+        kwargs["temperature"] = hyperparams.temperature
+
+    if hyperparams.top_p is not None:
+        kwargs["top_p"] = hyperparams.top_p
+
+    if hyperparams.reasoning_effort is not None:
+        kwargs["reasoning_effort"] = hyperparams.reasoning_effort
+
+    if hyperparams.verbosity is not None:
+        kwargs["verbosity"] = hyperparams.verbosity
+
+    if tools:
+        kwargs["tools"] = tools
+
+    if "openai" in hyperparams.extra:
+        kwargs.update(hyperparams.extra["openai"])
+
+    return kwargs
+
+
+def openai_response_to_completion(response: Any, provider: str) -> CompletionResponse:
+    """Convert OpenAI ChatCompletion response to canonical CompletionResponse."""
+    if hasattr(response, "error") and response.error:
+        error_msg = response.error.get("message", "Unknown error")
+        error_code = response.error.get("code", "unknown")
+        raise ProviderError(
+            provider,
+            Exception(f"API error (code={error_code}): {error_msg}"),
+        )
+
+    if not response.choices:
+        raise ProviderError(
+            provider,
+            Exception("No choices in response"),
+        )
+
+    choice = response.choices[0]
+    message = choice.message
+
+    text = message.content or ""
+
+    tool_calls = None
+    if message.tool_calls:
+        tool_calls = []
+        for tc in message.tool_calls:
+            try:
+                arguments = json.loads(tc.function.arguments)
+            except json.JSONDecodeError:
+                arguments = {}
+            tool_calls.append(
+                ToolCall(
+                    id=tc.id,
+                    name=tc.function.name,
+                    arguments=arguments,
+                )
+            )
+
+    usage = None
+    if response.usage:
+        usage = {
+            "prompt_tokens": response.usage.prompt_tokens,
+            "completion_tokens": response.usage.completion_tokens,
+            "total_tokens": response.usage.total_tokens,
+        }
+        cached = extract_cached_tokens(response.usage)
+        if cached:
+            usage["cached_tokens"] = cached
+
+    return CompletionResponse(
+        text=text,
+        tool_calls=tool_calls,
+        usage=usage,
+        raw=response.model_dump(),
+    )
+
+
+class ToolCallAccumulator:
+    """Accumulates streaming tool call deltas into complete ToolCalls.
+
+    During streaming, tool calls arrive as fragments (id, name, argument chunks).
+    OpenAI sends `id` only on the first chunk; subsequent chunks use `index`
+    to correlate. This accumulator uses index-based lookup to handle both.
+    """
+
+    def __init__(self) -> None:
+        self._pending: dict[int, ToolCallDelta] = {}
+
+    def add_delta(self, delta: ToolCallDelta) -> None:
+        """Add a streaming delta to the accumulator."""
+        idx = delta.index
+        if idx is None:
+            return
+
+        if idx not in self._pending:
+            self._pending[idx] = ToolCallDelta(index=idx)
+
+        existing = self._pending[idx]
+        if delta.id:
+            existing.id = delta.id
+        if delta.name:
+            existing.name = delta.name
+        existing.arguments_json += delta.arguments_json
+
+    def finalize(self) -> list[ToolCall]:
+        """Parse accumulated deltas into complete ToolCalls.
+
+        Returns tool calls sorted by index for deterministic ordering.
+        """
+        result = []
+        for idx in sorted(self._pending):
+            delta = self._pending[idx]
+            if delta.id and delta.name:
+                try:
+                    arguments = json.loads(delta.arguments_json) if delta.arguments_json else {}
+                except json.JSONDecodeError:
+                    logger.warning(f"Failed to parse tool call arguments: {delta.arguments_json}")
+                    arguments = {}
+
+                result.append(
+                    ToolCall(
+                        id=delta.id,
+                        name=delta.name,
+                        arguments=arguments,
+                    )
+                )
+            elif delta.arguments_json:
+                logger.warning(
+                    f"Dropping incomplete tool call at index {idx}: "
+                    f"id={delta.id}, name={delta.name}, args_len={len(delta.arguments_json)}"
+                )
+        return result
+
+    def has_pending(self) -> bool:
+        """Check if there are any pending tool calls."""
+        return len(self._pending) > 0
+
+    def default_index(self) -> int | None:
+        """Suggest an index for an unindexed single-tool delta.
+
+        Returns 0 when no calls are pending (first tool call), the sole
+        pending index when exactly one exists (continuation), or None
+        when multiple calls are pending (ambiguous -- caller should drop).
+        """
+        if len(self._pending) == 0:
+            return 0
+        if len(self._pending) == 1:
+            return next(iter(self._pending))
+        return None

forge/core/llm/clients/openrouter.py

@@ -0,0 +1,234 @@
+"""OpenRouter client implementation.
+
+Uses OpenAI SDK to call OpenRouter's API directly (no LiteLLM).
+OpenRouter is OpenAI-compatible, so this is a thin wrapper that adds
+OpenRouter-specific headers and translates parameters to OpenRouter's format.
+"""
+
+import json
+import logging
+from typing import Any, AsyncGenerator
+
+from openai import AsyncOpenAI
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+from ..credentials import CredentialManager
+from ..detection import ProviderType
+from ..errors import AuthenticationError, ProviderError
+from ..types import (
+    CompletionResponse,
+    Message,
+    ModelHyperparameters,
+    StreamEvent,
+    ToolCallDelta,
+)
+from .base import estimate_message_tokens, merge_hyperparams
+from .openai_compat import (
+    ToolCallAccumulator,
+    build_chat_completion_kwargs,
+    extract_cached_tokens,
+    is_retryable_error,
+    message_to_openai,
+    openai_response_to_completion,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class OpenRouterClient:
+    """OpenRouter client using OpenAI SDK.
+
+    Calls OpenRouter's API directly at https://openrouter.ai/api/v1.
+    Uses Chat Completions only (no Responses API).
+    """
+
+    def __init__(
+        self,
+        model: str,
+        provider: ProviderType,
+        credentials: CredentialManager | None = None,
+        default_hyperparams: ModelHyperparameters | None = None,
+    ) -> None:
+        self._model = model
+        self._provider = provider
+        self._credentials = credentials or CredentialManager.default()
+        self._default_hyperparams = default_hyperparams
+        self._client: AsyncOpenAI | None = None
+
+    @property
+    def model(self) -> str:
+        return self._model
+
+    async def _get_client(self) -> AsyncOpenAI:
+        if self._client is not None:
+            return self._client
+
+        creds = await self._credentials.get_credentials(self._provider)
+
+        self._client = AsyncOpenAI(
+            api_key=creds["api_key"],
+            base_url=creds["base_url"],
+            default_headers=creds.get("extra_headers", {}),
+        )
+        return self._client
+
+    @staticmethod
+    def _translate_params(kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Translate Forge params to OpenRouter's API format.
+
+        OpenRouter uses ``reasoning: {effort: ...}`` (object) and top-level
+        ``verbosity``, passed via ``extra_body`` since the OpenAI SDK does not
+        accept them as direct kwargs.
+        """
+        extra_body: dict[str, Any] = kwargs.pop("extra_body", None) or {}
+        effort = kwargs.pop("reasoning_effort", None)
+        if effort is not None:
+            extra_body["reasoning"] = {"effort": effort}
+        verbosity = kwargs.pop("verbosity", None)
+        if verbosity is not None:
+            extra_body["verbosity"] = verbosity
+        if extra_body:
+            kwargs["extra_body"] = extra_body
+        return kwargs
+
+    _is_retryable_error = staticmethod(is_retryable_error)
+
+    @retry(
+        retry=retry_if_exception(lambda e: isinstance(e, Exception) and is_retryable_error(e)),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+        reraise=True,
+    )
+    async def _make_completion_request(
+        self,
+        client: AsyncOpenAI,
+        messages: list[Message],
+        tools: list[dict[str, Any]] | None,
+        merged_params: ModelHyperparameters,
+    ) -> CompletionResponse:
+        kwargs = build_chat_completion_kwargs(self._model, messages, tools, merged_params)
+        kwargs = self._translate_params(kwargs)
+        response = await client.chat.completions.create(**kwargs)
+        return openai_response_to_completion(response, self._provider)
+
+    async def complete(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict[str, Any]] | None = None,
+        hyperparams: ModelHyperparameters | None = None,
+    ) -> CompletionResponse:
+        merged_params = merge_hyperparams(self._default_hyperparams, hyperparams)
+        client = await self._get_client()
+
+        try:
+            return await self._make_completion_request(client, messages, tools, merged_params)
+        except (ProviderError, AuthenticationError):
+            raise
+        except Exception as e:
+            error_str = str(e).lower()
+            if "authentication" in error_str or "unauthorized" in error_str:
+                await self._credentials.invalidate(self._provider)
+                await self._close_client()
+                raise AuthenticationError(self._provider, str(e)) from e
+            raise ProviderError(self._provider, e) from e
+
+    async def _close_client(self) -> None:
+        """Close and discard the cached HTTP client.
+
+        Forces credential re-resolution on next request, preventing
+        stale credentials from being reused after invalidation.
+        """
+        client = self._client
+        self._client = None
+        if client is not None:
+            try:
+                await client.close()
+            except Exception:
+                pass
+
+    async def stream(
+        self,
+        messages: list[Message],
+        *,
+        tools: list[dict[str, Any]] | None = None,
+        hyperparams: ModelHyperparameters | None = None,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        merged_params = merge_hyperparams(self._default_hyperparams, hyperparams)
+        client = await self._get_client()
+
+        accumulator = ToolCallAccumulator()
+        usage_data: dict[str, int] | None = None
+
+        try:
+            kwargs = build_chat_completion_kwargs(self._model, messages, tools, merged_params)
+            kwargs = self._translate_params(kwargs)
+            kwargs["stream"] = True
+            kwargs["stream_options"] = {"include_usage": True}
+
+            stream_resp = await client.chat.completions.create(**kwargs)
+
+            async for chunk in stream_resp:
+                if chunk.usage:
+                    usage_data = {
+                        "prompt_tokens": chunk.usage.prompt_tokens,
+                        "completion_tokens": chunk.usage.completion_tokens,
+                        "total_tokens": chunk.usage.total_tokens,
+                    }
+                    cached = extract_cached_tokens(chunk.usage)
+                    if cached:
+                        usage_data["cached_tokens"] = cached
+
+                if not chunk.choices:
+                    continue
+
+                choice = chunk.choices[0]
+                delta = choice.delta
+
+                if delta.content:
+                    yield StreamEvent(type="text_delta", text=delta.content)
+
+                if delta.tool_calls:
+                    for tc_delta in delta.tool_calls:
+                        idx = tc_delta.index
+                        if idx is None and len(delta.tool_calls) == 1:
+                            idx = accumulator.default_index()
+                        tool_delta = ToolCallDelta(
+                            index=idx,
+                            id=tc_delta.id,
+                            name=tc_delta.function.name if tc_delta.function else None,
+                            arguments_json=(tc_delta.function.arguments or "") if tc_delta.function else "",
+                        )
+                        accumulator.add_delta(tool_delta)
+                        yield StreamEvent(type="tool_call_delta", tool_call_delta=tool_delta)
+
+            if usage_data:
+                yield StreamEvent(type="usage", usage=usage_data)
+
+            final_tool_calls = accumulator.finalize() if accumulator.has_pending() else None
+            yield StreamEvent(
+                type="response_end",
+                tool_calls=final_tool_calls,
+                usage=usage_data,
+            )
+
+        except Exception as e:
+            error_str = str(e).lower()
+            if "authentication" in error_str or "unauthorized" in error_str:
+                await self._credentials.invalidate(self._provider)
+                await self._close_client()
+            yield StreamEvent(type="error", error=str(e))
+
+    async def count_tokens(
+        self,
+        messages: list[Message],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> int:
+        openai_messages = [message_to_openai(m) for m in messages]
+        total = estimate_message_tokens(openai_messages)
+
+        if tools:
+            tools_json = json.dumps(tools)
+            total += len(tools_json) // 4
+
+        return total