tests/unit/test_context_vectorize.py (1)

4-4: Unused pytest import.

The pytest import on line 4 is not used in any of the tests. Consider removing it.

🧹 Proposed fix

 # SPDX-License-Identifier: Apache-2.0
 """Tests for ModalContent and Vectorize media extension."""
-import pytest
 from openviking.core.context import ModalContent, Vectorize

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit/test_context_vectorize.py` at line 4, Remove the unused top-level
import of pytest (the "import pytest" statement) since it isn't referenced
anywhere in the test file; simply delete that import line to eliminate the
unused import warning and keep the test module clean.

examples/ov.conf.gemini.example (1)

5-6: Clarify dimension flexibility in comments.

The comment states dimension "MUST be 3072" and the server will reject other values, but according to the GeminiDenseEmbedder docstring, Gemini supports output dimensions 128–3072 (recommended: 768, 1536, 3072). Consider updating the comment to reflect that 3072 is the default and other values within the supported range are valid.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/ov.conf.gemini.example` around lines 5 - 6, Update the comment to
clarify that 3072 is the default/native output size for
gemini-embedding-2-preview but not strictly required; mention that Gemini
supports output dimensions in the range 128–3072 (with recommended values 768,
1536, 3072) and that reducing the dimension requires explicit truncation (the
server will reject dimensions outside that supported range). Reference the
gemini-embedding-2-preview model and the GeminiDenseEmbedder docstring so
readers know where the supported range is defined.

openviking/models/embedder/gemini_embedders.py (2)

145-151: Batch fallback may propagate RuntimeError from individual embeds.

The batch fallback catches APIError and falls back to individual embed() calls. However, embed() wraps APIError in RuntimeError (line 97). If a single text in the batch fails, the RuntimeError will propagate and the remaining texts in that batch won't be processed.

Consider wrapping the individual fallback calls to ensure all texts are attempted:

🔧 Proposed fix for robust batch fallback

             except APIError as e:
                 logger.warning(
                     f"Gemini batch embed failed (code={e.code}) for batch of {len(batch)}, "
                     "falling back to individual calls"
                 )
                 for text in batch:
-                    results.append(self.embed(text))
+                    try:
+                        results.append(self.embed(text))
+                    except RuntimeError as re:
+                        logger.error(f"Individual embed failed: {re}")
+                        # Return empty vector or re-raise depending on desired behavior
+                        raise

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@openviking/models/embedder/gemini_embedders.py` around lines 145 - 151, The
batch fallback currently calls self.embed(text) for each text inside the except
APIError block, but self.embed() wraps APIError into RuntimeError so a failure
on one text will abort the loop; update the fallback to call self.embed(text)
inside a per-item try/except (catch RuntimeError and APIError) so failures are
logged (use logger.warning/error with the text index or snippet and exception)
and the loop continues, ensuring you still append successful embeddings to
results and handle/append a placeholder (e.g., None) or skip entries
consistently for failed items; adjust code in the except APIError block around
variables batch, results, and self.embed to implement this per-item error
handling.

---

99-128: Consider handling empty text in multimodal fallback.

When media is None, has no data, or has an unsupported MIME type, the code falls back to self.embed(vectorize.text). If vectorize.text is empty or None, this will embed an empty string, which may produce unexpected results.

Consider adding a guard or logging when falling back with empty text:

🛡️ Proposed defensive check

         if (
             media is None
             or media.data is None
             or media.mime_type not in _SUPPORTED_MULTIMODAL_MIMES
         ):
+            if not vectorize.text:
+                logger.warning("embed_multimodal called with no usable media and empty text")
             return self.embed(vectorize.text)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@openviking/models/embedder/gemini_embedders.py` around lines 99 - 128, The
fallback in embed_multimodal currently calls self.embed(vectorize.text) even
when vectorize.text is empty/None; add a defensive check before any fallback
(both the early-return branch and the exception fallback) to detect falsy/blank
vectorize.text, log a warning via logger.warning mentioning embed_multimodal and
media.uri, and return a safe empty EmbedResult (e.g.,
EmbedResult(dense_vector=[])) instead of calling self.embed when text is empty;
update the logic inside embed_multimodal (references: embed_multimodal, embed,
EmbedResult, vectorize.text, media.uri) to perform this check in both fallback
paths.

