diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 1dc7ffcb8..3b940fe2e 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -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 diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index c42172f77..51907f141 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -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()