diff --git a/routstr/upstream/gemini.py b/routstr/upstream/gemini.py index e7c05799..54bf849b 100644 --- a/routstr/upstream/gemini.py +++ b/routstr/upstream/gemini.py @@ -1,18 +1,13 @@ from __future__ import annotations -import json -from collections.abc import AsyncGenerator from typing import TYPE_CHECKING, Any -from fastapi import Request -from fastapi.responses import Response, StreamingResponse - from . import gemini_messages from .base import BaseUpstreamProvider from .clients.gemini import GeminiClient if TYPE_CHECKING: - from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow + from ..core.db import UpstreamProviderRow from ..payment.models import Model from ..core.logging import get_logger @@ -21,6 +16,15 @@ logger = get_logger(__name__) class GeminiUpstreamProvider(BaseUpstreamProvider): + """Gemini provider — proxies through Gemini's OpenAI-compat surface. + + The chat-completions, embeddings, and models paths all flow through + :meth:`BaseUpstreamProvider.forward_request`; we only override + ``get_request_base_url`` to redirect to ``{base}/openai/...`` and + ``_dispatch_anthropic_messages`` to inject thought-signatures on the + /v1/messages path (see :mod:`gemini_messages` for that rationale). + """ + provider_type = "gemini" default_base_url = "https://generativelanguage.googleapis.com/v1beta" platform_url = "https://aistudio.google.com/app/apikey" @@ -41,7 +45,7 @@ class GeminiUpstreamProvider(BaseUpstreamProvider): @property def client(self) -> GeminiClient: - """Get or create the Gemini API client.""" + """Get or create the Gemini API client (used for the models listing).""" if self._client is None: self._client = GeminiClient(api_key=self.api_key) return self._client @@ -67,7 +71,41 @@ class GeminiUpstreamProvider(BaseUpstreamProvider): } def transform_model_name(self, model_id: str) -> str: - return model_id.removeprefix("gemini/") + """Reduce a routstr model id to the bare upstream Gemini name. + + Gemini's OpenAI-compat surface expects the literal model id + (e.g. ``gemini-3.1-flash-lite-preview``) — no ``gemini/`` provider + prefix and no ``google/`` vendor sub-prefix. Take the last path + segment so we tolerate any of: + + ``gemini-2.0-flash`` + ``gemini/gemini-2.0-flash`` + ``gemini/google/gemini-3.1-flash-lite-preview`` + """ + return model_id.rsplit("/", 1)[-1] + + @property + def compat_base_url(self) -> str: + """Gemini's OpenAI-compat surface, regardless of what's stored. + + Stored ``base_url`` may be ``.../v1beta`` (the native Gemini API + root) or ``.../v1beta/openai`` (already pointed at the compat + surface). Normalize to the latter. + """ + return self.base_url.rstrip("/").removesuffix("/openai") + "/openai" + + def get_request_base_url( + self, path: str, model_obj: "Model | None" = None + ) -> str: + """Route every proxied request to the OpenAI-compat surface. + + Required because the stored ``base_url`` typically points at the + native Gemini API (``/v1beta``), but :meth:`forward_request` + forwards OpenAI-shaped paths (``/chat/completions``, + ``/embeddings``, ``/models``) which only exist under the + ``/openai`` subtree. + """ + return self.compat_base_url async def _dispatch_anthropic_messages( self, @@ -76,284 +114,23 @@ class GeminiUpstreamProvider(BaseUpstreamProvider): *, log_extra: dict[str, Any] | None = None, ) -> tuple[bool, Any, str | None]: - """Dispatch /v1/messages through Gemini's OpenAI-compat endpoint - with thought-signature injection. + """Dispatch /v1/messages via the gemini-specific httpx path. - Two Gemini-specific problems make the default litellm path fail: - - 1. litellm's ``gemini/`` native route mishandles tool-use input - reassembly for some MCP tool schemas (Claude Code reports - ``Invalid tool parameters``). - 2. Gemini 2.5/3 thinking models reject inbound ``functionCall`` - parts that lack a ``thought_signature``, which Anthropic-Messages - clients (Claude Code) never produce. Setting - ``reasoning_effort="none"`` only suppresses *new* thinking; it - does not lift validation on prior tool calls. Litellm + the - openai SDK both drop unknown tool-call fields before the wire, - so we cannot inject the dummy signature - (``"skip_thought_signature_validator"``, see Google's - thought-signatures docs FAQ #1) via the litellm path. - - ``gemini_messages.dispatch_gemini_messages`` solves both by - translating the Anthropic body to OpenAI form via litellm's - translator, injecting - ``extra_content.google.thought_signature`` on every tool_call, - POSTing directly to ``{base}/openai/chat/completions`` via - ``httpx`` (preserves arbitrary fields), and translating the - OpenAI streaming response back to Anthropic SSE events. + See :mod:`routstr.upstream.gemini_messages` for the full rationale + (thought-signature injection, why litellm + the openai SDK can't + carry the required ``extra_content`` field). """ - compat_base_url = ( - self.base_url.rstrip("/").removesuffix("/openai") + "/openai" - ) return await gemini_messages.dispatch_gemini_messages( request_body=request_body, model_obj=model_obj, - base_url=compat_base_url, + base_url=self.compat_base_url, api_key=self.api_key, transform_model_name=self.transform_model_name, log_extra=log_extra, ) - async def forward_request( - self, - request: Request, - path: str, - headers: dict, - request_body: bytes | None, - key: ApiKey, - max_cost_for_model: int, - session: AsyncSession, - model_obj: Model, - ) -> Response | StreamingResponse: - # Remove provider prefix from model ID for Gemini API - if "/" in model_obj.id: - model_obj.id = model_obj.id.split("/", 1)[1] - - if not path.startswith("chat/completions"): - return await super().forward_request( - request, - path, - headers, - request_body, - key, - max_cost_for_model, - session, - model_obj, - ) - - if not request_body: - return await super().forward_request( - request, - path, - headers, - request_body, - key, - max_cost_for_model, - session, - model_obj, - ) - - try: - openai_data = json.loads(request_body) - messages = openai_data.get("messages", []) - temperature = openai_data.get("temperature") - max_tokens = openai_data.get("max_tokens") - top_p = openai_data.get("top_p") - is_streaming = openai_data.get("stream", False) - - logger.info( - "Processing Gemini request with client abstraction", - extra={ - "model": model_obj.id, - "is_streaming": is_streaming, - "message_count": len(messages), - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - if is_streaming: - final_usage_data: dict | None = None - - def usage_callback(usage_data: dict[str, Any]) -> None: - """Callback to capture usage data during streaming""" - nonlocal final_usage_data - final_usage_data = usage_data - - async def completion_callback( - model: str, usage_data: dict[str, Any] | None - ) -> None: - """Callback to handle payment when streaming completes""" - nonlocal final_usage_data - if usage_data: - final_usage_data = usage_data - - payment_data = { - "model": model, - "usage": final_usage_data, - } - - from ..auth import adjust_payment_for_tokens - from ..core.db import create_session - - async with create_session() as new_session: - fresh_key = await new_session.get(key.__class__, key.hashed_key) - if fresh_key: - try: - cost_data = await adjust_payment_for_tokens( - fresh_key, - payment_data, - new_session, - max_cost_for_model, - ) - - logger.info( - "Gemini streaming payment finalized", - extra={ - "cost_data": cost_data, - "usage_data": final_usage_data, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - except Exception as cost_error: - logger.error( - "Error finalizing Gemini streaming payment", - extra={ - "error": str(cost_error), - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - response_generator = self.client.generate_content_stream( - model=model_obj.id, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p, - usage_callback=usage_callback, - completion_callback=completion_callback, - ) - - async def stream_with_cost() -> AsyncGenerator[bytes, None]: - payment_finalized = False - - async def finalize_payment() -> None: - nonlocal payment_finalized - if payment_finalized: - return - from ..auth import adjust_payment_for_tokens - from ..core.db import create_session - - async with create_session() as new_session: - fresh_key = await new_session.get( - key.__class__, key.hashed_key - ) - if fresh_key: - try: - await adjust_payment_for_tokens( - fresh_key, - { - "model": model_obj.id, - "usage": final_usage_data, - }, - new_session, - max_cost_for_model, - ) - payment_finalized = True - except Exception as cost_error: - logger.error( - "Error finalizing Gemini streaming payment in fallback", - extra={ - "error": str(cost_error), - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - try: - async for chunk in response_generator: - sse_data = f"data: {json.dumps(chunk)}\n\n" - yield sse_data.encode() - except Exception as e: - logger.error( - "Error in Gemini streaming response", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - raise - finally: - if not payment_finalized: - await finalize_payment() - - return StreamingResponse( - stream_with_cost(), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, - ) - - else: - openai_format_response = await self.client.generate_content( - model=model_obj.id, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p, - ) - - from ..auth import adjust_payment_for_tokens - - cost_data = await adjust_payment_for_tokens( - key, openai_format_response, session, max_cost_for_model - ) - await session.refresh(key) - remaining_balance_msats = key.balance - openai_format_response["cost"] = cost_data - openai_format_response["cost"]["sats_cost"] = ( - cost_data.get("total_msats", 0) // 1000 - ) - openai_format_response["cost"]["remaining_balance_msats"] = ( - remaining_balance_msats - ) - - logger.info( - "Gemini non-streaming payment completed", - extra={ - "cost_data": cost_data, - "model": model_obj.id, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - return Response( - content=json.dumps(openai_format_response), - media_type="application/json", - headers={"Cache-Control": "no-cache"}, - ) - - except Exception as e: - logger.error( - "Error in Gemini forward_request", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "path": path, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - return await super().forward_request( - request, - path, - headers, - request_body, - key, - max_cost_for_model, - session, - model_obj, - ) - async def _fetch_provider_models(self) -> dict: - """Fetch models from Gemini API.""" + """Fetch models from Gemini API via the OpenAI-compat client.""" try: models_data = await self.client.list_models() diff --git a/routstr/upstream/gemini_messages.py b/routstr/upstream/gemini_messages.py index f603cfac..ffed1cdb 100644 --- a/routstr/upstream/gemini_messages.py +++ b/routstr/upstream/gemini_messages.py @@ -86,14 +86,12 @@ def inject_thought_signatures(messages: list[dict]) -> None: for tc in tool_calls: if not isinstance(tc, dict): continue - extra = tc.setdefault("extra_content", {}) + extra = tc.get("extra_content") if not isinstance(extra, dict): - extra = {} - tc["extra_content"] = extra - google_cfg = extra.setdefault("google", {}) + extra = tc["extra_content"] = {} + google_cfg = extra.get("google") if not isinstance(google_cfg, dict): - google_cfg = {} - extra["google"] = google_cfg + google_cfg = extra["google"] = {} google_cfg.setdefault("thought_signature", DUMMY_THOUGHT_SIGNATURE) @@ -142,6 +140,9 @@ async def _openai_chunks_to_anthropic_events( next_block_idx = 0 final_finish_reason: str | None = None final_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0} + chunks_seen = 0 + chunks_with_usage = 0 + last_chunk_preview: str | None = None def open_text_block() -> bytes: nonlocal text_block_idx, next_block_idx @@ -198,6 +199,9 @@ async def _openai_chunks_to_anthropic_events( if not isinstance(chunk, dict): continue + chunks_seen += 1 + last_chunk_preview = payload[:500] + if not started: started = True yield _sse_event( @@ -220,14 +224,18 @@ async def _openai_chunks_to_anthropic_events( }, ) + print(f"[gemini-compat upstream chunk] {payload[:1000]}", flush=True) + usage = chunk.get("usage") if isinstance(usage, dict): + chunks_with_usage += 1 in_tok = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 out_tok = usage.get("completion_tokens") or usage.get("output_tokens") or 0 if in_tok: final_usage["input_tokens"] = int(in_tok) if out_tok: final_usage["output_tokens"] = int(out_tok) + print(f"[gemini-compat usage chunk] {usage}", flush=True) choices = chunk.get("choices") or [] if not choices: @@ -286,6 +294,16 @@ async def _openai_chunks_to_anthropic_events( for block_idx in tool_block_indices.values(): yield close_block(block_idx) + print( + f"[gemini-compat stream done] chunks_seen={chunks_seen} " + f"chunks_with_usage={chunks_with_usage} " + f"input_tokens={final_usage['input_tokens']} " + f"output_tokens={final_usage['output_tokens']} " + f"finish_reason={final_finish_reason} " + f"last_chunk={last_chunk_preview}", + flush=True, + ) + # message_delta with stop_reason and usage stop_reason = _FINISH_TO_STOP.get(final_finish_reason or "", "end_turn") yield _sse_event( @@ -418,6 +436,18 @@ async def dispatch_gemini_messages( openai_kwargs.setdefault("reasoning_effort", "none") openai_kwargs["stream"] = True openai_kwargs["model"] = upstream_model + # OpenAI-compat backends (including Gemini's) only emit a final + # ``usage`` chunk when the request opts in via this flag. Without it + # the cost-calculation pipeline can't read real token counts and + # falls back to MaxCostData billing. + existing_stream_options = openai_kwargs.get("stream_options") + merged_stream_options = ( + dict(existing_stream_options) + if isinstance(existing_stream_options, dict) + else {} + ) + merged_stream_options.setdefault("include_usage", True) + openai_kwargs["stream_options"] = merged_stream_options logger.info( "Dispatching /v1/messages via gemini compat (httpx)", @@ -430,6 +460,12 @@ async def dispatch_gemini_messages( **(log_extra or {}), }, ) + print( + f"[gemini-compat outgoing] model={upstream_model} " + f"stream_options={openai_kwargs.get('stream_options')} " + f"reasoning_effort={openai_kwargs.get('reasoning_effort')}", + flush=True, + ) http_client, response = await _post_and_stream( base_url, api_key, openai_kwargs, log_extra