rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 12:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(python): cover OpenAI registry variable round-trip (#3863)

Browse Source 
## Summary

- replace the synthetic registry-variable metadata test with the OpenAI
embedding function reported in #2387
- verify the resolved API key survives table metadata reconstruction
- assert the OpenAI client receives the resolved key while serialized
metadata retains the variable reference

## Root cause

LanceDB 0.22.0 reconstructed embedding functions from table metadata
with the model constructor, bypassing EmbeddingFunction.create and
leaving the literal $var:api_key placeholder in OpenAI configuration.
The production path was corrected for duplicate #2181 by #2640; this
change gives that fix direct, network-free OpenAI regression coverage
for #2387.

## Validation

- uv run --extra tests pytest python/tests/test_embeddings.py -q (13
passed, 9 skipped)
- uv run --project python --extra dev ruff check .
- uv run --project python --extra dev ruff format --check
python/python/tests/test_embeddings.py
- git diff --check

Fixes #2387

<!-- lance-gatekeeper-fix:v1 agent=d453b1b9b2a298a776f2e4ea1b1449b5
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-08-06 16:36:05 +08:00
committed by GitHub
parent c3176a47ce
commit 9e2e711c7a
1 changed files with 14 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/tests/test_embeddings.py
+14 -27
View File
@@ -115,34 +115,16 @@ def test_embedding_function_variables():
assert func.safe_model_dump()["secret_key"] == "$var:secret"
def test_parse_functions_with_variables():
@register("variable-parsing-test")
class VariableParsingFunction(TextEmbeddingFunction):
api_key: str
base_url: Optional[str] = None
@staticmethod
def sensitive_keys():
return ["api_key"]
def ndims(self):
return 10
def generate_embeddings(self, texts):
# Mock implementation that just returns random embeddings
# In real usage, this would use the api_key to call an API
return [np.random.rand(self.ndims()).tolist() for _ in texts]
def test_openai_variables_survive_metadata_round_trip():
registry = EmbeddingFunctionRegistry.get_instance()
registry.set_var("test_api_key", "sk-test-key-12345")
registry.set_var("test_base_url", "https://api.example.com")
conf = EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=registry.get("variable-parsing-test").create(
api_key="$var:test_api_key", base_url="$var:test_base_url"
function=registry.get("openai").create(
api_key="$var:test_api_key", base_url="https://api.example.com"
),
)
@@ -150,7 +132,10 @@ def test_parse_functions_with_variables():
# Create a mock arrow table with the metadata
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))]
[
pa.field("text", pa.string()),
pa.field("vector", pa.list_(pa.float32(), 1536)),
]
)
table = pa.table({"text": [], "vector": []}, schema=schema)
table = table.replace_schema_metadata(metadata)
@@ -164,13 +149,15 @@ def test_parse_functions_with_variables():
assert parsed_func.api_key == "sk-test-key-12345"
assert parsed_func.base_url == "https://api.example.com"
embeddings = parsed_func.generate_embeddings(["test text"])
assert len(embeddings) == 1
assert len(embeddings[0]) == 10
assert parsed_func.safe_model_dump()["api_key"] == "$var:test_api_key"
with patch("lancedb.embeddings.openai.attempt_import_or_raise") as import_openai:
parsed_func._openai_client
import_openai.return_value.OpenAI.assert_called_once_with(
api_key="sk-test-key-12345", base_url="https://api.example.com"
)
def test_embedding_with_bad_results(tmp_path):
@register("null-embedding")