From 0bc28e79083d13df8383c9b665afeeb7d3d79f6c Mon Sep 17 00:00:00 2001 From: heaoxiang Date: Thu, 15 Oct 2026 20:05:15 +0800 Subject: [PATCH 1/6] feat(benchmark): accept multimodal content in AML Add --- benchmark/aml/README.md | 29 +++++ benchmark/aml/server.py | 88 ++++++++++++- tests/unit/test_aml_server.py | 239 +++++++++++++++++++++++++++++++++- 3 files changed, 349 insertions(+), 7 deletions(-) diff --git a/benchmark/aml/README.md b/benchmark/aml/README.md index c3667e79f1..cb72eeb932 100644 --- a/benchmark/aml/README.md +++ b/benchmark/aml/README.md @@ -62,6 +62,35 @@ 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` still accepts string queries +and returns text evidence; image-query and image-evidence support are not +included in this adapter yet. + 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..cd48c06d4e 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,8 @@ 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 RETRYABLE_OPENVIKING_CODES = { "ABORTED", "DEADLINE_EXCEEDED", @@ -103,14 +107,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,6 +199,20 @@ 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") @@ -298,7 +374,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) diff --git a/tests/unit/test_aml_server.py b/tests/unit/test_aml_server.py index 7d25a73d1f..548f4339fd 100644 --- a/tests/unit/test_aml_server.py +++ b/tests/unit/test_aml_server.py @@ -1,13 +1,22 @@ """Competition API contract through the real OpenViking SDK request builder.""" import asyncio +import base64 import json import httpx import pytest +from benchmark.aml import server from benchmark.aml.server import AMLSettings, create_app +IMAGE_URL = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a6lcAAAAASUVORK5CYII=" +) +IMAGE_PART = {"type": "image_url", "image_url": {"url": IMAGE_URL}} +IMAGE_SIZE = len(base64.b64decode(IMAGE_URL.partition(",")[2])) + @pytest.fixture def adapter_client(monkeypatch, tmp_path): @@ -26,6 +35,7 @@ def sdk_client(**kwargs): openviking_api_key="test-ov-key", aml_api_key="test-aml-key", retry_delay_seconds=0.001, + task_poll_seconds=0.001, ) ) return client_type( @@ -37,7 +47,230 @@ def sdk_client(**kwargs): return build -async def test_search_uses_native_find_and_preserves_user_scope_and_aml_response(adapter_client): +@pytest.fixture +def add_client(adapter_client): + def build(task_status="completed"): + requests = [] + + def handle(request): + requests.append(request) + if request.url.path == "/api/v1/sessions/session-1/messages/batch": + result = {"added": len(json.loads(request.content)["messages"])} + 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"} + else: + assert request.method == "GET" + assert request.url.path == "/api/v1/tasks/task-1" + result = {"status": task_status, "error": "extraction failed"} + return httpx.Response(200, json={"status": "ok", "result": result}) + + return adapter_client(handle), requests + + return build + + +@pytest.mark.parametrize( + "content", + [ + " Remember my preference for tea.\n", + "def cache_key(user_id):\n return f'cache:{user_id}'\n", + ], + ids=["textual", "coding"], +) +async def test_add_string_payload_and_commit_are_unchanged(add_client, content): + client, requests = add_client() + payload = { + "request_id": "request-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [ + {"role": "user", "content": content, "timestamp": 0}, + {"role": "assistant", "content": "Acknowledged."}, + ], + } + async with client: + response = await client.post("/add", json=payload) + assert response.status_code == 200 + assert response.json() == { + "success": True, + "request_id": "request-1", + "user_id": "alice", + "session_id": "session-1", + } + assert json.loads(requests[0].content) == { + "messages": [ + {"role": "user", "content": content, "created_at": "1970-01-01T00:00:00.000Z"}, + {"role": "assistant", "content": "Acknowledged."}, + ] + } + 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", + ] + assert all(request.headers["X-OpenViking-User"] == "alice" for request in requests) + + +@pytest.mark.parametrize( + "parts", + [ + [IMAGE_PART], + [ + {"type": "text", "text": " First image:\n"}, + IMAGE_PART, + {"type": "text", "text": "The same image again:"}, + IMAGE_PART, + ], + [{"type": "text", "text": "Text supplied as a content array."}], + ], + ids=["image-only", "ordered-multiple-images", "text-array"], +) +async def test_add_forwards_content_array_as_native_parts(add_client, parts): + client, requests = add_client() + async with client: + response = await client.post( + "/add", + json={ + "request_id": "multimodal-1", + "user_id": "bob", + "session_id": "session-1", + "messages": [ + {"role": "user", "content": "A plain message."}, + {"role": "assistant", "content": parts, "timestamp": 1704067200000}, + ], + }, + ) + assert response.status_code == 200 + assert response.json() == { + "success": True, + "request_id": "multimodal-1", + "user_id": "bob", + "session_id": "session-1", + } + assert json.loads(requests[0].content) == { + "messages": [ + {"role": "user", "content": "A plain message."}, + {"role": "assistant", "parts": parts, "created_at": "2024-01-01T00:00:00.000Z"}, + ] + } + assert len(requests) == 3 + assert all(request.headers["X-OpenViking-User"] == "bob" for request in requests) + assert all(request.headers["X-OpenViking-Account"] == "competition" for request in requests) + assert all(request.headers["X-API-Key"] == "test-ov-key" for request in requests) + + +@pytest.mark.parametrize( + "content", + [ + "", + " \n ", + [], + [{"type": "text", "text": " "}], + [{"type": "image_url", "image_url": {}}], + [{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}], + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%"}}], + [{"type": "image_url", "image_url": {"url": "data:image/gif;base64,YQ=="}}], + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,"}}], + [{"type": "tool", "tool_name": "find"}], + ], + ids=[ + "empty-string", + "blank-string", + "empty-array", + "blank-text-part", + "missing-image-url", + "remote-image", + "invalid-base64", + "unsupported-mime", + "empty-image", + "unsupported-part", + ], +) +async def test_add_rejects_invalid_content_before_writing(add_client, content): + client, requests = add_client() + async with client: + response = await client.post( + "/add", + json={ + "request_id": "invalid-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [{"role": "user", "content": content}], + }, + ) + assert response.status_code == 422 + assert requests == [] + + +@pytest.mark.parametrize( + ("limit", "max_bytes", "image_count", "expected_status"), + [ + ("MAX_IMAGE_BYTES", IMAGE_SIZE - 1, 1, 422), + ("MAX_IMAGE_BYTES", IMAGE_SIZE, 1, 200), + ("MAX_ADD_IMAGE_BYTES", IMAGE_SIZE * 2 - 1, 2, 422), + ("MAX_ADD_IMAGE_BYTES", IMAGE_SIZE * 2, 2, 200), + ], + ids=["image-over-limit", "image-at-limit", "request-over-limit", "request-at-limit"], +) +async def test_add_enforces_image_budget_across_messages( + add_client, monkeypatch, limit, max_bytes, image_count, expected_status +): + monkeypatch.setattr(server, limit, max_bytes) + client, requests = add_client() + async with client: + response = await client.post( + "/add", + json={ + "request_id": "oversized-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [{"role": "user", "content": [IMAGE_PART]} for _ in range(image_count)], + }, + ) + assert response.status_code == expected_status + if expected_status == 422: + assert requests == [] + + +async def test_add_multimodal_extraction_failure_is_not_reported_as_success(add_client): + client, requests = add_client(task_status="failed") + async with client: + response = await client.post( + "/add", + json={ + "request_id": "failed-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [{"role": "user", "content": [IMAGE_PART]}], + }, + ) + assert response.status_code == 503 + assert "success" not in response.json() + assert len(requests) == 3 + + +async def test_add_multimodal_requires_existing_authentication(add_client): + client, requests = add_client() + async with client: + response = await client.post( + "/add", + headers={"Authorization": "Token wrong-key"}, + json={ + "request_id": "unauthorized-1", + "user_id": "alice", + "session_id": "session-1", + "messages": [{"role": "user", "content": [IMAGE_PART]}], + }, + ) + assert response.status_code == 401 + assert requests == [] + + +@pytest.mark.parametrize("query", ["Where?", "Which function invalidates src/cache.py?"]) +async def test_search_uses_native_find_and_preserves_user_scope_and_aml_response( + adapter_client, query +): seen = [] async def handle(request): @@ -74,13 +307,13 @@ async def handle(request): async with adapter_client(handle) as client: responses = await asyncio.gather( - client.post("/search", json={"user_id": "alice", "query": "Where?", "top_k": 2}), + client.post("/search", json={"user_id": "alice", "query": query, "top_k": 2}), client.post( "/search", json={"user_id": "bob", "query": "Which?", "options": ["A", "B"], "top_k": 2}, ), ) - assert sorted(seen) == [("alice", "Where?"), ("bob", "Which?\n\nAnswer options:\n1. A\n2. B")] + assert sorted(seen) == [("alice", query), ("bob", "Which?\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() == { From 7d8e039fe438851a38a7029f3177d569cd21e4ff Mon Sep 17 00:00:00 2001 From: heaoxiang Date: Thu, 15 Oct 2026 21:02:18 +0800 Subject: [PATCH 2/6] feat(search): route hybrid image queries through multimodal dense embeddings --- benchmark/aml/README.md | 29 +++- benchmark/aml/server.py | 24 ++- docs/en/api/06-retrieval.md | 4 +- docs/en/guides/01-configuration.md | 2 + docs/zh/api/06-retrieval.md | 4 +- docs/zh/guides/01-configuration.md | 2 + openviking/models/embedder/base.py | 44 ++++-- .../models/embedder/volcengine_embedders.py | 141 ++++++++---------- .../retrieve/context_assembler/gather.py | 2 +- .../retrieve/context_assembler/params.py | 2 +- openviking/server/routers/search.py | 10 +- openviking/service/search_service.py | 12 +- openviking/storage/viking_fs/_base.py | 4 +- openviking/storage/viking_fs/_semantic.py | 4 +- openviking/utils/image_search.py | 13 +- sdk/python/openviking_sdk/client.py | 4 +- tests/unit/test_aml_server.py | 112 ++++++++++++++ .../test_composite_multimodal_embedder.py | 89 +++++++++++ tests/unit/test_image_search.py | 20 +++ tests/unit/test_search_service_keywords.py | 9 ++ tests/unit/test_volcengine_embedder_input.py | 105 ++++++++++++- 21 files changed, 508 insertions(+), 128 deletions(-) create mode 100644 tests/unit/test_composite_multimodal_embedder.py diff --git a/benchmark/aml/README.md b/benchmark/aml/README.md index cb72eeb932..9c5be5340f 100644 --- a/benchmark/aml/README.md +++ b/benchmark/aml/README.md @@ -87,9 +87,32 @@ Base64 JPEG, PNG, or WebP Data URIs, with at most 10 MiB decoded per image and 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` still accepts string queries -and returns text evidence; image-query and image-evidence support are not -included in this adapter yet. +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 diff --git a/benchmark/aml/server.py b/benchmark/aml/server.py index cd48c06d4e..e0e5e35ef1 100644 --- a/benchmark/aml/server.py +++ b/benchmark/aml/server.py @@ -28,6 +28,7 @@ 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", @@ -217,14 +218,23 @@ def validate_image_budget(cls, messages: list[AMLMessage]) -> list[AMLMessage]: 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 @@ -322,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", @@ -330,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, @@ -401,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) @@ -411,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..7b556769a1 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -72,9 +72,9 @@ def _discard_unsuccessful( self.pop(key, None) async def close(self) -> None: - pending = { - entry.task for entry in self.values() if not entry.task.done() - } | {task for task in self._pending_tasks if not task.done()} + pending = {entry.task for entry in self.values() if not entry.task.done()} | { + task for task in self._pending_tasks if not task.done() + } for task in pending: task.cancel() if pending: @@ -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..cfc89f7ef3 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -33,9 +33,7 @@ T = TypeVar("T") -def _record_failed_embedding_call( - embedder, *, duration_seconds: float, error: Exception -) -> None: +def _record_failed_embedding_call(embedder, *, duration_seconds: float, error: Exception) -> None: """Emit a failed Ark request attempt with the same model labels as a success.""" try: root_context = get_root_observability_context() @@ -281,9 +279,7 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) vector = response.data.embedding else: # Use text embeddings API (text-only) @@ -296,9 +292,7 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) vector = response.data[0].embedding vector = truncate_and_normalize(vector, self.dimension) @@ -314,9 +308,7 @@ def _embed_call(): raise RuntimeError(f"Volcengine embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get( - lambda: _create_async_ark(self._ark_kwargs) - ) + return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: @@ -332,9 +324,7 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) vector = response.data.embedding else: text = extract_text_from_content(content) @@ -346,9 +336,7 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) vector = response.data[0].embedding return EmbedResult(dense_vector=truncate_and_normalize(vector, self.dimension)) @@ -468,9 +456,7 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) item = response.data sparse_vector = getattr(item, "sparse_embedding", None) return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) @@ -485,9 +471,7 @@ def _embed_call(): raise RuntimeError(f"Volcengine sparse embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get( - lambda: _create_async_ark(self._ark_kwargs) - ) + return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: @@ -503,9 +487,7 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage( - response, duration_seconds=duration_seconds - ) + self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) item = response.data sparse_vector = getattr(item, "sparse_embedding", None) return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) @@ -536,8 +518,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 +597,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 +631,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( @@ -672,34 +659,30 @@ def _embed_call(): raise RuntimeError(f"Volcengine hybrid embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get( - lambda: _create_async_ark(self._ark_kwargs) - ) + return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: 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/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..ef25f2cb28 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", ): @@ -401,7 +401,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, 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/unit/test_aml_server.py b/tests/unit/test_aml_server.py index 548f4339fd..fea50e9e08 100644 --- a/tests/unit/test_aml_server.py +++ b/tests/unit/test_aml_server.py @@ -343,3 +343,115 @@ def handle(request): 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 + + +@pytest.mark.parametrize( + ("query", "text", "images"), + [ + ([IMAGE_PART], "", IMAGE_URL), + ([{"type": "text", "text": "Which receipt?"}, IMAGE_PART], "Which receipt?", IMAGE_URL), + ( + [IMAGE_PART, {"type": "text", "text": "Compare both"}, IMAGE_PART], + "Compare both", + [IMAGE_URL, IMAGE_URL], + ), + ( + [{"type": "text", "text": "First"}, {"type": "text", "text": "Second"}], + "First\nSecond", + None, + ), + ], + ids=["image-only", "mixed", "multiple-images", "text-array"], +) +async def test_search_forwards_images_to_native_find(adapter_client, query, text, images): + seen = [] + + def handle(request): + assert request.url.path == "/api/v1/search/find" + assert request.headers["X-OpenViking-User"] == "alice" + body = json.loads(request.content) + assert body["query"] == text + "\n\nAnswer options:\n1. A\n2. B" + assert body.get("image_url") == images + assert body["target_uri"] == "viking://~/memories" + assert body["context_type"] == ["memory"] + assert body["level"] == 2 + seen.append(body) + return httpx.Response( + 200, + json={ + "status": "ok", + "result": { + "memories": [{"uri": "viking://user/alice/memories/a", "content": "Evidence"}] + }, + }, + ) + + async with adapter_client(handle) as client: + response = await client.post( + "/search", json={"user_id": "alice", "query": query, "options": ["A", "B"], "top_k": 2} + ) + assert response.status_code == 200 + assert response.json() == { + "data": [{"id": "viking://user/alice/memories/a", "content": "Evidence"}] + } + assert len(seen) == 1 + + +@pytest.mark.parametrize( + "query", + [ + [], + " ", + [{"type": "text", "text": " "}], + [{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}], + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%"}}], + ], +) +async def test_search_rejects_invalid_queries_before_native_find(adapter_client, query): + def handle(request): + raise AssertionError("Invalid query must not reach the backend") + + async with adapter_client(handle) as client: + response = await client.post( + "/search", json={"user_id": "alice", "query": query, "top_k": 2} + ) + assert response.status_code == 422 + + +@pytest.mark.parametrize( + ("limit", "maximum", "count", "status"), + [ + ("MAX_IMAGE_BYTES", IMAGE_SIZE - 1, 1, 422), + ("MAX_QUERY_IMAGE_BYTES", 2 * IMAGE_SIZE - 1, 2, 422), + ("MAX_QUERY_IMAGE_BYTES", 2 * IMAGE_SIZE, 2, 200), + ], +) +async def test_search_validates_per_image_and_aggregate_size( + adapter_client, monkeypatch, limit, maximum, count, status +): + monkeypatch.setattr(server, limit, maximum) + seen = [] + + def handle(request): + seen.append(request) + return httpx.Response(200, json={"status": "ok", "result": {"memories": []}}) + + async with adapter_client(handle) as client: + response = await client.post( + "/search", json={"user_id": "alice", "query": [IMAGE_PART] * count, "top_k": 2} + ) + assert response.status_code == status + assert len(seen) == (status == 200) + + +async def test_image_search_requires_existing_authentication(adapter_client): + def handle(request): + raise AssertionError("Unauthorized query must not reach the backend") + + async with adapter_client(handle) as client: + response = await client.post( + "/search", + headers={"Authorization": "Token wrong-key"}, + json={"user_id": "alice", "query": [IMAGE_PART], "top_k": 2}, + ) + assert response.status_code == 401 diff --git a/tests/unit/test_composite_multimodal_embedder.py b/tests/unit/test_composite_multimodal_embedder.py new file mode 100644 index 0000000000..3d6d8b2655 --- /dev/null +++ b/tests/unit/test_composite_multimodal_embedder.py @@ -0,0 +1,89 @@ +"""Hybrid channels keep images in dense input and text in sparse input.""" + +import pytest + +from openviking.models.embedder.base import ( + CompositeHybridEmbedder, + DenseEmbedderBase, + EmbedResult, + SparseEmbedderBase, + embed_compat, +) + + +class Dense(DenseEmbedderBase): + def __init__(self, multimodal=True): + super().__init__("existing-dense") + self.multimodal = multimodal + self.seen = [] + + @property + def supports_multimodal(self): + return self.multimodal + + def embed(self, content, is_query=False): + self.seen.append((content, is_query)) + return EmbedResult(dense_vector=[1.0, 0.0]) + + async def embed_async(self, content, is_query=False): + return self.embed(content, is_query=is_query) + + def get_dimension(self): + return 2 + + +class Sparse(SparseEmbedderBase): + def __init__(self): + super().__init__("existing-sparse") + self.seen = [] + + def embed(self, content, is_query=False): + assert isinstance(content, str) + self.seen.append((content, is_query)) + return EmbedResult(sparse_vector={"word": 0.7}) + + async def embed_async(self, content, is_query=False): + return self.embed(content, is_query=is_query) + + def get_dimension(self): + return 0 + + +IMAGE = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YQ=="}} + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("content", "text"), + [ + ("original\ncode", "original\ncode"), + ( + [{"type": "text", "text": "first"}, IMAGE, {"type": "text", "text": "second"}, IMAGE], + "first\nsecond", + ), + ([IMAGE], None), + ([{"type": "text", "text": " "}, IMAGE], None), + ], + ids=["text", "mixed-multiple-images", "image-only", "blank-text"], +) +async def test_composite_routes_modalities_to_independent_channels(asynchronous, content, text): + dense, sparse = Dense(), Sparse() + embedder = CompositeHybridEmbedder(dense, sparse) + assert embedder.supports_multimodal is True + if asynchronous: + result = await embed_compat(embedder, content, is_query=True) + else: + result = embedder.embed(embedder.prepare_embedding_input(content), is_query=True) + assert dense.seen == [(content, True)] + assert sparse.seen == ([] if text is None else [(text, True)]) + assert result.dense_vector == [1.0, 0.0] + assert result.sparse_vector == (None if text is None else {"word": 0.7}) + assert embedder.get_dimension() == 2 + + +def test_text_only_dense_channel_does_not_advertise_image_support(): + embedder = CompositeHybridEmbedder(Dense(multimodal=False), Sparse()) + assert embedder.supports_multimodal is False + assert ( + embedder.prepare_embedding_input([{"type": "text", "text": "caption"}, IMAGE]) == "caption" + ) diff --git a/tests/unit/test_image_search.py b/tests/unit/test_image_search.py index ddcf4d7d54..fe6c09d8b1 100644 --- a/tests/unit/test_image_search.py +++ b/tests/unit/test_image_search.py @@ -45,6 +45,26 @@ def test_image_helpers_accept_base64_bytes_and_image_uris(tmp_path): assert normalize_client_image_input(image_path).startswith("data:image/jpeg;base64,") +def test_multiple_images_keep_their_order_in_embedding_input_and_sdk(): + from openviking_sdk.client import _normalize_image_input + + images = ["data:image/png;base64,YQ==", "viking://resources/images/b.png"] + assert build_multimodal_embedding_input("compare", images) == [ + {"type": "text", "text": "compare"}, + *[{"type": "image_url", "image_url": {"url": url}} for url in images], + ] + assert normalize_client_image_input(images) == images + assert _normalize_image_input(images) == images + + +def test_find_and_search_request_models_accept_multiple_images(): + from openviking.server.routers.search import FindRequest, SearchRequest + + images = ["data:image/png;base64,YQ==", "data:image/png;base64,Yg=="] + assert FindRequest(image_url=images).image_url == images + assert SearchRequest(image_url=images).image_url == images + + def test_model_image_input_downsamples_jpeg_incompatible_modes(): buf = io.BytesIO() Image.new("LA", (80, 220), color=(255, 255)).save(buf, format="PNG") diff --git a/tests/unit/test_search_service_keywords.py b/tests/unit/test_search_service_keywords.py index 97a63ecdc9..5929bcbd40 100644 --- a/tests/unit/test_search_service_keywords.py +++ b/tests/unit/test_search_service_keywords.py @@ -31,3 +31,12 @@ async def test_keywords_capability_check_propagates_metadata_failures(): supported_modes=("volcengine", "vikingdb"), raise_on_error=True, ) + + +async def test_multiple_query_images_are_resolved_under_same_request_context(): + read = AsyncMock(return_value=b"png") + service = SearchService(SimpleNamespace(read_file_bytes=read)) + ctx = _ctx() + urls = ["data:image/png;base64,YQ==", "viking://resources/photo.png"] + assert await service._resolve_image_url(urls, ctx) == [urls[0], "data:image/png;base64,cG5n"] + read.assert_awaited_once_with(urls[1], ctx=ctx) diff --git a/tests/unit/test_volcengine_embedder_input.py b/tests/unit/test_volcengine_embedder_input.py index 9cd5aad6b1..d07960e3e0 100644 --- a/tests/unit/test_volcengine_embedder_input.py +++ b/tests/unit/test_volcengine_embedder_input.py @@ -3,8 +3,11 @@ """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, patch +import pytest + +from openviking.models.embedder.base import embed_compat from openviking.models.embedder.volcengine_embedders import ( VolcengineDenseEmbedder, VolcengineHybridEmbedder, @@ -27,16 +30,109 @@ 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): +def test_hybrid_embedder_preserves_media_for_dense_channel(_mock_ark): embedder = VolcengineHybridEmbedder(model_name="test", api_key="test-key") + assert embedder.supports_multimodal is True + assert embedder.prepare_embedding_input(_MULTIMODAL_INPUT) == _MULTIMODAL_INPUT + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("content", "expected_requests"), + [ + ( + "original text", + [ + { + "input": [{"type": "text", "text": "original text"}], + "sparse_embedding": {"type": "enabled"}, + } + ], + ), + ( + [{"type": "text", "text": "original text"}], + [ + { + "input": [{"type": "text", "text": "original text"}], + "sparse_embedding": {"type": "enabled"}, + } + ], + ), + ( + _MULTIMODAL_INPUT, + [ + {"input": _MULTIMODAL_INPUT}, + { + "input": [{"type": "text", "text": "diagram of a heat exchanger"}], + "sparse_embedding": {"type": "enabled"}, + }, + ], + ), + (_MULTIMODAL_INPUT[1:], [{"input": _MULTIMODAL_INPUT[1:]}]), + ( + [{"type": "text", "text": " \n"}, *_MULTIMODAL_INPUT[1:]], + [{"input": [{"type": "text", "text": " \n"}, *_MULTIMODAL_INPUT[1:]]}], + ), + ], + ids=["text", "text-array", "mixed", "media-only", "blank-text-and-media"], +) +@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") +async def test_hybrid_routes_dense_media_and_text_sparse_independently( + mock_ark, asynchronous, content, expected_requests +): + responses = [ + SimpleNamespace( + data=SimpleNamespace( + embedding=[1.0, 0.0], sparse_embedding=[{"index": 7, "value": 0.5}] + ), + usage=None, + ), + # The text-only dense vector must never replace the multimodal vector. + SimpleNamespace( + data=SimpleNamespace( + embedding=[0.0, 1.0], sparse_embedding=[{"index": 9, "value": 0.8}] + ), + usage=None, + ), + ] + client = MagicMock() + create = AsyncMock(side_effect=responses) if asynchronous else MagicMock(side_effect=responses) + client.multimodal_embeddings.create = create + mock_ark.return_value = client + embedder = VolcengineHybridEmbedder( + model_name="existing-model", api_key="test-key", dimension=2 + ) + embedder._get_async_client = lambda: client + if asynchronous: + result = await embed_compat(embedder, content, is_query=True) + else: + result = embedder.embed(embedder.prepare_embedding_input(content), is_query=True) + assert result.dense_vector == [1.0, 0.0] + sparse_call = next( + (i for i, request in enumerate(expected_requests) if "sparse_embedding" in request), None + ) + assert result.sparse_vector == ( + None if sparse_call is None else {"7": 0.5} if sparse_call == 0 else {"9": 0.8} + ) + assert [call.kwargs for call in create.call_args_list] == [ + {"model": "existing-model", "extra_headers": embedder.extra_headers, **request} + for request in expected_requests + ] + + +@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") +def test_hybrid_text_input_mode_keeps_multimodal_guard(mock_ark): + embedder = VolcengineHybridEmbedder(model_name="test", api_key="test-key", input_type="text") assert embedder.supports_multimodal is False assert embedder.prepare_embedding_input(_MULTIMODAL_INPUT) == "diagram of a heat exchanger" @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 +141,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( From 940657c768cde5b67d325b95eb3cfdc065fea952 Mon Sep 17 00:00:00 2001 From: heaoxiang Date: Thu, 15 Oct 2026 21:07:49 +0800 Subject: [PATCH 3/6] fix(search): respect explicit scopes for image queries --- openviking/retrieve/hierarchical_retriever.py | 24 +++++++++++++- .../test_hierarchical_retriever_image.py | 31 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/openviking/retrieve/hierarchical_retriever.py b/openviking/retrieve/hierarchical_retriever.py index 6be9e25b8a..a486687eb2 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" @@ -177,7 +194,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 target_dirs + 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/tests/retrieve/test_hierarchical_retriever_image.py b/tests/retrieve/test_hierarchical_retriever_image.py index b02a8814b4..3b23b500ad 100644 --- a/tests/retrieve/test_hierarchical_retriever_image.py +++ b/tests/retrieve/test_hierarchical_retriever_image.py @@ -3,6 +3,7 @@ from openviking.models.embedder.base import EmbedResult from openviking.retrieve.hierarchical_retriever import HierarchicalRetriever from openviking.server.identity import RequestContext, Role +from openviking.storage.expr import And, Eq, RawDSL from openviking_cli.exceptions import InvalidArgumentError from openviking_cli.retrieve.types import TypedQuery from openviking_cli.session.user_id import UserIdentifier @@ -190,6 +191,36 @@ def _image_query(): ) +@pytest.mark.parametrize( + ("targets", "scope", "expected_type"), + [ + (["viking://user/user/memories"], None, None), + ([], {"op": "must", "field": "context_type", "conds": ["memory"]}, None), + ( + [], + {"op": "and", "conds": [{"op": "must", "field": "context_type", "conds": ["memory"]}]}, + None, + ), + ([], And([Eq("context_type", "memory")]), None), + ([], RawDSL({"op": "must", "field": "context_type", "conds": ["memory"]}), None), + ([], {"op": "must", "field": "tags", "conds": ["cat"]}, "resource"), + ], +) +async def test_image_query_respects_explicit_scope_before_resource_default( + monkeypatch, targets, scope, expected_type +): + monkeypatch.setattr( + "openviking.retrieve.hierarchical_retriever.VikingDBManagerProxy", FakeProxy + ) + query = _image_query() + query.target_directories = targets + retriever = HierarchicalRetriever(storage=object(), embedder=MultimodalEmbedder()) + await retriever.retrieve(query, ctx=_ctx(), scope_dsl=scope) + assert FakeProxy.captured["context_type"] == expected_type + assert FakeProxy.captured["target_directories"] == targets + assert FakeProxy.captured["extra_filter"] == scope + + @pytest.mark.asyncio async def test_image_query_through_account_bound_multimodal_embedder(account_embedding_provider): _manager, provider, clients = account_embedding_provider From bdfb558d0c57b14bd4d0ceb079896f99d29716bf Mon Sep 17 00:00:00 2001 From: heaoxiang Date: Fri, 16 Oct 2026 11:32:10 +0800 Subject: [PATCH 4/6] chore: remove unrelated formatting changes from multimodal PR --- openviking/models/embedder/base.py | 6 +-- .../models/embedder/volcengine_embedders.py | 40 ++++++++++++++----- 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index 7b556769a1..534d1530ba 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -72,9 +72,9 @@ def _discard_unsuccessful( self.pop(key, None) async def close(self) -> None: - pending = {entry.task for entry in self.values() if not entry.task.done()} | { - task for task in self._pending_tasks if not task.done() - } + pending = { + entry.task for entry in self.values() if not entry.task.done() + } | {task for task in self._pending_tasks if not task.done()} for task in pending: task.cancel() if pending: diff --git a/openviking/models/embedder/volcengine_embedders.py b/openviking/models/embedder/volcengine_embedders.py index cfc89f7ef3..37fdf2febc 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -33,7 +33,9 @@ T = TypeVar("T") -def _record_failed_embedding_call(embedder, *, duration_seconds: float, error: Exception) -> None: +def _record_failed_embedding_call( + embedder, *, duration_seconds: float, error: Exception +) -> None: """Emit a failed Ark request attempt with the same model labels as a success.""" try: root_context = get_root_observability_context() @@ -279,7 +281,9 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) vector = response.data.embedding else: # Use text embeddings API (text-only) @@ -292,7 +296,9 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) vector = response.data[0].embedding vector = truncate_and_normalize(vector, self.dimension) @@ -308,7 +314,9 @@ def _embed_call(): raise RuntimeError(f"Volcengine embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) + return self._async_client_cache.get( + lambda: _create_async_ark(self._ark_kwargs) + ) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: @@ -324,7 +332,9 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) vector = response.data.embedding else: text = extract_text_from_content(content) @@ -336,7 +346,9 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) vector = response.data[0].embedding return EmbedResult(dense_vector=truncate_and_normalize(vector, self.dimension)) @@ -456,7 +468,9 @@ def _embed_call(): extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) item = response.data sparse_vector = getattr(item, "sparse_embedding", None) return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) @@ -471,7 +485,9 @@ def _embed_call(): raise RuntimeError(f"Volcengine sparse embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) + return self._async_client_cache.get( + lambda: _create_async_ark(self._ark_kwargs) + ) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: @@ -487,7 +503,9 @@ async def _embed_call() -> EmbedResult: extra_headers=self.extra_headers, ), ) - self._update_telemetry_token_usage(response, duration_seconds=duration_seconds) + self._update_telemetry_token_usage( + response, duration_seconds=duration_seconds + ) item = response.data sparse_vector = getattr(item, "sparse_embedding", None) return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) @@ -659,7 +677,9 @@ def _embed_call(): raise RuntimeError(f"Volcengine hybrid embedding failed: {str(e)}") from e def _get_async_client(self): - return self._async_client_cache.get(lambda: _create_async_ark(self._ark_kwargs)) + return self._async_client_cache.get( + lambda: _create_async_ark(self._ark_kwargs) + ) @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: From d8283a541e831ecaf356b2d93df8a8f61ebb70b2 Mon Sep 17 00:00:00 2001 From: heaoxiang Date: Fri, 16 Oct 2026 11:46:02 +0800 Subject: [PATCH 5/6] fix(search): preserve default image scope after target expansion --- openviking/retrieve/hierarchical_retriever.py | 6 +- openviking/storage/viking_fs/_semantic.py | 2 + .../misc/test_vikingfs_find_without_rerank.py | 88 +++++++++++++++++++ 3 files changed, 95 insertions(+), 1 deletion(-) diff --git a/openviking/retrieve/hierarchical_retriever.py b/openviking/retrieve/hierarchical_retriever.py index a486687eb2..ed25d96170 100644 --- a/openviking/retrieve/hierarchical_retriever.py +++ b/openviking/retrieve/hierarchical_retriever.py @@ -119,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. @@ -131,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() @@ -157,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( @@ -197,7 +201,7 @@ async def retrieve( if ( image_query and context_type is None - and not target_dirs + and not has_explicit_target and not _has_context_type_filter(scope_dsl) ): context_type = ContextType.RESOURCE.value diff --git a/openviking/storage/viking_fs/_semantic.py b/openviking/storage/viking_fs/_semantic.py index ef25f2cb28..2c45873f58 100644 --- a/openviking/storage/viking_fs/_semantic.py +++ b/openviking/storage/viking_fs/_semantic.py @@ -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, @@ -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/tests/misc/test_vikingfs_find_without_rerank.py b/tests/misc/test_vikingfs_find_without_rerank.py index 448ad43b5b..96f69818bb 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 @@ -132,6 +134,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["mode"] = mode @@ -169,6 +172,91 @@ async def retrieve( assert captured["mode"] == RetrieverMode.QUICK +@pytest.mark.parametrize("method", ["find", "search"]) +@pytest.mark.parametrize("role", [Role.USER, Role.ADMIN, Role.ROOT]) +@pytest.mark.parametrize( + ("query", "image_url"), + [ + ("", "data:image/png;base64,abc"), + ("receipt", "data:image/png;base64,abc"), + ("remember tea", None), + ("def cache_key(user_id): return user_id", None), + ], + ids=["image", "mixed", "text", "code"], +) +@pytest.mark.parametrize( + ("target_uri", "scope", "expected_image_type"), + [ + ("", None, "resource"), + ("viking://user/user1/memories", None, None), + ("", {"op": "must", "field": "context_type", "conds": ["memory"]}, None), + ], + ids=["default", "memory-uri", "memory-type"], +) +async def test_find_and_search_preserve_default_and_explicit_image_scopes( + monkeypatch, method, role, query, image_url, target_uri, scope, expected_image_type +): + fs = _make_viking_fs() + fs._ctx_or_default.side_effect = lambda ctx: ctx + 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])) + captured = {} + + class FakeProxy: + def __init__(self, _storage, _ctx): + pass + + async def collection_exists_bound(self): + return True + + async def search_in_tenant(self, **kwargs): + captured.update(kwargs) + context_type = kwargs["context_type"] or ("memory" if scope else None) + directories = kwargs["target_directories"] + rows = [ + { + "uri": "viking://user/user1/memories/receipt.md", + "context_type": "memory", + "_score": 0.9, + }, + { + "uri": "viking://resources/receipt.png", + "context_type": "resource", + "_score": 0.8, + }, + ] + return [ + row + for row in rows + if (context_type is None or row["context_type"] == context_type) + and (not directories or any(row["uri"].startswith(d + "/") for d in directories)) + ][: kwargs["limit"]] + + monkeypatch.setattr( + "openviking.retrieve.hierarchical_retriever.VikingDBManagerProxy", FakeProxy + ) + result = await getattr(fs, method)( + query, + image_url=image_url, + target_uri=target_uri, + filter=scope, + ctx=RequestContext(user=UserIdentifier("acc1", "user1"), role=role), + limit=1, + ) + + assert captured["context_type"] == (expected_image_type if image_url else None) + assert captured["extra_filter"] == scope + assert result.total == 1 + if image_url and expected_image_type == "resource": + assert [hit.uri for hit in result.resources] == ["viking://resources/receipt.png"] + assert result.memories == [] + else: + assert [hit.uri for hit in result.memories] == ["viking://user/user1/memories/receipt.md"] + assert result.resources == [] + + @pytest.mark.asyncio async def test_find_uses_quick_mode_with_rerank_config(monkeypatch) -> None: fs = _make_viking_fs() From 4b39bb5fb98e109ac93e20d28568a244b0c8881d Mon Sep 17 00:00:00 2001 From: qin-ctx Date: Fri, 16 Oct 2026 14:44:55 +0800 Subject: [PATCH 6/6] test(search): consolidate multimodal contract coverage --- .../misc/test_vikingfs_find_without_rerank.py | 175 ++----- .../test_hierarchical_retriever_image.py | 31 -- tests/unit/test_aml_server.py | 438 ++++-------------- .../test_composite_multimodal_embedder.py | 89 ---- tests/unit/test_image_search.py | 27 +- tests/unit/test_retry_embedder.py | 18 +- tests/unit/test_search_service_keywords.py | 9 - tests/unit/test_volcengine_embedder_input.py | 143 ++---- 8 files changed, 197 insertions(+), 733 deletions(-) delete mode 100644 tests/unit/test_composite_multimodal_embedder.py diff --git a/tests/misc/test_vikingfs_find_without_rerank.py b/tests/misc/test_vikingfs_find_without_rerank.py index 96f69818bb..f5467e9b70 100644 --- a/tests/misc/test_vikingfs_find_without_rerank.py +++ b/tests/misc/test_vikingfs_find_without_rerank.py @@ -114,147 +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", - has_explicit_target=None, - ): - 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"], - ) - - monkeypatch.setattr( - "openviking.retrieve.hierarchical_retriever.HierarchicalRetriever", - FakeRetriever, - ) - - result = await fs.find( - "", - target_uri="viking://resources", - image_url="data:image/png;base64,abc", - ctx=_ctx(), - ) - - 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"}} - ] - assert captured["mode"] == RetrieverMode.QUICK - - -@pytest.mark.parametrize("method", ["find", "search"]) -@pytest.mark.parametrize("role", [Role.USER, Role.ADMIN, Role.ROOT]) -@pytest.mark.parametrize( - ("query", "image_url"), - [ - ("", "data:image/png;base64,abc"), - ("receipt", "data:image/png;base64,abc"), - ("remember tea", None), - ("def cache_key(user_id): return user_id", None), - ], - ids=["image", "mixed", "text", "code"], -) -@pytest.mark.parametrize( - ("target_uri", "scope", "expected_image_type"), - [ - ("", None, "resource"), - ("viking://user/user1/memories", None, None), - ("", {"op": "must", "field": "context_type", "conds": ["memory"]}, None), - ], - ids=["default", "memory-uri", "memory-type"], -) -async def test_find_and_search_preserve_default_and_explicit_image_scopes( - monkeypatch, method, role, query, image_url, target_uri, scope, expected_image_type -): - fs = _make_viking_fs() - fs._ctx_or_default.side_effect = lambda ctx: ctx 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])) - captured = {} - - class FakeProxy: - def __init__(self, _storage, _ctx): - pass - - async def collection_exists_bound(self): - return True - - async def search_in_tenant(self, **kwargs): - captured.update(kwargs) - context_type = kwargs["context_type"] or ("memory" if scope else None) - directories = kwargs["target_directories"] - rows = [ - { - "uri": "viking://user/user1/memories/receipt.md", - "context_type": "memory", - "_score": 0.9, - }, - { - "uri": "viking://resources/receipt.png", - "context_type": "resource", - "_score": 0.8, - }, - ] - return [ - row - for row in rows - if (context_type is None or row["context_type"] == context_type) - and (not directories or any(row["uri"].startswith(d + "/") for d in directories)) - ][: kwargs["limit"]] - + 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.VikingDBManagerProxy", FakeProxy + "openviking.retrieve.hierarchical_retriever.VikingDBManagerProxy", + MagicMock(return_value=proxy), ) - result = await getattr(fs, method)( - query, - image_url=image_url, - target_uri=target_uri, - filter=scope, - ctx=RequestContext(user=UserIdentifier("acc1", "user1"), role=role), - limit=1, + 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 captured["context_type"] == (expected_image_type if image_url else None) - assert captured["extra_filter"] == scope - assert result.total == 1 - if image_url and expected_image_type == "resource": - assert [hit.uri for hit in result.resources] == ["viking://resources/receipt.png"] - assert result.memories == [] - else: - assert [hit.uri for hit in result.memories] == ["viking://user/user1/memories/receipt.md"] - assert result.resources == [] + 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} + ] + 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/retrieve/test_hierarchical_retriever_image.py b/tests/retrieve/test_hierarchical_retriever_image.py index 3b23b500ad..b02a8814b4 100644 --- a/tests/retrieve/test_hierarchical_retriever_image.py +++ b/tests/retrieve/test_hierarchical_retriever_image.py @@ -3,7 +3,6 @@ from openviking.models.embedder.base import EmbedResult from openviking.retrieve.hierarchical_retriever import HierarchicalRetriever from openviking.server.identity import RequestContext, Role -from openviking.storage.expr import And, Eq, RawDSL from openviking_cli.exceptions import InvalidArgumentError from openviking_cli.retrieve.types import TypedQuery from openviking_cli.session.user_id import UserIdentifier @@ -191,36 +190,6 @@ def _image_query(): ) -@pytest.mark.parametrize( - ("targets", "scope", "expected_type"), - [ - (["viking://user/user/memories"], None, None), - ([], {"op": "must", "field": "context_type", "conds": ["memory"]}, None), - ( - [], - {"op": "and", "conds": [{"op": "must", "field": "context_type", "conds": ["memory"]}]}, - None, - ), - ([], And([Eq("context_type", "memory")]), None), - ([], RawDSL({"op": "must", "field": "context_type", "conds": ["memory"]}), None), - ([], {"op": "must", "field": "tags", "conds": ["cat"]}, "resource"), - ], -) -async def test_image_query_respects_explicit_scope_before_resource_default( - monkeypatch, targets, scope, expected_type -): - monkeypatch.setattr( - "openviking.retrieve.hierarchical_retriever.VikingDBManagerProxy", FakeProxy - ) - query = _image_query() - query.target_directories = targets - retriever = HierarchicalRetriever(storage=object(), embedder=MultimodalEmbedder()) - await retriever.retrieve(query, ctx=_ctx(), scope_dsl=scope) - assert FakeProxy.captured["context_type"] == expected_type - assert FakeProxy.captured["target_directories"] == targets - assert FakeProxy.captured["extra_filter"] == scope - - @pytest.mark.asyncio async def test_image_query_through_account_bound_multimodal_embedder(account_embedding_provider): _manager, provider, clients = account_embedding_provider diff --git a/tests/unit/test_aml_server.py b/tests/unit/test_aml_server.py index fea50e9e08..fddc8f48c5 100644 --- a/tests/unit/test_aml_server.py +++ b/tests/unit/test_aml_server.py @@ -1,7 +1,6 @@ """Competition API contract through the real OpenViking SDK request builder.""" import asyncio -import base64 import json import httpx @@ -10,12 +9,7 @@ from benchmark.aml import server from benchmark.aml.server import AMLSettings, create_app -IMAGE_URL = ( - "data:image/png;base64," - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a6lcAAAAASUVORK5CYII=" -) -IMAGE_PART = {"type": "image_url", "image_url": {"url": IMAGE_URL}} -IMAGE_SIZE = len(base64.b64decode(IMAGE_URL.partition(",")[2])) +IMAGES = ["data:image/png;base64,YQ==", "data:image/png;base64,Yg=="] @pytest.fixture @@ -35,7 +29,6 @@ def sdk_client(**kwargs): openviking_api_key="test-ov-key", aml_api_key="test-aml-key", retry_delay_seconds=0.001, - task_poll_seconds=0.001, ) ) return client_type( @@ -47,230 +40,7 @@ def sdk_client(**kwargs): return build -@pytest.fixture -def add_client(adapter_client): - def build(task_status="completed"): - requests = [] - - def handle(request): - requests.append(request) - if request.url.path == "/api/v1/sessions/session-1/messages/batch": - result = {"added": len(json.loads(request.content)["messages"])} - 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"} - else: - assert request.method == "GET" - assert request.url.path == "/api/v1/tasks/task-1" - result = {"status": task_status, "error": "extraction failed"} - return httpx.Response(200, json={"status": "ok", "result": result}) - - return adapter_client(handle), requests - - return build - - -@pytest.mark.parametrize( - "content", - [ - " Remember my preference for tea.\n", - "def cache_key(user_id):\n return f'cache:{user_id}'\n", - ], - ids=["textual", "coding"], -) -async def test_add_string_payload_and_commit_are_unchanged(add_client, content): - client, requests = add_client() - payload = { - "request_id": "request-1", - "user_id": "alice", - "session_id": "session-1", - "messages": [ - {"role": "user", "content": content, "timestamp": 0}, - {"role": "assistant", "content": "Acknowledged."}, - ], - } - async with client: - response = await client.post("/add", json=payload) - assert response.status_code == 200 - assert response.json() == { - "success": True, - "request_id": "request-1", - "user_id": "alice", - "session_id": "session-1", - } - assert json.loads(requests[0].content) == { - "messages": [ - {"role": "user", "content": content, "created_at": "1970-01-01T00:00:00.000Z"}, - {"role": "assistant", "content": "Acknowledged."}, - ] - } - 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", - ] - assert all(request.headers["X-OpenViking-User"] == "alice" for request in requests) - - -@pytest.mark.parametrize( - "parts", - [ - [IMAGE_PART], - [ - {"type": "text", "text": " First image:\n"}, - IMAGE_PART, - {"type": "text", "text": "The same image again:"}, - IMAGE_PART, - ], - [{"type": "text", "text": "Text supplied as a content array."}], - ], - ids=["image-only", "ordered-multiple-images", "text-array"], -) -async def test_add_forwards_content_array_as_native_parts(add_client, parts): - client, requests = add_client() - async with client: - response = await client.post( - "/add", - json={ - "request_id": "multimodal-1", - "user_id": "bob", - "session_id": "session-1", - "messages": [ - {"role": "user", "content": "A plain message."}, - {"role": "assistant", "content": parts, "timestamp": 1704067200000}, - ], - }, - ) - assert response.status_code == 200 - assert response.json() == { - "success": True, - "request_id": "multimodal-1", - "user_id": "bob", - "session_id": "session-1", - } - assert json.loads(requests[0].content) == { - "messages": [ - {"role": "user", "content": "A plain message."}, - {"role": "assistant", "parts": parts, "created_at": "2024-01-01T00:00:00.000Z"}, - ] - } - assert len(requests) == 3 - assert all(request.headers["X-OpenViking-User"] == "bob" for request in requests) - assert all(request.headers["X-OpenViking-Account"] == "competition" for request in requests) - assert all(request.headers["X-API-Key"] == "test-ov-key" for request in requests) - - -@pytest.mark.parametrize( - "content", - [ - "", - " \n ", - [], - [{"type": "text", "text": " "}], - [{"type": "image_url", "image_url": {}}], - [{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}], - [{"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%"}}], - [{"type": "image_url", "image_url": {"url": "data:image/gif;base64,YQ=="}}], - [{"type": "image_url", "image_url": {"url": "data:image/png;base64,"}}], - [{"type": "tool", "tool_name": "find"}], - ], - ids=[ - "empty-string", - "blank-string", - "empty-array", - "blank-text-part", - "missing-image-url", - "remote-image", - "invalid-base64", - "unsupported-mime", - "empty-image", - "unsupported-part", - ], -) -async def test_add_rejects_invalid_content_before_writing(add_client, content): - client, requests = add_client() - async with client: - response = await client.post( - "/add", - json={ - "request_id": "invalid-1", - "user_id": "alice", - "session_id": "session-1", - "messages": [{"role": "user", "content": content}], - }, - ) - assert response.status_code == 422 - assert requests == [] - - -@pytest.mark.parametrize( - ("limit", "max_bytes", "image_count", "expected_status"), - [ - ("MAX_IMAGE_BYTES", IMAGE_SIZE - 1, 1, 422), - ("MAX_IMAGE_BYTES", IMAGE_SIZE, 1, 200), - ("MAX_ADD_IMAGE_BYTES", IMAGE_SIZE * 2 - 1, 2, 422), - ("MAX_ADD_IMAGE_BYTES", IMAGE_SIZE * 2, 2, 200), - ], - ids=["image-over-limit", "image-at-limit", "request-over-limit", "request-at-limit"], -) -async def test_add_enforces_image_budget_across_messages( - add_client, monkeypatch, limit, max_bytes, image_count, expected_status -): - monkeypatch.setattr(server, limit, max_bytes) - client, requests = add_client() - async with client: - response = await client.post( - "/add", - json={ - "request_id": "oversized-1", - "user_id": "alice", - "session_id": "session-1", - "messages": [{"role": "user", "content": [IMAGE_PART]} for _ in range(image_count)], - }, - ) - assert response.status_code == expected_status - if expected_status == 422: - assert requests == [] - - -async def test_add_multimodal_extraction_failure_is_not_reported_as_success(add_client): - client, requests = add_client(task_status="failed") - async with client: - response = await client.post( - "/add", - json={ - "request_id": "failed-1", - "user_id": "alice", - "session_id": "session-1", - "messages": [{"role": "user", "content": [IMAGE_PART]}], - }, - ) - assert response.status_code == 503 - assert "success" not in response.json() - assert len(requests) == 3 - - -async def test_add_multimodal_requires_existing_authentication(add_client): - client, requests = add_client() - async with client: - response = await client.post( - "/add", - headers={"Authorization": "Token wrong-key"}, - json={ - "request_id": "unauthorized-1", - "user_id": "alice", - "session_id": "session-1", - "messages": [{"role": "user", "content": [IMAGE_PART]}], - }, - ) - assert response.status_code == 401 - assert requests == [] - - -@pytest.mark.parametrize("query", ["Where?", "Which function invalidates src/cache.py?"]) -async def test_search_uses_native_find_and_preserves_user_scope_and_aml_response( - adapter_client, query -): +async def test_search_uses_native_find_and_preserves_user_scope_and_aml_response(adapter_client): seen = [] async def handle(request): @@ -287,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( @@ -307,13 +78,26 @@ async def handle(request): async with adapter_client(handle) as client: responses = await asyncio.gather( - client.post("/search", json={"user_id": "alice", "query": query, "top_k": 2}), + 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", query), ("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() == { @@ -324,134 +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": []}}) - - async with adapter_client(handle) as client: - response = await client.post( - "/search", json={"user_id": "alice", "query": "Where?", "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 - - -@pytest.mark.parametrize( - ("query", "text", "images"), - [ - ([IMAGE_PART], "", IMAGE_URL), - ([{"type": "text", "text": "Which receipt?"}, IMAGE_PART], "Which receipt?", IMAGE_URL), - ( - [IMAGE_PART, {"type": "text", "text": "Compare both"}, IMAGE_PART], - "Compare both", - [IMAGE_URL, IMAGE_URL], - ), - ( - [{"type": "text", "text": "First"}, {"type": "text", "text": "Second"}], - "First\nSecond", - None, - ), - ], - ids=["image-only", "mixed", "multiple-images", "text-array"], -) -async def test_search_forwards_images_to_native_find(adapter_client, query, text, images): - seen = [] - - def handle(request): - assert request.url.path == "/api/v1/search/find" assert request.headers["X-OpenViking-User"] == "alice" - body = json.loads(request.content) - assert body["query"] == text + "\n\nAnswer options:\n1. A\n2. B" - assert body.get("image_url") == images - assert body["target_uri"] == "viking://~/memories" - assert body["context_type"] == ["memory"] - assert body["level"] == 2 - seen.append(body) - return httpx.Response( - 200, - json={ - "status": "ok", - "result": { - "memories": [{"uri": "viking://user/alice/memories/a", "content": "Evidence"}] - }, - }, - ) + 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": query, "options": ["A", "B"], "top_k": 2} + "/search", json={"user_id": "alice", "query": [parts[1]], "top_k": 2} ) assert response.status_code == 200 - assert response.json() == { - "data": [{"id": "viking://user/alice/memories/a", "content": "Evidence"}] - } - assert len(seen) == 1 - - -@pytest.mark.parametrize( - "query", - [ - [], - " ", - [{"type": "text", "text": " "}], - [{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}], - [{"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%"}}], - ], -) -async def test_search_rejects_invalid_queries_before_native_find(adapter_client, query): - def handle(request): - raise AssertionError("Invalid query must not reach the backend") - - async with adapter_client(handle) as client: - response = await client.post( - "/search", json={"user_id": "alice", "query": query, "top_k": 2} - ) - assert response.status_code == 422 - - -@pytest.mark.parametrize( - ("limit", "maximum", "count", "status"), - [ - ("MAX_IMAGE_BYTES", IMAGE_SIZE - 1, 1, 422), - ("MAX_QUERY_IMAGE_BYTES", 2 * IMAGE_SIZE - 1, 2, 422), - ("MAX_QUERY_IMAGE_BYTES", 2 * IMAGE_SIZE, 2, 200), - ], -) -async def test_search_validates_per_image_and_aggregate_size( - adapter_client, monkeypatch, limit, maximum, count, status -): - monkeypatch.setattr(server, limit, maximum) - seen = [] - - def handle(request): - seen.append(request) - return httpx.Response(200, json={"status": "ok", "result": {"memories": []}}) - - async with adapter_client(handle) as client: - response = await client.post( - "/search", json={"user_id": "alice", "query": [IMAGE_PART] * count, "top_k": 2} - ) - assert response.status_code == status - assert len(seen) == (status == 200) - - -async def test_image_search_requires_existing_authentication(adapter_client): - def handle(request): - raise AssertionError("Unauthorized query must not reach the backend") - - async with adapter_client(handle) as client: - response = await client.post( - "/search", - headers={"Authorization": "Token wrong-key"}, - json={"user_id": "alice", "query": [IMAGE_PART], "top_k": 2}, - ) - assert response.status_code == 401 + assert response.json() == {"data": []} + 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_composite_multimodal_embedder.py b/tests/unit/test_composite_multimodal_embedder.py deleted file mode 100644 index 3d6d8b2655..0000000000 --- a/tests/unit/test_composite_multimodal_embedder.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Hybrid channels keep images in dense input and text in sparse input.""" - -import pytest - -from openviking.models.embedder.base import ( - CompositeHybridEmbedder, - DenseEmbedderBase, - EmbedResult, - SparseEmbedderBase, - embed_compat, -) - - -class Dense(DenseEmbedderBase): - def __init__(self, multimodal=True): - super().__init__("existing-dense") - self.multimodal = multimodal - self.seen = [] - - @property - def supports_multimodal(self): - return self.multimodal - - def embed(self, content, is_query=False): - self.seen.append((content, is_query)) - return EmbedResult(dense_vector=[1.0, 0.0]) - - async def embed_async(self, content, is_query=False): - return self.embed(content, is_query=is_query) - - def get_dimension(self): - return 2 - - -class Sparse(SparseEmbedderBase): - def __init__(self): - super().__init__("existing-sparse") - self.seen = [] - - def embed(self, content, is_query=False): - assert isinstance(content, str) - self.seen.append((content, is_query)) - return EmbedResult(sparse_vector={"word": 0.7}) - - async def embed_async(self, content, is_query=False): - return self.embed(content, is_query=is_query) - - def get_dimension(self): - return 0 - - -IMAGE = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YQ=="}} - - -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize( - ("content", "text"), - [ - ("original\ncode", "original\ncode"), - ( - [{"type": "text", "text": "first"}, IMAGE, {"type": "text", "text": "second"}, IMAGE], - "first\nsecond", - ), - ([IMAGE], None), - ([{"type": "text", "text": " "}, IMAGE], None), - ], - ids=["text", "mixed-multiple-images", "image-only", "blank-text"], -) -async def test_composite_routes_modalities_to_independent_channels(asynchronous, content, text): - dense, sparse = Dense(), Sparse() - embedder = CompositeHybridEmbedder(dense, sparse) - assert embedder.supports_multimodal is True - if asynchronous: - result = await embed_compat(embedder, content, is_query=True) - else: - result = embedder.embed(embedder.prepare_embedding_input(content), is_query=True) - assert dense.seen == [(content, True)] - assert sparse.seen == ([] if text is None else [(text, True)]) - assert result.dense_vector == [1.0, 0.0] - assert result.sparse_vector == (None if text is None else {"word": 0.7}) - assert embedder.get_dimension() == 2 - - -def test_text_only_dense_channel_does_not_advertise_image_support(): - embedder = CompositeHybridEmbedder(Dense(multimodal=False), Sparse()) - assert embedder.supports_multimodal is False - assert ( - embedder.prepare_embedding_input([{"type": "text", "text": "caption"}, IMAGE]) == "caption" - ) diff --git a/tests/unit/test_image_search.py b/tests/unit/test_image_search.py index fe6c09d8b1..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 @@ -45,26 +50,6 @@ def test_image_helpers_accept_base64_bytes_and_image_uris(tmp_path): assert normalize_client_image_input(image_path).startswith("data:image/jpeg;base64,") -def test_multiple_images_keep_their_order_in_embedding_input_and_sdk(): - from openviking_sdk.client import _normalize_image_input - - images = ["data:image/png;base64,YQ==", "viking://resources/images/b.png"] - assert build_multimodal_embedding_input("compare", images) == [ - {"type": "text", "text": "compare"}, - *[{"type": "image_url", "image_url": {"url": url}} for url in images], - ] - assert normalize_client_image_input(images) == images - assert _normalize_image_input(images) == images - - -def test_find_and_search_request_models_accept_multiple_images(): - from openviking.server.routers.search import FindRequest, SearchRequest - - images = ["data:image/png;base64,YQ==", "data:image/png;base64,Yg=="] - assert FindRequest(image_url=images).image_url == images - assert SearchRequest(image_url=images).image_url == images - - def test_model_image_input_downsamples_jpeg_incompatible_modes(): buf = io.BytesIO() Image.new("LA", (80, 220), color=(255, 255)).save(buf, format="PNG") 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_search_service_keywords.py b/tests/unit/test_search_service_keywords.py index 5929bcbd40..97a63ecdc9 100644 --- a/tests/unit/test_search_service_keywords.py +++ b/tests/unit/test_search_service_keywords.py @@ -31,12 +31,3 @@ async def test_keywords_capability_check_propagates_metadata_failures(): supported_modes=("volcengine", "vikingdb"), raise_on_error=True, ) - - -async def test_multiple_query_images_are_resolved_under_same_request_context(): - read = AsyncMock(return_value=b"png") - service = SearchService(SimpleNamespace(read_file_bytes=read)) - ctx = _ctx() - urls = ["data:image/png;base64,YQ==", "viking://resources/photo.png"] - assert await service._resolve_image_url(urls, ctx) == [urls[0], "data:image/png;base64,cG5n"] - read.assert_awaited_once_with(urls[1], ctx=ctx) diff --git a/tests/unit/test_volcengine_embedder_input.py b/tests/unit/test_volcengine_embedder_input.py index d07960e3e0..bfac65c746 100644 --- a/tests/unit/test_volcengine_embedder_input.py +++ b/tests/unit/test_volcengine_embedder_input.py @@ -3,9 +3,7 @@ """Input-guard tests for Volcengine sparse-capable embedding modes.""" from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest +from unittest.mock import AsyncMock, MagicMock, call, patch from openviking.models.embedder.base import embed_compat from openviking.models.embedder.volcengine_embedders import ( @@ -30,102 +28,61 @@ def test_sparse_embedder_downgrades_multimodal_input_to_text(_mock_ark): @patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") -def test_hybrid_embedder_preserves_media_for_dense_channel(_mock_ark): - embedder = VolcengineHybridEmbedder(model_name="test", api_key="test-key") - - assert embedder.supports_multimodal is True - assert embedder.prepare_embedding_input(_MULTIMODAL_INPUT) == _MULTIMODAL_INPUT - - -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize( - ("content", "expected_requests"), - [ - ( - "original text", - [ - { - "input": [{"type": "text", "text": "original text"}], - "sparse_embedding": {"type": "enabled"}, - } - ], - ), - ( - [{"type": "text", "text": "original text"}], - [ - { - "input": [{"type": "text", "text": "original text"}], - "sparse_embedding": {"type": "enabled"}, - } - ], - ), - ( - _MULTIMODAL_INPUT, - [ - {"input": _MULTIMODAL_INPUT}, - { - "input": [{"type": "text", "text": "diagram of a heat exchanger"}], - "sparse_embedding": {"type": "enabled"}, - }, - ], - ), - (_MULTIMODAL_INPUT[1:], [{"input": _MULTIMODAL_INPUT[1:]}]), - ( - [{"type": "text", "text": " \n"}, *_MULTIMODAL_INPUT[1:]], - [{"input": [{"type": "text", "text": " \n"}, *_MULTIMODAL_INPUT[1:]]}], - ), - ], - ids=["text", "text-array", "mixed", "media-only", "blank-text-and-media"], -) -@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") -async def test_hybrid_routes_dense_media_and_text_sparse_independently( - mock_ark, asynchronous, content, expected_requests -): - responses = [ - SimpleNamespace( - data=SimpleNamespace( - embedding=[1.0, 0.0], sparse_embedding=[{"index": 7, "value": 0.5}] - ), - usage=None, - ), - # The text-only dense vector must never replace the multimodal vector. - SimpleNamespace( - data=SimpleNamespace( - embedding=[0.0, 1.0], sparse_embedding=[{"index": 9, "value": 0.8}] - ), - usage=None, - ), - ] - client = MagicMock() - create = AsyncMock(side_effect=responses) if asynchronous else MagicMock(side_effect=responses) - client.multimodal_embeddings.create = create - mock_ark.return_value = client +@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 ) - embedder._get_async_client = lambda: client - if asynchronous: - result = await embed_compat(embedder, content, is_query=True) - else: - result = embedder.embed(embedder.prepare_embedding_input(content), is_query=True) - assert result.dense_vector == [1.0, 0.0] - sparse_call = next( - (i for i, request in enumerate(expected_requests) if "sparse_embedding" in request), None - ) - assert result.sparse_vector == ( - None if sparse_call is None else {"7": 0.5} if sparse_call == 0 else {"9": 0.8} - ) - assert [call.kwargs for call in create.call_args_list] == [ - {"model": "existing-model", "extra_headers": embedder.extra_headers, **request} - for request in expected_requests + 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"}, + ), ] + 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:], + ) + ] -@patch("openviking.models.embedder.volcengine_embedders.volcenginesdkarkruntime.Ark") -def test_hybrid_text_input_mode_keeps_multimodal_guard(mock_ark): - embedder = VolcengineHybridEmbedder(model_name="test", api_key="test-key", input_type="text") - assert embedder.supports_multimodal is False - assert embedder.prepare_embedding_input(_MULTIMODAL_INPUT) == "diagram of a heat exchanger" + 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")