From 353ce4d80c9e66f8f1ac9e49dbcec0ec899f1f9d Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:07 -0400 Subject: [PATCH] fix(python): give each JinaEmbeddings instance its own session (#4547) `JinaEmbeddings` stores its `requests` session in a `ClassVar`. The first instance builds the session with its own API key, and every later instance reuses that session. So a second instance with a different `api_key` sends the first instance's key. ## Fix The session is now a per-instance `cached_property`, the same pattern `OpenAIEmbeddings` uses for `_openai_client`. `_init_client` and the `ClassVar` are gone. An instance with no `api_key` and no `JINA_API_KEY` now raises the missing-key error. Before, it reused the session of an earlier instance and sent that instance's key. ## Tests `test_jina_embeddings_use_their_own_api_key` creates two instances with different keys through the registry. It stubs `requests.Session.post` and records the `Authorization` header of each request. - On `main`, the test fails because both requests send `Bearer key-a`. With the fix, it passes. - `test_embeddings.py`: 18 passed, 8 skipped. - An instance still pickles after the session exists, and `safe_model_dump()` still returns only `api_key`. - On pydantic 2.7.4, the minimum version, a model with the same `cached_property` gets one session per instance and leaves it out of `model_dump()`. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `embeddings/jinaai.py` is byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/embeddings/jinaai.py | 26 +++++++++++----------- python/python/tests/test_embeddings.py | 25 +++++++++++++++++++++ 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/python/python/lancedb/embeddings/jinaai.py b/python/python/lancedb/embeddings/jinaai.py index f6ab601b3..25b5a940b 100644 --- a/python/python/lancedb/embeddings/jinaai.py +++ b/python/python/lancedb/embeddings/jinaai.py @@ -5,9 +5,10 @@ import os import io import base64 +from functools import cached_property from urllib.parse import urlparse from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, List, Union, Optional, Any, Dict +from typing import TYPE_CHECKING, List, Union, Optional, Any, Dict import numpy as np import pyarrow as pa @@ -51,7 +52,6 @@ class JinaEmbeddings(EmbeddingFunction): name: str = "jina-clip-v1" api_key: Optional[str] = None - _session: ClassVar = None def ndims(self): # TODO: fix hardcoding @@ -206,8 +206,7 @@ class JinaEmbeddings(EmbeddingFunction): texts: list[str] or np.ndarray (of str) The texts to embed """ - self._init_client() - resp = JinaEmbeddings._session.post( # type: ignore + resp = self._session.post( API_URL, json={"input": input, "model": self.name} ).json() if "data" not in resp: @@ -220,14 +219,15 @@ class JinaEmbeddings(EmbeddingFunction): return [result["embedding"] for result in sorted_embeddings] - def _init_client(self): + @cached_property + def _session(self): import requests - if JinaEmbeddings._session is None: - if self.api_key is None and os.environ.get("JINA_API_KEY") is None: - api_key_not_found_help("jina") - api_key = self.api_key or os.environ.get("JINA_API_KEY") - JinaEmbeddings._session = requests.Session() - JinaEmbeddings._session.headers.update( - {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} - ) + if self.api_key is None and os.environ.get("JINA_API_KEY") is None: + api_key_not_found_help("jina") + api_key = self.api_key or os.environ.get("JINA_API_KEY") + session = requests.Session() + session.headers.update( + {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} + ) + return session diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 129fa6742..e319d3048 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1053,3 +1053,28 @@ def test_jina_generate_image_input_dict_local_path(tmp_path): image_dict = JinaEmbeddings._generate_image_input_dict(image) assert "image" in image_dict assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0 + + +def test_jina_embeddings_use_their_own_api_key(monkeypatch): + """Each JinaEmbeddings instance must send its own API key.""" + pytest.importorskip("requests") + monkeypatch.delenv("JINA_API_KEY", raising=False) + registry = get_registry() + registry.set_var("jina_key_a", "key-a") + registry.set_var("jina_key_b", "key-b") + func_a = registry.get("jina").create(api_key="$var:jina_key_a") + func_b = registry.get("jina").create(api_key="$var:jina_key_b") + + sent = [] + + def fake_post(session, url, json=None, **kwargs): + sent.append(session.headers["Authorization"]) + response = MagicMock() + response.json.return_value = {"data": [{"index": 0, "embedding": [0.0]}]} + return response + + with patch("requests.Session.post", fake_post): + func_a.compute_source_embeddings(["hello"]) + func_b.compute_source_embeddings(["hello"]) + + assert sent == ["Bearer key-a", "Bearer key-b"]