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..d3e8c617 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 (text-only; multimodal planned) - Voyage AI: Dense only """ @@ -23,9 +24,9 @@ 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 from openviking.models.embedder.vikingdb_embedders import ( VikingDBDenseEmbedder, VikingDBHybridEmbedder, @@ -36,6 +37,7 @@ VolcengineHybridEmbedder, VolcengineSparseEmbedder, ) +from openviking.models.embedder.voyage_embedders import VoyageDenseEmbedder __all__ = [ # Base classes @@ -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..8e44ea9b --- /dev/null +++ b/openviking/models/embedder/gemini_embedders.py @@ -0,0 +1,317 @@ +# 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.""" + +from typing import Any, Dict, List, Optional + +from google import genai +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: + 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 _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 + supports_multimodal: bool = False # text-only; multimodal planned separately + + 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}") + 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 + + 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(): + 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: + result = self.client.models.embed_content( + model=self.model_name, + contents=text, + 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, + *, + 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()] + 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=non_empty_texts, + config=config, + ) + 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) + ) + 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", + e.code, + len(batch), + ) + for text in batch: + results.append(self.embed(text, is_query=is_query)) + 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._build_config() + ) + results[idx] = [ + EmbedResult( + dense_vector=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..08db7f41 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") + 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( + 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,10 +91,14 @@ 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() + 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") @@ -89,10 +113,18 @@ 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 +159,30 @@ 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") @@ -146,12 +202,17 @@ def get_effective_dimension(self) -> int: 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 +246,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 +275,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 +287,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 +366,26 @@ 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 +395,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 +454,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..05ce703f 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,7 @@ line-ending = "auto" [dependency-groups] dev = [ "pytest>=9.0.2", + "pytest-asyncio>=1.3.0", + "pytest-cov>=7.0.0", + "pytest-xdist>=3.0", ] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index fceae36b..628c305a 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,118 @@ 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 (ImportError, ModuleNotFoundError, AttributeError): + 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..c3e3507d --- /dev/null +++ b/tests/integration/test_gemini_e2e.py @@ -0,0 +1,112 @@ +# 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 pytest + +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 _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 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..b828f2af --- /dev/null +++ b/tests/integration/test_gemini_embedding_it.py @@ -0,0 +1,123 @@ +# 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 + + _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") + + +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..15d19d0e --- /dev/null +++ b/tests/integration/test_gemini_openviking_it.py @@ -0,0 +1,209 @@ +# 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) + 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}", + 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 return the indexed pytest doc" diff --git a/tests/unit/test_embedding_config_gemini.py b/tests/unit/test_embedding_config_gemini.py new file mode 100644 index 00000000..f8761ae2 --- /dev/null +++ b/tests/unit/test_embedding_config_gemini.py @@ -0,0 +1,120 @@ +# 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 + + 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 new file mode 100644 index 00000000..a6dbc06d --- /dev/null +++ b/tests/unit/test_gemini_embedder.py @@ -0,0 +1,498 @@ +# 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() + + @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 ( + _MODEL_TOKEN_LIMITS, + GeminiDenseEmbedder, + ) + + 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") + + @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") + 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_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 + + 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 ( + _VALID_TASK_TYPES, + GeminiDenseEmbedder, + ) + + 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") + + +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 ( + _HTTP_RETRY_AVAILABLE, + GeminiDenseEmbedder, + ) + + 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 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"