tests/integration/test_gemini_e2e.py (2)

92-92: Timing assertion may cause flaky tests.

assert elapsed < 15 is network-dependent and could fail due to API latency spikes, CI resource contention, or regional network conditions. Consider either:

1. Using a more generous threshold (e.g., 60s)
2. Removing the timing assertion and validating only correctness
3. Marking this test with @pytest.mark.flaky if available

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_e2e.py` at line 92, The timing assertion assert
elapsed < 15 is flaky; update the test in tests/integration/test_gemini_e2e.py
to avoid brittle timing checks by either increasing the threshold (e.g., replace
the 15s limit with 60s for elapsed), removing the timing assertion altogether
and only asserting correctness, or decorating the test with a flaky marker
(e.g., `@pytest.mark.flaky`) so transient API/CI latency won't fail the test;
modify the assertion referencing elapsed accordingly and add a brief comment
explaining the reason.

---

56-63: Batch vs individual comparison may have false positives.

The test compares batch results to individual results expecting sim > 0.99. However, if the Gemini API is non-deterministic (e.g., different request batching), this could occasionally fail. The threshold of 0.99 is quite strict.

Consider loosening to 0.95 or documenting that exact reproducibility is expected.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_e2e.py` around lines 56 - 63, The test
test_embed_batch_matches_individual is too strict and can false-fail if
embeddings vary; update the assertion threshold from 0.99 to a looser value
(e.g., 0.95) in the loop that compares batch_results from embed_batch to
individual_results from embed using _cosine_similarity, or add a clear comment
documenting that exact reproducibility is required; change the numeric check
(sim > 0.99) to sim > 0.95 (or configurable) so the test tolerates minor
non-determinism.

openviking_cli/utils/config/embedding_config.py (2)

155-159: Duplicate task type constants across files.

_GEMINI_TASK_TYPES is defined here and _VALID_TASK_TYPES is defined in gemini_embedders.py (line 40-49). These should be consolidated to a single source of truth to prevent drift.

♻️ Proposed approach

Import from the embedder module:

+from openviking.models.embedder.gemini_embedders import _VALID_TASK_TYPES as _GEMINI_TASK_TYPES
-            _GEMINI_TASK_TYPES = {
-                "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY",
-                "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING",
-                "FACT_VERIFICATION", "CODE_RETRIEVAL_QUERY",
-            }

Or extract to a shared constants module.

🤖 Prompt for AI Agents

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

In `@openviking_cli/utils/config/embedding_config.py` around lines 155 - 159, The
duplicate task-type sets (_GEMINI_TASK_TYPES in embedding_config.py and
_VALID_TASK_TYPES in gemini_embedders.py) should be consolidated to a single
source of truth: either export the canonical set from gemini_embedders.py (e.g.,
rename or re-export _VALID_TASK_TYPES) and import it into embedding_config.py,
or move the set into a new shared constants module and import it from both
places; update references to use the single exported symbol and remove the
duplicate definition in embedding_config.py so both modules rely on the same
constant.

---

438-468: New embedder methods not integrated into service layer.

The get_query_embedder() and get_document_embedder() methods are defined but unused. The service layer (openviking/service/core.py, lines 101 and 247) continues to call config.embedding.get_embedder(), so context-aware embeddings for non-symmetric providers (Gemini, Jina) won't be utilized without explicit service layer updates.

Consider integrating these methods into the service layer or documenting whether this is intentional (opt-in behavior to be used elsewhere).

🤖 Prompt for AI Agents

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

In `@openviking_cli/utils/config/embedding_config.py` around lines 438 - 468, The
new methods get_query_embedder and get_document_embedder are never used; update
the service layer code that currently calls config.embedding.get_embedder() to
instead call get_query_embedder when producing/querying embeddings and
get_document_embedder when creating/indexing document embeddings (identify the
call sites by the functions that currently invoke
config.embedding.get_embedder()); ensure the service logic preserves the
fallback behavior (call get_embedder if contextual methods are unavailable) and
keep existing behavior for symmetric providers by relying on
_get_contextual_embedder's effective_context logic.

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

186-192: Fallback logic may re-raise error for persistent API failures.

When embed_batch falls back to individual embed() calls on line 192, if the API error is persistent (e.g., 401 auth error), each individual call will also fail and raise a RuntimeError. This will abort processing mid-batch rather than gracefully handling all items.

