diff --git a/benchmark/aml/README.md b/benchmark/aml/README.md index c3667e79f1..9c5be5340f 100644 --- a/benchmark/aml/README.md +++ b/benchmark/aml/README.md @@ -62,6 +62,58 @@ default; use `AML_RETRY_ATTEMPTS` and `AML_RETRY_DELAY_SECONDS` to adjust that behavior. `Token`, `Bearer`, and `X-Api-Key` authentication are accepted when `AML_API_KEY` is set. +`/add` accepts either a content string (Textual/Coding) or an ordered array of +`text` and `image_url` parts in `messages[].content` (Multimodal). Strings are +forwarded as native `content`; arrays are forwarded as native Session `parts`, +preserving part order, message roles, and timestamps. Images must be inline +Base64 JPEG, PNG, or WebP Data URIs, with at most 10 MiB decoded per image and +30 MiB decoded across one Add request. For example: + +```json +{ + "request_id": "example:add:1", + "user_id": "example:user", + "session_id": "example:session", + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "Remember this receipt."}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,"}} + ] + }] +} +``` + +The example's image payload is a placeholder; supply actual Base64 image bytes. +Image understanding and memory extraction use OpenViking's existing Session +commit pipeline and require a configured vision-capable VLM. The adapter does +not perform separate image captioning. + +`/search` accepts a query string or an array of the same `text` and `image_url` +parts, including image-only and multiple-image queries. Text parts are joined +with newlines and image URLs retain their relative order. The adapter forwards +them to native `find` using `query` and `image_url`, without an extra VLM call. +String queries keep their original request and response behavior. For example: + +```json +{ + "user_id": "example:user", + "query": [ + {"type": "text", "text": "Which workshop was this receipt for?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,"}} + ], + "top_k": 5 +} +``` + +The same 10 MiB per-image limit applies, with a 30 MiB decoded total per Search +request. Image queries require a multimodal dense embedding model compatible +with the stored vectors. In hybrid mode, OV embeds the complete image/text +input through dense and only the text through sparse; image-only queries skip +sparse. Volcengine text-only queries retain their combined dense/sparse request. +Keep the configured dense model and dimension consistent with the existing +index. Search responses remain text memory evidence rather than raw images. + The native `find` call uses quick retrieval without reranking, even when the OpenViking server has a reranker configured. The adapter returns each hit's `content` when present, otherwise its `abstract`, without an additional diff --git a/benchmark/aml/server.py b/benchmark/aml/server.py index 5372132610..e0e5e35ef1 100644 --- a/benchmark/aml/server.py +++ b/benchmark/aml/server.py @@ -5,6 +5,8 @@ import argparse import asyncio +import base64 +import binascii import hmac import math import os @@ -12,7 +14,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, AsyncIterator, Awaitable, Callable, Literal, TypeVar +from typing import Annotated, Any, AsyncIterator, Awaitable, Callable, Literal, TypeVar import httpx import uvicorn @@ -24,6 +26,9 @@ DEFAULT_PORT = 8088 DEFAULT_OPENVIKING_URL = "http://127.0.0.1:1933" MEMORY_ROOT_URI = "viking://~/memories" +MAX_IMAGE_BYTES = 10 * 1024 * 1024 +MAX_ADD_IMAGE_BYTES = 30 * 1024 * 1024 +MAX_QUERY_IMAGE_BYTES = 30 * 1024 * 1024 RETRYABLE_OPENVIKING_CODES = { "ABORTED", "DEADLINE_EXCEEDED", @@ -103,14 +108,72 @@ def _not_blank(value: str) -> str: return value +class AMLTextPart(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["text"] + text: str = Field(min_length=1) + + _text_not_blank = field_validator("text")(_not_blank) + + +class AMLImageURL(BaseModel): + model_config = ConfigDict(extra="forbid") + + url: str + detail: Literal["auto", "low", "high"] | None = None + + @field_validator("url") + @classmethod + def validate_inline_image(cls, value: str) -> str: + header, _, encoded = value.partition(",") + if ( + header + not in { + "data:image/jpeg;base64", + "data:image/png;base64", + "data:image/webp;base64", + } + or not encoded + ): + raise ValueError("image_url must be an inline Base64 JPEG, PNG, or WebP Data URI") + if len(encoded) > 4 * ((MAX_IMAGE_BYTES + 2) // 3): + raise ValueError("image exceeds the 10 MiB decoded size limit") + try: + size = len(base64.b64decode(encoded, validate=True)) + except (binascii.Error, ValueError) as exc: + raise ValueError("image_url contains invalid Base64 data") from exc + if size > MAX_IMAGE_BYTES: + raise ValueError("image exceeds the 10 MiB decoded size limit") + return value + + @property + def decoded_size(self) -> int: + # Base64 was validated above; count bytes without decoding images again. + return len(self.url.partition(",")[2].rstrip("=")) * 3 // 4 + + +class AMLImagePart(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["image_url"] + image_url: AMLImageURL + + +AMLContentPart = Annotated[AMLTextPart | AMLImagePart, Field(discriminator="type")] + + class AMLMessage(BaseModel): model_config = ConfigDict(extra="forbid") role: Literal["user", "assistant"] - content: str = Field(min_length=1) + content: str | list[AMLContentPart] = Field(min_length=1) timestamp: int | None = Field(default=None, ge=0, strict=True) - _content_not_blank = field_validator("content")(_not_blank) + @field_validator("content") + @classmethod + def content_not_blank(cls, value: str | list[AMLContentPart]) -> str | list[AMLContentPart]: + return _not_blank(value) if isinstance(value, str) else value @field_validator("timestamp") @classmethod @@ -137,18 +200,41 @@ class AMLAddRequest(BaseModel): _user_id_not_blank = field_validator("user_id")(_not_blank) _session_id_not_blank = field_validator("session_id")(_not_blank) + @field_validator("messages") + @classmethod + def validate_image_budget(cls, messages: list[AMLMessage]) -> list[AMLMessage]: + size = sum( + part.image_url.decoded_size + for message in messages + if isinstance(message.content, list) + for part in message.content + if isinstance(part, AMLImagePart) + ) + if size > MAX_ADD_IMAGE_BYTES: + raise ValueError("images exceed the 30 MiB decoded size limit per Add request") + return messages + class AMLSearchRequest(BaseModel): model_config = ConfigDict(extra="forbid") - query: str = Field(min_length=1) + query: str | list[AMLContentPart] = Field(min_length=1) options: list[str] | None = None user_id: str = Field(min_length=1) top_k: int = Field(ge=1, strict=True) - _query_not_blank = field_validator("query")(_not_blank) _user_id_not_blank = field_validator("user_id")(_not_blank) + @field_validator("query") + @classmethod + def validate_query(cls, value: str | list[AMLContentPart]) -> str | list[AMLContentPart]: + if isinstance(value, str): + return _not_blank(value) + size = sum(part.image_url.decoded_size for part in value if isinstance(part, AMLImagePart)) + if size > MAX_QUERY_IMAGE_BYTES: + raise ValueError("images exceed the 30 MiB decoded size limit per Search request") + return value + class OpenVikingBackendError(RuntimeError): pass @@ -246,7 +332,9 @@ async def add_and_commit( raise OpenVikingBackendError("OpenViking commit timed out") await asyncio.sleep(min(self.settings.task_poll_seconds, remaining)) - async def find(self, *, user_id: str, query: str, limit: int) -> list[dict[str, Any]]: + async def find( + self, *, user_id: str, query: str, limit: int, image: str | list[str] | None = None + ) -> list[dict[str, Any]]: async with self._client(user_id) as client: result = await self._retry( "find", @@ -254,6 +342,7 @@ async def find(self, *, user_id: str, query: str, limit: int) -> list[dict[str, query=query, target_uri=MEMORY_ROOT_URI, limit=limit, + **({"image": image} if image else {}), options={ "context_type": ["memory"], "level": 2, @@ -298,7 +387,11 @@ def authorize(self, request: Request) -> None: async def add(self, request: AMLAddRequest) -> dict[str, Any]: messages: list[dict[str, Any]] = [] for message in request.messages: - item: dict[str, Any] = {"role": message.role, "content": message.content} + item: dict[str, Any] = {"role": message.role} + if isinstance(message.content, str): + item["content"] = message.content + else: + item["parts"] = [part.model_dump(exclude_none=True) for part in message.content] if message.timestamp is not None: item["created_at"] = _format_timestamp(message.timestamp) messages.append(item) @@ -321,6 +414,10 @@ async def add(self, request: AMLAddRequest) -> dict[str, Any]: async def search(self, request: AMLSearchRequest) -> dict[str, list[dict[str, Any]]]: query = request.query + images = [] + if isinstance(query, list): + images = [part.image_url.url for part in query if isinstance(part, AMLImagePart)] + query = "\n".join(part.text for part in query if isinstance(part, AMLTextPart)) if request.options: options = "\n".join( f"{index}. {option}" for index, option in enumerate(request.options, start=1) @@ -331,6 +428,7 @@ async def search(self, request: AMLSearchRequest) -> dict[str, list[dict[str, An user_id=request.user_id, query=query, limit=request.top_k, + **({"image": images[0] if len(images) == 1 else images} if images else {}), ) except OpenVikingBackendError as exc: raise HTTPException( diff --git a/docs/en/api/06-retrieval.md b/docs/en/api/06-retrieval.md index 83cdb16304..759033e26f 100644 --- a/docs/en/api/06-retrieval.md +++ b/docs/en/api/06-retrieval.md @@ -54,7 +54,7 @@ The `find()` method runs one global vector similarity search in QUICK mode for s | Parameter | Type | Required | Default | Description | |-----------|------|----------|---------|-------------| | query | str | No | "" | Search query string. May be omitted for an image query or a filter-only lookup | -| image_url | str | No | None | Image query as a `data:image/...;base64,...`, `http(s)://`, or `viking://` URI. Requires a multimodal embedding model | +| image_url | str or list[str] | No | None | One image URI or an ordered list of image URIs, each as a `data:image/...;base64,...`, `http(s)://`, or `viking://` URI. Requires a multimodal embedding model | | target_uri | str \| List[str] | No | "" | Limit search to specific URI prefix | | events_time_decay_protection | str \| null | No | null | Omit or pass `null` to disable decay. Pass `"0"` to decay immediately, or a duration such as `"7d"` to preserve the original score during that period and decay afterward. Supports non-negative integer `Xm`/`Xh`/`Xd` | | context_type | str \| List[str] | No | None | Limit results to one or more `ContextType` values: `memory`, `resource`, or `skill` | @@ -394,7 +394,7 @@ Semantic retrieval with optional session-aware query planning. | Parameter | Type | Required | Default | Description | |-----------|------|----------|---------|-------------| | query | str | No | "" | Search query string. Required unless `image_url` is provided | -| image_url | str | No | None | Image query as a `data:image/...;base64,...`, `http(s)://`, or `viking://` URI. Requires a multimodal embedding model | +| image_url | str or list[str] | No | None | One image URI or an ordered list of image URIs, each as a `data:image/...;base64,...`, `http(s)://`, or `viking://` URI. Requires a multimodal embedding model | | search_type | `"semantic"` \| `"keywords"` | No | `"semantic"` | Retrieval type. `semantic` uses embeddings and vector recall; `keywords` uses BM25 keyword recall | | mode | `"list"` \| `"context"` | No | `"list"` | Return ranked hits or assembled context. CLI `--mode` maps to this field; context-specific parameters are described below | | target_uri | str \| List[str] | No | "" | Limit search to specific URI prefix | diff --git a/docs/en/guides/01-configuration.md b/docs/en/guides/01-configuration.md index fc78e2fa39..da749f3acb 100644 --- a/docs/en/guides/01-configuration.md +++ b/docs/en/guides/01-configuration.md @@ -388,6 +388,8 @@ When the embedding provider experiences consecutive transient failures (e.g. `42 With `input: "multimodal"`, OpenViking can embed text, images (PNG, JPG, etc.), and mixed content. Image-to-image search requires this mode; text-only embedding models continue to index image summaries but cannot accept image queries. +Hybrid image support depends on the dense channel. The complete mixed input goes to dense, while sparse receives only text; image-only inputs skip sparse. Volcengine hybrid keeps one combined request for text-only input and uses separate multimodal dense and text-only sparse requests for mixed input. Keep the dense model and dimension compatible with the vectors already stored in the index. + **Supported providers:** - `openai`: OpenAI Embedding API - `azure`: Azure OpenAI Embedding API diff --git a/docs/zh/api/06-retrieval.md b/docs/zh/api/06-retrieval.md index a1f999ce53..41aaa8386a 100644 --- a/docs/zh/api/06-retrieval.md +++ b/docs/zh/api/06-retrieval.md @@ -54,7 +54,7 @@ OpenViking 提供多种检索方法,包括简单的向量相似度搜索、关 | 参数 | 类型 | 必填 | 默认值 | 说明 | |------|------|------|--------|------| | query | str | 否 | "" | 查询文本;图片查询或仅按过滤条件查询时可省略 | -| image_url | str | 否 | None | 图片查询,支持 `data:image/...;base64,...`、`http(s)://` 或 `viking://` URI;需要 multimodal embedding 模型 | +| image_url | str 或 list[str] | 否 | None | 单个图片 URI 或有序图片 URI 数组,每项支持 `data:image/...;base64,...`、`http(s)://` 或 `viking://` URI;需要 multimodal embedding 模型 | | target_uri | str \| List[str] | 否 | "" | 限制搜索范围到指定的 URI 前缀 | | events_time_decay_protection | str \| null | 否 | null | 不传或传 `null` 关闭衰减;传 `"0"` 立即衰减;传 `"7d"` 等时长则在保护期内保持原分,之后衰减。支持非负整数 `Xm`/`Xh`/`Xd` | | context_type | str \| List[str] | 否 | None | 限定一个或多个 `ContextType` 取值:`memory`、`resource` 或 `skill` | @@ -394,7 +394,7 @@ ov find "红色海报风格" --image ./poster.png --uri "viking://resources/imag | 参数 | 类型 | 必填 | 默认值 | 说明 | |------|------|------|--------|------| | query | str | 否 | "" | 搜索查询字符串;未提供 `image_url` 时必填 | -| image_url | str | 否 | None | 图片查询,支持 `data:image/...;base64,...`、`http(s)://` 或 `viking://` URI;需要 multimodal embedding 模型 | +| image_url | str 或 list[str] | 否 | None | 单个图片 URI 或有序图片 URI 数组,每项支持 `data:image/...;base64,...`、`http(s)://` 或 `viking://` URI;需要 multimodal embedding 模型 | | search_type | `"semantic"` \| `"keywords"` | 否 | `"semantic"` | 检索类型。`semantic` 使用 Embedding 和向量召回;`keywords` 使用 BM25 关键词召回 | | mode | `"list"` \| `"context"` | 否 | `"list"` | 返回排序命中列表或组装好的上下文。CLI `--mode` 会映射到这个字段;context 专用参数见下文 | | target_uri | str \| List[str] | 否 | "" | 限制搜索范围到指定的 URI 前缀 | diff --git a/docs/zh/guides/01-configuration.md b/docs/zh/guides/01-configuration.md index 7a219cc789..8fee90c8ec 100644 --- a/docs/zh/guides/01-configuration.md +++ b/docs/zh/guides/01-configuration.md @@ -385,6 +385,8 @@ OpenAI 已于 2026 年 8 月 31 日[停止在 ChatGPT 登录的 Codex 中提供 使用 `input: "multimodal"` 时,OpenViking 可以嵌入文本、图片(PNG、JPG 等)和混合内容。以图搜图需要该模式;纯文本 embedding 模型仍会索引图片 summary,但不能接收图片查询。 +Hybrid 是否支持图片取决于 dense 通道。图文完整输入进入 dense,sparse 只接收文字;纯图输入跳过 sparse。Volcengine hybrid 对纯文本保留一次合并请求,对图文分别请求多模态 dense 和文本 sparse。dense 模型和维度应与索引中已有的向量保持兼容。 + **支持的 provider:** - `openai`: OpenAI Embedding API - `azure`: Azure OpenAI Embedding API diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index a8e21c2350..534d1530ba 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -656,28 +656,38 @@ def __init__(self, dense_embedder: DenseEmbedderBase, sparse_embedder: SparseEmb @property def supports_multimodal(self) -> bool: - """Supports multimodal input only if both sub-embedders do.""" - return self.dense_embedder.supports_multimodal and self.sparse_embedder.supports_multimodal + """Images are handled by the dense channel; sparse consumes only text.""" + return self.dense_embedder.supports_multimodal + + def _sparse_input(self, content: "EmbeddingInput") -> Optional[str]: + text = extract_text_from_content(content) + if isinstance(content, list) and not text.strip(): + return None + return self.sparse_embedder.prepare_embedding_input(text) def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: """Combine results from both embedders""" dense_input = self.dense_embedder.prepare_embedding_input(content) - sparse_input = self.sparse_embedder.prepare_embedding_input(content) + sparse_input = self._sparse_input(content) dense_res = self.dense_embedder.embed(dense_input, is_query=is_query) - sparse_res = self.sparse_embedder.embed(sparse_input, is_query=is_query) + sparse_res = ( + self.sparse_embedder.embed(sparse_input, is_query=is_query) + if sparse_input is not None + else None + ) return EmbedResult( - dense_vector=dense_res.dense_vector, sparse_vector=sparse_res.sparse_vector + dense_vector=dense_res.dense_vector, + sparse_vector=sparse_res.sparse_vector if sparse_res is not None else None, ) async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: dense_input = self.dense_embedder.prepare_embedding_input(content) - sparse_input = self.sparse_embedder.prepare_embedding_input(content) - calls = asyncio.gather( - self.dense_embedder.embed_async(dense_input, is_query=is_query), - self.sparse_embedder.embed_async(sparse_input, is_query=is_query), - return_exceptions=True, - ) + sparse_input = self._sparse_input(content) + providers = [self.dense_embedder.embed_async(dense_input, is_query=is_query)] + if sparse_input is not None: + providers.append(self.sparse_embedder.embed_async(sparse_input, is_query=is_query)) + calls = asyncio.gather(*providers, return_exceptions=True) # Keep the parent borrow alive until every child has stopped. Shielding # also lets cancellation wait for providers with asynchronous cleanup. from openviking.service.task_tracker_concurrency import run_to_completion @@ -686,9 +696,11 @@ async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) - for result in results: if isinstance(result, BaseException): raise result - dense_res, sparse_res = results + dense_res = results[0] + sparse_res = results[1] if sparse_input is not None else None return EmbedResult( - dense_vector=dense_res.dense_vector, sparse_vector=sparse_res.sparse_vector + dense_vector=dense_res.dense_vector, + sparse_vector=sparse_res.sparse_vector if sparse_res is not None else None, ) def get_dimension(self) -> int: diff --git a/openviking/models/embedder/volcengine_embedders.py b/openviking/models/embedder/volcengine_embedders.py index b13b360729..37fdf2febc 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -536,8 +536,9 @@ def supports_multimodal(self) -> bool: class VolcengineHybridEmbedder(HybridEmbedderBase): """Volcengine Hybrid Embedder Implementation - Generates both dense and sparse embeddings simultaneously using Volcengine's - multimodal embedding API. + Text uses one combined dense/sparse request. Images use the same dense + model without sparse output, with a separate text-only sparse request when + text is present. """ def __init__( @@ -614,14 +615,22 @@ def _usage_value(key: str, default: int = 0) -> int: @property def supports_multimodal(self) -> bool: - """Hybrid vectors include sparse output and therefore require text input. + """The dense channel accepts images independently of sparse support.""" + return self.input_type == "multimodal" - A hybrid provider request enables sparse embedding, for which image and - video parts are unsupported. The base guard performs a deterministic - text-only downgrade for mixed content. Image-only search remains - unsupported because dropping its only input would produce an empty query. - """ - return False + def _embedding_requests(self, content: "EmbeddingInput") -> List[Dict[str, Any]]: + parts = to_multimodal_input(content) + if all(part.get("type") == "text" for part in parts): + return [{"input": parts, "sparse_embedding": {"type": "enabled"}}] + # Ark rejects media when sparse output is enabled. Keep every media + # part in the dense request, and derive sparse weights from text only. + requests = [{"input": parts}] + text = extract_text_from_content(content) + if text.strip(): + requests.append( + {"input": to_multimodal_input(text), "sparse_embedding": {"type": "enabled"}} + ) + return requests @model_error_boundary def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: @@ -640,27 +649,23 @@ def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResul """ def _embed_call(): - # Always use multimodal for hybrid to get both - - response, duration_seconds = _measure_embedding_call( - self, - lambda: self.client.multimodal_embeddings.create( - input=to_multimodal_input(content), - model=self.model_name, - sparse_embedding={"type": "enabled"}, - extra_headers=self.extra_headers, - ), - ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) - item = response.data - dense_vector = truncate_and_normalize(item.embedding, self.dimension) - sparse_vector = getattr(item, "sparse_embedding", None) - - return EmbedResult( - dense_vector=dense_vector, sparse_vector=process_sparse_embedding(sparse_vector) - ) + dense_vector = None + sparse_vector = None + for request in self._embedding_requests(content): + response, duration_seconds = _measure_embedding_call( + self, + lambda request=request: self.client.multimodal_embeddings.create( + model=self.model_name, extra_headers=self.extra_headers, **request + ), + ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + if dense_vector is None: + dense_vector = truncate_and_normalize(response.data.embedding, self.dimension) + if "sparse_embedding" in request: + sparse_vector = process_sparse_embedding( + getattr(response.data, "sparse_embedding", None) + ) + return EmbedResult(dense_vector=dense_vector, sparse_vector=sparse_vector) try: return self._run_once( @@ -681,25 +686,23 @@ async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) - client = self._get_async_client() async def _embed_call() -> EmbedResult: - response, duration_seconds = await _measure_embedding_call_async( - self, - lambda: client.multimodal_embeddings.create( - input=to_multimodal_input(content), - model=self.model_name, - sparse_embedding={"type": "enabled"}, - extra_headers=self.extra_headers, - ), - ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) - item = response.data - dense_vector = truncate_and_normalize(item.embedding, self.dimension) - sparse_vector = getattr(item, "sparse_embedding", None) - return EmbedResult( - dense_vector=dense_vector, - sparse_vector=process_sparse_embedding(sparse_vector), - ) + dense_vector = None + sparse_vector = None + for request in self._embedding_requests(content): + response, duration_seconds = await _measure_embedding_call_async( + self, + lambda request=request: client.multimodal_embeddings.create( + model=self.model_name, extra_headers=self.extra_headers, **request + ), + ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + if dense_vector is None: + dense_vector = truncate_and_normalize(response.data.embedding, self.dimension) + if "sparse_embedding" in request: + sparse_vector = process_sparse_embedding( + getattr(response.data, "sparse_embedding", None) + ) + return EmbedResult(dense_vector=dense_vector, sparse_vector=sparse_vector) try: return await self._run_once_async( diff --git a/openviking/retrieve/context_assembler/gather.py b/openviking/retrieve/context_assembler/gather.py index f21ee374da..ee5ca32cae 100644 --- a/openviking/retrieve/context_assembler/gather.py +++ b/openviking/retrieve/context_assembler/gather.py @@ -211,7 +211,7 @@ async def gather_candidates( limit: int, score_threshold: Optional[float], filter: Optional[Dict[str, Any]] = None, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, events_time_decay_protection: Optional[str] = None, search_type: SearchType = "semantic", peer_scope: str = "all", diff --git a/openviking/retrieve/context_assembler/params.py b/openviking/retrieve/context_assembler/params.py index 656496ce36..3eb03fc30f 100644 --- a/openviking/retrieve/context_assembler/params.py +++ b/openviking/retrieve/context_assembler/params.py @@ -124,7 +124,7 @@ class AssembleParams: query: str = "" search_type: SearchType = "semantic" - image_url: Optional[str] = None + image_url: str | list[str] | None = None limit: int = DEFAULT_LIMIT score_threshold: Optional[float] = None filter: Optional[Dict[str, Any]] = None diff --git a/openviking/retrieve/hierarchical_retriever.py b/openviking/retrieve/hierarchical_retriever.py index 6be9e25b8a..ed25d96170 100644 --- a/openviking/retrieve/hierarchical_retriever.py +++ b/openviking/retrieve/hierarchical_retriever.py @@ -44,6 +44,23 @@ logger = get_logger(__name__) +def _has_context_type_filter(scope: Any) -> bool: + """Recognize explicit type constraints in public DSL and internal filter ASTs.""" + if isinstance(scope, dict): + return scope.get("field") == "context_type" or any( + _has_context_type_filter(value) + for value in scope.values() + if isinstance(value, (dict, list)) + ) + if isinstance(scope, list): + return any(_has_context_type_filter(value) for value in scope) + if hasattr(scope, "payload"): + return _has_context_type_filter(scope.payload) + return getattr(scope, "field", None) == "context_type" or any( + _has_context_type_filter(value) for value in getattr(scope, "conds", []) + ) + + class RetrieverMode(str): THINKING = "thinking" QUICK = "quick" @@ -102,6 +119,7 @@ async def retrieve( events_time_decay_protection: Optional[str] = None, request_now: Optional[datetime] = None, search_type: SearchType = "semantic", + has_explicit_target: Optional[bool] = None, ) -> QueryResult: """ Run one global vector search, then optionally rerank its candidates. @@ -114,6 +132,7 @@ async def retrieve( score_gte: True uses >=, False uses > scope_dsl: Additional scope constraints passed from public find/search filter level: Optional result level filter (0=L0, 1=L1, 2=L2) + has_explicit_target: Whether target_uri was supplied before default expansion """ t0 = time.monotonic() telemetry = get_current_telemetry() @@ -140,6 +159,8 @@ async def retrieve( vector_proxy = VikingDBManagerProxy(self.vector_store, ctx) target_dirs = [d for d in (query.target_directories or []) if d] + if has_explicit_target is None: + has_explicit_target = bool(target_dirs) if not await vector_proxy.collection_exists_bound(): logger.warning( @@ -177,7 +198,12 @@ async def retrieve( root_uris = default_target_directories(ctx, context_type=query.context_type) context_type = query.context_type.value if query.context_type else None - if image_query and context_type is None: + if ( + image_query + and context_type is None + and not has_explicit_target + and not _has_context_type_filter(scope_dsl) + ): context_type = ContextType.RESOURCE.value search_limit = limit * self.RERANK_CANDIDATE_MULTIPLIER if use_rerank else limit diff --git a/openviking/server/routers/search.py b/openviking/server/routers/search.py index 26e4fe9d6d..59546d1ba0 100644 --- a/openviking/server/routers/search.py +++ b/openviking/server/routers/search.py @@ -112,7 +112,11 @@ def _resolve_uri_list(uris: Sequence[str], ctx: RequestContext) -> list[str]: return [validate_request_viking_uri(resolve_path_variables(uri), ctx) for uri in uris] -def _resolve_image_url(image_url: Optional[str], ctx: RequestContext) -> Optional[str]: +def _resolve_image_url( + image_url: str | list[str] | None, ctx: RequestContext +) -> str | list[str] | None: + if isinstance(image_url, list): + return [_resolve_image_url(url, ctx) for url in image_url] if not image_url: return image_url resolved = resolve_path_variables(image_url) @@ -127,7 +131,7 @@ class FindRequest(BaseModel): model_config = ConfigDict(extra="forbid") query: str = "" - image_url: Optional[str] = None + image_url: str | list[str] | None = None target_uri: Union[str, List[str]] = "" context_type: Optional[Union[str, List[str]]] = None limit: int = Field(default=DEFAULT_LIMIT, ge=1) @@ -220,7 +224,7 @@ class SearchRequest(BaseModel): query: str = "" search_type: SearchType = "semantic" - image_url: Optional[str] = None + image_url: str | list[str] | None = None target_uri: Union[str, List[str]] = "" context_type: Optional[Union[str, List[str]]] = None session_id: Optional[str] = None diff --git a/openviking/service/search_service.py b/openviking/service/search_service.py index 59e7e3d50e..321c416527 100644 --- a/openviking/service/search_service.py +++ b/openviking/service/search_service.py @@ -30,7 +30,7 @@ def _validate_query( query: str, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, filter: Optional[Dict] = None, ) -> None: """Reject malformed query text or a request with nothing to search for. @@ -95,9 +95,11 @@ async def ensure_keywords_search_supported(self, ctx: RequestContext) -> None: async def _resolve_image_url( self, - image_url: Optional[str], + image_url: str | list[str] | None, ctx: RequestContext, - ) -> Optional[str]: + ) -> str | list[str] | None: + if isinstance(image_url, list): + return [await self._resolve_image_url(url, ctx) for url in image_url] if not image_url: return None if is_viking_uri(image_url): @@ -121,7 +123,7 @@ async def search( score_threshold: Optional[float] = None, filter: Optional[Dict] = None, level: Optional[List[int]] = None, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, events_time_decay_protection: Optional[str] = None, search_type: SearchType = "semantic", context_types: Optional[List[ContextType]] = None, @@ -179,7 +181,7 @@ async def find( score_threshold: Optional[float] = None, filter: Optional[Dict] = None, level: Optional[List[int]] = None, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, events_time_decay_protection: Optional[str] = None, search_type: SearchType = "semantic", ) -> Any: diff --git a/openviking/storage/viking_fs/_base.py b/openviking/storage/viking_fs/_base.py index 4607b6f020..139be3d2f5 100644 --- a/openviking/storage/viking_fs/_base.py +++ b/openviking/storage/viking_fs/_base.py @@ -87,12 +87,12 @@ def _prepare_snapshot_diff( return change_type, before, after -def _ensure_non_empty_search_query(query: str, image_url: Optional[str] = None) -> None: +def _ensure_non_empty_search_query(query: str, image_url: str | list[str] | None = None) -> None: if not query.strip() and not image_url: raise InvalidArgumentError("Search query or image_url must not be empty.") -def is_filter_only_query(query: str, image_url: Optional[str] = None) -> bool: +def is_filter_only_query(query: str, image_url: str | list[str] | None = None) -> bool: """Return True when the caller supplied no query and no image.""" return not query.strip() and not image_url diff --git a/openviking/storage/viking_fs/_semantic.py b/openviking/storage/viking_fs/_semantic.py index ff019b8917..2c45873f58 100644 --- a/openviking/storage/viking_fs/_semantic.py +++ b/openviking/storage/viking_fs/_semantic.py @@ -203,7 +203,7 @@ async def find( filter: Optional[Dict] = None, ctx: Optional[RequestContext] = None, level: Optional[List[int]] = None, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, events_time_decay_protection: Optional[str] = None, search_type: SearchType = "semantic", ): @@ -306,6 +306,7 @@ async def find( limit=limit, mode=RetrieverMode.QUICK, score_threshold=score_threshold, + has_explicit_target=bool(retrieval_targets.first_explicit_directory), scope_dsl=filter, level=level, events_time_decay_protection=events_time_decay_protection, @@ -401,7 +402,7 @@ async def search( filter: Optional[Dict] = None, ctx: Optional[RequestContext] = None, level: Optional[List[int]] = None, - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, events_time_decay_protection: Optional[str] = None, search_type: SearchType = "semantic", context_types: Optional[List[ContextType]] = None, @@ -542,6 +543,7 @@ async def _execute(tq: TypedQuery): ctx=real_ctx, limit=limit, score_threshold=score_threshold, + has_explicit_target=bool(retrieval_targets.first_explicit_directory), scope_dsl=filter, level=level, events_time_decay_protection=events_time_decay_protection, diff --git a/openviking/utils/image_search.py b/openviking/utils/image_search.py index d5a8b7d1f9..269efd3464 100644 --- a/openviking/utils/image_search.py +++ b/openviking/utils/image_search.py @@ -9,7 +9,7 @@ import mimetypes import os from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List from PIL import Image, UnidentifiedImageError @@ -94,19 +94,22 @@ def is_viking_uri(value: str) -> bool: def build_multimodal_embedding_input( text: str = "", - image_url: Optional[str] = None, + image_url: str | list[str] | None = None, ) -> List[Dict[str, Any]]: parts: List[Dict[str, Any]] = [] if text.strip(): parts.append({"type": "text", "text": text}) - if image_url: - parts.append({"type": "image_url", "image_url": {"url": image_url}}) + for url in image_url if isinstance(image_url, list) else [image_url]: + if url: + parts.append({"type": "image_url", "image_url": {"url": url}}) return parts -def normalize_client_image_input(image: Any) -> Optional[str]: +def normalize_client_image_input(image: Any) -> str | list[str] | None: if image is None: return None + if isinstance(image, list): + return [normalize_client_image_input(item) for item in image] if isinstance(image, (bytes, bytearray, memoryview)): return image_bytes_to_data_uri(image) diff --git a/sdk/python/openviking_sdk/client.py b/sdk/python/openviking_sdk/client.py index c01669d003..a6ca49014c 100644 --- a/sdk/python/openviking_sdk/client.py +++ b/sdk/python/openviking_sdk/client.py @@ -136,9 +136,11 @@ def _atomic_write_bytes(path: Path, data: bytes) -> None: raise -def _normalize_image_input(image: Any) -> Optional[str]: +def _normalize_image_input(image: Any) -> str | list[str] | None: if image is None: return None + if isinstance(image, list): + return [_normalize_image_input(item) for item in image] if isinstance(image, (bytes, bytearray, memoryview)): return _image_to_data_uri(image) value = os.fspath(image) if isinstance(image, os.PathLike) else str(image) diff --git a/tests/misc/test_vikingfs_find_without_rerank.py b/tests/misc/test_vikingfs_find_without_rerank.py index 448ad43b5b..f5467e9b70 100644 --- a/tests/misc/test_vikingfs_find_without_rerank.py +++ b/tests/misc/test_vikingfs_find_without_rerank.py @@ -7,6 +7,7 @@ import pytest +from openviking.models.embedder.base import EmbedResult from openviking.retrieve.hierarchical_retriever import RetrieverMode from openviking.server.identity import RequestContext, Role from openviking.storage.viking_fs import VikingFS @@ -58,6 +59,7 @@ async def retrieve( events_time_decay_protection=None, request_now=None, search_type="semantic", + has_explicit_target=None, ): captured["typed_query"] = typed_query captured["ctx"] = ctx @@ -112,61 +114,58 @@ async def retrieve( @pytest.mark.asyncio -async def test_find_accepts_image_url_without_text_query(monkeypatch) -> None: +async def test_image_retrieval_preserves_default_and_explicit_scopes(monkeypatch) -> None: fs = _make_viking_fs() - captured = {} - - class FakeRetriever: - def __init__(self, storage, embedder, rerank_config): - pass - - async def retrieve( - self, - typed_query, - ctx, - limit, - mode, - score_threshold, - scope_dsl, - level, - events_time_decay_protection=None, - request_now=None, - search_type="semantic", - ): - captured["typed_query"] = typed_query - captured["mode"] = mode - return QueryResult( - query=typed_query, - matched_contexts=[ - MatchedContext( - uri="viking://resources/photos/cat.png", - context_type=ContextType.RESOURCE, - score=0.9, - ) - ], - searched_directories=["viking://resources"], - ) - + fs.abstract = AsyncMock(return_value="") + fs.query_embedder.supports_multimodal = True + fs.query_embedder.prepare_embedding_input.side_effect = lambda content: content + fs.query_embedder.embed_async = AsyncMock(return_value=EmbedResult(dense_vector=[1.0])) + proxy = MagicMock() + proxy.collection_exists_bound = AsyncMock(return_value=True) + proxy.search_in_tenant = AsyncMock( + return_value=[ + {"uri": "viking://resources/cat.png", "context_type": "resource", "_score": 0.9} + ] + ) monkeypatch.setattr( - "openviking.retrieve.hierarchical_retriever.HierarchicalRetriever", - FakeRetriever, + "openviking.retrieve.hierarchical_retriever.VikingDBManagerProxy", + MagicMock(return_value=proxy), ) - - result = await fs.find( - "", - target_uri="viking://resources", - image_url="data:image/png;base64,abc", - ctx=_ctx(), + images = ["data:image/png;base64,YQ==", "data:image/png;base64,Yg=="] + + found = await fs.find("", image_url=images, ctx=_ctx()) + assert [hit.uri for hit in found.resources] == ["viking://resources/cat.png"] + assert proxy.search_in_tenant.call_args.kwargs["context_type"] == "resource" + fs.query_embedder.embed_async.assert_awaited_once_with( + [ + {"type": "image_url", "image_url": {"url": images[0]}}, + {"type": "image_url", "image_url": {"url": images[1]}}, + ], + is_query=True, ) - assert result.total == 1 - typed_query = captured["typed_query"] - assert typed_query.query == "" - assert typed_query.image_query is True - assert typed_query.embedding_input == [ - {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}} + searched = await fs.search("", image_url=images, ctx=_ctx()) + assert [hit.uri for hit in searched.resources] == ["viking://resources/cat.png"] + assert proxy.search_in_tenant.call_args.kwargs["context_type"] == "resource" + + # An explicit memory URI or type filter overrides the image resource default. + proxy.search_in_tenant.return_value = [ + {"uri": "viking://user/user1/memories/cat.md", "context_type": "memory", "_score": 0.9} ] - assert captured["mode"] == RetrieverMode.QUICK + found = await fs.find( + "", image_url=images, target_uri="viking://user/user1/memories", ctx=_ctx() + ) + assert [hit.uri for hit in found.memories] == ["viking://user/user1/memories/cat.md"] + assert proxy.search_in_tenant.call_args.kwargs["context_type"] is None + assert proxy.search_in_tenant.call_args.kwargs["target_directories"] == [ + "viking://user/user1/memories" + ] + + scope = {"op": "must", "field": "context_type", "conds": ["memory"]} + searched = await fs.search("cat", image_url=images, filter=scope, ctx=_ctx()) + assert [hit.uri for hit in searched.memories] == ["viking://user/user1/memories/cat.md"] + assert proxy.search_in_tenant.call_args.kwargs["context_type"] is None + assert proxy.search_in_tenant.call_args.kwargs["extra_filter"] == scope @pytest.mark.asyncio diff --git a/tests/unit/test_aml_server.py b/tests/unit/test_aml_server.py index 7d25a73d1f..fddc8f48c5 100644 --- a/tests/unit/test_aml_server.py +++ b/tests/unit/test_aml_server.py @@ -6,8 +6,11 @@ import httpx import pytest +from benchmark.aml import server from benchmark.aml.server import AMLSettings, create_app +IMAGES = ["data:image/png;base64,YQ==", "data:image/png;base64,Yg=="] + @pytest.fixture def adapter_client(monkeypatch, tmp_path): @@ -54,6 +57,7 @@ async def handle(request): assert "mode" not in body assert not body.get("session_id") assert not body.get("read_content") + assert body.get("image_url") == (IMAGES if user == "bob" else None) seen.append((user, body["query"])) await asyncio.sleep(0) return httpx.Response( @@ -77,10 +81,23 @@ async def handle(request): client.post("/search", json={"user_id": "alice", "query": "Where?", "top_k": 2}), client.post( "/search", - json={"user_id": "bob", "query": "Which?", "options": ["A", "B"], "top_k": 2}, + json={ + "user_id": "bob", + "query": [ + {"type": "image_url", "image_url": {"url": IMAGES[0]}}, + {"type": "text", "text": "Which?"}, + {"type": "image_url", "image_url": {"url": IMAGES[1]}}, + {"type": "text", "text": "Compare both."}, + ], + "options": ["A", "B"], + "top_k": 2, + }, ), ) - assert sorted(seen) == [("alice", "Where?"), ("bob", "Which?\n\nAnswer options:\n1. A\n2. B")] + assert sorted(seen) == [ + ("alice", "Where?"), + ("bob", "Which?\nCompare both.\n\nAnswer options:\n1. A\n2. B"), + ] for user, response in zip(("alice", "bob"), responses, strict=True): assert response.status_code == 200 assert response.json() == { @@ -91,22 +108,76 @@ async def handle(request): } -async def test_search_retries_transient_native_find_failure(adapter_client): +async def test_add_commits_native_parts_and_image_search_retries(adapter_client, monkeypatch): requests = [] + parts = [ + {"type": "text", "text": "Remember these receipts."}, + {"type": "image_url", "image_url": {"url": IMAGES[0]}}, + {"type": "image_url", "image_url": {"url": IMAGES[1]}}, + ] + payload = { + "request_id": "request-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [ + {"role": "user", "content": parts, "timestamp": 0}, + {"role": "assistant", "content": " Saved.\n"}, + ], + } def handle(request): requests.append(request) - if len(requests) == 1: - raise httpx.ReadTimeout("temporary failure", request=request) - return httpx.Response(200, json={"status": "ok", "result": {"memories": []}}) + assert request.headers["X-OpenViking-User"] == "alice" + assert request.headers["X-OpenViking-Account"] == "competition" + assert request.headers["X-API-Key"] == "test-ov-key" + if request.url.path == "/api/v1/sessions/session-1/messages/batch": + assert json.loads(request.content) == { + "messages": [ + {"role": "user", "parts": parts, "created_at": "1970-01-01T00:00:00.000Z"}, + {"role": "assistant", "content": " Saved.\n"}, + ] + } + result = {"added": 2} + elif request.url.path == "/api/v1/sessions/session-1/commit": + assert json.loads(request.content) == {"keep_recent_count": 0} + result = {"status": "accepted", "task_id": "task-1"} + elif request.url.path == "/api/v1/tasks/task-1": + assert request.method == "GET" + result = {"status": "completed"} + else: + assert request.url.path == "/api/v1/search/find" + if len(requests) == 4: + raise httpx.ReadTimeout("temporary failure", request=request) + result = {"memories": []} + return httpx.Response(200, json={"status": "ok", "result": result}) async with adapter_client(handle) as client: + # Each image fits; the aggregate budget is enforced before any write. + monkeypatch.setattr(server, "MAX_ADD_IMAGE_BYTES", 1) + rejected = await client.post("/add", json=payload) + assert rejected.status_code == 422 + assert requests == [] + monkeypatch.setattr(server, "MAX_ADD_IMAGE_BYTES", 2) + added = await client.post("/add", json=payload) + assert added.status_code == 200 + assert added.json() == { + "success": True, + "request_id": "request-1", + "user_id": "alice", + "session_id": "session-1", + } + assert [request.url.path for request in requests] == [ + "/api/v1/sessions/session-1/messages/batch", + "/api/v1/sessions/session-1/commit", + "/api/v1/tasks/task-1", + ] response = await client.post( - "/search", json={"user_id": "alice", "query": "Where?", "top_k": 2} + "/search", json={"user_id": "alice", "query": [parts[1]], "top_k": 2} ) assert response.status_code == 200 assert response.json() == {"data": []} - assert len(requests) == 2 - assert all(request.url.path == "/api/v1/search/find" for request in requests) - assert requests[0].content == requests[1].content - assert requests[0].headers == requests[1].headers + assert len(requests) == 5 + assert json.loads(requests[-1].content)["query"] == "" + assert json.loads(requests[-1].content)["image_url"] == IMAGES[0] + assert requests[-2].content == requests[-1].content + assert requests[-2].headers == requests[-1].headers diff --git a/tests/unit/test_image_search.py b/tests/unit/test_image_search.py index ddcf4d7d54..0e3679f64e 100644 --- a/tests/unit/test_image_search.py +++ b/tests/unit/test_image_search.py @@ -30,9 +30,14 @@ def supports_multimodal(self) -> bool: def test_text_only_embedder_drops_image_parts(): parts = build_multimodal_embedding_input( "cat photo", - "data:image/png;base64,abc", + ["data:image/png;base64,YQ==", "viking://resources/photos/cat.png"], ) + assert parts == [ + {"type": "text", "text": "cat photo"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,YQ=="}}, + {"type": "image_url", "image_url": {"url": "viking://resources/photos/cat.png"}}, + ] assert TextOnlyEmbedder("test").prepare_embedding_input(parts) == "cat photo" assert MultimodalEmbedder("test").prepare_embedding_input(parts) == parts diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index 9851e7138c..9b2c1e1184 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +from unittest.mock import AsyncMock import pytest @@ -214,17 +215,30 @@ async def no_sleep(_delay): return None monkeypatch.setattr("openviking.utils.model_retry.asyncio.sleep", no_sleep) - dense = FakeDenseEmbedder([[0.1, 0.2]]) + dense = FakeDenseEmbedder([[0.1, 0.2], [0.3, 0.4]]) sparse = FakeSparseEmbedder([ProviderError(503), {"term": 1.0}], max_retries=1) + monkeypatch.setattr(dense, "embed_async", AsyncMock(wraps=dense.embed_async)) + monkeypatch.setattr(sparse, "embed_async", AsyncMock(wraps=sparse.embed_async)) composite = CompositeHybridEmbedder(RetryEmbedder(dense), RetryEmbedder(sparse)) + image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YQ=="}} + content = [{"type": "text", "text": "query"}, image] - result = await composite.embed_async("query") + result = await composite.embed_async(composite.prepare_embedding_input(content), is_query=True) + dense.embed_async.assert_awaited_once_with(content, is_query=True) + sparse.embed_async.assert_awaited_with("query", is_query=True) assert result.dense_vector == [0.1, 0.2] assert result.sparse_vector == {"term": 1.0} assert dense.calls == 1 assert sparse.calls == 2 + # Removing text must skip sparse entirely, including its retry machinery. + image_result = await composite.embed_async([image], is_query=True) + assert image_result.dense_vector == [0.3, 0.4] + assert image_result.sparse_vector is None + assert dense.calls == 2 + assert sparse.calls == 2 + def test_failover_runs_after_current_credential_retry_exhaustion(): primary = FakeDenseEmbedder([ProviderError(503), ProviderError(503)], max_retries=1) diff --git a/tests/unit/test_volcengine_embedder_input.py b/tests/unit/test_volcengine_embedder_input.py index 9cd5aad6b1..bfac65c746 100644 --- a/tests/unit/test_volcengine_embedder_input.py +++ b/tests/unit/test_volcengine_embedder_input.py @@ -3,8 +3,9 @@ """Input-guard tests for Volcengine sparse-capable embedding modes.""" from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch +from openviking.models.embedder.base import embed_compat from openviking.models.embedder.volcengine_embedders import ( VolcengineDenseEmbedder, VolcengineHybridEmbedder, @@ -27,16 +28,68 @@ def test_sparse_embedder_downgrades_multimodal_input_to_text(_mock_ark): @patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") -def test_hybrid_embedder_downgrades_multimodal_input_to_text(_mock_ark): - embedder = VolcengineHybridEmbedder(model_name="test", api_key="test-key") +@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.AsyncArk") +async def test_hybrid_embedder_routes_media_to_dense_and_text_to_sparse(mock_async_ark, mock_ark): + dense_response = SimpleNamespace(data=SimpleNamespace(embedding=[1.0, 0.0]), usage=None) + sparse_response = SimpleNamespace( + data=SimpleNamespace(embedding=[0.0, 1.0], sparse_embedding=[{"index": 9, "value": 0.8}]), + usage=None, + ) + create = mock_ark.return_value.multimodal_embeddings.create + create.side_effect = [dense_response, sparse_response, sparse_response] + async_create = AsyncMock(side_effect=[dense_response, sparse_response, dense_response]) + mock_async_ark.return_value.multimodal_embeddings.create = async_create + embedder = VolcengineHybridEmbedder( + model_name="existing-model", api_key="test-key", dimension=2 + ) + expected_requests = [ + call(model="existing-model", extra_headers=embedder.extra_headers, input=_MULTIMODAL_INPUT), + call( + model="existing-model", + extra_headers=embedder.extra_headers, + input=_MULTIMODAL_INPUT[:1], + sparse_embedding={"type": "enabled"}, + ), + ] - assert embedder.supports_multimodal is False - assert embedder.prepare_embedding_input(_MULTIMODAL_INPUT) == "diagram of a heat exchanger" + result = embedder.embed(embedder.prepare_embedding_input(_MULTIMODAL_INPUT), is_query=True) + assert result.dense_vector == [1.0, 0.0] + assert result.sparse_vector == {"9": 0.8} + assert create.call_args_list == expected_requests + + async_result = await embed_compat(embedder, _MULTIMODAL_INPUT, is_query=True) + assert async_result == result + assert async_create.await_args_list == expected_requests + + image_result = await embed_compat(embedder, _MULTIMODAL_INPUT[1:], is_query=True) + assert image_result.dense_vector == [1.0, 0.0] + assert image_result.sparse_vector is None + assert async_create.await_args_list[2:] == [ + call( + model="existing-model", + extra_headers=embedder.extra_headers, + input=_MULTIMODAL_INPUT[1:], + ) + ] + + text_result = embedder.embed("hello", is_query=True) + assert text_result.dense_vector == [0.0, 1.0] + assert text_result.sparse_vector == {"9": 0.8} + assert create.call_args_list[2:] == [ + call( + model="existing-model", + extra_headers=embedder.extra_headers, + input=[{"type": "text", "text": "hello"}], + sparse_embedding={"type": "enabled"}, + ) + ] @patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") @patch("openviking.models.embedder.volcengine_embedders.time.perf_counter") -def test_dense_embedder_passes_local_provider_duration_to_metrics(mock_perf_counter, mock_ark): +def test_dense_embedder_passes_local_provider_duration_to_metrics( + mock_perf_counter, mock_ark, monkeypatch +): """Embedding call duration must come from this SDK request, not shared embedder state.""" response = SimpleNamespace( usage=SimpleNamespace(prompt_tokens=3, total_tokens=3), @@ -45,6 +98,9 @@ def test_dense_embedder_passes_local_provider_duration_to_metrics(mock_perf_coun client = MagicMock() client.embeddings.create.return_value = response mock_ark.return_value = client + monkeypatch.setattr( + "openviking.models.embedder.volcengine_embedders.get_current_telemetry", MagicMock() + ) mock_perf_counter.side_effect = [100.0, 100.25] embedder = VolcengineDenseEmbedder(