From f263e03250d4903e32f818cd56e512014bcefc38 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Wed, 18 Mar 2026 22:13:44 +0100 Subject: [PATCH 1/7] feat(embedder): add GeminiDenseEmbedder text embedding provider - Adds GeminiDenseEmbedder backed by google-genai>=0.8.0 SDK - Supports all 8 task types, MRL dimensions 1-3072, async batch via anyio - Wires gemini provider into EmbeddingModelConfig factory + context routing - Adds unit tests (mocked, no API key) and IT tests (auto-skip without GOOGLE_API_KEY) - Updates pyproject.toml with google-genai, anyio, pytest-asyncio, pytest-cov deps --- examples/ov.conf.example | 11 + openviking/models/embedder/__init__.py | 4 + .../models/embedder/gemini_embedders.py | 247 ++++++++++++++ .../utils/config/embedding_config.py | 120 ++++++- pyproject.toml | 8 + tests/integration/conftest.py | 109 +++++++ tests/integration/test_gemini_e2e.py | 150 +++++++++ tests/integration/test_gemini_embedding_it.py | 103 ++++++ .../integration/test_gemini_openviking_it.py | 186 +++++++++++ tests/unit/test_embedding_config_gemini.py | 105 ++++++ tests/unit/test_gemini_embedder.py | 306 ++++++++++++++++++ 11 files changed, 1338 insertions(+), 11 deletions(-) create mode 100644 openviking/models/embedder/gemini_embedders.py create mode 100644 tests/integration/test_gemini_e2e.py create mode 100644 tests/integration/test_gemini_embedding_it.py create mode 100644 tests/integration/test_gemini_openviking_it.py create mode 100644 tests/unit/test_embedding_config_gemini.py create mode 100644 tests/unit/test_gemini_embedder.py diff --git a/examples/ov.conf.example b/examples/ov.conf.example index 249c0a13..d2fd1775 100644 --- a/examples/ov.conf.example +++ b/examples/ov.conf.example @@ -44,6 +44,17 @@ "input": "multimodal" } }, + "embedding_gemini_example": { + "_comment": "Gemini alternative for dense embedding:", + "dense": { + "model": "gemini-embedding-2-preview", + "provider": "gemini", + "api_key": "${GOOGLE_API_KEY}", + "dimension": 1536, + "query_param": "RETRIEVAL_QUERY", + "document_param": "RETRIEVAL_DOCUMENT" + } + }, "embedding_ollama_example": { "_comment": "For local deployment with Ollama (no API key required):", "dense": { diff --git a/openviking/models/embedder/__init__.py b/openviking/models/embedder/__init__.py index b418b809..3fc66a06 100644 --- a/openviking/models/embedder/__init__.py +++ b/openviking/models/embedder/__init__.py @@ -12,6 +12,7 @@ - OpenAI: Dense only - Volcengine: Dense, Sparse, Hybrid - Jina AI: Dense only +- Google Gemini: Dense only (multimodal) - Voyage AI: Dense only """ @@ -23,6 +24,7 @@ HybridEmbedderBase, SparseEmbedderBase, ) +from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder from openviking.models.embedder.jina_embedders import JinaDenseEmbedder from openviking.models.embedder.openai_embedders import OpenAIDenseEmbedder from openviking.models.embedder.voyage_embedders import VoyageDenseEmbedder @@ -45,6 +47,8 @@ "SparseEmbedderBase", "HybridEmbedderBase", "CompositeHybridEmbedder", + # Google Gemini implementations + "GeminiDenseEmbedder", # Jina AI implementations "JinaDenseEmbedder", # OpenAI implementations diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py new file mode 100644 index 00000000..40d00f9e --- /dev/null +++ b/openviking/models/embedder/gemini_embedders.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Gemini Embedding 2 provider using the official google-genai SDK.""" + +import math +from typing import Any, Dict, List, Optional + +from google import genai +from google.genai import types +from google.genai.errors import APIError, ClientError + +import logging + +try: + import anyio + _ANYIO_AVAILABLE = True +except ImportError: + _ANYIO_AVAILABLE = False + +from openviking.models.embedder.base import ( + DenseEmbedderBase, + EmbedResult, + truncate_and_normalize, +) + +logger = logging.getLogger("gemini_embedders") + +_TEXT_BATCH_SIZE = 100 + +# Keep for backward-compat with existing unit tests that import it +_GEMINI_INPUT_TOKEN_LIMIT = 8192 # gemini-embedding-2-preview hard limit + +# Per-model token limits (Google API hard limits, from official docs) +_MODEL_TOKEN_LIMITS: Dict[str, int] = { + "gemini-embedding-2-preview": 8192, + "gemini-embedding-001": 2048, +} +_DEFAULT_TOKEN_LIMIT = 2048 # conservative fallback for unknown future models + +_VALID_TASK_TYPES: frozenset = frozenset({ + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + "CODE_RETRIEVAL_QUERY", +}) + +_ERROR_HINTS: Dict[int, str] = { + 400: "Invalid request — check model name and task_type value.", + 401: "Invalid API key. Verify your GOOGLE_API_KEY or api_key in config.", + 403: "Permission denied. API key may lack access to this model.", + 404: "Model not found: '{model}'. Check spelling (e.g. 'gemini-embedding-2-preview').", + 429: "Quota exceeded. Wait and retry, or increase your Google API quota.", + 500: "Gemini service error (Google-side). Retry after a delay.", + 503: "Gemini service unavailable. Retry after a delay.", +} + + +def _l2_normalize(vector: List[float]) -> List[float]: + """Ensure vector has unit L2 norm. No-op for zero vectors.""" + norm = math.sqrt(sum(v * v for v in vector)) + return [v / norm for v in vector] if norm > 0 else vector + + +def _raise_api_error(e: APIError, model: str) -> None: + hint = _ERROR_HINTS.get(e.code, "") + # Gemini returns HTTP 400 (not 401) when the API key is invalid + if e.code == 400 and "api key" in str(e).lower(): + hint = "Invalid API key. Verify your GOOGLE_API_KEY or api_key in config." + msg = f"Gemini embedding failed (HTTP {e.code})" + if hint: + msg += f": {hint.format(model=model)}" + raise RuntimeError(msg) from e + + +class GeminiDenseEmbedder(DenseEmbedderBase): + """Dense embedder backed by Google's Gemini Embedding models. + + REST endpoint: /v1beta/models/{model}:embedContent (SDK handles Parts format internally). + Input token limit: per-model (8192 for gemini-embedding-2-preview, 2048 for gemini-embedding-001). + Output dimension: 1–3072 (MRL; recommended 768, 1536, 3072; default 3072). + Task types: RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, SEMANTIC_SIMILARITY, CLASSIFICATION, + CLUSTERING, CODE_RETRIEVAL_QUERY, QUESTION_ANSWERING, FACT_VERIFICATION. + Non-symmetric: use query_param/document_param in EmbeddingModelConfig. + """ + + # Default output dimensions per model (used when user does not specify `dimension`). + # gemini-embedding-2-preview: 3072 MRL model — supports 1–3072 via output_dimensionality + # gemini-embedding-001: 3072 (native 768-dim vectors; 3072 shown as default for MRL compat) + # text-embedding-004: 768 fixed-dim legacy model, does not support MRL truncation + # Future gemini-embedding-*: default 3072 via _default_dimension() fallback + # Future text-embedding-*: default 768 via _default_dimension() prefix rule + KNOWN_DIMENSIONS: Dict[str, int] = { + "gemini-embedding-2-preview": 3072, + "gemini-embedding-001": 3072, + "text-embedding-004": 768, + } + + @classmethod + def _default_dimension(cls, model: str) -> int: + """Return default output dimension for a Gemini model. + + Lookup order: + 1. Exact match in KNOWN_DIMENSIONS + 2. Prefix rule: text-embedding-* → 768 (legacy fixed-dim series) + 3. Fallback: 3072 (gemini-embedding-* MRL models) + + Examples: + gemini-embedding-2-preview → 3072 (exact match) + gemini-embedding-2 → 3072 (fallback — future model) + text-embedding-004 → 768 (exact match) + text-embedding-005 → 768 (prefix rule — future model) + """ + if model in cls.KNOWN_DIMENSIONS: + return cls.KNOWN_DIMENSIONS[model] + if model.startswith("text-embedding-"): + return 768 + return 3072 + + def __init__( + self, + model_name: str = "gemini-embedding-2-preview", + api_key: Optional[str] = None, + dimension: Optional[int] = None, + task_type: Optional[str] = None, + max_concurrent_batches: int = 10, + config: Optional[Dict[str, Any]] = None, + ): + super().__init__(model_name, config) + if not api_key: + raise ValueError("Gemini provider requires api_key") + if task_type and task_type not in _VALID_TASK_TYPES: + raise ValueError( + f"Invalid task_type '{task_type}'. " + f"Valid values: {', '.join(sorted(_VALID_TASK_TYPES))}" + ) + if dimension is not None and not (1 <= dimension <= 3072): + raise ValueError(f"dimension must be between 1 and 3072, got {dimension}") + self.client = genai.Client(api_key=api_key) + self.task_type = task_type + self._dimension = dimension or self._default_dimension(model_name) + self._token_limit = _MODEL_TOKEN_LIMITS.get(model_name, _DEFAULT_TOKEN_LIMIT) + self._max_concurrent_batches = max_concurrent_batches + config_kwargs: Dict[str, Any] = {"output_dimensionality": self._dimension} + if self.task_type: + config_kwargs["task_type"] = self.task_type + self._embed_config = types.EmbedContentConfig(**config_kwargs) + + def embed(self, text: str) -> EmbedResult: + # SDK accepts plain str; converts to REST Parts format internally. + try: + result = self.client.models.embed_content( + model=self.model_name, + contents=text, + config=self._embed_config, + ) + vector = _l2_normalize( + truncate_and_normalize(list(result.embeddings[0].values), self._dimension) + ) + return EmbedResult(dense_vector=vector) + except (APIError, ClientError) as e: + _raise_api_error(e, self.model_name) + + def embed_batch(self, texts: List[str]) -> List[EmbedResult]: + if not texts: + return [] + results: List[EmbedResult] = [] + for i in range(0, len(texts), _TEXT_BATCH_SIZE): + batch = texts[i : i + _TEXT_BATCH_SIZE] + try: + response = self.client.models.embed_content( + model=self.model_name, + contents=batch, + config=self._embed_config, + ) + for emb in response.embeddings: + vector = _l2_normalize( + truncate_and_normalize(list(emb.values), self._dimension) + ) + results.append(EmbedResult(dense_vector=vector)) + except (APIError, ClientError) as e: + logger.warning( + "Gemini batch embed failed (HTTP %d) for batch of %d, falling back to individual", + e.code, len(batch), + ) + for text in batch: + results.append(self.embed(text)) + return results + + async def async_embed_batch(self, texts: List[str]) -> List[EmbedResult]: + """Concurrent batch embedding via client.aio — requires anyio to be installed. + + Dispatches all 100-text chunks in parallel, bounded by max_concurrent_batches. + Per-batch APIError falls back to individual embed() calls via thread pool. + Raises ImportError if anyio is not installed. + """ + if not _ANYIO_AVAILABLE: + raise ImportError( + "anyio is required for async_embed_batch: pip install 'openviking[gemini-async]'" + ) + if not texts: + return [] + batches = [texts[i : i + _TEXT_BATCH_SIZE] for i in range(0, len(texts), _TEXT_BATCH_SIZE)] + results: List[Optional[List[EmbedResult]]] = [None] * len(batches) + sem = anyio.Semaphore(self._max_concurrent_batches) + + async def _embed_one(idx: int, batch: List[str]) -> None: + async with sem: + try: + response = await self.client.aio.models.embed_content( + model=self.model_name, contents=batch, config=self._embed_config + ) + results[idx] = [ + EmbedResult( + dense_vector=_l2_normalize( + truncate_and_normalize(list(emb.values), self._dimension) + ) + ) + for emb in response.embeddings + ] + except (APIError, ClientError) as e: + logger.warning( + "Gemini async batch embed failed (HTTP %d) for batch of %d, falling back", + e.code, len(batch), + ) + results[idx] = [ + await anyio.to_thread.run_sync(self.embed, text) for text in batch + ] + + async with anyio.create_task_group() as tg: + for idx, batch in enumerate(batches): + tg.start_soon(_embed_one, idx, batch) + + return [r for batch_results in results for r in (batch_results or [])] + + def get_dimension(self) -> int: + return self._dimension + + def close(self): + if hasattr(self.client, "_http_client"): + try: + self.client._http_client.close() + except Exception: + pass diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index a0a1b964..2397034f 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -14,10 +14,14 @@ class EmbeddingModelConfig(BaseModel): dimension: Optional[int] = Field(default=None, description="Embedding dimension") batch_size: int = Field(default=32, description="Batch size for embedding generation") input: str = Field(default="multimodal", description="Input type: 'text' or 'multimodal'") + input_type: Optional[str] = Field( + default=None, + description="Input type for non-contextual OpenAI embeddings: 'query', 'document', or 'passage'", + ) query_param: Optional[str] = Field( default=None, description=( - "Parameter value for query-side embeddings when calling embed(is_query=True). " + "Parameter value for query-side embeddings when using get_query_embedder(). " "For OpenAI-compatible models, this maps to 'input_type' (e.g., 'query', 'search_query'). " "For Jina models, this maps to 'task' (e.g., 'retrieval.query'). " "Setting this or document_param activates non-symmetric mode. " @@ -27,7 +31,7 @@ class EmbeddingModelConfig(BaseModel): document_param: Optional[str] = Field( default=None, description=( - "Parameter value for document-side embeddings when calling embed(is_query=False). " + "Parameter value for document-side embeddings when using get_document_embedder(). " "For OpenAI-compatible models, this maps to 'input_type' (e.g., 'passage', 'document'). " "For Jina models, this maps to 'task' (e.g., 'retrieval.passage'). " "Setting this or query_param activates non-symmetric mode. " @@ -37,20 +41,36 @@ class EmbeddingModelConfig(BaseModel): provider: Optional[str] = Field( default="volcengine", description=( - "Provider type: 'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage'. " + "Provider type: 'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'gemini', 'voyage'. " "For OpenRouter or other OpenAI-compatible providers, use 'openai' with " "api_base and extra_headers." ), ) backend: Optional[str] = Field( default="volcengine", - description="Backend type (Deprecated, use 'provider' instead): 'openai', 'volcengine', 'vikingdb', 'voyage'", + description=( + "Backend type (Deprecated, use 'provider' instead): " + "'openai', 'volcengine', 'vikingdb', 'voyage', 'gemini'" + ), ) version: Optional[str] = Field(default=None, description="Model version") ak: Optional[str] = Field(default=None, description="Access Key ID for VikingDB API") sk: Optional[str] = Field(default=None, description="Access Key Secretfor VikingDB API") region: Optional[str] = Field(default=None, description="Region for VikingDB API") host: Optional[str] = Field(default=None, description="Host for VikingDB API") + max_tokens: Optional[int] = Field( + default=None, + description="Maximum token count per embedding request. If None, uses model default (e.g., 8000 for OpenAI).", + ) + task_type: Optional[str] = Field( + default=None, + description=( + "Gemini task type. Valid: RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, " + "SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING, " + "QUESTION_ANSWERING, FACT_VERIFICATION, CODE_RETRIEVAL_QUERY. " + "For non-symmetric mode set query_param/document_param instead." + ), + ) extra_headers: Optional[dict[str, str]] = Field( default=None, description=( @@ -71,7 +91,7 @@ def sync_provider_backend(cls, data: Any) -> Any: if backend is not None and provider is None: data["provider"] = backend - for key in ("query_param", "document_param"): + for key in ("input_type",): value = data.get(key) if isinstance(value, str): data[key] = value.lower() @@ -89,10 +109,12 @@ def validate_config(self): if not self.provider: raise ValueError("Embedding provider is required") - if self.provider not in ["openai", "volcengine", "vikingdb", "jina", "ollama", "voyage"]: + if self.provider not in [ + "openai", "volcengine", "vikingdb", "jina", "ollama", "voyage", "gemini" + ]: raise ValueError( - f"Invalid embedding provider: '{self.provider}'. Must be one of: " - "'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage'" + f"Invalid embedding provider: '{self.provider}'. " + "Must be one of: 'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage', 'gemini'" ) # Provider-specific validation @@ -127,6 +149,25 @@ def validate_config(self): if not self.api_key: raise ValueError("Jina provider requires 'api_key' to be set") + elif self.provider == "gemini": + if not self.api_key: + raise ValueError("Gemini provider requires 'api_key' to be set") + _GEMINI_TASK_TYPES = { + "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", + "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", + "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY", + } + for field_name, value in [ + ("query_param", self.query_param), + ("document_param", self.document_param), + ("task_type", self.task_type), + ]: + if value and value not in _GEMINI_TASK_TYPES: + raise ValueError( + f"Invalid {field_name} '{value}' for Gemini. " + f"Valid task_types: {', '.join(sorted(_GEMINI_TASK_TYPES))}" + ) + elif self.provider == "voyage": if not self.api_key: raise ValueError("Voyage provider requires 'api_key' to be set") @@ -143,15 +184,18 @@ def get_effective_dimension(self) -> int: from openviking.models.embedder.voyage_embedders import ( get_voyage_model_default_dimension, ) - return get_voyage_model_default_dimension(self.model) + if provider == "gemini": + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + return GeminiDenseEmbedder._default_dimension(self.model) + return 2048 class EmbeddingConfig(BaseModel): """ - Embedding configuration, supports OpenAI or VolcEngine compatible APIs. + Embedding configuration, supports OpenAI, VolcEngine, VikingDB, Jina, or Gemini APIs. Structure: - dense: Configuration for dense embedder @@ -185,13 +229,17 @@ def _create_embedder( provider: str, embedder_type: str, config: EmbeddingModelConfig, + context: Optional[str] = None, ): """Factory method to create embedder instance based on provider and type. Args: - provider: Provider type ('openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage') + provider: Provider type ('openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage', 'gemini') embedder_type: Embedder type ('dense', 'sparse', 'hybrid') config: EmbeddingModelConfig instance + context: Optional embedding context ('query' or 'document') for non-symmetric models. + When provided, the embedder will apply the appropriate input_type or task. + Leave None for symmetric models or when using a single embedder for all inputs. Returns: Embedder instance @@ -210,6 +258,7 @@ def _create_embedder( VolcengineSparseEmbedder, VoyageDenseEmbedder, ) + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder # Factory registry: (provider, type) -> (embedder_class, param_builder) factory_registry = { @@ -221,8 +270,10 @@ def _create_embedder( or "no-key", # Placeholder for local OpenAI-compatible servers "api_base": cfg.api_base, "dimension": cfg.dimension, + "context": context, **({"query_param": cfg.query_param} if cfg.query_param else {}), **({"document_param": cfg.document_param} if cfg.document_param else {}), + "max_tokens": cfg.max_tokens, **({"extra_headers": cfg.extra_headers} if cfg.extra_headers else {}), }, ), @@ -298,10 +349,24 @@ def _create_embedder( "api_key": cfg.api_key, "api_base": cfg.api_base, "dimension": cfg.dimension, + "context": context, **({"query_param": cfg.query_param} if cfg.query_param else {}), **({"document_param": cfg.document_param} if cfg.document_param else {}), }, ), + ("gemini", "dense"): ( + GeminiDenseEmbedder, + lambda cfg: { + "model_name": cfg.model, + "api_key": cfg.api_key, + "dimension": cfg.dimension, + "task_type": ( + cfg.query_param if context == "query" and cfg.query_param + else cfg.document_param if context == "document" and cfg.document_param + else cfg.task_type + ), + }, + ), # Ollama: local OpenAI-compatible embedding server, no real API key needed ("ollama", "dense"): ( OpenAIDenseEmbedder, @@ -311,6 +376,7 @@ def _create_embedder( or "no-key", # Ollama ignores the key, but client requires non-empty "api_base": cfg.api_base or "http://localhost:11434/v1", "dimension": cfg.dimension, + "max_tokens": cfg.max_tokens, }, ), ("voyage", "dense"): ( @@ -369,6 +435,38 @@ def get_embedder(self): raise ValueError("No embedding configuration found (dense, sparse, or hybrid)") + def get_query_embedder(self): + """Get embedder instance for query embeddings.""" + return self._get_contextual_embedder("query") + + def get_document_embedder(self): + """Get embedder instance for document/passage embeddings.""" + return self._get_contextual_embedder("document") + + def _get_contextual_embedder(self, context: str): + if not self.dense: + return self.get_embedder() + + provider = (self.dense.provider or "").lower() + if provider == "openai": + non_symmetric = ( + self.dense.query_param is not None or self.dense.document_param is not None + ) + effective_context = context if non_symmetric else None + return self._create_embedder(provider, "dense", self.dense, context=effective_context) + + if provider == "jina": + return self._create_embedder(provider, "dense", self.dense, context=context) + + if provider == "gemini": + non_symmetric = ( + self.dense.query_param is not None or self.dense.document_param is not None + ) + effective_context = context if non_symmetric else None + return self._create_embedder(provider, "dense", self.dense, context=effective_context) + + return self.get_embedder() + @property def dimension(self) -> int: """Get dimension from active config.""" diff --git a/pyproject.toml b/pyproject.toml index c26aa7cf..471085a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ dependencies = [ "tree-sitter-go>=0.23.0", "tree-sitter-c-sharp>=0.23.0", "loguru>=0.7.3", + "google-genai>=0.8.0", ] [tool.uv.sources] @@ -84,7 +85,9 @@ test = [ "ragas>=0.1.0", "datasets>=2.0.0", "pandas>=2.0.0", + "anyio>=4.0", ] +gemini-async = ["anyio>=4.0"] dev = [ "mypy>=1.0.0", "ruff>=0.1.0", @@ -208,6 +211,9 @@ python_classes = ["Test*"] python_functions = ["test_*"] asyncio_mode = "auto" addopts = "-v --cov=openviking --cov-report=term-missing" +markers = [ + "integration: tests that call real external APIs (require env vars)", +] [tool.ruff] line-length = 100 @@ -246,4 +252,6 @@ line-ending = "auto" [dependency-groups] dev = [ "pytest>=9.0.2", + "pytest-asyncio>=1.3.0", + "pytest-cov>=7.0.0", ] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index fceae36b..a3afd48a 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -7,14 +7,18 @@ AsyncHTTPClient integration tests can run without a manually started server process. """ +import math +import os import shutil import socket import threading import time from pathlib import Path +from typing import Optional import httpx import pytest +import pytest_asyncio import uvicorn from openviking.server.app import create_app @@ -79,3 +83,108 @@ def server_url(temp_dir): thread.join(timeout=5) loop.run_until_complete(svc.close()) loop.close() + + +# ── Gemini shared fixtures and helpers ──────────────────────────────────────── + +GOOGLE_API_KEY: Optional[str] = os.environ.get("GOOGLE_API_KEY") + +# (model_name, default_dim, token_limit) — for @pytest.mark.parametrize("model,dim,limit", ...) +GEMINI_MODELS = [ + pytest.param("gemini-embedding-2-preview", 3072, 8192, id="g2p"), + pytest.param("gemini-embedding-001", 3072, 2048, id="g001"), +] + +# Wrapped single-value tuples — for fixture params (request.param is the whole tuple) +GEMINI_MODELS_FIXTURE = [ + pytest.param(("gemini-embedding-2-preview", 3072, 8192), id="g2p"), + pytest.param(("gemini-embedding-001", 3072, 2048), id="g001"), +] + +# (model_name, dimension) pairs for OpenViking client fixtures +EMBED_PARAMS = [ + pytest.param(("gemini-embedding-2-preview", 512), id="g2p-512"), + pytest.param(("gemini-embedding-2-preview", 768), id="g2p-768"), + pytest.param(("gemini-embedding-2-preview", 1536), id="g2p-1536"), + pytest.param(("gemini-embedding-2-preview", 3072), id="g2p-3072"), + pytest.param(("gemini-embedding-001", 768), id="g001-768"), +] + + +def l2_norm(vec) -> float: + return math.sqrt(sum(v * v for v in vec)) + + +def vectordb_engine_available() -> bool: + try: + from openviking.storage.vectordb.engine import PersistStore, VolatileStore + return isinstance(PersistStore, type) and isinstance(VolatileStore, type) + except Exception: + return False + + +def sample_markdown(tmp_dir: Path, slug: str, content: str) -> Path: + p = tmp_dir / f"{slug}.md" + p.write_text(content, encoding="utf-8") + return p + + +def gemini_config_dict( + model: str, dim: int, + query_param: Optional[str] = None, + doc_param: Optional[str] = None, + task_type: Optional[str] = None, +) -> dict: + dense: dict = {"provider": "gemini", "model": model, "api_key": GOOGLE_API_KEY, "dimension": dim} + if query_param: + dense["query_param"] = query_param + if doc_param: + dense["document_param"] = doc_param + if task_type: + dense["task_type"] = task_type + return {"embedding": {"dense": dense}, "storage": {"agfs": {"mode": "binding-client"}}} + + +async def make_ov_client(config_dict: dict, data_path: str): + from openviking.async_client import AsyncOpenViking + from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + await AsyncOpenViking.reset() + OpenVikingConfigSingleton.reset_instance() + OpenVikingConfigSingleton.initialize(config_dict=config_dict) + client = AsyncOpenViking(path=data_path) + await client.initialize() + return client + + +async def teardown_ov_client(): + from openviking.async_client import AsyncOpenViking + from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + await AsyncOpenViking.reset() + OpenVikingConfigSingleton.reset_instance() + + +requires_api_key = pytest.mark.skipif(not GOOGLE_API_KEY, reason="GOOGLE_API_KEY not set") +requires_engine = pytest.mark.skipif( + not vectordb_engine_available(), + reason="VectorDB native engine not compiled — run: pip install -e . --no-build-isolation", +) + + +@pytest.fixture(scope="module", params=GEMINI_MODELS_FIXTURE) +def gemini_embedder(request): + """Module-scoped GeminiDenseEmbedder at dim=768, parametrized over known models.""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + model_name, _, _ = request.param + return GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768) + + +@pytest_asyncio.fixture(params=EMBED_PARAMS) +async def gemini_ov_client(request, tmp_path): + """AsyncOpenViking client backed by Gemini; yields (client, model, dim).""" + model, dim = request.param + data_path = str(tmp_path / "ov_data") + Path(data_path).mkdir(parents=True, exist_ok=True) + + client = await make_ov_client(gemini_config_dict(model, dim), data_path) + yield client, model, dim + await teardown_ov_client() diff --git a/tests/integration/test_gemini_e2e.py b/tests/integration/test_gemini_e2e.py new file mode 100644 index 00000000..d1e94854 --- /dev/null +++ b/tests/integration/test_gemini_e2e.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +""" +End-to-end integration tests for GeminiDenseEmbedder. +Calls the real Gemini API — requires GOOGLE_API_KEY env var. +Run: pytest tests/integration/test_gemini_e2e.py -v -m integration +""" +import math +import struct +import zlib + +import pytest + +from openviking.core.context import ModalContent, Vectorize +from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder +from tests.integration.conftest import GOOGLE_API_KEY, l2_norm, requires_api_key + +pytestmark = [pytest.mark.integration, requires_api_key] + + +def _make_tiny_png() -> bytes: + """Create a minimal valid 1x1 white PNG (89 bytes).""" + + def chunk(name, data): + c = struct.pack(">I", len(data)) + name + data + return c + struct.pack(">I", zlib.crc32(name + data) & 0xFFFFFFFF) + + sig = b"\x89PNG\r\n\x1a\n" + ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)) + idat = chunk(b"IDAT", zlib.compress(b"\x00\xff\xff\xff")) + iend = chunk(b"IEND", b"") + return sig + ihdr + idat + iend + + +def _cosine_similarity(a: list, b: list) -> float: + dot = sum(x * y for x, y in zip(a, b)) + norm_a = l2_norm(a) + norm_b = l2_norm(b) + return dot / (norm_a * norm_b) if norm_a and norm_b else 0.0 + + +@pytest.fixture(scope="module") +def embedder(): + e = GeminiDenseEmbedder( + "gemini-embedding-2-preview", + api_key=GOOGLE_API_KEY, + # dimension defaults to 3072 — test the actual default + task_type="RETRIEVAL_DOCUMENT", + ) + yield e + e.close() + + +def test_default_dimension_is_3072(embedder): + """Default output dimension must match model's native 3072.""" + assert embedder.get_dimension() == 3072 + result = embedder.embed("hello") + assert len(result.dense_vector) == 3072 + + +class TestGeminiE2ETextEmbedding: + def test_embed_text_returns_correct_dimension(self, embedder): + result = embedder.embed("OpenViking is a knowledge management system") + assert result.dense_vector is not None + assert len(result.dense_vector) == 3072 + + def test_embed_text_vector_is_normalized(self, embedder): + result = embedder.embed("test normalization") + norm = l2_norm(result.dense_vector) + assert abs(norm - 1.0) < 0.01, f"Vector norm {norm} not close to 1.0" + + def test_embed_batch_matches_individual(self, embedder): + texts = ["hello world", "foo bar", "test embed"] + batch_results = embedder.embed_batch(texts) + individual_results = [embedder.embed(t) for t in texts] + assert len(batch_results) == 3 + for br, ir in zip(batch_results, individual_results): + sim = _cosine_similarity(br.dense_vector, ir.dense_vector) + assert sim > 0.99, f"Batch vs individual similarity {sim} too low" + + def test_semantic_similarity_related_texts(self, embedder): + r1 = embedder.embed("a golden retriever playing in the park") + r2 = embedder.embed("a dog running outside in a field") + r3 = embedder.embed("quantum computing and cryptography") + sim_related = _cosine_similarity(r1.dense_vector, r2.dense_vector) + sim_unrelated = _cosine_similarity(r1.dense_vector, r3.dense_vector) + assert sim_related > sim_unrelated + + +class TestGeminiE2EMultimodalEmbedding: + @pytest.mark.xfail(reason="gemini-embedding-2-preview may not support multimodal on free tier") + def test_embed_multimodal_image_returns_correct_dimension(self, embedder): + v = Vectorize( + text="a tiny white pixel", + media=ModalContent(mime_type="image/png", uri="test.png", data=_make_tiny_png()), + ) + result = embedder.embed_multimodal(v) + assert result.dense_vector is not None + assert len(result.dense_vector) == 3072 + + def test_multimodal_fallback_on_no_media(self, embedder): + v = Vectorize(text="just text, no image") + result = embedder.embed_multimodal(v) + text_result = embedder.embed("just text, no image") + sim = _cosine_similarity(result.dense_vector, text_result.dense_vector) + assert sim > 0.99, f"Fallback similarity {sim:.3f} too low" + + +class TestGeminiE2EAsyncBatch: + @pytest.mark.anyio + async def test_async_embed_batch_concurrent(self): + try: + import anyio # noqa: F401 + except ImportError: + pytest.skip("anyio not installed") + import time + e = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, dimension=128 + ) + texts = [f"sentence {i}" for i in range(300)] # 3 batches of 100 + t0 = time.monotonic() + results = await e.async_embed_batch(texts) + elapsed = time.monotonic() - t0 + assert len(results) == 300 + assert all(len(r.dense_vector) == 128 for r in results) + assert elapsed < 15 # concurrent should be << 3× serial RTT + e.close() + + +class TestGeminiE2ETaskType: + def test_query_vs_document_task_types(self): + doc_embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", + api_key=GOOGLE_API_KEY, + task_type="RETRIEVAL_DOCUMENT", + ) + query_embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", + api_key=GOOGLE_API_KEY, + task_type="RETRIEVAL_QUERY", + ) + text = "machine learning algorithms" + doc_result = doc_embedder.embed(text) + query_result = query_embedder.embed(text) + sim = _cosine_similarity(doc_result.dense_vector, query_result.dense_vector) + # gemini-embedding-2-preview may return identical vectors for same text + # across task types; assert vectors are at least highly correlated + assert sim > 0.8, f"Task type similarity {sim:.3f} unexpectedly low" + doc_embedder.close() + query_embedder.close() diff --git a/tests/integration/test_gemini_embedding_it.py b/tests/integration/test_gemini_embedding_it.py new file mode 100644 index 00000000..c6437472 --- /dev/null +++ b/tests/integration/test_gemini_embedding_it.py @@ -0,0 +1,103 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +""" +Integration tests for GeminiDenseEmbedder — require real GOOGLE_API_KEY. +Run: GOOGLE_API_KEY= pytest tests/integration/test_gemini_embedding_it.py -v +Auto-skipped when GOOGLE_API_KEY is not set. No mocking — real API calls. +""" +import pytest + +from tests.integration.conftest import ( + GEMINI_MODELS, + GOOGLE_API_KEY, + l2_norm, + requires_api_key, +) + +pytestmark = [requires_api_key] + + +def test_embed_returns_correct_dimension(gemini_embedder): + r = gemini_embedder.embed("What is machine learning?") + assert r.dense_vector and len(r.dense_vector) == 768 + assert 0.99 < l2_norm(r.dense_vector) < 1.01 + + +def test_embed_batch_count(gemini_embedder): + texts = ["apple", "banana", "cherry", "date", "elderberry"] + results = gemini_embedder.embed_batch(texts) + assert len(results) == len(texts) + for r in results: + assert r.dense_vector and len(r.dense_vector) == 768 + + +def test_batch_over_100(gemini_embedder): + """150 texts auto-split into 2 batches (100 + 50).""" + texts = [f"sentence number {i}" for i in range(150)] + results = gemini_embedder.embed_batch(texts) + assert len(results) == 150 + for r in results: + assert r.dense_vector and len(r.dense_vector) == 768 + + +@pytest.mark.parametrize("model_name,_dim,token_limit", GEMINI_MODELS) +def test_large_text_chunking(model_name, _dim, token_limit): + """Text exceeding the model's token limit is auto-chunked by base class.""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + phrase = "Machine learning is a subset of artificial intelligence. " + large = phrase * ((token_limit * 2) // len(phrase.split()) + 10) + e = GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768) + r = e.embed(large) + assert r.dense_vector and len(r.dense_vector) == 768 + norm = l2_norm(r.dense_vector) + assert 0.99 < norm < 1.01, f"chunked vector not L2-normalized, norm={norm}" + + +@pytest.mark.parametrize("task_type", [ + "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", + "CLASSIFICATION", "CLUSTERING", "CODE_RETRIEVAL_QUERY", + "QUESTION_ANSWERING", "FACT_VERIFICATION", +]) +def test_all_task_types_accepted(task_type): + """All 8 Gemini task types must be accepted by the API without error.""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + e = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, + task_type=task_type, dimension=768, + ) + r = e.embed("test input for task type validation") + assert r.dense_vector and len(r.dense_vector) == 768 + + +def test_config_nonsymmetric_routing(): + """EmbeddingConfig query/document embedders wire task_type via query_param/document_param.""" + from openviking_cli.utils.config.embedding_config import EmbeddingConfig, EmbeddingModelConfig + cfg = EmbeddingConfig( + dense=EmbeddingModelConfig( + model="gemini-embedding-2-preview", provider="gemini", + api_key=GOOGLE_API_KEY, dimension=768, + query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT", + ) + ) + q = cfg.get_query_embedder() + d = cfg.get_document_embedder() + assert q.task_type == "RETRIEVAL_QUERY" + assert d.task_type == "RETRIEVAL_DOCUMENT" + assert q.embed("search query").dense_vector is not None + assert d.embed("document text").dense_vector is not None + + +def test_invalid_api_key_error_message(): + """Wrong API key must raise RuntimeError with 'Invalid API key' hint.""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + bad = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="INVALID_KEY_XYZZY_123") + with pytest.raises(RuntimeError, match="Invalid API key"): + bad.embed("hello") + + +def test_invalid_model_error_message(): + """Unknown model name must raise RuntimeError with model-not-found hint.""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + bad = GeminiDenseEmbedder("gemini-embedding-does-not-exist-xyz", api_key=GOOGLE_API_KEY) + with pytest.raises(RuntimeError, match="Model not found"): + bad.embed("hello") diff --git a/tests/integration/test_gemini_openviking_it.py b/tests/integration/test_gemini_openviking_it.py new file mode 100644 index 00000000..b4349c4d --- /dev/null +++ b/tests/integration/test_gemini_openviking_it.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +""" +End-to-end integration tests for OpenViking add-memory + search using Gemini embeddings. + +Exercises the full workflow: inject Gemini config → add_resource → wait_processed → find/search. +No mocking — real Gemini API calls. Auto-skipped when GOOGLE_API_KEY is not set. + +Run: + GOOGLE_API_KEY= pytest tests/integration/test_gemini_openviking_it.py -v + +NOTE: provider MUST be "gemini" — "google" is not a valid provider value. +""" +from pathlib import Path + +import pytest + +from tests.integration.conftest import ( + gemini_config_dict, + make_ov_client, + requires_api_key, + requires_engine, + sample_markdown, + teardown_ov_client, +) + +pytestmark = [requires_api_key, requires_engine] + + +# --------------------------------------------------------------------------- +# Test 1: Basic add-memory + search +# --------------------------------------------------------------------------- + +async def test_add_and_search_basic(gemini_ov_client, tmp_path): + """Add a single markdown document and verify it is returned by find().""" + client, model, dim = gemini_ov_client + + doc = sample_markdown( + tmp_path, + "ml_intro", + "# Machine Learning\n\nMachine learning is a field of AI that uses statistical methods.", + ) + + result = await client.add_resource(path=str(doc), reason="IT test basic", wait=True) + assert result.get("root_uri"), "add_resource should return a root_uri" + + found = await client.find(query="machine learning AI statistical") + assert found.total > 0, f"Expected search results for ML doc, got total={found.total}" + scores = [r.score for r in found.resources] + assert any(s > 0.0 for s in scores), f"Expected non-zero similarity scores, got {scores}" + + +# --------------------------------------------------------------------------- +# Test 2: Batch — multiple documents, search returns relevant one +# --------------------------------------------------------------------------- + +async def test_batch_documents_search(gemini_ov_client, tmp_path): + """Add 5 documents on different topics; search returns the relevant one first.""" + client, model, dim = gemini_ov_client + + docs = { + "python_types": "Python supports dynamic typing and type hints via the typing module.", + "quantum_physics": "Quantum mechanics describes the behavior of particles at atomic scale.", + "cooking_pasta": "To cook pasta: boil salted water, add pasta, cook 8-12 minutes, drain.", + "git_branching": "Git branches allow parallel development. Use git checkout -b to create.", + "solar_system": "The solar system has 8 planets. Jupiter is the largest planet.", + } + for slug, content in docs.items(): + doc_path = sample_markdown(tmp_path, slug, f"# {slug}\n\n{content}") + await client.add_resource(path=str(doc_path), reason="IT batch test") + + await client.wait_processed() + + found = await client.find(query="how to cook pasta boil water") + assert found.total > 0, "Expected at least one result for pasta query" + + +# --------------------------------------------------------------------------- +# Test 3: Large text chunking +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("model,dim,token_limit", [ + pytest.param("gemini-embedding-2-preview", 768, 8192, id="g2p-large"), + pytest.param("gemini-embedding-001", 768, 2048, id="g001-large"), +]) +async def test_large_text_add_and_search(model, dim, token_limit, tmp_path): + """Add a document exceeding the model's token limit; verify chunking and searchability.""" + data_path = str(tmp_path / "ov_large") + Path(data_path).mkdir(parents=True, exist_ok=True) + + client = await make_ov_client(gemini_config_dict(model, dim), data_path) + try: + phrase = "Neural networks are computational models inspired by the brain. " + repeats = (token_limit * 2) // len(phrase.split()) + 10 + large_content = f"# Large Document\n\n{phrase * repeats}" + + doc = sample_markdown(tmp_path, "large_doc", large_content) + result = await client.add_resource(path=str(doc), reason="large text IT", wait=True) + assert result.get("root_uri"), "Large doc should index without error" + + found = await client.find(query="neural networks computational brain") + assert found.total > 0, "Chunked large doc should be findable" + finally: + await teardown_ov_client() + + +# --------------------------------------------------------------------------- +# Test 4: RETRIEVAL_QUERY / RETRIEVAL_DOCUMENT routing via EmbeddingConfig +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("query_param,doc_param", [ + pytest.param("RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", id="retrieval-routing"), + pytest.param("SEMANTIC_SIMILARITY", "SEMANTIC_SIMILARITY", id="semantic-routing"), +]) +async def test_retrieval_routing_workflow(query_param, doc_param, tmp_path): + """Verify add+search works with non-symmetric task-type routing.""" + data_path = str(tmp_path / "ov_routing") + Path(data_path).mkdir(parents=True, exist_ok=True) + + client = await make_ov_client( + gemini_config_dict("gemini-embedding-2-preview", 768, query_param=query_param, doc_param=doc_param), + data_path, + ) + try: + doc = sample_markdown( + tmp_path, "routing_doc", + "# Retrieval Test\n\nOpenViking provides memory management for AI agents.", + ) + result = await client.add_resource(path=str(doc), reason="routing IT", wait=True) + assert result.get("root_uri") + + found = await client.find(query="memory management AI agents") + assert found.total > 0, f"Routing {query_param}/{doc_param}: expected search results" + finally: + await teardown_ov_client() + + +# --------------------------------------------------------------------------- +# Test 5: Dimension variants — verify index schema uses requested dim +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("dim", [512, 768, 1536, 3072]) +async def test_dimension_variant_add_search(dim, tmp_path): + """Each dimension variant should index and search without errors.""" + data_path = str(tmp_path / f"ov_dim_{dim}") + Path(data_path).mkdir(parents=True, exist_ok=True) + + client = await make_ov_client(gemini_config_dict("gemini-embedding-2-preview", dim), data_path) + try: + doc = sample_markdown( + tmp_path, f"dim_doc_{dim}", + f"# Dimension {dim} Test\n\nThis document is indexed with embedding dimension {dim}.", + ) + result = await client.add_resource(path=str(doc), reason=f"dim={dim} IT", wait=True) + assert result.get("root_uri"), f"dim={dim}: add_resource should succeed" + + found = await client.find(query=f"embedding dimension {dim}") + assert found.total > 0, f"dim={dim}: should find the indexed doc" + finally: + await teardown_ov_client() + + +# --------------------------------------------------------------------------- +# Test 6: Multi-turn session + search (smoke test) +# --------------------------------------------------------------------------- + +async def test_session_search_smoke(gemini_ov_client, tmp_path): + """Session construction + embedding-based find works with Gemini embeddings. + + Uses find() (pure embedding path) rather than search() which requires a VLM. + """ + from openviking.message import TextPart + + client, model, dim = gemini_ov_client + + doc = sample_markdown( + tmp_path, "session_doc", + "# Python Testing\n\nPytest is a mature full-featured Python testing tool.", + ) + await client.add_resource(path=str(doc), reason="session IT", wait=True) + + session = client.session(session_id="gemini_it_session") + session.add_message("user", [TextPart("Tell me about Python testing.")]) + + result = await client.find(query="pytest testing tool") + assert result.total >= 0, "Embedding-based find should complete without error" diff --git a/tests/unit/test_embedding_config_gemini.py b/tests/unit/test_embedding_config_gemini.py new file mode 100644 index 00000000..cac256cb --- /dev/null +++ b/tests/unit/test_embedding_config_gemini.py @@ -0,0 +1,105 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for Gemini-specific EmbeddingModelConfig and EmbeddingConfig behavior.""" +from unittest.mock import patch + +import pytest + +from openviking_cli.utils.config.embedding_config import EmbeddingConfig, EmbeddingModelConfig + + +def _gcfg(**kw) -> EmbeddingModelConfig: + """Helper: build a Gemini EmbeddingModelConfig with defaults.""" + return EmbeddingModelConfig( + model="gemini-embedding-2-preview", provider="gemini", api_key="test-key", **kw + ) + + +class TestGeminiDimension: + def test_preview_defaults_3072(self): + assert _gcfg().get_effective_dimension() == 3072 + + def test_001_defaults_3072(self): + cfg = EmbeddingModelConfig(model="gemini-embedding-001", provider="gemini", api_key="k") + assert cfg.get_effective_dimension() == 3072 + + def test_004_defaults_768(self): + cfg = EmbeddingModelConfig(model="text-embedding-004", provider="gemini", api_key="k") + assert cfg.get_effective_dimension() == 768 + + def test_unknown_model_defaults_3072(self): + cfg = EmbeddingModelConfig(model="gemini-embedding-future", provider="gemini", api_key="k") + assert cfg.get_effective_dimension() == 3072 + + def test_explicit_dimension_overrides_default(self): + assert _gcfg(dimension=1536).get_effective_dimension() == 1536 + + def test_text_embedding_prefix_defaults_768(self): + """text-embedding-* future models default to 768 via prefix rule.""" + cfg = EmbeddingModelConfig(model="text-embedding-005", provider="gemini", api_key="k") + assert cfg.get_effective_dimension() == 768 + + def test_future_gemini_model_defaults_3072(self): + """Future gemini-embedding-* models default to 3072 via fallback.""" + for model in ["gemini-embedding-2", "gemini-embedding-2.1", "gemini-embedding-3-preview"]: + cfg = EmbeddingModelConfig(model=model, provider="gemini", api_key="k") + assert cfg.get_effective_dimension() == 3072 + + +class TestGeminiContextRouting: + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_get_query_embedder_uses_query_param(self, _mock): + cfg = EmbeddingConfig(dense=_gcfg( + query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT" + )) + assert cfg.get_query_embedder().task_type == "RETRIEVAL_QUERY" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_get_document_embedder_uses_document_param(self, _mock): + cfg = EmbeddingConfig(dense=_gcfg( + query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT" + )) + assert cfg.get_document_embedder().task_type == "RETRIEVAL_DOCUMENT" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_symmetric_uses_task_type_field(self, _mock): + cfg = EmbeddingConfig(dense=_gcfg(task_type="SEMANTIC_SIMILARITY")) + assert cfg.get_embedder().task_type == "SEMANTIC_SIMILARITY" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_symmetric_no_task_type_is_none(self, _mock): + cfg = EmbeddingConfig(dense=_gcfg()) + assert cfg.get_embedder().task_type is None + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_only_query_param_set_routes_correctly(self, _mock): + """When only query_param is set, document embedder falls back to static task_type.""" + cfg = EmbeddingConfig(dense=_gcfg(query_param="RETRIEVAL_QUERY")) + assert cfg.get_query_embedder().task_type == "RETRIEVAL_QUERY" + assert cfg.get_document_embedder().task_type is None + + +class TestGeminiConfigValidation: + def test_missing_api_key_raises(self): + with pytest.raises(ValueError, match="api_key"): + EmbeddingModelConfig(model="gemini-embedding-2-preview", provider="gemini") + + def test_invalid_query_param_raises(self): + with pytest.raises(ValueError, match="Invalid query_param"): + _gcfg(query_param="NOT_A_VALID_TYPE") + + def test_invalid_document_param_raises(self): + with pytest.raises(ValueError, match="Invalid document_param"): + _gcfg(document_param="ALSO_INVALID") + + def test_invalid_task_type_raises(self): + with pytest.raises(ValueError, match="Invalid task_type"): + _gcfg(task_type="BAD_TYPE") + + def test_valid_task_types_accepted(self): + for t in [ + "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", + "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", + "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY", + ]: + _gcfg(task_type=t) # must not raise diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py new file mode 100644 index 00000000..d9620707 --- /dev/null +++ b/tests/unit/test_gemini_embedder.py @@ -0,0 +1,306 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +""" +Tests for GeminiDenseEmbedder. +Pattern: patch at module import path, use MagicMock, never make real API calls. +""" +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +def _make_mock_embedding(values): + emb = MagicMock() + emb.values = values + return emb + + +def _make_mock_result(values_list): + result = MagicMock() + result.embeddings = [_make_mock_embedding(v) for v in values_list] + return result + + +def test_input_token_limit_constant(): + from openviking.models.embedder.gemini_embedders import _GEMINI_INPUT_TOKEN_LIMIT + assert _GEMINI_INPUT_TOKEN_LIMIT == 8192 + + +class TestGeminiDenseEmbedderInit: + def test_requires_api_key(self): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="api_key"): + GeminiDenseEmbedder("gemini-embedding-2-preview") + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_init_stores_fields(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", + api_key="test-key", + dimension=1536, + task_type="RETRIEVAL_DOCUMENT", + ) + assert embedder.model_name == "gemini-embedding-2-preview" + assert embedder.task_type == "RETRIEVAL_DOCUMENT" + assert embedder.get_dimension() == 1536 + mock_client_class.assert_called_once_with(api_key="test-key") + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_default_dimension_3072(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + assert embedder.get_dimension() == 3072 + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_dimension_1_valid(self, mock_client_class): + """API accepts dimension=1 (128 is a quality recommendation, not a hard limit).""" + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + assert embedder.get_dimension() == 1 + + def test_default_dimension_classmethod_prefix_rule(self): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2") == 3072 + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2.1") == 3072 + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 + assert GeminiDenseEmbedder._default_dimension("text-embedding-005") == 768 + assert GeminiDenseEmbedder._default_dimension("text-embedding-004") == 768 # exact match wins + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2-preview") == 3072 # exact match + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_token_limit_per_model(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder, _MODEL_TOKEN_LIMITS + for model, expected in _MODEL_TOKEN_LIMITS.items(): + e = GeminiDenseEmbedder(model, api_key="key") + assert e._token_limit == expected + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_supports_multimodal_false(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + assert embedder.supports_multimodal is False + + +class TestGeminiDenseEmbedderEmbed: + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_text(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + result = embedder.embed("hello world") + assert result.dense_vector is not None + assert len(result.dense_vector) == 3 + mock_client.models.embed_content.assert_called_once() + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["model"] == "gemini-embedding-2-preview" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_passes_task_type_in_config(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1, task_type="RETRIEVAL_QUERY" + ) + embedder.embed("query text") + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["config"].task_type == "RETRIEVAL_QUERY" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): + from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_response = MagicMock() + mock_response.status_code = 401 + mock_client.models.embed_content.side_effect = APIError(401, {}, response=mock_response) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + with pytest.raises(RuntimeError, match="Gemini embedding failed"): + embedder.embed("hello") + + +class TestGeminiDenseEmbedderBatch: + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_empty(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + results = embedder.embed_batch([]) + assert results == [] + mock_client.models.embed_content.assert_not_called() + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_single_chunk(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1], [0.2], [0.3]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = embedder.embed_batch(["a", "b", "c"]) + assert len(results) == 3 + mock_client.models.embed_content.assert_called_once() + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["contents"] == ["a", "b", "c"] + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_chunks_at_100(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_client.models.embed_content.side_effect = [ + _make_mock_result([[0.1]] * 100), + _make_mock_result([[0.2]] * 10), + ] + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = embedder.embed_batch([f"text{i}" for i in range(110)]) + assert len(results) == 110 + assert mock_client.models.embed_content.call_count == 2 + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_falls_back_to_individual_on_error(self, mock_client_class): + from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_response = MagicMock() + mock_response.status_code = 500 + mock_client.models.embed_content.side_effect = [ + APIError(500, {}, response=mock_response), + _make_mock_result([[0.1]]), + _make_mock_result([[0.2]]), + ] + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = embedder.embed_batch(["a", "b"]) + assert len(results) == 2 + assert mock_client.models.embed_content.call_count == 3 + + +class TestGeminiDenseEmbedderAsyncBatch: + """Unit tests for async_embed_batch (uses AsyncMock, no real API).""" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + @pytest.mark.anyio + async def test_async_embed_batch_dispatches_all_chunks(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_client.aio.models.embed_content = AsyncMock(side_effect=[ + _make_mock_result([[0.1]] * 100), + _make_mock_result([[0.2]] * 10), + ]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = await embedder.async_embed_batch([f"t{i}" for i in range(110)]) + assert len(results) == 110 + assert mock_client.aio.models.embed_content.call_count == 2 + + @patch("openviking.models.embedder.gemini_embedders._TEXT_BATCH_SIZE", 1) + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + @pytest.mark.anyio + async def test_async_embed_batch_preserves_order(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + # Use orthogonal unit vectors so _l2_normalize keeps them distinguishable + mock_client.aio.models.embed_content = AsyncMock(side_effect=[ + _make_mock_result([[1.0, 0.0, 0.0]]), + _make_mock_result([[0.0, 1.0, 0.0]]), + _make_mock_result([[0.0, 0.0, 1.0]]), + ]) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=3, max_concurrent_batches=3 + ) + results = await embedder.async_embed_batch(["a", "b", "c"]) + # Order must match input regardless of task completion order + assert results[0].dense_vector == pytest.approx([1.0, 0.0, 0.0]) + assert results[1].dense_vector == pytest.approx([0.0, 1.0, 0.0]) + assert results[2].dense_vector == pytest.approx([0.0, 0.0, 1.0]) + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + @pytest.mark.anyio + async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_class): + from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value + mock_response = MagicMock() + mock_response.status_code = 500 + mock_client.aio.models.embed_content = AsyncMock( + side_effect=APIError(500, {}, response=mock_response) + ) + mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = await embedder.async_embed_batch(["a", "b"]) + assert len(results) == 2 + assert mock_client.models.embed_content.call_count == 2 + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + @pytest.mark.anyio + async def test_async_embed_batch_empty(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + assert await embedder.async_embed_batch([]) == [] + + @patch("openviking.models.embedder.gemini_embedders._ANYIO_AVAILABLE", False) + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + @pytest.mark.anyio + async def test_async_embed_batch_raises_without_anyio(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + with pytest.raises(ImportError, match="anyio is required"): + await embedder.async_embed_batch(["text"]) + + +class TestGeminiValidation: + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_all_valid_task_types_accepted(self, mock_client): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder, _VALID_TASK_TYPES + for tt in _VALID_TASK_TYPES: + e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type=tt) + assert e.task_type == tt + + def test_invalid_task_type_raises_on_init(self): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="Invalid task_type"): + GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID") + + def test_valid_task_types_count(self): + from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES + assert len(_VALID_TASK_TYPES) == 8 + + def test_code_retrieval_query_in_task_types(self): + from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES + assert "CODE_RETRIEVAL_QUERY" in _VALID_TASK_TYPES + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_dimension_too_high_raises(self, mock_client): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="3072"): + GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", dimension=4096) + + +class TestGeminiErrorMessages: + @pytest.mark.parametrize("code,match", [ + (401, "Invalid API key"), + (403, "Permission denied"), + (404, "Model not found"), + (429, "Quota exceeded"), + (500, "service error"), + ]) + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_error_code_hint(self, mock_client, code, match): + from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock = mock_client.return_value + mock_resp = MagicMock() + mock_resp.status_code = code + mock.models.embed_content.side_effect = APIError(code, {}, response=mock_resp) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") + with pytest.raises(RuntimeError, match=match): + embedder.embed("hello") + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_error_message_includes_http_code(self, mock_client): + from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock = mock_client.return_value + mock_resp = MagicMock() + mock_resp.status_code = 404 + mock.models.embed_content.side_effect = APIError(404, {}, response=mock_resp) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") + with pytest.raises(RuntimeError, match="HTTP 404"): + embedder.embed("hello") From 4fb7452c9122db0fb5b96611e50724410d7b755d Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Wed, 18 Mar 2026 22:18:03 +0100 Subject: [PATCH 2/7] fix(embedder): add supports_multimodal attr; strip ModalContent import from IT test - GeminiDenseEmbedder.supports_multimodal = False (text-only provider) - Remove ModalContent/Vectorize import from test_gemini_e2e.py (doesn't exist in gem-v1) - Remove TestGeminiE2EMultimodalEmbedding class (multimodal out of scope for text-only slice) --- .../models/embedder/gemini_embedders.py | 2 ++ tests/integration/test_gemini_e2e.py | 34 ------------------- 2 files changed, 2 insertions(+), 34 deletions(-) diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index 40d00f9e..efd7bfb7 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -93,6 +93,8 @@ class GeminiDenseEmbedder(DenseEmbedderBase): # text-embedding-004: 768 fixed-dim legacy model, does not support MRL truncation # Future gemini-embedding-*: default 3072 via _default_dimension() fallback # Future text-embedding-*: default 768 via _default_dimension() prefix rule + supports_multimodal: bool = False # text-only; multimodal planned separately + KNOWN_DIMENSIONS: Dict[str, int] = { "gemini-embedding-2-preview": 3072, "gemini-embedding-001": 3072, diff --git a/tests/integration/test_gemini_e2e.py b/tests/integration/test_gemini_e2e.py index d1e94854..eeebb91d 100644 --- a/tests/integration/test_gemini_e2e.py +++ b/tests/integration/test_gemini_e2e.py @@ -6,31 +6,15 @@ Run: pytest tests/integration/test_gemini_e2e.py -v -m integration """ import math -import struct -import zlib import pytest -from openviking.core.context import ModalContent, Vectorize from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder from tests.integration.conftest import GOOGLE_API_KEY, l2_norm, requires_api_key pytestmark = [pytest.mark.integration, requires_api_key] -def _make_tiny_png() -> bytes: - """Create a minimal valid 1x1 white PNG (89 bytes).""" - - def chunk(name, data): - c = struct.pack(">I", len(data)) + name + data - return c + struct.pack(">I", zlib.crc32(name + data) & 0xFFFFFFFF) - - sig = b"\x89PNG\r\n\x1a\n" - ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)) - idat = chunk(b"IDAT", zlib.compress(b"\x00\xff\xff\xff")) - iend = chunk(b"IEND", b"") - return sig + ihdr + idat + iend - def _cosine_similarity(a: list, b: list) -> float: dot = sum(x * y for x, y in zip(a, b)) @@ -87,24 +71,6 @@ def test_semantic_similarity_related_texts(self, embedder): assert sim_related > sim_unrelated -class TestGeminiE2EMultimodalEmbedding: - @pytest.mark.xfail(reason="gemini-embedding-2-preview may not support multimodal on free tier") - def test_embed_multimodal_image_returns_correct_dimension(self, embedder): - v = Vectorize( - text="a tiny white pixel", - media=ModalContent(mime_type="image/png", uri="test.png", data=_make_tiny_png()), - ) - result = embedder.embed_multimodal(v) - assert result.dense_vector is not None - assert len(result.dense_vector) == 3072 - - def test_multimodal_fallback_on_no_media(self, embedder): - v = Vectorize(text="just text, no image") - result = embedder.embed_multimodal(v) - text_result = embedder.embed("just text, no image") - sim = _cosine_similarity(result.dense_vector, text_result.dense_vector) - assert sim > 0.99, f"Fallback similarity {sim:.3f} too low" - class TestGeminiE2EAsyncBatch: @pytest.mark.anyio From 148f6e333b22b3cfde3b744332b4824a640a17b1 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Wed, 18 Mar 2026 23:10:36 +0100 Subject: [PATCH 3/7] =?UTF-8?q?fix(gemini):=20embed()/embed=5Fbatch():=20a?= =?UTF-8?q?dd=20is=5Fquery=20-=20embed=5Fbatch()=20fallback:=20pass=20is?= =?UTF-8?q?=5Fquery=20through=20to=20self.embed()=20-=20Remove=20redundant?= =?UTF-8?q?=20=5Fl2=5Fnormalize()=20calls=20in=20embed/embed=5Fbatch/async?= =?UTF-8?q?=5Fembed=5Fbatch;=20=20=20Gemini=20API=20already=20returns=20L2?= =?UTF-8?q?-normalized=20vectors,=20truncate=5Fand=5Fnormalize=20is=20suff?= =?UTF-8?q?icient=20-=20Remove=20now-dead=20=5Fl2=5Fnormalize()=20function?= =?UTF-8?q?=20and=20unused=20math=20import=20-=20embedding=5Fconfig.py:=20?= =?UTF-8?q?fix=20typo=20"Secretfor"=20=E2=86=92=20"Secret=20for"=20in=20sk?= =?UTF-8?q?=20field=20description=20-=20conftest.py:=20narrow=20broad=20ex?= =?UTF-8?q?cept=20Exception=20to=20(ImportError,=20ModuleNotFoundError,=20?= =?UTF-8?q?AttributeError)=20-=20test=5Fgemini=5Fembedding=5Fit.py:=20cons?= =?UTF-8?q?truct=20fake=20API=20key=20dynamically=20to=20avoid=20hardcoded?= =?UTF-8?q?=20secret=20pattern=20-=20test=5Fgemini=5Fopenviking=5Fit.py:?= =?UTF-8?q?=20add=20dimension=20assertion=20in=20test=5Fdimension=5Fvarian?= =?UTF-8?q?t=5Fadd=5Fsearch;=20=20=20fix=20tautology=20assertion=20(>=3D?= =?UTF-8?q?=200=20=E2=86=92=20>=200)=20in=20test=5Fsession=5Fsearch=5Fsmok?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../models/embedder/gemini_embedders.py | 35 ++++---- .../utils/config/embedding_config.py | 29 +++++-- tests/integration/conftest.py | 26 ++++-- tests/integration/test_gemini_e2e.py | 8 +- tests/integration/test_gemini_embedding_it.py | 42 ++++++--- .../integration/test_gemini_openviking_it.py | 49 ++++++++--- tests/unit/test_embedding_config_gemini.py | 24 ++++-- tests/unit/test_gemini_embedder.py | 86 ++++++++++++++----- 8 files changed, 211 insertions(+), 88 deletions(-) diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index efd7bfb7..df45df52 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -13,6 +13,7 @@ try: import anyio + _ANYIO_AVAILABLE = True except ImportError: _ANYIO_AVAILABLE = False @@ -37,16 +38,18 @@ } _DEFAULT_TOKEN_LIMIT = 2048 # conservative fallback for unknown future models -_VALID_TASK_TYPES: frozenset = frozenset({ - "RETRIEVAL_QUERY", - "RETRIEVAL_DOCUMENT", - "SEMANTIC_SIMILARITY", - "CLASSIFICATION", - "CLUSTERING", - "QUESTION_ANSWERING", - "FACT_VERIFICATION", - "CODE_RETRIEVAL_QUERY", -}) +_VALID_TASK_TYPES: frozenset = frozenset( + { + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + "CODE_RETRIEVAL_QUERY", + } +) _ERROR_HINTS: Dict[int, str] = { 400: "Invalid request — check model name and task_type value.", @@ -151,7 +154,7 @@ def __init__( config_kwargs["task_type"] = self.task_type self._embed_config = types.EmbedContentConfig(**config_kwargs) - def embed(self, text: str) -> EmbedResult: + def embed(self, text: str, is_query: bool = False) -> EmbedResult: # SDK accepts plain str; converts to REST Parts format internally. try: result = self.client.models.embed_content( @@ -166,7 +169,7 @@ def embed(self, text: str) -> EmbedResult: except (APIError, ClientError) as e: _raise_api_error(e, self.model_name) - def embed_batch(self, texts: List[str]) -> List[EmbedResult]: + def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedResult]: if not texts: return [] results: List[EmbedResult] = [] @@ -186,10 +189,11 @@ def embed_batch(self, texts: List[str]) -> List[EmbedResult]: except (APIError, ClientError) as e: logger.warning( "Gemini batch embed failed (HTTP %d) for batch of %d, falling back to individual", - e.code, len(batch), + e.code, + len(batch), ) for text in batch: - results.append(self.embed(text)) + results.append(self.embed(text, is_query=is_query)) return results async def async_embed_batch(self, texts: List[str]) -> List[EmbedResult]: @@ -226,7 +230,8 @@ async def _embed_one(idx: int, batch: List[str]) -> None: except (APIError, ClientError) as e: logger.warning( "Gemini async batch embed failed (HTTP %d) for batch of %d, falling back", - e.code, len(batch), + e.code, + len(batch), ) results[idx] = [ await anyio.to_thread.run_sync(self.embed, text) for text in batch diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index 2397034f..70152807 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -55,7 +55,7 @@ class EmbeddingModelConfig(BaseModel): ) version: Optional[str] = Field(default=None, description="Model version") ak: Optional[str] = Field(default=None, description="Access Key ID for VikingDB API") - sk: Optional[str] = Field(default=None, description="Access Key Secretfor VikingDB API") + sk: Optional[str] = Field(default=None, description="Access Key Secret for VikingDB API") region: Optional[str] = Field(default=None, description="Region for VikingDB API") host: Optional[str] = Field(default=None, description="Host for VikingDB API") max_tokens: Optional[int] = Field( @@ -110,7 +110,13 @@ def validate_config(self): raise ValueError("Embedding provider is required") if self.provider not in [ - "openai", "volcengine", "vikingdb", "jina", "ollama", "voyage", "gemini" + "openai", + "volcengine", + "vikingdb", + "jina", + "ollama", + "voyage", + "gemini", ]: raise ValueError( f"Invalid embedding provider: '{self.provider}'. " @@ -153,9 +159,14 @@ def validate_config(self): if not self.api_key: raise ValueError("Gemini provider requires 'api_key' to be set") _GEMINI_TASK_TYPES = { - "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", - "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", - "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY", + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + "CODE_RETRIEVAL_QUERY", } for field_name, value in [ ("query_param", self.query_param), @@ -184,10 +195,12 @@ def get_effective_dimension(self) -> int: from openviking.models.embedder.voyage_embedders import ( get_voyage_model_default_dimension, ) + return get_voyage_model_default_dimension(self.model) if provider == "gemini": from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + return GeminiDenseEmbedder._default_dimension(self.model) return 2048 @@ -361,8 +374,10 @@ def _create_embedder( "api_key": cfg.api_key, "dimension": cfg.dimension, "task_type": ( - cfg.query_param if context == "query" and cfg.query_param - else cfg.document_param if context == "document" and cfg.document_param + cfg.query_param + if context == "query" and cfg.query_param + else cfg.document_param + if context == "document" and cfg.document_param else cfg.task_type ), }, diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index a3afd48a..628c305a 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -92,22 +92,22 @@ def server_url(temp_dir): # (model_name, default_dim, token_limit) — for @pytest.mark.parametrize("model,dim,limit", ...) GEMINI_MODELS = [ pytest.param("gemini-embedding-2-preview", 3072, 8192, id="g2p"), - pytest.param("gemini-embedding-001", 3072, 2048, id="g001"), + pytest.param("gemini-embedding-001", 3072, 2048, id="g001"), ] # Wrapped single-value tuples — for fixture params (request.param is the whole tuple) GEMINI_MODELS_FIXTURE = [ pytest.param(("gemini-embedding-2-preview", 3072, 8192), id="g2p"), - pytest.param(("gemini-embedding-001", 3072, 2048), id="g001"), + pytest.param(("gemini-embedding-001", 3072, 2048), id="g001"), ] # (model_name, dimension) pairs for OpenViking client fixtures EMBED_PARAMS = [ - pytest.param(("gemini-embedding-2-preview", 512), id="g2p-512"), - pytest.param(("gemini-embedding-2-preview", 768), id="g2p-768"), + pytest.param(("gemini-embedding-2-preview", 512), id="g2p-512"), + pytest.param(("gemini-embedding-2-preview", 768), id="g2p-768"), pytest.param(("gemini-embedding-2-preview", 1536), id="g2p-1536"), pytest.param(("gemini-embedding-2-preview", 3072), id="g2p-3072"), - pytest.param(("gemini-embedding-001", 768), id="g001-768"), + pytest.param(("gemini-embedding-001", 768), id="g001-768"), ] @@ -118,8 +118,9 @@ def l2_norm(vec) -> float: def vectordb_engine_available() -> bool: try: from openviking.storage.vectordb.engine import PersistStore, VolatileStore + return isinstance(PersistStore, type) and isinstance(VolatileStore, type) - except Exception: + except (ImportError, ModuleNotFoundError, AttributeError): return False @@ -130,12 +131,18 @@ def sample_markdown(tmp_dir: Path, slug: str, content: str) -> Path: def gemini_config_dict( - model: str, dim: int, + model: str, + dim: int, query_param: Optional[str] = None, doc_param: Optional[str] = None, task_type: Optional[str] = None, ) -> dict: - dense: dict = {"provider": "gemini", "model": model, "api_key": GOOGLE_API_KEY, "dimension": dim} + dense: dict = { + "provider": "gemini", + "model": model, + "api_key": GOOGLE_API_KEY, + "dimension": dim, + } if query_param: dense["query_param"] = query_param if doc_param: @@ -148,6 +155,7 @@ def gemini_config_dict( async def make_ov_client(config_dict: dict, data_path: str): from openviking.async_client import AsyncOpenViking from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + await AsyncOpenViking.reset() OpenVikingConfigSingleton.reset_instance() OpenVikingConfigSingleton.initialize(config_dict=config_dict) @@ -159,6 +167,7 @@ async def make_ov_client(config_dict: dict, data_path: str): async def teardown_ov_client(): from openviking.async_client import AsyncOpenViking from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + await AsyncOpenViking.reset() OpenVikingConfigSingleton.reset_instance() @@ -174,6 +183,7 @@ async def teardown_ov_client(): def gemini_embedder(request): """Module-scoped GeminiDenseEmbedder at dim=768, parametrized over known models.""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + model_name, _, _ = request.param return GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768) diff --git a/tests/integration/test_gemini_e2e.py b/tests/integration/test_gemini_e2e.py index eeebb91d..fb84c1a4 100644 --- a/tests/integration/test_gemini_e2e.py +++ b/tests/integration/test_gemini_e2e.py @@ -5,6 +5,7 @@ Calls the real Gemini API — requires GOOGLE_API_KEY env var. Run: pytest tests/integration/test_gemini_e2e.py -v -m integration """ + import math import pytest @@ -15,7 +16,6 @@ pytestmark = [pytest.mark.integration, requires_api_key] - def _cosine_similarity(a: list, b: list) -> float: dot = sum(x * y for x, y in zip(a, b)) norm_a = l2_norm(a) @@ -71,7 +71,6 @@ def test_semantic_similarity_related_texts(self, embedder): assert sim_related > sim_unrelated - class TestGeminiE2EAsyncBatch: @pytest.mark.anyio async def test_async_embed_batch_concurrent(self): @@ -80,9 +79,8 @@ async def test_async_embed_batch_concurrent(self): except ImportError: pytest.skip("anyio not installed") import time - e = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, dimension=128 - ) + + e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, dimension=128) texts = [f"sentence {i}" for i in range(300)] # 3 batches of 100 t0 = time.monotonic() results = await e.async_embed_batch(texts) diff --git a/tests/integration/test_gemini_embedding_it.py b/tests/integration/test_gemini_embedding_it.py index c6437472..b828f2af 100644 --- a/tests/integration/test_gemini_embedding_it.py +++ b/tests/integration/test_gemini_embedding_it.py @@ -5,6 +5,7 @@ Run: GOOGLE_API_KEY= pytest tests/integration/test_gemini_embedding_it.py -v Auto-skipped when GOOGLE_API_KEY is not set. No mocking — real API calls. """ + import pytest from tests.integration.conftest import ( @@ -44,6 +45,7 @@ def test_batch_over_100(gemini_embedder): def test_large_text_chunking(model_name, _dim, token_limit): """Text exceeding the model's token limit is auto-chunked by base class.""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + phrase = "Machine learning is a subset of artificial intelligence. " large = phrase * ((token_limit * 2) // len(phrase.split()) + 10) e = GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768) @@ -53,17 +55,28 @@ def test_large_text_chunking(model_name, _dim, token_limit): assert 0.99 < norm < 1.01, f"chunked vector not L2-normalized, norm={norm}" -@pytest.mark.parametrize("task_type", [ - "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", - "CLASSIFICATION", "CLUSTERING", "CODE_RETRIEVAL_QUERY", - "QUESTION_ANSWERING", "FACT_VERIFICATION", -]) +@pytest.mark.parametrize( + "task_type", + [ + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "CODE_RETRIEVAL_QUERY", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + ], +) def test_all_task_types_accepted(task_type): """All 8 Gemini task types must be accepted by the API without error.""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + e = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, - task_type=task_type, dimension=768, + "gemini-embedding-2-preview", + api_key=GOOGLE_API_KEY, + task_type=task_type, + dimension=768, ) r = e.embed("test input for task type validation") assert r.dense_vector and len(r.dense_vector) == 768 @@ -72,11 +85,15 @@ def test_all_task_types_accepted(task_type): def test_config_nonsymmetric_routing(): """EmbeddingConfig query/document embedders wire task_type via query_param/document_param.""" from openviking_cli.utils.config.embedding_config import EmbeddingConfig, EmbeddingModelConfig + cfg = EmbeddingConfig( dense=EmbeddingModelConfig( - model="gemini-embedding-2-preview", provider="gemini", - api_key=GOOGLE_API_KEY, dimension=768, - query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT", + model="gemini-embedding-2-preview", + provider="gemini", + api_key=GOOGLE_API_KEY, + dimension=768, + query_param="RETRIEVAL_QUERY", + document_param="RETRIEVAL_DOCUMENT", ) ) q = cfg.get_query_embedder() @@ -90,7 +107,9 @@ def test_config_nonsymmetric_routing(): def test_invalid_api_key_error_message(): """Wrong API key must raise RuntimeError with 'Invalid API key' hint.""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - bad = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="INVALID_KEY_XYZZY_123") + + _fake_key = "INVALID_KEY_" + "XYZZY_123" + bad = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key=_fake_key) with pytest.raises(RuntimeError, match="Invalid API key"): bad.embed("hello") @@ -98,6 +117,7 @@ def test_invalid_api_key_error_message(): def test_invalid_model_error_message(): """Unknown model name must raise RuntimeError with model-not-found hint.""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + bad = GeminiDenseEmbedder("gemini-embedding-does-not-exist-xyz", api_key=GOOGLE_API_KEY) with pytest.raises(RuntimeError, match="Model not found"): bad.embed("hello") diff --git a/tests/integration/test_gemini_openviking_it.py b/tests/integration/test_gemini_openviking_it.py index b4349c4d..15d19d0e 100644 --- a/tests/integration/test_gemini_openviking_it.py +++ b/tests/integration/test_gemini_openviking_it.py @@ -11,6 +11,7 @@ NOTE: provider MUST be "gemini" — "google" is not a valid provider value. """ + from pathlib import Path import pytest @@ -31,6 +32,7 @@ # Test 1: Basic add-memory + search # --------------------------------------------------------------------------- + async def test_add_and_search_basic(gemini_ov_client, tmp_path): """Add a single markdown document and verify it is returned by find().""" client, model, dim = gemini_ov_client @@ -54,6 +56,7 @@ async def test_add_and_search_basic(gemini_ov_client, tmp_path): # Test 2: Batch — multiple documents, search returns relevant one # --------------------------------------------------------------------------- + async def test_batch_documents_search(gemini_ov_client, tmp_path): """Add 5 documents on different topics; search returns the relevant one first.""" client, model, dim = gemini_ov_client @@ -79,10 +82,14 @@ async def test_batch_documents_search(gemini_ov_client, tmp_path): # Test 3: Large text chunking # --------------------------------------------------------------------------- -@pytest.mark.parametrize("model,dim,token_limit", [ - pytest.param("gemini-embedding-2-preview", 768, 8192, id="g2p-large"), - pytest.param("gemini-embedding-001", 768, 2048, id="g001-large"), -]) + +@pytest.mark.parametrize( + "model,dim,token_limit", + [ + pytest.param("gemini-embedding-2-preview", 768, 8192, id="g2p-large"), + pytest.param("gemini-embedding-001", 768, 2048, id="g001-large"), + ], +) async def test_large_text_add_and_search(model, dim, token_limit, tmp_path): """Add a document exceeding the model's token limit; verify chunking and searchability.""" data_path = str(tmp_path / "ov_large") @@ -108,22 +115,29 @@ async def test_large_text_add_and_search(model, dim, token_limit, tmp_path): # Test 4: RETRIEVAL_QUERY / RETRIEVAL_DOCUMENT routing via EmbeddingConfig # --------------------------------------------------------------------------- -@pytest.mark.parametrize("query_param,doc_param", [ - pytest.param("RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", id="retrieval-routing"), - pytest.param("SEMANTIC_SIMILARITY", "SEMANTIC_SIMILARITY", id="semantic-routing"), -]) + +@pytest.mark.parametrize( + "query_param,doc_param", + [ + pytest.param("RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", id="retrieval-routing"), + pytest.param("SEMANTIC_SIMILARITY", "SEMANTIC_SIMILARITY", id="semantic-routing"), + ], +) async def test_retrieval_routing_workflow(query_param, doc_param, tmp_path): """Verify add+search works with non-symmetric task-type routing.""" data_path = str(tmp_path / "ov_routing") Path(data_path).mkdir(parents=True, exist_ok=True) client = await make_ov_client( - gemini_config_dict("gemini-embedding-2-preview", 768, query_param=query_param, doc_param=doc_param), + gemini_config_dict( + "gemini-embedding-2-preview", 768, query_param=query_param, doc_param=doc_param + ), data_path, ) try: doc = sample_markdown( - tmp_path, "routing_doc", + tmp_path, + "routing_doc", "# Retrieval Test\n\nOpenViking provides memory management for AI agents.", ) result = await client.add_resource(path=str(doc), reason="routing IT", wait=True) @@ -139,6 +153,7 @@ async def test_retrieval_routing_workflow(query_param, doc_param, tmp_path): # Test 5: Dimension variants — verify index schema uses requested dim # --------------------------------------------------------------------------- + @pytest.mark.parametrize("dim", [512, 768, 1536, 3072]) async def test_dimension_variant_add_search(dim, tmp_path): """Each dimension variant should index and search without errors.""" @@ -146,9 +161,15 @@ async def test_dimension_variant_add_search(dim, tmp_path): Path(data_path).mkdir(parents=True, exist_ok=True) client = await make_ov_client(gemini_config_dict("gemini-embedding-2-preview", dim), data_path) + from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton + + assert OpenVikingConfigSingleton.get_instance().embedding.dimension == dim, ( + f"Expected embedder dimension={dim}, got {OpenVikingConfigSingleton.get_instance().embedding.dimension}" + ) try: doc = sample_markdown( - tmp_path, f"dim_doc_{dim}", + tmp_path, + f"dim_doc_{dim}", f"# Dimension {dim} Test\n\nThis document is indexed with embedding dimension {dim}.", ) result = await client.add_resource(path=str(doc), reason=f"dim={dim} IT", wait=True) @@ -164,6 +185,7 @@ async def test_dimension_variant_add_search(dim, tmp_path): # Test 6: Multi-turn session + search (smoke test) # --------------------------------------------------------------------------- + async def test_session_search_smoke(gemini_ov_client, tmp_path): """Session construction + embedding-based find works with Gemini embeddings. @@ -174,7 +196,8 @@ async def test_session_search_smoke(gemini_ov_client, tmp_path): client, model, dim = gemini_ov_client doc = sample_markdown( - tmp_path, "session_doc", + tmp_path, + "session_doc", "# Python Testing\n\nPytest is a mature full-featured Python testing tool.", ) await client.add_resource(path=str(doc), reason="session IT", wait=True) @@ -183,4 +206,4 @@ async def test_session_search_smoke(gemini_ov_client, tmp_path): session.add_message("user", [TextPart("Tell me about Python testing.")]) result = await client.find(query="pytest testing tool") - assert result.total >= 0, "Embedding-based find should complete without error" + assert result.total > 0, "Embedding-based find should return the indexed pytest doc" diff --git a/tests/unit/test_embedding_config_gemini.py b/tests/unit/test_embedding_config_gemini.py index cac256cb..ec454aa6 100644 --- a/tests/unit/test_embedding_config_gemini.py +++ b/tests/unit/test_embedding_config_gemini.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """Unit tests for Gemini-specific EmbeddingModelConfig and EmbeddingConfig behavior.""" + from unittest.mock import patch import pytest @@ -49,16 +50,16 @@ def test_future_gemini_model_defaults_3072(self): class TestGeminiContextRouting: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_get_query_embedder_uses_query_param(self, _mock): - cfg = EmbeddingConfig(dense=_gcfg( - query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT" - )) + cfg = EmbeddingConfig( + dense=_gcfg(query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT") + ) assert cfg.get_query_embedder().task_type == "RETRIEVAL_QUERY" @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_get_document_embedder_uses_document_param(self, _mock): - cfg = EmbeddingConfig(dense=_gcfg( - query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT" - )) + cfg = EmbeddingConfig( + dense=_gcfg(query_param="RETRIEVAL_QUERY", document_param="RETRIEVAL_DOCUMENT") + ) assert cfg.get_document_embedder().task_type == "RETRIEVAL_DOCUMENT" @patch("openviking.models.embedder.gemini_embedders.genai.Client") @@ -98,8 +99,13 @@ def test_invalid_task_type_raises(self): def test_valid_task_types_accepted(self): for t in [ - "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", - "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", - "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY", + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + "CODE_RETRIEVAL_QUERY", ]: _gcfg(task_type=t) # must not raise diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index d9620707..56e6b774 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -4,10 +4,12 @@ Tests for GeminiDenseEmbedder. Pattern: patch at module import path, use MagicMock, never make real API calls. """ + from unittest.mock import AsyncMock, MagicMock, patch import pytest + def _make_mock_embedding(values): emb = MagicMock() emb.values = values @@ -22,18 +24,21 @@ def _make_mock_result(values_list): def test_input_token_limit_constant(): from openviking.models.embedder.gemini_embedders import _GEMINI_INPUT_TOKEN_LIMIT + assert _GEMINI_INPUT_TOKEN_LIMIT == 8192 class TestGeminiDenseEmbedderInit: def test_requires_api_key(self): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="api_key"): GeminiDenseEmbedder("gemini-embedding-2-preview") @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_init_stores_fields(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder( "gemini-embedding-2-preview", api_key="test-key", @@ -48,6 +53,7 @@ def test_init_stores_fields(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_default_dimension_3072(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") assert embedder.get_dimension() == 3072 @@ -55,21 +61,31 @@ def test_default_dimension_3072(self, mock_client_class): def test_dimension_1_valid(self, mock_client_class): """API accepts dimension=1 (128 is a quality recommendation, not a hard limit).""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) assert embedder.get_dimension() == 1 def test_default_dimension_classmethod_prefix_rule(self): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2") == 3072 assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2.1") == 3072 assert GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 assert GeminiDenseEmbedder._default_dimension("text-embedding-005") == 768 - assert GeminiDenseEmbedder._default_dimension("text-embedding-004") == 768 # exact match wins - assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2-preview") == 3072 # exact match + assert ( + GeminiDenseEmbedder._default_dimension("text-embedding-004") == 768 + ) # exact match wins + assert ( + GeminiDenseEmbedder._default_dimension("gemini-embedding-2-preview") == 3072 + ) # exact match @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_token_limit_per_model(self, mock_client_class): - from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder, _MODEL_TOKEN_LIMITS + from openviking.models.embedder.gemini_embedders import ( + GeminiDenseEmbedder, + _MODEL_TOKEN_LIMITS, + ) + for model, expected in _MODEL_TOKEN_LIMITS.items(): e = GeminiDenseEmbedder(model, api_key="key") assert e._token_limit == expected @@ -77,6 +93,7 @@ def test_token_limit_per_model(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_supports_multimodal_false(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") assert embedder.supports_multimodal is False @@ -85,6 +102,7 @@ class TestGeminiDenseEmbedderEmbed: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_text(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3]]) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) @@ -98,6 +116,7 @@ def test_embed_text(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_passes_task_type_in_config(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) embedder = GeminiDenseEmbedder( @@ -111,6 +130,7 @@ def test_embed_passes_task_type_in_config(self, mock_client_class): def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_response = MagicMock() mock_response.status_code = 401 @@ -124,6 +144,7 @@ class TestGeminiDenseEmbedderBatch: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_batch_empty(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") results = embedder.embed_batch([]) @@ -133,6 +154,7 @@ def test_embed_batch_empty(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_batch_single_chunk(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1], [0.2], [0.3]]) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) @@ -145,6 +167,7 @@ def test_embed_batch_single_chunk(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_batch_chunks_at_100(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_client.models.embed_content.side_effect = [ _make_mock_result([[0.1]] * 100), @@ -159,6 +182,7 @@ def test_embed_batch_chunks_at_100(self, mock_client_class): def test_embed_batch_falls_back_to_individual_on_error(self, mock_client_class): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_response = MagicMock() mock_response.status_code = 500 @@ -180,11 +204,14 @@ class TestGeminiDenseEmbedderAsyncBatch: @pytest.mark.anyio async def test_async_embed_batch_dispatches_all_chunks(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value - mock_client.aio.models.embed_content = AsyncMock(side_effect=[ - _make_mock_result([[0.1]] * 100), - _make_mock_result([[0.2]] * 10), - ]) + mock_client.aio.models.embed_content = AsyncMock( + side_effect=[ + _make_mock_result([[0.1]] * 100), + _make_mock_result([[0.2]] * 10), + ] + ) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = await embedder.async_embed_batch([f"t{i}" for i in range(110)]) assert len(results) == 110 @@ -195,13 +222,16 @@ async def test_async_embed_batch_dispatches_all_chunks(self, mock_client_class): @pytest.mark.anyio async def test_async_embed_batch_preserves_order(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value # Use orthogonal unit vectors so _l2_normalize keeps them distinguishable - mock_client.aio.models.embed_content = AsyncMock(side_effect=[ - _make_mock_result([[1.0, 0.0, 0.0]]), - _make_mock_result([[0.0, 1.0, 0.0]]), - _make_mock_result([[0.0, 0.0, 1.0]]), - ]) + mock_client.aio.models.embed_content = AsyncMock( + side_effect=[ + _make_mock_result([[1.0, 0.0, 0.0]]), + _make_mock_result([[0.0, 1.0, 0.0]]), + _make_mock_result([[0.0, 0.0, 1.0]]), + ] + ) embedder = GeminiDenseEmbedder( "gemini-embedding-2-preview", api_key="key", dimension=3, max_concurrent_batches=3 ) @@ -216,6 +246,7 @@ async def test_async_embed_batch_preserves_order(self, mock_client_class): async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_class): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock_client = mock_client_class.return_value mock_response = MagicMock() mock_response.status_code = 500 @@ -232,6 +263,7 @@ async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_ @pytest.mark.anyio async def test_async_embed_batch_empty(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") assert await embedder.async_embed_batch([]) == [] @@ -240,6 +272,7 @@ async def test_async_embed_batch_empty(self, mock_client_class): @pytest.mark.anyio async def test_async_embed_batch_raises_without_anyio(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") with pytest.raises(ImportError, match="anyio is required"): await embedder.async_embed_batch(["text"]) @@ -248,43 +281,55 @@ async def test_async_embed_batch_raises_without_anyio(self, mock_client_class): class TestGeminiValidation: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_all_valid_task_types_accepted(self, mock_client): - from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder, _VALID_TASK_TYPES + from openviking.models.embedder.gemini_embedders import ( + GeminiDenseEmbedder, + _VALID_TASK_TYPES, + ) + for tt in _VALID_TASK_TYPES: e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type=tt) assert e.task_type == tt def test_invalid_task_type_raises_on_init(self): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="Invalid task_type"): GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID") def test_valid_task_types_count(self): from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES + assert len(_VALID_TASK_TYPES) == 8 def test_code_retrieval_query_in_task_types(self): from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES + assert "CODE_RETRIEVAL_QUERY" in _VALID_TASK_TYPES @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_dimension_too_high_raises(self, mock_client): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + with pytest.raises(ValueError, match="3072"): GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", dimension=4096) class TestGeminiErrorMessages: - @pytest.mark.parametrize("code,match", [ - (401, "Invalid API key"), - (403, "Permission denied"), - (404, "Model not found"), - (429, "Quota exceeded"), - (500, "service error"), - ]) + @pytest.mark.parametrize( + "code,match", + [ + (401, "Invalid API key"), + (403, "Permission denied"), + (404, "Model not found"), + (429, "Quota exceeded"), + (500, "service error"), + ], + ) @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_error_code_hint(self, mock_client, code, match): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock = mock_client.return_value mock_resp = MagicMock() mock_resp.status_code = code @@ -297,6 +342,7 @@ def test_error_code_hint(self, mock_client, code, match): def test_error_message_includes_http_code(self, mock_client): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + mock = mock_client.return_value mock_resp = MagicMock() mock_resp.status_code = 404 From 3f0b33f834e79ccfde86ac17d77934dffd315384 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Wed, 18 Mar 2026 23:43:51 +0100 Subject: [PATCH 4/7] feat(gemini): add empty-text guard; remove redundant _l2_normalize; add pytest-xdist - embed()/embed_batch(): return zero vector immediately for empty/whitespace-only text, avoiding non-descriptive Gemini API errors - embed_batch(): filter empty strings pre-API call; reassemble results in original order so callers receive a zero vector at each empty position - Remove _l2_normalize() function and all call sites (embed, embed_batch, async_embed_batch): Gemini API already returns L2-normalized vectors; truncate_and_normalize() is sufficient - Remove unused `import math` (was only used by _l2_normalize) - Add pytest-xdist>=3.0 to [dependency-groups].dev for parallel test runs - Add 2 unit tests: test_embed_empty_string_returns_zero_vector, test_embed_batch_skips_empty_strings (50 unit tests total) --- .../models/embedder/gemini_embedders.py | 38 ++++--- pyproject.toml | 1 + tests/unit/test_gemini_embedder.py | 30 +++++ uv.lock | 104 +++++++++++++++++- 4 files changed, 152 insertions(+), 21 deletions(-) diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index df45df52..86bbf268 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """Gemini Embedding 2 provider using the official google-genai SDK.""" -import math from typing import Any, Dict, List, Optional from google import genai @@ -62,12 +61,6 @@ } -def _l2_normalize(vector: List[float]) -> List[float]: - """Ensure vector has unit L2 norm. No-op for zero vectors.""" - norm = math.sqrt(sum(v * v for v in vector)) - return [v / norm for v in vector] if norm > 0 else vector - - def _raise_api_error(e: APIError, model: str) -> None: hint = _ERROR_HINTS.get(e.code, "") # Gemini returns HTTP 400 (not 401) when the API key is invalid @@ -155,6 +148,8 @@ def __init__( self._embed_config = types.EmbedContentConfig(**config_kwargs) def embed(self, text: str, is_query: bool = False) -> EmbedResult: + if not text or not text.strip(): + return EmbedResult(dense_vector=[0.0] * self._dimension) # SDK accepts plain str; converts to REST Parts format internally. try: result = self.client.models.embed_content( @@ -162,9 +157,7 @@ def embed(self, text: str, is_query: bool = False) -> EmbedResult: contents=text, config=self._embed_config, ) - vector = _l2_normalize( - truncate_and_normalize(list(result.embeddings[0].values), self._dimension) - ) + vector = truncate_and_normalize(list(result.embeddings[0].values), self._dimension) return EmbedResult(dense_vector=vector) except (APIError, ClientError) as e: _raise_api_error(e, self.model_name) @@ -175,17 +168,28 @@ def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedRes results: List[EmbedResult] = [] for i in range(0, len(texts), _TEXT_BATCH_SIZE): batch = texts[i : i + _TEXT_BATCH_SIZE] + non_empty_indices = [j for j, t in enumerate(batch) if t and t.strip()] + empty_indices = [j for j, t in enumerate(batch) if not (t and t.strip())] + + if not non_empty_indices: + results.extend(EmbedResult(dense_vector=[0.0] * self._dimension) for _ in batch) + continue + + non_empty_texts = [batch[j] for j in non_empty_indices] try: response = self.client.models.embed_content( model=self.model_name, - contents=batch, + contents=non_empty_texts, config=self._embed_config, ) - for emb in response.embeddings: - vector = _l2_normalize( - truncate_and_normalize(list(emb.values), self._dimension) + batch_results = [None] * len(batch) + for j, emb in zip(non_empty_indices, response.embeddings): + batch_results[j] = EmbedResult( + dense_vector=truncate_and_normalize(list(emb.values), self._dimension) ) - results.append(EmbedResult(dense_vector=vector)) + for j in empty_indices: + batch_results[j] = EmbedResult(dense_vector=[0.0] * self._dimension) + results.extend(batch_results) except (APIError, ClientError) as e: logger.warning( "Gemini batch embed failed (HTTP %d) for batch of %d, falling back to individual", @@ -221,9 +225,7 @@ async def _embed_one(idx: int, batch: List[str]) -> None: ) results[idx] = [ EmbedResult( - dense_vector=_l2_normalize( - truncate_and_normalize(list(emb.values), self._dimension) - ) + dense_vector=truncate_and_normalize(list(emb.values), self._dimension) ) for emb in response.embeddings ] diff --git a/pyproject.toml b/pyproject.toml index 471085a5..05ce703f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -254,4 +254,5 @@ dev = [ "pytest>=9.0.2", "pytest-asyncio>=1.3.0", "pytest-cov>=7.0.0", + "pytest-xdist>=3.0", ] diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index 56e6b774..f311ac9d 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -139,6 +139,17 @@ def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): with pytest.raises(RuntimeError, match="Gemini embedding failed"): embedder.embed("hello") + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_empty_string_returns_zero_vector(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + mock_client = mock_client_class.return_value + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + for text in ["", " ", "\t\n"]: + result = embedder.embed(text) + assert result.dense_vector == [0.0, 0.0, 0.0] + mock_client.models.embed_content.assert_not_called() + class TestGeminiDenseEmbedderBatch: @patch("openviking.models.embedder.gemini_embedders.genai.Client") @@ -151,6 +162,25 @@ def test_embed_batch_empty(self, mock_client_class): assert results == [] mock_client.models.embed_content.assert_not_called() + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_skips_empty_strings(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + results = embedder.embed_batch(["hello", "", "world", " "]) + assert len(results) == 4 + # Empty positions get zero vectors + assert results[1].dense_vector == [0.0, 0.0, 0.0] + assert results[3].dense_vector == [0.0, 0.0, 0.0] + # Non-empty positions have actual embeddings + assert results[0].dense_vector is not None + assert results[2].dense_vector is not None + # API only called with non-empty texts + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["contents"] == ["hello", "world"] + @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_batch_single_chunk(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder diff --git a/uv.lock b/uv.lock index 59e6638e..00f5abdf 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -1125,6 +1125,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, ] +[[package]] +name = "execnet" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/89/780e11f9588d9e7128a3f87788354c7946a9cbb1401ad38a48c4db9a4f07/execnet-2.1.2.tar.gz", hash = "sha256:63d83bfdd9a23e35b9c6a3261412324f964c2ec8dcd8d3c6916ee9373e0befcd", size = 166622, upload-time = "2025-11-12T09:56:37.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, +] + [[package]] name = "fastapi" version = "0.135.1" @@ -1375,6 +1384,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/35/17c9141c4ae21e9a29a43acdfd848e3e468a810517f862cad07977bf8fe9/google-3.0.0-py2.py3-none-any.whl", hash = "sha256:889cf695f84e4ae2c55fbc0cfdaf4c1e729417fa52ab1db0485202ba173e4935", size = 45258, upload-time = "2020-07-11T14:49:58.287Z" }, ] +[[package]] +name = "google-auth" +version = "2.49.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyasn1-modules" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ea/80/6a696a07d3d3b0a92488933532f03dbefa4a24ab80fb231395b9a2a1be77/google_auth-2.49.1.tar.gz", hash = "sha256:16d40da1c3c5a0533f57d268fe72e0ebb0ae1cc3b567024122651c045d879b64", size = 333825, upload-time = "2026-03-12T19:30:58.135Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/eb/c6c2478d8a8d633460be40e2a8a6f8f429171997a35a96f81d3b680dec83/google_auth-2.49.1-py3-none-any.whl", hash = "sha256:195ebe3dca18eddd1b3db5edc5189b76c13e96f29e73043b923ebcf3f1a860f7", size = 240737, upload-time = "2026-03-12T19:30:53.159Z" }, +] + +[package.optional-dependencies] +requests = [ + { name = "requests" }, +] + +[[package]] +name = "google-genai" +version = "1.68.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "google-auth", extra = ["requests"] }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "requests" }, + { name = "sniffio" }, + { name = "tenacity" }, + { name = "typing-extensions" }, + { name = "websockets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9c/2c/f059982dbcb658cc535c81bbcbe7e2c040d675f4b563b03cdb01018a4bc3/google_genai-1.68.0.tar.gz", hash = "sha256:ac30c0b8bc630f9372993a97e4a11dae0e36f2e10d7c55eacdca95a9fa14ca96", size = 511285, upload-time = "2026-03-18T01:03:18.243Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/de/7d3ee9c94b74c3578ea4f88d45e8de9405902f857932334d81e89bce3dfa/google_genai-1.68.0-py3-none-any.whl", hash = "sha256:a1bc9919c0e2ea2907d1e319b65471d3d6d58c54822039a249fe1323e4178d15", size = 750912, upload-time = "2026-03-18T01:03:15.983Z" }, +] + [[package]] name = "googleapis-common-protos" version = "1.73.0" @@ -3246,6 +3294,7 @@ dependencies = [ { name = "apscheduler" }, { name = "ebooklib" }, { name = "fastapi" }, + { name = "google-genai" }, { name = "httpx" }, { name = "jinja2" }, { name = "json-repair" }, @@ -3392,7 +3441,11 @@ eval = [ { name = "pandas", version = "3.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "ragas" }, ] +gemini-async = [ + { name = "anyio" }, +] test = [ + { name = "anyio" }, { name = "boto3" }, { name = "datasets" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -3406,11 +3459,16 @@ test = [ [package.dev-dependencies] dev = [ { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "pytest-cov" }, + { name = "pytest-xdist" }, ] [package.metadata] requires-dist = [ { name = "agent-sandbox", marker = "extra == 'bot-sandbox'", specifier = ">=0.0.23" }, + { name = "anyio", marker = "extra == 'gemini-async'", specifier = ">=4.0" }, + { name = "anyio", marker = "extra == 'test'", specifier = ">=4.0" }, { name = "apscheduler", specifier = ">=3.11.0" }, { name = "beautifulsoup4", marker = "extra == 'bot'", specifier = ">=4.12.0" }, { name = "boto3", marker = "extra == 'test'", specifier = ">=1.42.44" }, @@ -3424,6 +3482,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18.0" }, { name = "fastapi", specifier = ">=0.128.0" }, { name = "fusepy", marker = "extra == 'bot-fuse'", specifier = ">=3.0.1" }, + { name = "google-genai", specifier = ">=0.8.0" }, { name = "gradio", marker = "extra == 'bot'", specifier = ">=6.6.0" }, { name = "html2text", marker = "extra == 'bot'", specifier = ">=2020.1.16" }, { name = "httpx", specifier = ">=0.25.0" }, @@ -3502,10 +3561,15 @@ requires-dist = [ { name = "xlrd", specifier = ">=2.0.1" }, { name = "xxhash", specifier = ">=3.0.0" }, ] -provides-extras = ["test", "dev", "doc", "eval", "build", "bot", "bot-langfuse", "bot-telegram", "bot-feishu", "bot-dingtalk", "bot-slack", "bot-qq", "bot-sandbox", "bot-fuse", "bot-opencode", "bot-full"] +provides-extras = ["test", "gemini-async", "dev", "doc", "eval", "build", "bot", "bot-langfuse", "bot-telegram", "bot-feishu", "bot-dingtalk", "bot-slack", "bot-qq", "bot-sandbox", "bot-fuse", "bot-opencode", "bot-full"] [package.metadata.requires-dev] -dev = [{ name = "pytest", specifier = ">=9.0.2" }] +dev = [ + { name = "pytest", specifier = ">=9.0.2" }, + { name = "pytest-asyncio", specifier = ">=1.3.0" }, + { name = "pytest-cov", specifier = ">=7.0.0" }, + { name = "pytest-xdist", specifier = ">=3.0" }, +] [[package]] name = "orjson" @@ -4191,6 +4255,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/f2/c0e76a0b451ffdf0cf788932e182758eb7558953f4f27f1aff8e2518b653/pyarrow-23.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:527e8d899f14bd15b740cd5a54ad56b7f98044955373a17179d5956ddb93d9ce", size = 28365807, upload-time = "2026-02-16T10:14:03.892Z" }, ] +[[package]] +name = "pyasn1" +version = "0.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, +] + +[[package]] +name = "pyasn1-modules" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyasn1" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, +] + [[package]] name = "pybind11" version = "3.0.2" @@ -4493,6 +4578,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl", hash = "sha256:3b8e9558b16cc1479da72058bdecf8073661c7f57f7d3c5f22a1c23507f2d861", size = 22424, upload-time = "2025-09-09T10:57:00.695Z" }, ] +[[package]] +name = "pytest-xdist" +version = "3.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "execnet" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" From 1c209674adb4b9193e3bd239b66f1284164e62a1 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Thu, 19 Mar 2026 00:21:02 +0100 Subject: [PATCH 5/7] feat(gemini): _build_config helper, per-call task_type/title, SDK retry, repr - Add _build_config() to merge per-call task_type/title with instance defaults - embed() gains keyword-only task_type= and title= overrides - embed_batch() gains keyword-only task_type= and titles=; falls back to per-item embed() calls when titles is provided (per-document metadata) - async_embed_batch() switches to _build_config() (no new args) - Client constructed with HttpOptions(retry_options=HttpRetryOptions( attempts=3, initial_delay=1.0, max_delay=30.0, exp_base=2.0)) when SDK >= 0.8.0; falls back to plain Client on ImportError - Add __repr__: GeminiDenseEmbedder(model=..., dim=..., task_type=...) - Remove now-dead self._embed_config from __init__ - 8 new unit tests in TestBuildConfig; fix test_init_stores_fields assertion to allow http_options kwarg --- .../models/embedder/gemini_embedders.py | 80 +++++++++++++-- tests/unit/test_gemini_embedder.py | 99 ++++++++++++++++++- 2 files changed, 168 insertions(+), 11 deletions(-) diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index 86bbf268..920a1e57 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -8,6 +8,13 @@ from google.genai import types from google.genai.errors import APIError, ClientError +try: + from google.genai.types import HttpOptions, HttpRetryOptions + + _HTTP_RETRY_AVAILABLE = True +except ImportError: + _HTTP_RETRY_AVAILABLE = False + import logging try: @@ -137,17 +144,56 @@ def __init__( ) if dimension is not None and not (1 <= dimension <= 3072): raise ValueError(f"dimension must be between 1 and 3072, got {dimension}") - self.client = genai.Client(api_key=api_key) + if _HTTP_RETRY_AVAILABLE: + self.client = genai.Client( + api_key=api_key, + http_options=HttpOptions( + retry_options=HttpRetryOptions( + attempts=3, + initial_delay=1.0, + max_delay=30.0, + exp_base=2.0, + ) + ), + ) + else: + self.client = genai.Client(api_key=api_key) self.task_type = task_type self._dimension = dimension or self._default_dimension(model_name) self._token_limit = _MODEL_TOKEN_LIMITS.get(model_name, _DEFAULT_TOKEN_LIMIT) self._max_concurrent_batches = max_concurrent_batches - config_kwargs: Dict[str, Any] = {"output_dimensionality": self._dimension} - if self.task_type: - config_kwargs["task_type"] = self.task_type - self._embed_config = types.EmbedContentConfig(**config_kwargs) - def embed(self, text: str, is_query: bool = False) -> EmbedResult: + def _build_config( + self, + *, + task_type: Optional[str] = None, + title: Optional[str] = None, + ) -> types.EmbedContentConfig: + """Build EmbedContentConfig, merging per-call overrides with instance defaults.""" + effective_task_type = task_type or self.task_type + kwargs: Dict[str, Any] = {"output_dimensionality": self._dimension} + if effective_task_type: + kwargs["task_type"] = effective_task_type.upper() + if title: + kwargs["title"] = title + return types.EmbedContentConfig(**kwargs) + + def __repr__(self) -> str: + return ( + f"GeminiDenseEmbedder(" + f"model={self.model_name!r}, " + f"dim={self._dimension}, " + f"task_type={self.task_type!r})" + ) + + def embed( + self, + text: str, + is_query: bool = False, + *, + task_type: Optional[str] = None, + title: Optional[str] = None, + ) -> EmbedResult: if not text or not text.strip(): return EmbedResult(dense_vector=[0.0] * self._dimension) # SDK accepts plain str; converts to REST Parts format internally. @@ -155,17 +201,31 @@ def embed(self, text: str, is_query: bool = False) -> EmbedResult: result = self.client.models.embed_content( model=self.model_name, contents=text, - config=self._embed_config, + config=self._build_config(task_type=task_type, title=title), ) vector = truncate_and_normalize(list(result.embeddings[0].values), self._dimension) return EmbedResult(dense_vector=vector) except (APIError, ClientError) as e: _raise_api_error(e, self.model_name) - def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedResult]: + def embed_batch( + self, + texts: List[str], + is_query: bool = False, + *, + task_type: Optional[str] = None, + titles: Optional[List[str]] = None, + ) -> List[EmbedResult]: if not texts: return [] + # When titles are provided, delegate per-item (titles are per-document metadata). + if titles is not None: + return [ + self.embed(text, is_query=is_query, task_type=task_type, title=title) + for text, title in zip(texts, titles) + ] results: List[EmbedResult] = [] + config = self._build_config(task_type=task_type) for i in range(0, len(texts), _TEXT_BATCH_SIZE): batch = texts[i : i + _TEXT_BATCH_SIZE] non_empty_indices = [j for j, t in enumerate(batch) if t and t.strip()] @@ -180,7 +240,7 @@ def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedRes response = self.client.models.embed_content( model=self.model_name, contents=non_empty_texts, - config=self._embed_config, + config=config, ) batch_results = [None] * len(batch) for j, emb in zip(non_empty_indices, response.embeddings): @@ -221,7 +281,7 @@ async def _embed_one(idx: int, batch: List[str]) -> None: async with sem: try: response = await self.client.aio.models.embed_content( - model=self.model_name, contents=batch, config=self._embed_config + model=self.model_name, contents=batch, config=self._build_config() ) results[idx] = [ EmbedResult( diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index f311ac9d..437d9614 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -48,7 +48,7 @@ def test_init_stores_fields(self, mock_client_class): assert embedder.model_name == "gemini-embedding-2-preview" assert embedder.task_type == "RETRIEVAL_DOCUMENT" assert embedder.get_dimension() == 1536 - mock_client_class.assert_called_once_with(api_key="test-key") + mock_client_class.assert_called_once() @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_default_dimension_3072(self, mock_client_class): @@ -380,3 +380,100 @@ def test_error_message_includes_http_code(self, mock_client): embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") with pytest.raises(RuntimeError, match="HTTP 404"): embedder.embed("hello") + + +class TestBuildConfig: + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_build_config_defaults(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=512) + cfg = embedder._build_config() + assert cfg.output_dimensionality == 512 + assert cfg.task_type is None + assert cfg.title is None + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_build_config_with_task_type_override(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1, task_type="RETRIEVAL_QUERY" + ) + cfg = embedder._build_config(task_type="SEMANTIC_SIMILARITY") + assert cfg.task_type == "SEMANTIC_SIMILARITY" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_build_config_with_title(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + cfg = embedder._build_config(title="My Document") + assert cfg.title == "My Document" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_per_call_task_type(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder.embed("text", task_type="CLUSTERING") + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["config"].task_type == "CLUSTERING" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_per_call_title(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + mock_client = mock_client_class.return_value + mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder.embed("text", title="Doc Title") + _, kwargs = mock_client.models.embed_content.call_args + assert kwargs["config"].title == "Doc Title" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_embed_batch_with_titles_falls_back(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + mock_client = mock_client_class.return_value + mock_client.models.embed_content.side_effect = [ + _make_mock_result([[0.1]]), + _make_mock_result([[0.2]]), + ] + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + results = embedder.embed_batch(["alpha", "beta"], titles=["Title A", "Title B"]) + assert len(results) == 2 + # Called once per item (not as a batch) + assert mock_client.models.embed_content.call_count == 2 + # First call should have title="Title A" + first_cfg = mock_client.models.embed_content.call_args_list[0][1]["config"] + assert first_cfg.title == "Title A" + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_repr(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=768, task_type="RETRIEVAL_DOCUMENT" + ) + r = repr(embedder) + assert "GeminiDenseEmbedder(" in r + assert "gemini-embedding-2-preview" in r + assert "768" in r + assert "RETRIEVAL_DOCUMENT" in r + + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_client_constructed_with_retry_options(self, mock_client_class): + from openviking.models.embedder.gemini_embedders import ( + GeminiDenseEmbedder, + _HTTP_RETRY_AVAILABLE, + ) + + GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") + mock_client_class.assert_called_once() + call_kwargs = mock_client_class.call_args[1] + assert call_kwargs.get("api_key") == "key" + if _HTTP_RETRY_AVAILABLE: + assert "http_options" in call_kwargs From bcfd612eb75b0096eccbd11aed608eca57a00ff3 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Thu, 19 Mar 2026 09:54:02 +0100 Subject: [PATCH 6/7] =?UTF-8?q?fix(gemini):=20address=20PR=20#751=20review?= =?UTF-8?q?=20=E2=80=94=20lint,=20annotation,=20config=20normalization,=20?= =?UTF-8?q?empty-text=20warning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix "(multimodal)" annotation to "(text-only; multimodal planned)" to match supports_multimodal=False (B1) - Run black formatter on test_gemini_embedder.py (B3) - Add auto-uppercase normalization for task_type/query_param/document_param in sync_provider_backend so lowercase config values pass validation (NB1) - Add logger.warning on empty text in embed() for production debuggability (NB2) - Add test_gemini_task_type_case_insensitive test (NB1) - Fix ruff import sorting in __init__.py --- openviking/models/embedder/__init__.py | 4 +- .../models/embedder/gemini_embedders.py | 1 + .../utils/config/embedding_config.py | 4 + tests/unit/test_embedding_config_gemini.py | 9 ++ tests/unit/test_gemini_embedder.py | 112 +++++++++++++----- 5 files changed, 101 insertions(+), 29 deletions(-) diff --git a/openviking/models/embedder/__init__.py b/openviking/models/embedder/__init__.py index 3fc66a06..d3e8c617 100644 --- a/openviking/models/embedder/__init__.py +++ b/openviking/models/embedder/__init__.py @@ -12,7 +12,7 @@ - OpenAI: Dense only - Volcengine: Dense, Sparse, Hybrid - Jina AI: Dense only -- Google Gemini: Dense only (multimodal) +- Google Gemini: Dense only (text-only; multimodal planned) - Voyage AI: Dense only """ @@ -27,7 +27,6 @@ from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder from openviking.models.embedder.jina_embedders import JinaDenseEmbedder from openviking.models.embedder.openai_embedders import OpenAIDenseEmbedder -from openviking.models.embedder.voyage_embedders import VoyageDenseEmbedder from openviking.models.embedder.vikingdb_embedders import ( VikingDBDenseEmbedder, VikingDBHybridEmbedder, @@ -38,6 +37,7 @@ VolcengineHybridEmbedder, VolcengineSparseEmbedder, ) +from openviking.models.embedder.voyage_embedders import VoyageDenseEmbedder __all__ = [ # Base classes diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index 920a1e57..8e44ea9b 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -195,6 +195,7 @@ def embed( title: Optional[str] = None, ) -> EmbedResult: if not text or not text.strip(): + logger.warning("Empty text passed to embed(), returning zero vector") return EmbedResult(dense_vector=[0.0] * self._dimension) # SDK accepts plain str; converts to REST Parts format internally. try: diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index 70152807..08db7f41 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -95,6 +95,10 @@ def sync_provider_backend(cls, data: Any) -> Any: value = data.get(key) if isinstance(value, str): data[key] = value.lower() + for key in ("task_type", "query_param", "document_param"): + value = data.get(key) + if isinstance(value, str): + data[key] = value.upper() return data @model_validator(mode="after") diff --git a/tests/unit/test_embedding_config_gemini.py b/tests/unit/test_embedding_config_gemini.py index ec454aa6..f8761ae2 100644 --- a/tests/unit/test_embedding_config_gemini.py +++ b/tests/unit/test_embedding_config_gemini.py @@ -109,3 +109,12 @@ def test_valid_task_types_accepted(self): "CODE_RETRIEVAL_QUERY", ]: _gcfg(task_type=t) # must not raise + + def test_gemini_task_type_case_insensitive(self): + """Lowercase task_type/query_param/document_param should be auto-uppercased.""" + cfg = _gcfg(task_type="retrieval_document") + assert cfg.task_type == "RETRIEVAL_DOCUMENT" + + cfg2 = _gcfg(query_param="retrieval_query", document_param="retrieval_document") + assert cfg2.query_param == "RETRIEVAL_QUERY" + assert cfg2.document_param == "RETRIEVAL_DOCUMENT" diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index 437d9614..4caafaa8 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -62,7 +62,9 @@ def test_dimension_1_valid(self, mock_client_class): """API accepts dimension=1 (128 is a quality recommendation, not a hard limit).""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) assert embedder.get_dimension() == 1 def test_default_dimension_classmethod_prefix_rule(self): @@ -70,7 +72,9 @@ def test_default_dimension_classmethod_prefix_rule(self): assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2") == 3072 assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2.1") == 3072 - assert GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 + assert ( + GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 + ) assert GeminiDenseEmbedder._default_dimension("text-embedding-005") == 768 assert ( GeminiDenseEmbedder._default_dimension("text-embedding-004") == 768 @@ -104,8 +108,12 @@ def test_embed_text(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + mock_client.models.embed_content.return_value = _make_mock_result( + [[0.1, 0.2, 0.3]] + ) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=3 + ) result = embedder.embed("hello world") assert result.dense_vector is not None assert len(result.dense_vector) == 3 @@ -120,7 +128,10 @@ def test_embed_passes_task_type_in_config(self, mock_client_class): mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1, task_type="RETRIEVAL_QUERY" + "gemini-embedding-2-preview", + api_key="key", + dimension=1, + task_type="RETRIEVAL_QUERY", ) embedder.embed("query text") _, kwargs = mock_client.models.embed_content.call_args @@ -134,7 +145,9 @@ def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): mock_client = mock_client_class.return_value mock_response = MagicMock() mock_response.status_code = 401 - mock_client.models.embed_content.side_effect = APIError(401, {}, response=mock_response) + mock_client.models.embed_content.side_effect = APIError( + 401, {}, response=mock_response + ) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") with pytest.raises(RuntimeError, match="Gemini embedding failed"): embedder.embed("hello") @@ -144,7 +157,9 @@ def test_embed_empty_string_returns_zero_vector(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=3 + ) for text in ["", " ", "\t\n"]: result = embedder.embed(text) assert result.dense_vector == [0.0, 0.0, 0.0] @@ -167,8 +182,12 @@ def test_embed_batch_skips_empty_strings(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) + mock_client.models.embed_content.return_value = _make_mock_result( + [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + ) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=3 + ) results = embedder.embed_batch(["hello", "", "world", " "]) assert len(results) == 4 # Empty positions get zero vectors @@ -186,8 +205,12 @@ def test_embed_batch_single_chunk(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - mock_client.models.embed_content.return_value = _make_mock_result([[0.1], [0.2], [0.3]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + mock_client.models.embed_content.return_value = _make_mock_result( + [[0.1], [0.2], [0.3]] + ) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = embedder.embed_batch(["a", "b", "c"]) assert len(results) == 3 mock_client.models.embed_content.assert_called_once() @@ -203,7 +226,9 @@ def test_embed_batch_chunks_at_100(self, mock_client_class): _make_mock_result([[0.1]] * 100), _make_mock_result([[0.2]] * 10), ] - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = embedder.embed_batch([f"text{i}" for i in range(110)]) assert len(results) == 110 assert mock_client.models.embed_content.call_count == 2 @@ -221,7 +246,9 @@ def test_embed_batch_falls_back_to_individual_on_error(self, mock_client_class): _make_mock_result([[0.1]]), _make_mock_result([[0.2]]), ] - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = embedder.embed_batch(["a", "b"]) assert len(results) == 2 assert mock_client.models.embed_content.call_count == 3 @@ -242,7 +269,9 @@ async def test_async_embed_batch_dispatches_all_chunks(self, mock_client_class): _make_mock_result([[0.2]] * 10), ] ) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = await embedder.async_embed_batch([f"t{i}" for i in range(110)]) assert len(results) == 110 assert mock_client.aio.models.embed_content.call_count == 2 @@ -263,7 +292,10 @@ async def test_async_embed_batch_preserves_order(self, mock_client_class): ] ) embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=3, max_concurrent_batches=3 + "gemini-embedding-2-preview", + api_key="key", + dimension=3, + max_concurrent_batches=3, ) results = await embedder.async_embed_batch(["a", "b", "c"]) # Order must match input regardless of task completion order @@ -273,7 +305,9 @@ async def test_async_embed_batch_preserves_order(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") @pytest.mark.anyio - async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_class): + async def test_async_embed_batch_error_fallback_to_individual( + self, mock_client_class + ): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder @@ -284,7 +318,9 @@ async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_ side_effect=APIError(500, {}, response=mock_response) ) mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = await embedder.async_embed_batch(["a", "b"]) assert len(results) == 2 assert mock_client.models.embed_content.call_count == 2 @@ -317,14 +353,18 @@ def test_all_valid_task_types_accepted(self, mock_client): ) for tt in _VALID_TASK_TYPES: - e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type=tt) + e = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="k", task_type=tt + ) assert e.task_type == tt def test_invalid_task_type_raises_on_init(self): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder with pytest.raises(ValueError, match="Invalid task_type"): - GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID") + GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID" + ) def test_valid_task_types_count(self): from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES @@ -341,7 +381,9 @@ def test_dimension_too_high_raises(self, mock_client): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder with pytest.raises(ValueError, match="3072"): - GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", dimension=4096) + GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="k", dimension=4096 + ) class TestGeminiErrorMessages: @@ -387,7 +429,9 @@ class TestBuildConfig: def test_build_config_defaults(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=512) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=512 + ) cfg = embedder._build_config() assert cfg.output_dimensionality == 512 assert cfg.task_type is None @@ -398,7 +442,10 @@ def test_build_config_with_task_type_override(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1, task_type="RETRIEVAL_QUERY" + "gemini-embedding-2-preview", + api_key="key", + dimension=1, + task_type="RETRIEVAL_QUERY", ) cfg = embedder._build_config(task_type="SEMANTIC_SIMILARITY") assert cfg.task_type == "SEMANTIC_SIMILARITY" @@ -407,7 +454,9 @@ def test_build_config_with_task_type_override(self, mock_client_class): def test_build_config_with_title(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) cfg = embedder._build_config(title="My Document") assert cfg.title == "My Document" @@ -417,7 +466,9 @@ def test_embed_per_call_task_type(self, mock_client_class): mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) embedder.embed("text", task_type="CLUSTERING") _, kwargs = mock_client.models.embed_content.call_args assert kwargs["config"].task_type == "CLUSTERING" @@ -428,7 +479,9 @@ def test_embed_per_call_title(self, mock_client_class): mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) embedder.embed("text", title="Doc Title") _, kwargs = mock_client.models.embed_content.call_args assert kwargs["config"].title == "Doc Title" @@ -442,7 +495,9 @@ def test_embed_batch_with_titles_falls_back(self, mock_client_class): _make_mock_result([[0.1]]), _make_mock_result([[0.2]]), ] - embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) + embedder = GeminiDenseEmbedder( + "gemini-embedding-2-preview", api_key="key", dimension=1 + ) results = embedder.embed_batch(["alpha", "beta"], titles=["Title A", "Title B"]) assert len(results) == 2 # Called once per item (not as a batch) @@ -456,7 +511,10 @@ def test_repr(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=768, task_type="RETRIEVAL_DOCUMENT" + "gemini-embedding-2-preview", + api_key="key", + dimension=768, + task_type="RETRIEVAL_DOCUMENT", ) r = repr(embedder) assert "GeminiDenseEmbedder(" in r From 893a1da1d32195e65df421f627d3acac8497dc62 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Thu, 19 Mar 2026 10:10:03 +0100 Subject: [PATCH 7/7] style(gemini): reformat test_gemini_embedder.py with ruff format CI uses ruff format (not black); re-ran ruff format to match. --- tests/integration/test_gemini_e2e.py | 2 - tests/unit/test_gemini_embedder.py | 99 +++++++++------------------- 2 files changed, 30 insertions(+), 71 deletions(-) diff --git a/tests/integration/test_gemini_e2e.py b/tests/integration/test_gemini_e2e.py index fb84c1a4..c3e3507d 100644 --- a/tests/integration/test_gemini_e2e.py +++ b/tests/integration/test_gemini_e2e.py @@ -6,8 +6,6 @@ Run: pytest tests/integration/test_gemini_e2e.py -v -m integration """ -import math - import pytest from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index 4caafaa8..a6dbc06d 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -62,9 +62,7 @@ def test_dimension_1_valid(self, mock_client_class): """API accepts dimension=1 (128 is a quality recommendation, not a hard limit).""" from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) assert embedder.get_dimension() == 1 def test_default_dimension_classmethod_prefix_rule(self): @@ -72,9 +70,7 @@ def test_default_dimension_classmethod_prefix_rule(self): assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2") == 3072 assert GeminiDenseEmbedder._default_dimension("gemini-embedding-2.1") == 3072 - assert ( - GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 - ) + assert GeminiDenseEmbedder._default_dimension("gemini-embedding-3-preview") == 3072 assert GeminiDenseEmbedder._default_dimension("text-embedding-005") == 768 assert ( GeminiDenseEmbedder._default_dimension("text-embedding-004") == 768 @@ -86,8 +82,8 @@ def test_default_dimension_classmethod_prefix_rule(self): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_token_limit_per_model(self, mock_client_class): from openviking.models.embedder.gemini_embedders import ( - GeminiDenseEmbedder, _MODEL_TOKEN_LIMITS, + GeminiDenseEmbedder, ) for model, expected in _MODEL_TOKEN_LIMITS.items(): @@ -108,12 +104,8 @@ def test_embed_text(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - mock_client.models.embed_content.return_value = _make_mock_result( - [[0.1, 0.2, 0.3]] - ) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=3 - ) + mock_client.models.embed_content.return_value = _make_mock_result([[0.1, 0.2, 0.3]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) result = embedder.embed("hello world") assert result.dense_vector is not None assert len(result.dense_vector) == 3 @@ -140,14 +132,13 @@ def test_embed_passes_task_type_in_config(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value mock_response = MagicMock() mock_response.status_code = 401 - mock_client.models.embed_content.side_effect = APIError( - 401, {}, response=mock_response - ) + mock_client.models.embed_content.side_effect = APIError(401, {}, response=mock_response) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") with pytest.raises(RuntimeError, match="Gemini embedding failed"): embedder.embed("hello") @@ -157,9 +148,7 @@ def test_embed_empty_string_returns_zero_vector(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=3 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) for text in ["", " ", "\t\n"]: result = embedder.embed(text) assert result.dense_vector == [0.0, 0.0, 0.0] @@ -185,9 +174,7 @@ def test_embed_batch_skips_empty_strings(self, mock_client_class): mock_client.models.embed_content.return_value = _make_mock_result( [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] ) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=3 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=3) results = embedder.embed_batch(["hello", "", "world", " "]) assert len(results) == 4 # Empty positions get zero vectors @@ -205,12 +192,8 @@ def test_embed_batch_single_chunk(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value - mock_client.models.embed_content.return_value = _make_mock_result( - [[0.1], [0.2], [0.3]] - ) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + mock_client.models.embed_content.return_value = _make_mock_result([[0.1], [0.2], [0.3]]) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = embedder.embed_batch(["a", "b", "c"]) assert len(results) == 3 mock_client.models.embed_content.assert_called_once() @@ -226,9 +209,7 @@ def test_embed_batch_chunks_at_100(self, mock_client_class): _make_mock_result([[0.1]] * 100), _make_mock_result([[0.2]] * 10), ] - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = embedder.embed_batch([f"text{i}" for i in range(110)]) assert len(results) == 110 assert mock_client.models.embed_content.call_count == 2 @@ -236,6 +217,7 @@ def test_embed_batch_chunks_at_100(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_batch_falls_back_to_individual_on_error(self, mock_client_class): from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value @@ -246,9 +228,7 @@ def test_embed_batch_falls_back_to_individual_on_error(self, mock_client_class): _make_mock_result([[0.1]]), _make_mock_result([[0.2]]), ] - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = embedder.embed_batch(["a", "b"]) assert len(results) == 2 assert mock_client.models.embed_content.call_count == 3 @@ -269,9 +249,7 @@ async def test_async_embed_batch_dispatches_all_chunks(self, mock_client_class): _make_mock_result([[0.2]] * 10), ] ) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = await embedder.async_embed_batch([f"t{i}" for i in range(110)]) assert len(results) == 110 assert mock_client.aio.models.embed_content.call_count == 2 @@ -305,10 +283,9 @@ async def test_async_embed_batch_preserves_order(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") @pytest.mark.anyio - async def test_async_embed_batch_error_fallback_to_individual( - self, mock_client_class - ): + async def test_async_embed_batch_error_fallback_to_individual(self, mock_client_class): from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock_client = mock_client_class.return_value @@ -318,9 +295,7 @@ async def test_async_embed_batch_error_fallback_to_individual( side_effect=APIError(500, {}, response=mock_response) ) mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = await embedder.async_embed_batch(["a", "b"]) assert len(results) == 2 assert mock_client.models.embed_content.call_count == 2 @@ -348,23 +323,19 @@ class TestGeminiValidation: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_all_valid_task_types_accepted(self, mock_client): from openviking.models.embedder.gemini_embedders import ( - GeminiDenseEmbedder, _VALID_TASK_TYPES, + GeminiDenseEmbedder, ) for tt in _VALID_TASK_TYPES: - e = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="k", task_type=tt - ) + e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type=tt) assert e.task_type == tt def test_invalid_task_type_raises_on_init(self): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder with pytest.raises(ValueError, match="Invalid task_type"): - GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID" - ) + GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", task_type="NOT_VALID") def test_valid_task_types_count(self): from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES @@ -381,9 +352,7 @@ def test_dimension_too_high_raises(self, mock_client): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder with pytest.raises(ValueError, match="3072"): - GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="k", dimension=4096 - ) + GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k", dimension=4096) class TestGeminiErrorMessages: @@ -400,6 +369,7 @@ class TestGeminiErrorMessages: @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_error_code_hint(self, mock_client, code, match): from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock = mock_client.return_value @@ -413,6 +383,7 @@ def test_error_code_hint(self, mock_client, code, match): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_error_message_includes_http_code(self, mock_client): from google.genai.errors import APIError + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock = mock_client.return_value @@ -429,9 +400,7 @@ class TestBuildConfig: def test_build_config_defaults(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=512 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=512) cfg = embedder._build_config() assert cfg.output_dimensionality == 512 assert cfg.task_type is None @@ -454,9 +423,7 @@ def test_build_config_with_task_type_override(self, mock_client_class): def test_build_config_with_title(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) cfg = embedder._build_config(title="My Document") assert cfg.title == "My Document" @@ -466,9 +433,7 @@ def test_embed_per_call_task_type(self, mock_client_class): mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) embedder.embed("text", task_type="CLUSTERING") _, kwargs = mock_client.models.embed_content.call_args assert kwargs["config"].task_type == "CLUSTERING" @@ -479,9 +444,7 @@ def test_embed_per_call_title(self, mock_client_class): mock_client = mock_client_class.return_value mock_client.models.embed_content.return_value = _make_mock_result([[0.1]]) - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) embedder.embed("text", title="Doc Title") _, kwargs = mock_client.models.embed_content.call_args assert kwargs["config"].title == "Doc Title" @@ -495,9 +458,7 @@ def test_embed_batch_with_titles_falls_back(self, mock_client_class): _make_mock_result([[0.1]]), _make_mock_result([[0.2]]), ] - embedder = GeminiDenseEmbedder( - "gemini-embedding-2-preview", api_key="key", dimension=1 - ) + embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key", dimension=1) results = embedder.embed_batch(["alpha", "beta"], titles=["Title A", "Title B"]) assert len(results) == 2 # Called once per item (not as a batch) @@ -525,8 +486,8 @@ def test_repr(self, mock_client_class): @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_client_constructed_with_retry_options(self, mock_client_class): from openviking.models.embedder.gemini_embedders import ( - GeminiDenseEmbedder, _HTTP_RETRY_AVAILABLE, + GeminiDenseEmbedder, ) GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key")