mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 12:12:27 +00:00
fix(python): keep embedding functions registered without an alias readable (#4356)
`EmbeddingFunctionRegistry.register()` documents that the class name is
used when no alias is given, and it does register the class under that
name. But it stored the alias itself (`None`) as
`__embedding_function_registry_alias__`, which is what
`function_to_metadata()` writes into the table schema. A table using
such a function gets `"name": null` in its `embedding_functions`
metadata, and any later read of the functions fails:
```python
@get_registry().register()
class MyEmbed(TextEmbeddingFunction): ...
f = MyEmbed.create()
class Schema(LanceModel):
text: str = f.SourceField()
vector: Vector(4) = f.VectorField()
tbl = db.create_table("t", schema=Schema)
tbl.add([{"text": "hi"}]) # KeyError: None
tbl.search("hi") # KeyError: None
```
This stores the key the class is registered under instead, so the class
name round-trips. For aliased classes the key is the alias, so their
metadata is unchanged. The TypeScript registry already falls back to
`ctor.name` in the same way.
Added `test_embedding_function_registered_without_alias`, which creates
a table with such a function, reopens it, and adds a row. It fails with
`KeyError` before the change. `test_embeddings.py`, `test_pydantic.py`,
`test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff
check` are clean.
This commit is contained in:
@@ -73,7 +73,9 @@ class EmbeddingFunctionRegistry:
|
||||
raise KeyError(f"{cls.__name__} was already registered")
|
||||
key = alias or cls.__name__
|
||||
self._functions[key] = cls
|
||||
cls.__embedding_function_registry_alias__ = alias
|
||||
# Store the name the class is registered under, so the table
|
||||
# metadata written for it can be looked up again when reading.
|
||||
cls.__embedding_function_registry_alias__ = key
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
@@ -64,6 +64,37 @@ def test_embedding_function(tmp_path):
|
||||
assert np.allclose(actual, expected)
|
||||
|
||||
|
||||
def test_embedding_function_registered_without_alias(tmp_path):
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
|
||||
@registry.register()
|
||||
class NoAliasEmbeddingFunction(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 2
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
return [np.array([1.0, 2.0]) for _ in texts]
|
||||
|
||||
assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction
|
||||
|
||||
func = NoAliasEmbeddingFunction.create()
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(2) = func.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
db.create_table("test", schema=Schema)
|
||||
|
||||
# The class name is written to the table metadata, so the function can be
|
||||
# resolved again when the table is reopened.
|
||||
table = lancedb.connect(tmp_path).open_table("test")
|
||||
table.add([{"text": "hello"}])
|
||||
conf = table.embedding_functions["vector"]
|
||||
assert isinstance(conf.function, NoAliasEmbeddingFunction)
|
||||
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]]
|
||||
|
||||
|
||||
def test_instructor_ndims_uses_instruction():
|
||||
instructor = get_registry().get("instructor").create()
|
||||
model = MagicMock()
|
||||
|
||||
Reference in New Issue
Block a user