diff --git a/openviking/models/embedder/volcengine_embedders.py b/openviking/models/embedder/volcengine_embedders.py index 86b9c0fc09..30c8d0dc14 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -172,6 +172,7 @@ def __init__( input_type: str = "multimodal", extra_headers: Optional[Dict[str, str]] = None, config: Optional[Dict[str, Any]] = None, + query_param: Optional[str] = None, ): """Initialize Volcengine Dense Embedder @@ -182,12 +183,15 @@ def __init__( dimension: Target dimension for truncation (optional) input_type: Input type - "text" or "multimodal" (default: "multimodal") config: Additional configuration dict + query_param: Instruction prefix for text queries (``is_query=True``). + Documents are embedded unchanged, so setting it needs no reindex. Raises: ValueError: If api_key is not provided """ super().__init__(model_name, config) self.provider = "volcengine" + self.query_param = query_param self.api_key = api_key self.api_base = api_base or "https://ark.cn-beijing.volces.com/api/v3" @@ -254,6 +258,12 @@ def _usage_value(key: str, default: int = 0) -> int: duration_seconds=duration_seconds, ) + def _query_input(self, content: "EmbeddingInput", is_query: bool) -> "EmbeddingInput": + # Text queries only; image/multimodal queries keep their plain input. + if is_query and self.query_param and isinstance(content, str): + return self.query_param + content + return content + def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: """Perform dense embedding on text or multimodal content @@ -268,6 +278,7 @@ def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResul Raises: RuntimeError: When API call fails """ + content = self._query_input(content, is_query) def _embed_call(): if self.input_type == "multimodal": @@ -319,6 +330,7 @@ def _get_async_client(self): async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: client = self._get_async_client() + content = self._query_input(content, is_query) async def _embed_call() -> EmbedResult: if self.input_type == "multimodal": diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index 5f773c86d5..55e6279423 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -52,6 +52,7 @@ class EmbeddingModelConfig(BaseModel): "Parameter value for query-side embeddings when calling embed(is_query=True). " "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'). " + "For Volcengine dense models, it is an instruction prefixed to text queries. " "Setting this or document_param activates non-symmetric mode. " "Leave both unset for symmetric models." ), @@ -158,10 +159,13 @@ 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"): - value = data.get(key) - if isinstance(value, str): - data[key] = value.lower() + # Volcengine takes query_param as instruction text, where case matters; + # the other providers take parameter values such as "query". + if str(data.get("provider") or "volcengine").lower() != "volcengine": + for key in ("query_param", "document_param"): + value = data.get(key) + if isinstance(value, str): + data[key] = value.lower() return data _VALID_PROVIDERS: ClassVar[Tuple[str, ...]] = ( @@ -824,6 +828,7 @@ def _create_embedder( "dimension": cfg.dimension, "input_type": cfg.input, "config": dict(runtime_config), + **({"query_param": cfg.query_param} if cfg.query_param else {}), **({"extra_headers": cfg.extra_headers} if cfg.extra_headers else {}), }, ), diff --git a/tests/unit/test_volcengine_embedder_input.py b/tests/unit/test_volcengine_embedder_input.py index 9cd5aad6b1..d9ae26f3e8 100644 --- a/tests/unit/test_volcengine_embedder_input.py +++ b/tests/unit/test_volcengine_embedder_input.py @@ -55,3 +55,47 @@ def test_dense_embedder_passes_local_provider_duration_to_metrics(mock_perf_coun embedder.embed("hello") assert embedder.update_token_usage.call_args.kwargs["duration_seconds"] == 0.25 + + +_INSTRUCTION = "Target_modality: text.\nInstruction: Retrieve the memory notes relevant to the question\nQuery:" + + +def test_dense_embedder_prefixes_text_queries_only(): + embedder = VolcengineDenseEmbedder( + model_name="test", api_key="test-key", dimension=2, query_param=_INSTRUCTION + ) + sent = [] + + def create(input, model, extra_headers): + sent.append(input) + return SimpleNamespace(data=SimpleNamespace(embedding=[1.0, 0.0]), usage=None) + + embedder.client = SimpleNamespace(multimodal_embeddings=SimpleNamespace(create=create)) + image = [{"type": "image_url", "image_url": {"url": "https://example.test/a.png"}}] + + embedder.embed("Which plan did we pick?", is_query=True) + embedder.embed("document text") + embedder.embed(image, is_query=True) + + assert sent == [ + [{"type": "text", "text": _INSTRUCTION + "Which plan did we pick?"}], + [{"type": "text", "text": "document text"}], + image, + ] + + +@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") +def test_embedding_config_passes_volcengine_query_param_unchanged(_mock_ark): + from openviking_cli.utils.config.embedding_config import EmbeddingConfig + + config = EmbeddingConfig( + dense={ + "provider": "volcengine", + "model": "test", + "api_key": "test-key", + "dimension": 2, + "query_param": _INSTRUCTION, + } + ) + + assert config.get_embedder().query_param == _INSTRUCTION