Consider whether this is the intended behavior or if the fallback should only apply to transient errors (5xx).

🤖 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 186 - 192, The
fallback in embed_batch that iterates calling self.embed(text) for each item can
re-raise persistent API errors (e.g., 401) and abort processing; change
embed_batch (and the except block that catches (APIError, ClientError) as e) to
only perform per-item fallback for transient server errors (HTTP 5xx) and
re-raise for persistent client/auth errors (4xx). Concretely, inspect e.code (or
error.kind) inside the except in embed_batch: if it's a 5xx, loop and attempt
per-item embedding but wrap each self.embed(text) in its own try/except to
catch/record individual failures (log and append a sentinel or result
placeholder); if e.code is 4xx (or other non-transient), re-raise the original
exception immediately. This targets the embed_batch method and the call to
self.embed to avoid aborting mid-batch on persistent errors.

---

226-233: Async fallback runs sync embed() sequentially, negating concurrency benefits.

When async batch embedding fails, the fallback at line 232 runs self.embed(text) via anyio.to_thread.run_sync for each text sequentially within the list comprehension. This makes the fallback path significantly slower than the primary async path.

Consider parallelizing the fallback or documenting this as expected degraded behavior.

🤖 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 226 - 233, The
fallback currently runs self.embed(text) sequentially inside a list
comprehension (via anyio.to_thread.run_sync) which loses concurrency; update the
except handler to run the per-text fallbacks in parallel by spawning concurrent
tasks (e.g., using anyio.create_task_group or anyio.from_thread.run) that each
call anyio.to_thread.run_sync(self.embed, text) and then await/join and collect
results into results[idx], so each embed call executes concurrently rather than
sequentially in the except block handling APIError/ClientError.

---

244-249: Use the public client.close() method instead of accessing the private _http_client attribute.

The google-genai SDK provides an official close() method on the Client class for cleanup. Directly accessing the private _http_client attribute is fragile and can break across SDK versions. Replace self.client._http_client.close() with self.client.close().

🤖 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 244 - 249,
Replace the private attribute access in the close routine: instead of checking
for and calling self.client._http_client.close(), call the public
self.client.close() method; update the close method in the Gemini embedder class
(the method named close that currently references self.client and _http_client)
to invoke self.client.close() inside a try/except to handle errors and remove
reliance on the private _http_client attribute.

tests/integration/test_gemini_openviking_it.py (2)

185-189: ⚠️ Potential issue | 🟠 Major

Use the session in the search you assert.

The stronger > 0 check helps, but client.find(...) still ignores the session created on Lines 185-186. This remains a plain find smoke test, so a regression in session-scoped retrieval would still slip through.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_openviking_it.py` around lines 185 - 189, The
test creates a session object (session =
client.session(session_id="gemini_it_session") and adds a message, but then
calls client.find(...) without scoping to that session; update the find call to
use the created session (pass the session object or its session_id) so the
search is performed against the session-scoped data (e.g., client.find(...,
session=session) or client.find(..., session_id="gemini_it_session") instead of
the current call to client.find(query="pytest testing tool")).

---

149-151: ⚠️ Potential issue | 🟠 Major

Assert the runtime dimension, not the config echo.

OpenVikingConfigSingleton.instance().embedding.dimension is still derived from config, so this test stays green even if the factory drops dimension or the active embedder/index schema falls back to a different size. Please inspect the created embedder or the index schema directly here.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_openviking_it.py` around lines 149 - 151, The
test currently asserts OpenVikingConfigSingleton.instance().embedding.dimension
(a config echo) so change it to assert the runtime embedder/index schema
dimension instead: after creating the embedder or index in the test, obtain the
actual runtime object (e.g. the created embedder instance or the index schema
from the component under test) and assert its reported dimension equals dim
rather than checking OpenVikingConfigSingleton.instance().embedding.dimension;
locate uses of OpenVikingConfigSingleton.instance() in the test and replace the
assertion with a check against the embedder.dimension or index.schema.dimension
on the actual created object.

tests/integration/conftest.py (1)

173-178: Close the module-scoped Gemini client in fixture teardown.

This fixture creates a real GeminiDenseEmbedder and never calls close(), so the underlying HTTP client stays open for the life of the test process. Switching to yield keeps the live suite cleaner and less flaky.

