diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 52323daef..135bca8fd 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -608,6 +608,9 @@ class RemoteDBConnection(DBConnection): One of "error", "drop", "fill". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + embedding_functions: list of EmbeddingFunctionConfig, optional + The embedding functions to store in the table schema and use to + generate vectors when creating the table or adding data. Returns ------- @@ -682,12 +685,6 @@ class RemoteDBConnection(DBConnection): if namespace_path is None: namespace_path = [] validate_table_name(name) - if embedding_functions is not None: - logging.warning( - "embedding_functions is not yet supported on LanceDB Cloud." - "Please vote https://github.com/lancedb/lancedb/issues/626 " - "for this feature." - ) from .table import RemoteTable @@ -700,6 +697,7 @@ class RemoteDBConnection(DBConnection): schema=schema, on_bad_vectors=on_bad_vectors, fill_value=fill_value, + embedding_functions=embedding_functions, ) ) return RemoteTable( diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index a9e411a3b..36c318591 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -17,7 +17,8 @@ from packaging.version import Version import lancedb import numpy as np -from lancedb.conftest import MockTextEmbeddingFunction +from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction +from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry from lancedb.query import AsyncQuery, ColumnOrdering from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError @@ -596,6 +597,106 @@ def test_remote_permutation_is_picklable(): ] +@pytest.mark.parametrize("with_schema", [True, False]) +def test_create_table_embedding_functions(with_schema): + func = MockNonNormTextEmbeddingFunction.create() + config = EmbeddingFunctionConfig( + source_column="text", vector_column="vector", function=func + ) + schema = pa.schema( + [pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))] + ) + received = {} + + def handler(request): + if request.path == "/v1/table/test/describe/": + # Echo the metadata actually sent by create_table, as the server does. + metadata = received["create"].schema.metadata or {} + send_json( + request, + { + "version": 1, + "schema": { + "fields": [ + { + "name": "text", + "type": {"type": "string"}, + "nullable": True, + }, + { + "name": "vector", + "type": { + "type": "fixed_size_list", + "fields": [ + { + "name": "item", + "type": {"type": "float"}, + "nullable": True, + } + ], + "length": 10, + }, + "nullable": True, + }, + ], + "metadata": { + k.decode(): v.decode() for k, v in metadata.items() + }, + }, + }, + ) + elif request.path in ( + "/v1/table/test/create/?mode=create", + "/v1/table/test/insert/", + ): + if request.headers.get("Transfer-Encoding") == "chunked": + body = bytearray() + while True: + size = int(request.rfile.readline(), 16) + if size == 0: + request.rfile.readline() + break + body.extend(request.rfile.read(size)) + request.rfile.read(2) + else: + body = request.rfile.read(int(request.headers["Content-Length"])) + operation = "create" if "/create/" in request.path else "insert" + received[operation] = pa.ipc.open_stream(body).read_all() + send_json(request, {}) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + table = db.create_table( + "test", + schema=schema if with_schema else None, + data=None if with_schema else [{"text": "hello world"}], + embedding_functions=[config], + ) + metadata = pa.schema( + [], + metadata=EmbeddingFunctionRegistry.get_instance().get_table_metadata( + [config] + ), + ).metadata + assert received["create"].schema.metadata == metadata + assert table.schema.metadata == metadata + assert table.schema.field("vector").type == pa.list_(pa.float32(), 10) + assert table.embedding_functions["vector"].source_column == "text" + if with_schema: + assert received["create"].num_rows == 0 + else: + assert received["create"]["vector"].to_pylist() == [ + func.compute_source_embeddings(["hello world"])[0].tolist() + ] + + table.add([{"text": "goodbye world"}]) + assert received["insert"]["vector"].to_pylist() == [ + func.compute_source_embeddings(["goodbye world"])[0].tolist() + ] + + def test_create_table_exist_ok(): def handler(request): if request.path == "/v1/table/test/create/?mode=exist_ok":