diff --git a/src/rotator_library/client/executor.py b/src/rotator_library/client/executor.py index 6c88ff039..84e730761 100644 --- a/src/rotator_library/client/executor.py +++ b/src/rotator_library/client/executor.py @@ -619,9 +619,16 @@ async def _execute_non_streaming( await self._run_pre_request_callback(context, kwargs) # Make the API call + is_embedding = context.request_type == "embedding" + if plugin and plugin.has_custom_logic(): kwargs["credential_identifier"] = cred - response = await plugin.acompletion( + call_fn = ( + plugin.aembedding + if is_embedding + else plugin.acompletion + ) + response = await call_fn( self._http_client, **kwargs ) else: @@ -630,7 +637,12 @@ async def _execute_non_streaming( self._apply_litellm_logger(kwargs) # Remove internal context before litellm call kwargs.pop("transaction_context", None) - response = await litellm.acompletion(**kwargs) + call_fn = ( + litellm.aembedding + if is_embedding + else litellm.acompletion + ) + response = await call_fn(**kwargs) # Success! Extract token usage if available ( diff --git a/src/rotator_library/client/rotating_client.py b/src/rotator_library/client/rotating_client.py index ffcfbd9c1..61b577953 100644 --- a/src/rotator_library/client/rotating_client.py +++ b/src/rotator_library/client/rotating_client.py @@ -426,7 +426,7 @@ async def acompletion( return await self._executor.execute(context) - def aembedding( + async def aembedding( self, request: Optional[Any] = None, pre_request_callback: Optional[callable] = None, @@ -449,13 +449,14 @@ def aembedding( provider=provider, kwargs=kwargs, streaming=False, + request_type="embedding", credentials=self.all_credentials.get(provider, []), deadline=time.time() + self.global_timeout, request=request, pre_request_callback=pre_request_callback, ) - return self._executor.execute(context) + return await self._executor.execute(context) def token_count(self, **kwargs) -> int: """Calculate token count for text or messages.""" diff --git a/src/rotator_library/core/types.py b/src/rotator_library/core/types.py index 6da9cd013..bef96dc8f 100644 --- a/src/rotator_library/core/types.py +++ b/src/rotator_library/core/types.py @@ -76,6 +76,7 @@ class RequestContext: streaming: bool credentials: List[str] deadline: float + request_type: Literal["completion", "embedding"] = "completion" session_id: Optional[str] = None request: Optional[Any] = None # FastAPI Request object pre_request_callback: Optional[Callable] = None diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 6a4f755e0..0fabb9c94 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -1824,6 +1824,181 @@ async def logging_stream_wrapper(): raise last_error raise ValueError("No fallback models available") + @staticmethod + def _normalize_embedding_inputs(raw_input: Any) -> List[str]: + if raw_input is None: + return [] + if isinstance(raw_input, str): + return [raw_input] + texts: List[str] = [] + for item in raw_input: + if item is None: + continue + if isinstance(item, str): + texts.append(item) + else: + texts.append(str(item)) + return texts + + @staticmethod + def _parse_embed_content_response( + data: Dict[str, Any], + ) -> Tuple[List[float], int]: + payload = data.get("response", data) + embedding = payload.get("embedding") + values = None + if isinstance(embedding, dict): + values = embedding.get("values") + elif isinstance(embedding, list): + values = embedding + if values is None: + embeddings = payload.get("embeddings") + if isinstance(embeddings, list) and embeddings: + first = embeddings[0] + if isinstance(first, dict): + values = first.get("values") + elif isinstance(first, list): + values = first + if not values: + raise ValueError( + f"Gemini CLI embedContent returned no embedding values: {data}" + ) + + usage = payload.get("usageMetadata") or data.get("usageMetadata") or {} + tokens = ( + usage.get("promptTokenCount") + or usage.get("totalTokenCount") + or usage.get("totalTokens") + or 0 + ) + return list(values), int(tokens) + + async def aembedding( + self, client: httpx.AsyncClient, **kwargs + ) -> litellm.EmbeddingResponse: + model = kwargs["model"] + credential_path = kwargs.pop("credential_identifier") + kwargs.pop("transaction_context", None) + + auth_header = await self.get_auth_header(credential_path) + project_id = self.project_id_cache.get(credential_path) + if not project_id: + access_token = auth_header["Authorization"].split(" ")[1] + project_id = await self._discover_project_id( + credential_path, access_token, kwargs.get("litellm_params", {}) + ) + + model_name = model.split("/")[-1] + texts = self._normalize_embedding_inputs(kwargs.get("input")) + + headers = auth_header.copy() + headers.update(self._get_gemini_cli_request_headers(model_name)) + + data_items: List[Dict[str, Any]] = [] + total_tokens = 0 + for index, text in enumerate(texts): + request_body: Dict[str, Any] = { + "content": {"parts": [{"text": text}]}, + } + dimensions = kwargs.get("dimensions") + if dimensions: + request_body["outputDimensionality"] = dimensions + # HTTP EmbeddingRequest exposes input_type; accept task_type/taskType too. + task_type = ( + kwargs.get("task_type") + or kwargs.get("taskType") + or kwargs.get("input_type") + ) + if task_type: + request_body["taskType"] = task_type + + request_payload = { + "model": model_name, + "project": project_id, + "request": request_body, + } + + last_endpoint_error = None + response_data = None + for endpoint_idx, base_endpoint in enumerate(GEMINI_CLI_ENDPOINT_FALLBACKS): + url = f"{base_endpoint}:embedContent" + try: + response = await client.post( + url, + headers=headers, + json=request_payload, + timeout=TimeoutConfig.non_streaming(), + ) + response.raise_for_status() + response_data = response.json() + last_endpoint_error = None + break + except httpx.HTTPStatusError as e: + error_body = None + if e.response is not None: + try: + error_body = e.response.text + except Exception: + pass + if e.response is not None and e.response.status_code == 429: + retry_after = extract_retry_after_from_body(error_body) + retry_info = ( + f" (retry after {retry_after}s)" if retry_after else "" + ) + error_msg = f"Gemini CLI rate limit exceeded{retry_info}" + if error_body: + error_msg = f"{error_msg} | {error_body}" + raise RateLimitError( + message=error_msg, + llm_provider="gemini_cli", + model=model, + response=e.response, + ) + if ( + e.response is not None + and e.response.status_code >= 500 + and endpoint_idx < len(GEMINI_CLI_ENDPOINT_FALLBACKS) - 1 + ): + last_endpoint_error = e + lib_logger.warning( + f"embedContent: endpoint {base_endpoint} returned {e.response.status_code}, trying fallback" + ) + continue + raise + except (httpx.ConnectError, httpx.TimeoutException) as e: + last_endpoint_error = e + if endpoint_idx < len(GEMINI_CLI_ENDPOINT_FALLBACKS) - 1: + lib_logger.warning( + f"embedContent: connection error to {base_endpoint}, trying fallback" + ) + continue + raise + + if response_data is None: + if last_endpoint_error: + raise last_endpoint_error + raise ValueError("Gemini CLI embedContent failed with no response") + + values, tokens = self._parse_embed_content_response(response_data) + total_tokens += tokens + data_items.append( + { + "object": "embedding", + "index": index, + "embedding": values, + } + ) + + return litellm.EmbeddingResponse( + model=model, + data=data_items, + usage=litellm.Usage( + prompt_tokens=total_tokens, + completion_tokens=0, + total_tokens=total_tokens, + ), + ) + async def count_tokens( self, client: httpx.AsyncClient, diff --git a/src/rotator_library/request_sanitizer.py b/src/rotator_library/request_sanitizer.py index 083ae366e..536ba027a 100644 --- a/src/rotator_library/request_sanitizer.py +++ b/src/rotator_library/request_sanitizer.py @@ -3,11 +3,22 @@ from typing import Dict, Any + +def _supports_dimensions(model: str) -> bool: + """Models that accept an OpenAI-style `dimensions` embedding parameter.""" + if model.startswith("openai/text-embedding-3"): + return True + # Gemini Code Assist :embedContent maps dimensions -> outputDimensionality. + if model.startswith("gemini_cli/") and "embedding" in model.lower(): + return True + return False + + def sanitize_request_payload(payload: Dict[str, Any], model: str) -> Dict[str, Any]: """ Removes unsupported parameters from the request payload based on the model. """ - if "dimensions" in payload and not model.startswith("openai/text-embedding-3"): + if "dimensions" in payload and not _supports_dimensions(model): del payload["dimensions"] if payload.get("thinking") == {"type": "enabled", "budget_tokens": -1}: diff --git a/tests/test_gemini_cli_embeddings.py b/tests/test_gemini_cli_embeddings.py new file mode 100644 index 000000000..af481297a --- /dev/null +++ b/tests/test_gemini_cli_embeddings.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""Embedding requests for gemini_cli must hit aembedding / :embedContent, not acompletion.""" + +import sys +import time +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +SRC = Path(__file__).resolve().parents[1] / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +from rotator_library.client.executor import RequestExecutor +from rotator_library.client.rotating_client import RotatingClient +from rotator_library.core.types import FilterResult, RequestContext +from litellm.exceptions import RateLimitError +from rotator_library.providers.gemini_cli_provider import GeminiCliProvider + + +class _FakeCredContext: + def __init__(self, credential: str): + self.credential = credential + self.stable_id = "stable-id" + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def mark_success(self, **kwargs): + return None + + +class _FakeUsageManager: + initialized = True + states = {} + + def get_model_quota_group(self, model): + return None + + async def get_availability_stats(self, model, quota_group): + return { + "available": 1, + "total": 1, + "rotation_mode": "sequential", + "blocked_by": {}, + } + + async def acquire_credential(self, **kwargs): + return _FakeCredContext("cred.json") + + +class _FakeFilter: + def filter_by_tier(self, credentials, model, provider): + return FilterResult(compatible=list(credentials)) + + +class _FakeTransforms: + async def apply(self, provider, model, cred, kwargs): + return kwargs + + +def _embedding_response(): + return SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=1, completion_tokens=0), + model_dump=lambda: {}, + ) + + +class TestExecutorEmbeddingDispatch(unittest.IsolatedAsyncioTestCase): + async def test_custom_provider_embedding_calls_aembedding_not_acompletion(self): + plugin = MagicMock() + plugin.has_custom_logic.return_value = True + plugin.skip_cost_calculation = True + plugin.aembedding = AsyncMock(return_value=_embedding_response()) + plugin.acompletion = AsyncMock(return_value=_embedding_response()) + + executor = RequestExecutor( + usage_managers={"gemini_cli": _FakeUsageManager()}, + cooldown_manager=None, + credential_filter=_FakeFilter(), + provider_transforms=_FakeTransforms(), + provider_plugins={"gemini_cli": plugin}, + http_client=MagicMock(), + ) + + context = RequestContext( + model="gemini_cli/gemini-embedding-001", + provider="gemini_cli", + kwargs={ + "model": "gemini_cli/gemini-embedding-001", + "input": ["test"], + }, + streaming=False, + credentials=["cred.json"], + deadline=time.time() + 30, + request_type="embedding", + ) + + await executor.execute(context) + + plugin.aembedding.assert_awaited_once() + plugin.acompletion.assert_not_awaited() + + async def test_custom_provider_completion_still_calls_acompletion(self): + plugin = MagicMock() + plugin.has_custom_logic.return_value = True + plugin.skip_cost_calculation = True + plugin.aembedding = AsyncMock(return_value=_embedding_response()) + plugin.acompletion = AsyncMock(return_value=_embedding_response()) + + executor = RequestExecutor( + usage_managers={"gemini_cli": _FakeUsageManager()}, + cooldown_manager=None, + credential_filter=_FakeFilter(), + provider_transforms=_FakeTransforms(), + provider_plugins={"gemini_cli": plugin}, + http_client=MagicMock(), + ) + + context = RequestContext( + model="gemini_cli/gemini-2.5-flash", + provider="gemini_cli", + kwargs={ + "model": "gemini_cli/gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + }, + streaming=False, + credentials=["cred.json"], + deadline=time.time() + 30, + ) + + await executor.execute(context) + + plugin.acompletion.assert_awaited_once() + plugin.aembedding.assert_not_awaited() + + async def test_rotating_client_aembedding_sets_request_type(self): + captured = {} + + class _FakeExecutor: + async def execute(self, context): + captured["request_type"] = context.request_type + return _embedding_response() + + client = RotatingClient.__new__(RotatingClient) + client.all_credentials = {"gemini_cli": ["cred.json"]} + client.global_timeout = 30 + client._executor = _FakeExecutor() + + await client.aembedding( + model="gemini_cli/gemini-embedding-001", input=["test"] + ) + + self.assertEqual(captured["request_type"], "embedding") + + +class TestGeminiCliProviderEmbedding(unittest.IsolatedAsyncioTestCase): + async def test_aembedding_posts_to_embed_content_not_stream_generate(self): + provider = GeminiCliProvider() + provider.project_id_cache["cred.json"] = "test-project" + provider.get_auth_header = AsyncMock( + return_value={"Authorization": "Bearer fake-token"} + ) + + posted = {} + + class _FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return {"embedding": {"values": [0.1, 0.2, 0.3]}} + + async def fake_post(url, **kwargs): + posted["url"] = url + posted["json"] = kwargs.get("json") + return _FakeResponse() + + client = MagicMock() + client.post = fake_post + + response = await provider.aembedding( + client, + model="gemini_cli/gemini-embedding-001", + input=["test"], + credential_identifier="cred.json", + ) + + self.assertIn(":embedContent", posted["url"]) + self.assertNotIn("streamGenerateContent", posted["url"]) + self.assertEqual(posted["json"]["model"], "gemini-embedding-001") + self.assertEqual( + posted["json"]["request"]["content"]["parts"][0]["text"], "test" + ) + + data = response.data + first = data[0] + values = first["embedding"] if isinstance(first, dict) else first.embedding + self.assertEqual(list(values), [0.1, 0.2, 0.3]) + + async def test_aembedding_accepts_code_assist_wrapped_response(self): + provider = GeminiCliProvider() + provider.project_id_cache["cred.json"] = "test-project" + provider.get_auth_header = AsyncMock( + return_value={"Authorization": "Bearer fake-token"} + ) + + class _FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return { + "response": { + "embedding": {"values": [0.4, 0.5]}, + "usageMetadata": {"promptTokenCount": 2}, + } + } + + client = MagicMock() + client.post = AsyncMock(return_value=_FakeResponse()) + + response = await provider.aembedding( + client, + model="gemini_cli/gemini-embedding-001", + input="hello", + credential_identifier="cred.json", + ) + first = response.data[0] + values = first["embedding"] if isinstance(first, dict) else first.embedding + self.assertEqual(list(values), [0.4, 0.5]) + self.assertEqual(response.usage.prompt_tokens, 2) + + async def test_aembedding_maps_dimensions_and_input_type(self): + provider = GeminiCliProvider() + provider.project_id_cache["cred.json"] = "test-project" + provider.get_auth_header = AsyncMock( + return_value={"Authorization": "Bearer fake-token"} + ) + + posted = {} + + class _FakeResponse: + def raise_for_status(self): + return None + + def json(self): + return {"embedding": {"values": [0.1, 0.2]}} + + async def fake_post(url, **kwargs): + posted["json"] = kwargs.get("json") + return _FakeResponse() + + client = MagicMock() + client.post = fake_post + + await provider.aembedding( + client, + model="gemini_cli/gemini-embedding-001", + input=["doc"], + dimensions=768, + input_type="RETRIEVAL_DOCUMENT", + credential_identifier="cred.json", + ) + + req = posted["json"]["request"] + self.assertEqual(req["outputDimensionality"], 768) + self.assertEqual(req["taskType"], "RETRIEVAL_DOCUMENT") + + async def test_aembedding_429_raises_rate_limit_error(self): + import httpx + + provider = GeminiCliProvider() + provider.project_id_cache["cred.json"] = "test-project" + provider.get_auth_header = AsyncMock( + return_value={"Authorization": "Bearer fake-token"} + ) + + class _FakeResponse: + status_code = 429 + text = '{"error":"rate limited"}' + + def raise_for_status(self): + raise httpx.HTTPStatusError( + "429", + request=MagicMock(), + response=self, + ) + + client = MagicMock() + client.post = AsyncMock(return_value=_FakeResponse()) + + with self.assertRaises(RateLimitError): + await provider.aembedding( + client, + model="gemini_cli/gemini-embedding-001", + input=["x"], + credential_identifier="cred.json", + ) + + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_request_sanitizer_dimensions.py b/tests/test_request_sanitizer_dimensions.py new file mode 100644 index 000000000..a2950df9b --- /dev/null +++ b/tests/test_request_sanitizer_dimensions.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: LGPL-3.0-only +# Copyright (c) 2026 Mirrowel + +"""dimensions must reach gemini_cli embeddings (review on #169).""" + +import sys +import unittest +from pathlib import Path + +SRC = Path(__file__).resolve().parents[1] / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +from rotator_library.request_sanitizer import sanitize_request_payload + + +class TestSanitizeDimensions(unittest.TestCase): + def test_keeps_dimensions_for_openai_text_embedding_3(self): + out = sanitize_request_payload( + {"input": ["a"], "dimensions": 256}, + "openai/text-embedding-3-small", + ) + self.assertEqual(out.get("dimensions"), 256) + + def test_keeps_dimensions_for_gemini_cli_embedding_models(self): + out = sanitize_request_payload( + {"input": ["a"], "dimensions": 768, "taskType": "RETRIEVAL_DOCUMENT"}, + "gemini_cli/gemini-embedding-001", + ) + self.assertEqual(out.get("dimensions"), 768) + self.assertEqual(out.get("taskType"), "RETRIEVAL_DOCUMENT") + + def test_strips_dimensions_for_unrelated_chat_models(self): + out = sanitize_request_payload( + {"messages": [], "dimensions": 256}, + "gemini_cli/gemini-2.5-flash", + ) + self.assertNotIn("dimensions", out) + + +if __name__ == "__main__": + unittest.main()