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 <noreply@anthropic.com>
This commit is contained in:
Yash Raj Pandey
2026-10-23 03:23:07 +08:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 5fa30ae404
commit 353ce4d80c
2 changed files with 38 additions and 13 deletions
+13 -13
View File
@@ -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
+25
View File
@@ -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"]