openviking/utils/embedding_utils.py (1)

32-34: Move the shared MIME allowlist out of the Gemini private module.

embedding_utils is generic ingestion code, but this now imports _SUPPORTED_MULTIMODAL_MIMES, a private implementation detail of openviking.models.embedder.gemini_embedders, at module import time. That makes unrelated callers sensitive to provider-internal refactors. Prefer exposing the allowlist from a neutral/public module instead.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@openviking/utils/embedding_utils.py` around lines 32 - 34, The module
embedding_utils imports the private symbol _SUPPORTED_MULTIMODAL_MIMES from
openviking.models.embedder.gemini_embedders and derives
_GEMINI_SUPPORTED_IMAGE_MIMES from it, coupling generic ingestion code to a
provider-internal detail; instead move or expose the MIME allowlist as a
neutral/public constant (e.g. SUPPORTED_MULTIMODAL_MIMES or
SUPPORTED_IMAGE_MIMES) in a shared module (for example a new
openviking.models.embedder.mime or openviking.utils.mime) and update
embedding_utils to import that public symbol and derive the frozenset there
(replace references to _SUPPORTED_MULTIMODAL_MIMES and
_GEMINI_SUPPORTED_IMAGE_MIMES with the new public names).

openviking/storage/queuefs/embedding_msg_converter.py (1)

26-29: Keep media-only messages textless instead of injecting the URI.

The dequeue path later passes embedding_msg.message into Vectorize(text=...) for embed_multimodal(), so this fallback becomes part of the model input, not just queue bookkeeping. For media-only items that means embedding viking://.../file.ext alongside the bytes and also using the URI as the text-only fallback payload.

Possible direction

         vectorization_text = context.get_vectorization_text()
-        # If text is empty but media is present, use URI as fallback so the message isn't dropped
-        if not vectorization_text and getattr(context.vectorize, "media", None) is not None:
-            vectorization_text = context.uri or ""
-        if not vectorization_text:
+        has_media = getattr(context.vectorize, "media", None) is not None
+        if not vectorization_text and not has_media:
             return None

-            message=vectorization_text,
+            message=vectorization_text or "",

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@openviking/storage/queuefs/embedding_msg_converter.py` around lines 26 - 29,
Remove the fallback that injects context.uri into vectorization_text for
media-only items so media-only messages remain textless; specifically, in
embedding_msg_converter.py do not set vectorization_text = context.uri when
getattr(context.vectorize, "media", None) is not None — leave vectorization_text
empty for media-only cases so embedding_msg.message (and subsequently
Vectorize(text=...) used by embed_multimodal()) does not receive the URI as
model input; if bookkeeping needs a non-empty identifier, handle that separately
(e.g., store URI in metadata) rather than in vectorization_text.

tests/storage/test_embedding_msg_converter_tenant.py (1)

43-45: Consolidate imports at the top of the file.

The imports at lines 45 and 75-76 are placed mid-file. Additionally, EmbeddingMsgConverter at line 76 duplicates the import at line 9. Moving all imports to the top improves readability and follows PEP 8 conventions.

♻️ Proposed refactor

Move these imports to the top of the file with the other imports:

 from openviking.core.context import Context
+from openviking.core.context import ModalContent, Vectorize
 from openviking.storage.queuefs.embedding_msg_converter import EmbeddingMsgConverter
+from openviking.storage.queuefs.embedding_msg import EmbeddingMsg
 from openviking_cli.session.user_id import UserIdentifier

Then remove the mid-file imports at lines 45 and 75-76.