♻️ Suggested change

 `@pytest.fixture`(scope="module", params=GEMINI_MODELS_FIXTURE)
 def gemini_embedder(request):
     """Module-scoped GeminiDenseEmbedder at dim=768, parametrized over known models."""
     from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder
     model_name, _, _ = request.param
-    return GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768)
+    embedder = GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768)
+    yield embedder
+    embedder.close()

🤖 Prompt for AI Agents

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

In `@tests/integration/conftest.py` around lines 173 - 178, The gemini_embedder
fixture creates a GeminiDenseEmbedder but never closes it; change the
module-scoped fixture "gemini_embedder" to use yield instead of return:
instantiate GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY,
dimension=768), yield the instance, and after yield call instance.close() in the
teardown so the underlying HTTP client is properly closed; reference the fixture
name gemini_embedder and the class GeminiDenseEmbedder to locate the code.

tests/integration/test_gemini_openviking_it.py (2)

205-209: ⚠️ Potential issue | 🟠 Major

The session setup is still outside the asserted behavior.

client.find() on Lines 208-209 does not consume the session or the message added on Line 206, so this test only covers plain embedding search. Use the session-aware retrieval path here or drop the unused session scaffolding.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_openviking_it.py` around lines 205 - 209, The
test creates a session via client.session(...) and adds a message with
session.add_message(...) but then calls client.find(...) without using that
session; either remove the unused session setup or switch to the session-aware
retrieval API. Fix by either deleting the session/session.add_message lines, or
call the session-aware finder (e.g., pass the session into the find call like
client.find(session=session, query="pytest testing tool") or use whatever
session.find(...) helper exists) so the added message is actually consumed by
the retrieval assertion.

---

163-168: ⚠️ Potential issue | 🟠 Major

This still only verifies the configured dimension, not the runtime one.

OpenVikingConfigSingleton.instance().embedding.dimension on Lines 166-168 just re-reads the requested config value. The test still goes green if the live embedder or index schema later falls back to 3072, so please assert the runtime dimension instead.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_openviking_it.py` around lines 163 - 168,
Replace the config-only assertion with a runtime check: after creating client
via make_ov_client(...), inspect the actual embedder/index created by the client
(e.g., client.embedder.dimension or client._embedder.dimension, or the index
schema such as client.index.schema.embedding_dimension) and assert that that
runtime dimension equals dim instead of using
OpenVikingConfigSingleton.instance().embedding.dimension; update the test to
pull the dimension from the live client/embedder/index and assert equality to
dim.

openviking_cli/utils/config/embedding_config.py (1)

370-383: ⚠️ Potential issue | 🟠 Major

Gemini still ignores the configured concurrency limit.

EmbeddingConfig.max_concurrent on Line 225 is the public throttle, but this lambda always instantiates GeminiDenseEmbedder with its constructor default. Any override is ignored, which makes rate-limit tuning and throughput control ineffective.

♻️ Suggested change

             ("gemini", "dense"): (
                 GeminiDenseEmbedder,
                 lambda cfg: {
                     "model_name": cfg.model,
                     "api_key": cfg.api_key,
                     "dimension": cfg.dimension,
+                    "max_concurrent_batches": self.max_concurrent,
                     "task_type": (
                         cfg.query_param
                         if context == "query" and cfg.query_param
                         else cfg.document_param
                         if context == "document" and cfg.document_param
                         else cfg.task_type
                     ),
                 },
             ),

🤖 Prompt for AI Agents

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

In `@openviking_cli/utils/config/embedding_config.py` around lines 370 - 383, The
mapping for ("gemini", "dense") ignores EmbeddingConfig.max_concurrent by not
passing it into GeminiDenseEmbedder; update the lambda that constructs
GeminiDenseEmbedder to include the concurrency override (use cfg.max_concurrent
or EmbeddingConfig.max_concurrent as exposed on cfg) in the args passed to
GeminiDenseEmbedder so the instance respects the configured throttle, e.g., add
a "max_concurrent" (or the exact constructor param name used by
GeminiDenseEmbedder) entry to the returned dict alongside model_name, api_key,
dimension and task_type.

tests/integration/test_gemini_e2e.py (1)

81-90: ⚠️ Potential issue | 🟠 Major

Remove the hard latency SLA from the real-API test.

elapsed < 15 will flap with runner load, network jitter, or transient throttling even when batching is correct. Keep the functional assertions here and move timing checks to an opt-in perf/benchmark test.

