diff --git a/docs/en/guides/01-configuration.md b/docs/en/guides/01-configuration.md index dacb665c..f2791981 100644 --- a/docs/en/guides/01-configuration.md +++ b/docs/en/guides/01-configuration.md @@ -115,7 +115,7 @@ Embedding model configuration for vector search, supporting dense, sparse, and h | Parameter | Type | Description | |-----------|------|-------------| | `max_concurrent` | int | Maximum concurrent embedding requests (`embedding.max_concurrent`, default: `10`) | -| `provider` | str | `"volcengine"`, `"openai"`, `"vikingdb"`, `"jina"`, or `"voyage"` | +| `provider` | str | `"volcengine"`, `"openai"`, `"vikingdb"`, `"jina"`, `"voyage"`, or `"google"` | | `api_key` | str | API key | | `model` | str | Model name | | `dimension` | int | Vector dimension. For Voyage, this maps to `output_dimension` | @@ -124,10 +124,11 @@ Embedding model configuration for vector search, supporting dense, sparse, and h **Available Models** -| Model | Dimension | Input Type | Notes | -|-------|-----------|------------|-------| -| `doubao-embedding-vision-250615` | 1024 | multimodal | Recommended | -| `doubao-embedding-250615` | 1024 | text | Text only | +| Provider | Model | Dimension | Input Type | Notes | +|----------|-------|-----------|------------|-------| +| `volcengine` | `doubao-embedding-vision-250615` | 1024 | multimodal | Recommended | +| `volcengine` | `doubao-embedding-250615` | 1024 | text | Text only | +| `google` | `gemini-embedding-2-preview` | 3072 | text | Google Gemini Embedding 2 with MRL | With `input: "multimodal"`, OpenViking can embed text, images (PNG, JPG, etc.), and mixed content. @@ -137,6 +138,7 @@ With `input: "multimodal"`, OpenViking can embed text, images (PNG, JPG, etc.), - `vikingdb`: VikingDB Embedding API - `jina`: Jina AI Embedding API - `voyage`: Voyage AI Embedding API +- `google`: Google/Gemini AI Embedding API **vikingdb provider example:** @@ -192,6 +194,29 @@ Get your API key at https://jina.ai } ``` +**google provider example:** + +```json +{ + "embedding": { + "dense": { + "provider": "google", + "api_key": "your-google-api-key", + "model": "gemini-embedding-2-preview", + "dimension": 1024, + "query_param": "RETRIEVAL_QUERY", + "document_param": "RETRIEVAL_DOCUMENT" + } + } +} +``` + +For Google/Gemini embeddings: +- `query_param` and `document_param` support task-specific embeddings +- Valid task types: `RETRIEVAL_QUERY`, `RETRIEVAL_DOCUMENT`, `SEMANTIC_SIMILARITY`, `CLASSIFICATION`, `CLUSTERING` +- Enhanced format: `"task_type=RETRIEVAL_QUERY,output_dimensionality=1024"` +- Get your API key at https://aistudio.google.com/app/apikey + Supported Voyage text embedding models include: - `voyage-4-lite` - `voyage-4` diff --git a/docs/zh/guides/01-configuration.md b/docs/zh/guides/01-configuration.md index bf6a8ea2..61dd7261 100644 --- a/docs/zh/guides/01-configuration.md +++ b/docs/zh/guides/01-configuration.md @@ -121,7 +121,7 @@ OpenViking 使用 JSON 配置文件（`ov.conf`）进行设置。配置文件支 | 参数 | 类型 | 说明 | |------|------|------| | `max_concurrent` | int | 最大并发 Embedding 请求数（`embedding.max_concurrent`，默认：`10`） | -| `provider` | str | `"volcengine"`、`"openai"`、`"vikingdb"` 或 `"jina"` | +| `provider` | str | `"volcengine"`、`"openai"`、`"vikingdb"`、`"jina"`、`"voyage"` 或 `"google"` | | `api_key` | str | API Key | | `model` | str | 模型名称 | | `dimension` | int | 向量维度 | diff --git a/openviking/models/embedder/__init__.py b/openviking/models/embedder/__init__.py index b418b809..d17e4a31 100644 --- a/openviking/models/embedder/__init__.py +++ b/openviking/models/embedder/__init__.py @@ -13,6 +13,7 @@ - Volcengine: Dense, Sparse, Hybrid - Jina AI: Dense only - Voyage AI: Dense only +- Google/Gemini: Dense only """ from openviking.models.embedder.base import ( @@ -23,9 +24,9 @@ HybridEmbedderBase, SparseEmbedderBase, ) +from openviking.models.embedder.google_embedders import GoogleDenseEmbedder 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 + "GoogleDenseEmbedder", # Jina AI implementations "JinaDenseEmbedder", # OpenAI implementations diff --git a/openviking/models/embedder/google_embedders.py b/openviking/models/embedder/google_embedders.py new file mode 100644 index 00000000..04789913 --- /dev/null +++ b/openviking/models/embedder/google_embedders.py @@ -0,0 +1,313 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Google/Gemini AI Embedder Implementation""" + +import logging +import re +from typing import Any, Dict, List, Optional + +import requests + +from openviking.models.embedder.base import ( + DenseEmbedderBase, + EmbedResult, + exponential_backoff_retry, +) + +logger = logging.getLogger(__name__) + +# Default dimensions for Google/Gemini embedding models +GOOGLE_MODEL_DIMENSIONS = { + "gemini-embedding-2-preview": 3072, # Gemini Embedding 2 with MRL support +} + + +class GoogleDenseEmbedder(DenseEmbedderBase): + """Google Gemini Embedding 2 Dense Embedder Implementation + + Uses native Google Gemini embedding API with Parts format. + Supports Gemini Embedding 2 (gemini-embedding-2-preview) only. + Supports Matryoshka dimension reduction via output_dimensionality. + + ## Note: taskType not supported by gemini-embedding-2-preview + + Tested 2026-03-19 against the live API at full 3072 dimensions: + the taskType parameter is accepted without error but produces bit-for-bit + identical vectors regardless of which task type is specified. All eight + documented task types (RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, + SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING, CODE_RETRIEVAL_QUERY, + QUESTION_ANSWERING, FACT_VERIFICATION) return the same embedding as the + default (no taskType). The parameter is therefore not sent. + + By contrast, gemini-embedding-001 does produce distinct vectors per task + type. This is because taskType in Gemini embedding models is implemented as + an instruction prefix injected into the embedding input — effectively + "task: {task_type}, content: {text}" — rather than a separate model head or + fine-tuned adapter. gemini-embedding-2-preview appears to have dropped this + instruction-following behaviour. + + Example: + >>> embedder = GoogleDenseEmbedder( + ... api_key="your-gemini-api-key", + ... dimension=1024, + ... ) + >>> result = embedder.embed("Hello world") + >>> print(len(result.dense_vector)) + 1024 + """ + + def __init__( + self, + model_name: str = "gemini-embedding-2-preview", + api_key: Optional[str] = None, + api_base: Optional[str] = None, + dimension: Optional[int] = None, + config: Optional[Dict[str, Any]] = None, + max_tokens: Optional[int] = None, + extra_headers: Optional[Dict[str, str]] = None, + ): + """Initialize Google Gemini Embedding 2 Dense Embedder + + Args: + model_name: Must be "gemini-embedding-2-preview" (default and only supported model) + api_key: Google API key, required + api_base: API base URL, defaults to https://generativelanguage.googleapis.com/v1beta + dimension: Dimension for Matryoshka reduction, optional (max 3072) + config: Additional configuration dict + max_tokens: Maximum token count per embedding request, None to use default (8192) + extra_headers: Extra HTTP headers to include in API requests + + Raises: + ValueError: If api_key is not provided or unsupported model is specified + """ + super().__init__(model_name, config) + self.api_key = api_key + self.api_base = api_base or "https://generativelanguage.googleapis.com/v1beta" + self.dimension = dimension + self._max_tokens = max_tokens or 8192 + self.extra_headers = extra_headers or {} + + if not self.api_key: + raise ValueError("api_key is required") + + if model_name not in GOOGLE_MODEL_DIMENSIONS: + raise ValueError( + f"Unsupported model '{model_name}'. Only 'gemini-embedding-2-preview' is supported." + ) + + max_dim = GOOGLE_MODEL_DIMENSIONS[model_name] + if dimension is not None and dimension > max_dim: + raise ValueError( + f"Requested dimension {dimension} exceeds maximum {max_dim} for model '{model_name}'. " + f"Gemini Embedding 2 supports Matryoshka dimension reduction up to {max_dim}." + ) + self._dimension = dimension if dimension is not None else max_dim + + @property + def max_tokens(self) -> int: + """Maximum token count per embedding request.""" + return self._max_tokens + + def _estimate_tokens(self, text: str) -> int: + """Estimate token count. Falls back to character-based heuristic if tiktoken unavailable.""" + try: + import tiktoken + + enc = tiktoken.encoding_for_model(self.model_name) + return len(enc.encode(text)) + except Exception: + return max(len(text) // 3, len(text.encode("utf-8")) // 4) + + def _chunk_text(self, text: str) -> List[str]: + """Split text into chunks each within max_tokens. + + Splitting priority: paragraphs (\\n\\n) > sentences (。.!?\\n) > fixed length. + """ + max_tok = self.max_tokens + if self._estimate_tokens(text) <= max_tok: + return [text] + + paragraphs = text.split("\n\n") + if len(paragraphs) > 1: + chunks = self._merge_segments(paragraphs, max_tok, "\n\n") + if all(self._estimate_tokens(c) <= max_tok for c in chunks): + return chunks + + sentences = re.split(r"(?<=[。.!?\n])", text) + sentences = [s for s in sentences if s] + if len(sentences) > 1: + chunks = self._merge_segments(sentences, max_tok, "") + if all(self._estimate_tokens(c) <= max_tok for c in chunks): + return chunks + + return self._fixed_length_split(text, max_tok) + + def _merge_segments(self, segments: List[str], max_tok: int, separator: str) -> List[str]: + chunks: List[str] = [] + current = "" + for seg in segments: + candidate = (current + separator + seg) if current else seg + if self._estimate_tokens(candidate) <= max_tok: + current = candidate + else: + if current: + chunks.append(current) + current = seg + if current: + chunks.append(current) + return chunks + + def _fixed_length_split(self, text: str, max_tok: int) -> List[str]: + total_tokens = self._estimate_tokens(text) + chars_per_token = len(text) / max(total_tokens, 1) + chunk_size = max(int(max_tok * chars_per_token * 0.9), 100) + + chunks: List[str] = [] + start = 0 + while start < len(text): + end = start + chunk_size + if end < len(text): + boundary = text.rfind(" ", start, end) + if boundary > start: + end = boundary + chunks.append(text[start:end]) + start = end + return chunks + + def _update_telemetry_token_usage(self, response_data: Dict[str, Any]) -> None: + """Update telemetry with token usage from API response""" + pass + + def _embed_single(self, text: str) -> EmbedResult: + """Perform raw embedding without chunking logic. + + Args: + text: Input text + + Returns: + EmbedResult: Result containing only dense_vector + + Raises: + RuntimeError: When API call fails + """ + try: + url = f"{self.api_base}/models/{self.model_name}:embedContent" + + headers = { + "Content-Type": "application/json", + "x-goog-api-key": self.api_key, + **self.extra_headers, + } + + request_body: Dict[str, Any] = {"content": {"parts": [{"text": text}]}} + if self.dimension: + request_body["output_dimensionality"] = self.dimension + + def _do_request(): + resp = requests.post(url, json=request_body, headers=headers, timeout=30) + resp.raise_for_status() + return resp + + response = exponential_backoff_retry( + _do_request, + is_retryable=lambda e: isinstance(e, requests.exceptions.ConnectionError) + or isinstance(e, requests.exceptions.Timeout), + logger=logger, + ) + + response_data = response.json() + + if "embedding" in response_data and "values" in response_data["embedding"]: + vector = response_data["embedding"]["values"] + else: + raise RuntimeError(f"Unexpected response format: {response_data}") + + self._update_telemetry_token_usage(response_data) + + return EmbedResult(dense_vector=vector) + + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Google/Gemini API request error: {str(e)}") from e + except Exception as e: + raise RuntimeError(f"Embedding failed: {str(e)}") from e + + def embed(self, text: str, is_query: bool = False) -> EmbedResult: + """Embed single text, with automatic chunking for oversized input. + + Args: + text: Input text + is_query: Ignored. gemini-embedding-2-preview does not support taskType. + + Returns: + EmbedResult: Result containing only dense_vector + + Raises: + RuntimeError: When API call fails + """ + if not text or not text.strip(): + return EmbedResult() + + if self._estimate_tokens(text) > self.max_tokens: + return self._chunk_and_embed(text) + + return self._embed_single(text) + + def _chunk_and_embed(self, text: str) -> EmbedResult: + """Chunk oversized text and average the embeddings.""" + chunks = self._chunk_text(text) + chunk_vectors: List[List[float]] = [] + + for chunk in chunks: + result = self._embed_single(chunk) + chunk_vectors.append(result.dense_vector) + + if not chunk_vectors: + return EmbedResult(dense_vector=[0.0] * self._dimension) + + avg_vector = [ + sum(v[i] for v in chunk_vectors) / len(chunk_vectors) + for i in range(len(chunk_vectors[0])) + ] + + return EmbedResult(dense_vector=avg_vector) + + def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedResult]: + """Batch embedding with automatic chunking for oversized inputs. + + Individual texts are processed sequentially since Google's native API + does not support batch requests. + + Args: + texts: List of texts + is_query: Ignored. gemini-embedding-2-preview does not support taskType. + + Returns: + List[EmbedResult]: List of embedding results + + Raises: + RuntimeError: When API call fails + """ + if not texts: + return [] + + results: List[EmbedResult] = [] + + for text in texts: + if not text or not text.strip(): + results.append(EmbedResult()) + continue + if self._estimate_tokens(text) <= self.max_tokens: + result = self._embed_single(text) + else: + result = self._chunk_and_embed(text) + results.append(result) + + return results + + def get_dimension(self) -> int: + """Get embedding dimension + + Returns: + int: Vector dimension + """ + return self._dimension diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index a0a1b964..72d61297 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -37,7 +37,7 @@ 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', 'voyage', 'google'. " "For OpenRouter or other OpenAI-compatible providers, use 'openai' with " "api_base and extra_headers." ), @@ -51,6 +51,10 @@ class EmbeddingModelConfig(BaseModel): 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., 8192 for Google).", + ) extra_headers: Optional[dict[str, str]] = Field( default=None, description=( @@ -89,10 +93,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", + "google", + ]: raise ValueError( f"Invalid embedding provider: '{self.provider}'. Must be one of: " - "'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage'" + "'openai', 'volcengine', 'vikingdb', 'jina', 'ollama', 'voyage', 'google'" ) # Provider-specific validation @@ -131,6 +143,10 @@ def validate_config(self): if not self.api_key: raise ValueError("Voyage provider requires 'api_key' to be set") + elif self.provider == "google": + if not self.api_key: + raise ValueError("Google provider requires 'api_key' to be set") + return self def get_effective_dimension(self) -> int: @@ -146,6 +162,11 @@ def get_effective_dimension(self) -> int: return get_voyage_model_default_dimension(self.model) + if provider == "google": + from openviking.models.embedder.google_embedders import GOOGLE_MODEL_DIMENSIONS + + return GOOGLE_MODEL_DIMENSIONS.get(self.model, 3072) + return 2048 @@ -189,7 +210,7 @@ def _create_embedder( """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', 'google') embedder_type: Embedder type ('dense', 'sparse', 'hybrid') config: EmbeddingModelConfig instance @@ -200,6 +221,7 @@ def _create_embedder( ValueError: If provider/type combination is not supported """ from openviking.models.embedder import ( + GoogleDenseEmbedder, JinaDenseEmbedder, OpenAIDenseEmbedder, VikingDBDenseEmbedder, @@ -322,6 +344,17 @@ def _create_embedder( "dimension": cfg.dimension, }, ), + ("google", "dense"): ( + GoogleDenseEmbedder, + lambda cfg: { + "model_name": cfg.model, + "api_key": cfg.api_key, + "api_base": cfg.api_base, + "dimension": cfg.dimension, + "max_tokens": cfg.max_tokens, + **({"extra_headers": cfg.extra_headers} if cfg.extra_headers else {}), + }, + ), } key = (provider, embedder_type) diff --git a/pyproject.toml b/pyproject.toml index c26aa7cf..60e55714 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -246,4 +246,5 @@ line-ending = "auto" [dependency-groups] dev = [ "pytest>=9.0.2", + "pytest-asyncio>=1.3.0", ] diff --git a/tests/unit/test_google_embedder.py b/tests/unit/test_google_embedder.py new file mode 100644 index 00000000..50bd2229 --- /dev/null +++ b/tests/unit/test_google_embedder.py @@ -0,0 +1,143 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for Google/Gemini Embedder""" + +from unittest.mock import MagicMock, patch + +import pytest +import requests + +from openviking.models.embedder import GoogleDenseEmbedder + + +def _make_response(values: list) -> MagicMock: + mock_resp = MagicMock() + mock_resp.json.return_value = {"embedding": {"values": values}} + return mock_resp + + +def _embedder(**kwargs) -> GoogleDenseEmbedder: + return GoogleDenseEmbedder( + model_name="gemini-embedding-2-preview", api_key="test-key", **kwargs + ) + + +class TestGoogleDenseEmbedderInit: + def test_requires_api_key(self): + with pytest.raises(ValueError, match="api_key is required"): + GoogleDenseEmbedder(model_name="gemini-embedding-2-preview") + + def test_rejects_unsupported_model(self): + with pytest.raises(ValueError, match="Unsupported model"): + GoogleDenseEmbedder(model_name="unknown-model", api_key="key") + + def test_rejects_dimension_exceeding_max(self): + with pytest.raises(ValueError, match="exceeds maximum"): + _embedder(dimension=9999) + + def test_defaults(self): + e = _embedder() + assert e.get_dimension() == 3072 + assert e.max_tokens == 8192 + assert e.api_base == "https://generativelanguage.googleapis.com/v1beta" + + def test_custom_values(self): + e = _embedder(dimension=1024, max_tokens=4096, api_base="https://custom/v1") + assert e.get_dimension() == 1024 + assert e.max_tokens == 4096 + assert e.api_base == "https://custom/v1" + + +class TestGoogleDenseEmbedderEmbed: + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_embed_request_structure(self, mock_post): + """Single embed call — verify URL, auth header, body, and return value.""" + mock_post.return_value = _make_response([0.1] * 3072) + + result = _embedder().embed("Hello world") + + assert result.dense_vector is not None + assert len(result.dense_vector) == 3072 + mock_post.assert_called_once() + url = mock_post.call_args[0][0] + assert "gemini-embedding-2-preview:embedContent" in url + headers = mock_post.call_args[1]["headers"] + assert headers["x-goog-api-key"] == "test-key" + body = mock_post.call_args[1]["json"] + assert body["content"]["parts"][0]["text"] == "Hello world" + assert "taskType" not in body + assert "task_type" not in body + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_dimension_sent_as_output_dimensionality(self, mock_post): + mock_post.return_value = _make_response([0.1] * 1024) + _embedder(dimension=1024).embed("Hello world") + body = mock_post.call_args[1]["json"] + assert body["output_dimensionality"] == 1024 + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_no_dimension_omits_output_dimensionality(self, mock_post): + mock_post.return_value = _make_response([0.1] * 3072) + _embedder().embed("Hello world") + assert "output_dimensionality" not in mock_post.call_args[1]["json"] + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_extra_headers_forwarded(self, mock_post): + mock_post.return_value = _make_response([0.1] * 3072) + _embedder(extra_headers={"X-Custom": "value"}).embed("Hello world") + assert mock_post.call_args[1]["headers"]["X-Custom"] == "value" + + @pytest.mark.parametrize("text", ["", " "]) + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_blank_text_returns_empty_without_request(self, mock_post, text): + result = _embedder().embed(text) + assert result.dense_vector is None + mock_post.assert_not_called() + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_api_error_raises_runtime_error(self, mock_post): + mock_resp = MagicMock() + mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError() + mock_post.return_value = mock_resp + with pytest.raises(RuntimeError): + _embedder().embed("Hello world") + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_unexpected_response_raises(self, mock_post): + mock_resp = MagicMock() + mock_resp.json.return_value = {"unexpected": "format"} + mock_post.return_value = mock_resp + with pytest.raises(RuntimeError, match="Unexpected response format"): + _embedder().embed("Hello world") + + +class TestGoogleDenseEmbedderBatch: + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_batch_results_and_empty_skipped(self, mock_post): + mock_post.return_value = _make_response([0.1] * 3072) + results = _embedder().embed_batch(["Hello", "", "World"]) + assert len(results) == 3 + assert results[0].dense_vector is not None + assert results[1].dense_vector is None + assert results[2].dense_vector is not None + assert mock_post.call_count == 2 + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_batch_empty_list(self, mock_post): + assert _embedder().embed_batch([]) == [] + mock_post.assert_not_called() + + +class TestGoogleDenseEmbedderChunking: + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_oversized_text_chunked_and_averaged(self, mock_post): + mock_post.return_value = _make_response([0.5] * 3072) + result = _embedder(max_tokens=5).embed("word " * 100) + assert result.dense_vector is not None + assert mock_post.call_count > 1 + + @patch("openviking.models.embedder.google_embedders.requests.post") + def test_normal_text_single_request(self, mock_post): + mock_post.return_value = _make_response([0.1] * 3072) + _embedder().embed("Hello world") + assert mock_post.call_count == 1 diff --git a/uv.lock b/uv.lock index 59e6638e..2e02b1eb 100644 --- a/uv.lock +++ b/uv.lock @@ -3406,6 +3406,7 @@ test = [ [package.dev-dependencies] dev = [ { name = "pytest" }, + { name = "pytest-asyncio" }, ] [package.metadata] @@ -3505,7 +3506,10 @@ requires-dist = [ 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"] [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" }, +] [[package]] name = "orjson"