Also applies to: 75-76

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/storage/test_embedding_msg_converter_tenant.py` around lines 43 - 45,
Imports for EmbeddingMsg and EmbeddingMsgConverter are currently placed
mid-file; move the import statements for EmbeddingMsg (from
openviking.storage.queuefs.embedding_msg) and EmbeddingMsgConverter to the top
alongside the other imports, remove the duplicate mid-file import(s), and ensure
there's only one canonical import per symbol (EmbeddingMsg,
EmbeddingMsgConverter) at the module-level to satisfy PEP8 and avoid
duplication.

tests/storage/test_collection_schemas.py (1)

220-250: Consider patching VikingFS in the security test for robustness.

The test relies on the URI mismatch check happening before the file read. While this matches the current implementation, adding a _FakeVikingFS that raises an error would make the test more defensive against future code changes that might reorder the checks.

♻️ Suggested defensive patch

+    class _FailingVikingFS:
+        async def read_file_bytes(self, uri, ctx=None):
+            raise AssertionError("Should not attempt to read mismatched URI")
+
-    handler = TextEmbeddingHandler(_FakeVikingDB())
-    handler.set_callbacks(on_success=lambda: None, on_error=lambda *_: None)
-    await handler.on_dequeue(payload)
+    with patch("openviking.storage.viking_fs.get_viking_fs", return_value=_FailingVikingFS()):
+        handler = TextEmbeddingHandler(_FakeVikingDB())
+        handler.set_callbacks(on_success=lambda: None, on_error=lambda *_: None)
+        await handler.on_dequeue(payload)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/storage/test_collection_schemas.py` around lines 220 - 250, The test
test_handler_rejects_media_uri_not_matching_context_uri should also patch
VikingFS to ensure no file read occurs if the media_uri mismatches; add a
_FakeVikingFS class that raises an exception on read/open and monkeypatch it
into the code path the TextEmbeddingHandler uses (so that any accidental reorder
of checks that tries to access the file will fail the test), keep the existing
_FakeVikingDB and assertions that embedder.embed_multimodal_calls == 0 and
embedder.embed_calls == 1 to verify fallback-to-text behavior.

docs/zh/guides/01-configuration.md (1)

102-106: Consider documenting task_type in Gemini examples.

task_type is now supported and useful for retrieval tuning; adding it to both Gemini snippets will reduce under-configuration by users.

📝 Suggested doc tweak

 {
   "embedding": {
     "dense": {
       "provider": "gemini",
       "api_key": "your-google-api-key",
       "model": "gemini-embedding-2-preview",
-      "dimension": 3072
+      "dimension": 3072,
+      "task_type": "RETRIEVAL_DOCUMENT"
     }
   }
 }

Also applies to: 214-218

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@docs/zh/guides/01-configuration.md` around lines 102 - 106, The Gemini
provider examples currently lack the optional task_type setting; update both
Gemini configuration snippets that include "provider": "gemini" (the JSON blocks
with "api_key" and "model" such as "gemini-embedding-2-preview") to add a
"task_type" field (e.g., "task_type": "text-embedding" or the appropriate Gemini
task like "embedding") so users can tune retrieval behavior; ensure the new key
appears alongside "model" and "dimension" in each Gemini snippet and use
consistent example values across both occurrences.

tests/storage/test_embedding_msg_converter_tenant.py (1)

43-46: Consolidate duplicate imports at file top.

Context and EmbeddingMsgConverter are imported twice (lines 8-9 and lines 75-76). Consider moving all imports to the top of the file for cleaner organization.

 from openviking.core.context import Context
+from openviking.core.context import ModalContent, Vectorize
 from openviking.storage.queuefs.embedding_msg_converter import EmbeddingMsgConverter
+from openviking.storage.queuefs.embedding_msg import EmbeddingMsg
 from openviking_cli.session.user_id import UserIdentifier
 
-...
-
-from openviking.storage.queuefs.embedding_msg import EmbeddingMsg
-
-...
-
-from openviking.core.context import Context, ModalContent, Vectorize
-from openviking.storage.queuefs.embedding_msg_converter import EmbeddingMsgConverter

Also applies to: 75-77

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/storage/test_embedding_msg_converter_tenant.py` around lines 43 - 46,
The file has duplicate imports of Context and EmbeddingMsgConverter (and
EmbeddingMsg) scattered later in the test; consolidate by moving all imports for
Context, EmbeddingMsgConverter, and EmbeddingMsg to the top import block and
remove the redundant import statements (references: Context,
EmbeddingMsgConverter, EmbeddingMsg) so each symbol is imported once at file
head for cleaner organization.