♻️ Suggested fix

-        import time
-
         e = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key=GOOGLE_API_KEY, dimension=128)
         texts = [f"sentence {i}" for i in range(300)]  # 3 batches of 100
-        t0 = time.monotonic()
         results = await e.async_embed_batch(texts)
-        elapsed = time.monotonic() - t0
         assert len(results) == 300
         assert all(len(r.dense_vector) == 128 for r in results)
-        assert elapsed < 15  # concurrent should be << 3× serial RTT
         e.close()

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_e2e.py` around lines 81 - 90, The test is
asserting a hard latency SLA which can flap; remove the timing assertion and
keep only the functional checks: instantiate GeminiDenseEmbedder, call
async_embed_batch, assert len(results) == 300 and each r.dense_vector has length
128, and drop the elapsed < 15 assertion (or move any timing assertions into a
separate, opt-in perf/benchmark test). Locate the timing logic around
GeminiDenseEmbedder and async_embed_batch in the test and delete the elapsed
measurement assertion while preserving the setup and functional assertions.

pyproject.toml (1)

82-89: ⚠️ Potential issue | 🟠 Major

pytest-asyncio auto mode still conflicts with the new AnyIO tests.

Line 212 still sets asyncio_mode = "auto", while this PR adds anyio support and @pytest.mark.anyio tests. AnyIO’s testing docs explicitly call out pytest-asyncio auto mode as incompatible with the AnyIO plugin, and pytest-asyncio documents strict mode as the coexistence mode when other async test plugins are in play. Either switch to strict mode and mark asyncio tests explicitly, or drop the AnyIO markers. (anyio.readthedocs.io)

Also applies to: 214-216

🤖 Prompt for AI Agents

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

In `@pyproject.toml` around lines 82 - 89, The test configuration currently sets
pytest-asyncio asyncio_mode = "auto" while the PR adds AnyIO-marked tests
(`@pytest.mark.anyio`), which conflicts; update the test config to use
asyncio_mode = "strict" (so pytest-asyncio coexists with the AnyIO plugin) and
then explicitly mark asyncio tests with `@pytest.mark.asyncio` where needed, or
alternatively remove the AnyIO markers and keep auto mode—change the
asyncio_mode setting and adjust test decorators (`@pytest.mark.anyio` /
`@pytest.mark.asyncio`) accordingly.

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

143-143: ⚠️ Potential issue | 🟠 Major

_token_limit is still never enforced before Gemini calls.

Line 143 records the per-model ceiling, but embed(), embed_batch(), and async_embed_batch() never check it before sending contents to the SDK. Oversized inputs will still fail remotely instead of being chunked, truncated, or rejected locally with a clear error.

Also applies to: 150-160, 165-184, 223-225

🤖 Prompt for AI Agents

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

In `@openviking/models/embedder/gemini_embedders.py` at line 143, The per-model
ceiling stored in self._token_limit is never enforced before calling Gemini;
update embed, embed_batch, and async_embed_batch to validate inputs against
self._token_limit (locations around the current embed() implementation and the
batch methods referenced) by tokenizing each content (use your existing
tokenizer utility) and either chunking long inputs into acceptable-sized pieces,
truncating with a clear warning, or raising a descriptive local error instead of
sending oversized payloads to the SDK; ensure the chosen behavior is applied
consistently in embed(), embed_batch(), and async_embed_batch(), and reference
_MODEL_TOKEN_LIMITS/_DEFAULT_TOKEN_LIMIT where needed to determine limits.

tests/integration/conftest.py (1)

182-188: ⚠️ Potential issue | 🟡 Minor

Close the module-scoped Gemini fixture.

gemini_embedder returns a live GeminiDenseEmbedder but never closes it. With multiple parametrized models, the suite can keep SDK clients/sockets open until process exit.

♻️ Suggested fix

 `@pytest.fixture`(scope="module", params=GEMINI_MODELS_FIXTURE)
 def gemini_embedder(request):
     """Module-scoped GeminiDenseEmbedder at dim=768, parametrized over known models."""
     from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder
 
     model_name, _, _ = request.param
-    return GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768)
+    embedder = GeminiDenseEmbedder(model_name, api_key=GOOGLE_API_KEY, dimension=768)
+    try:
+        yield embedder
+    finally:
+        embedder.close()

🤖 Prompt for AI Agents

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

In `@tests/integration/conftest.py` around lines 182 - 188, The fixture
gemini_embedder returns a GeminiDenseEmbedder without closing it; change the
fixture to assign the instance to a local variable, register a teardown via
request.addfinalizer that calls the embedder's cleanup method (e.g.,
embedder.close() or embedder.shutdown(), or await embedder.aclose() if async)
and then return the instance; reference the fixture name gemini_embedder and the
class GeminiDenseEmbedder so the finalizer ensures SDK clients/sockets are
closed after each module param.

openviking_cli/utils/config/embedding_config.py (1)

225-227: ⚠️ Potential issue | 🟡 Minor

EmbeddingConfig.max_concurrent still does nothing for Gemini.

Lines 225-227 expose the limiter, but the Gemini factory never passes it as max_concurrent_batches, so tuning max_concurrent has no effect on async_embed_batch().

♻️ Suggested fix

             ("gemini", "dense"): (
                 GeminiDenseEmbedder,
                 lambda cfg: {
                     "model_name": cfg.model,
                     "api_key": cfg.api_key,
                     "dimension": cfg.dimension,
+                    "max_concurrent_batches": self.max_concurrent,
                     "task_type": (
                         cfg.query_param
                         if context == "query" and cfg.query_param
                         else cfg.document_param
                         if context == "document" and cfg.document_param

Also applies to: 370-383

🤖 Prompt for AI Agents

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

In `@openviking_cli/utils/config/embedding_config.py` around lines 225 - 227,
EmbeddingConfig.max_concurrent is defined but not forwarded to Gemini's batching
limiter; update the Gemini embedding provider/factory to pass
EmbeddingConfig.max_concurrent as the max_concurrent_batches argument so
async_embed_batch actually uses it. Locate where the Gemini provider/factory
(the function/class that instantiates the Gemini embeddings client) constructs
or calls async_embed_batch and modify the call/constructor to accept and forward
embedding_config.max_concurrent (also update the second occurrence around the
other factory/initialization block noted in the review). Ensure the parameter
name matches the Gemini client API (max_concurrent_batches) and propagate the
value from EmbeddingConfig.max_concurrent through the factory to the
async_embed_batch invocations.

tests/integration/test_gemini_openviking_it.py (1)

163-168: ⚠️ Potential issue | 🟠 Major

This still checks config, not the active embedding dimension.

OpenVikingConfigSingleton.get_instance().embedding.dimension just echoes the configured value. If the Gemini factory silently falls back to 3072 and both indexing/search use that default, this test still passes, so it won't catch the regression it is meant to guard against. Assert the live embedder or index schema dimension instead.

🤖 Prompt for AI Agents

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

In `@tests/integration/test_gemini_openviking_it.py` around lines 163 - 168, The
test is asserting the configured dimension via OpenVikingConfigSingleton rather
than the runtime/embedder/index dimension, so replace the config check with an
assertion against the live component returned by make_ov_client: after creating
client = await make_ov_client(...), query the actual embedder or index schema on
that client (e.g., inspect client.embedder.dimension or
client.index.schema.dimension or similar runtime property provided by the
OpenViking client) and assert it equals dim to ensure the active embedder/index
uses the expected dimension instead of relying on the config value.

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

311-316: close() relies on SDK private attribute _http_client.

Accessing self.client._http_client is fragile since _http_client is an internal implementation detail of the google-genai SDK and may change or be removed without notice.

💡 Consider a more defensive approach

     def close(self):
-        if hasattr(self.client, "_http_client"):
-            try:
-                self.client._http_client.close()
-            except Exception:
-                pass
+        # The google-genai SDK manages its own connection lifecycle.
+        # If explicit cleanup is needed in the future, check SDK docs for a public API.
+        pass

Alternatively, if cleanup is critical, document that this relies on SDK internals and may need updates when upgrading the SDK version.

🤖 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 311 - 316, The
close() method is accessing the private SDK attribute self.client._http_client
which is fragile; update close() to avoid relying on that internal name by first
checking for public/official cleanup methods (e.g., a close/async_close method
or transport attribute) via getattr/hasattr on self.client and invoking those if
present, otherwise fall back to a safe getattr(self.client, "_http_client",
None) check and then calling close on it inside a try/except; also add a short
comment documenting the fallback to the private attribute so future maintainers
know this is an SDK-internal fallback that may need updating when upgrading
google-genai.