code: #  Copyright 2023 LanceDB Developers
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.
  import sys
  from typing import List, Union
  
  import lance
  import lancedb
  import numpy as np
  import pyarrow as pa
  import pytest
  from lancedb.conftest import MockTextEmbeddingFunction
  from lancedb.embeddings import (
      EmbeddingFunctionConfig,
      EmbeddingFunctionRegistry,
      with_embeddings,
  )
  from lancedb.embeddings.base import TextEmbeddingFunction
  from lancedb.embeddings.registry import get_registry, register
  from lancedb.pydantic import LanceModel, Vector
  
  
  def mock_embed_func(input_data):
      return [np.random.randn(128).tolist() for _ in range(len(input_data))]
  
  
  def test_with_embeddings():
      for wrap_api in [True, False]:
          if wrap_api and sys.version_info.minor >= 11:
              # ratelimiter package doesn't work on 3.11
              continue
          data = pa.Table.from_arrays(
              [
                  pa.array(["foo", "bar"]),
                  pa.array([10.0, 20.0]),
              ],
              names=["text", "price"],
          )
          data = with_embeddings(mock_embed_func, data, wrap_api=wrap_api)
          assert data.num_columns == 3
          assert data.num_rows == 2
          assert data.column_names == ["text", "price", "vector"]
          assert data.column("text").to_pylist() == ["foo", "bar"]
          assert data.column("price").to_pylist() == [10.0, 20.0]
  
  
  def test_embedding_function(tmp_path):
      registry = EmbeddingFunctionRegistry.get_instance()
  
      # let's create a table
      table = pa.table(
          {
              "text": pa.array(["hello world", "goodbye world"]),
              "vector": [np.random.randn(10), np.random.randn(10)],
          }
      )
      conf = EmbeddingFunctionConfig(
          source_column="text",
          vector_column="vector",
          function=MockTextEmbeddingFunction(),
      )
      metadata = registry.get_table_metadata([conf])
      table = table.replace_schema_metadata(metadata)
  
      # Write it to disk
      lance.write_dataset(table, tmp_path / "test.lance")
  
      # Load this back
      ds = lance.dataset(tmp_path / "test.lance")
  
      # can we get the serialized version back out?
      configs = registry.parse_functions(ds.schema.metadata)
  
      conf = configs["vector"]
      func = conf.function
      actual = func.compute_query_embeddings("hello world")
  
      # And we make sure we can call it
      expected = func.compute_query_embeddings("hello world")
  
      assert np.allclose(actual, expected)
  
  
  @pytest.mark.slow
  def test_embedding_function_rate_limit(tmp_path):
      def _get_schema_from_model(model):
          class Schema(LanceModel):
              text: str = model.SourceField()
              vector: Vector(model.ndims()) = model.VectorField()
  
          return Schema
  
      db = lancedb.connect(tmp_path)
      registry = EmbeddingFunctionRegistry.get_instance()
      model = registry.get("test-rate-limited").create(max_retries=0)
      schema = _get_schema_from_model(model)
      table = db.create_table("test", schema=schema, mode="overwrite")
      table.add([{"text": "hello world"}])
      with pytest.raises(Exception):
          table.add([{"text": "hello world"}])
      assert len(table) == 1
  
      model = registry.get("test-rate-limited").create()
      schema = _get_schema_from_model(model)
      table = db.create_table("test", schema=schema, mode="overwrite")
      table.add([{"text": "hello world"}])
      table.add([{"text": "hello world"}])
      assert len(table) == 2
  
  
  def test_add_optional_vector(tmp_path):
      @register("mock-embedding")
      class MockEmbeddingFunction(TextEmbeddingFunction):
          def ndims(self):
              return 128
  
          def generate_embeddings(
              self, texts: Union[List[str], np.ndarray]
          ) -> List[np.array]:
              """
              Generate the embeddings for the given texts
              """
              return [np.random.randn(self.ndims()).tolist() for _ in range(len(texts))]
  
      registry = get_registry()
      model = registry.get("mock-embedding").create()
  
      class LanceSchema(LanceModel):
          id: str
          vector: Vector(model.ndims()) = model.VectorField(default=None)
          text: str = model.SourceField()
  
      db = lancedb.connect(tmp_path)
      tbl = db.create_table("optional_vector", schema=LanceSchema)
  
      # add works
      expected = LanceSchema(id="id", text="text")
      tbl.add([expected])
      assert not (np.abs(tbl.to_pandas()["vector"][0]) < 1e-6).all()
  
apis: lancedb.embeddings.EmbeddingFunctionRegistry.get_instance lancedb.embeddings.registry.get_registry lancedb.conftest.MockTextEmbeddingFunction lancedb.embeddings.registry.register lancedb.embeddings.with_embeddings lancedb.connect
extract_api: [((1948, 1988), 'lancedb.embeddings.EmbeddingFunctionRegistry.get_instance', 'EmbeddingFunctionRegistry.get_instance', ([], {}), '()\n', (1986, 1988), False, 'from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry, with_embeddings\n'), ((2476, 2527), 'lance.write_dataset', 'lance.write_dataset', (['table', "(tmp_path / 'test.lance')"], {}), "(table, tmp_path / 'test.lance')\n", (2495, 2527), False, 'import lance\n'), ((2559, 2597), 'lance.dataset', 'lance.dataset', (["(tmp_path / 'test.lance')"], {}), "(tmp_path / 'test.lance')\n", (2572, 2597), False, 'import lance\n'), ((2932, 2961), 'numpy.allclose', 'np.allclose', (['actual', 'expected'], {}), '(actual, expected)\n', (2943, 2961), True, 'import numpy as np\n'), ((3246, 3271), 'lancedb.connect', 'lancedb.connect', (['tmp_path'], {}), '(tmp_path)\n', (3261, 3271), False, 'import lancedb\n'), ((3287, 3327), 'lancedb.embeddings.EmbeddingFunctionRegistry.get_instance', 'EmbeddingFunctionRegistry.get_instance', ([], {}), '()\n', (3325, 3327), False, 'from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry, with_embeddings\n'), ((3980, 4006), 'lancedb.embeddings.registry.register', 'register', (['"""mock-embedding"""'], {}), "('mock-embedding')\n", (3988, 4006), False, 'from lancedb.embeddings.registry import get_registry, register\n'), ((4419, 4433), 'lancedb.embeddings.registry.get_registry', 'get_registry', ([], {}), '()\n', (4431, 4433), False, 'from lancedb.embeddings.registry import get_registry, register\n'), ((4660, 4685), 'lancedb.connect', 'lancedb.connect', (['tmp_path'], {}), '(tmp_path)\n', (4675, 4685), False, 'import lancedb\n'), ((1570, 1627), 'lancedb.embeddings.with_embeddings', 'with_embeddings', (['mock_embed_func', 'data'], {'wrap_api': 'wrap_api'}), '(mock_embed_func, data, wrap_api=wrap_api)\n', (1585, 1627), False, 'from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry, with_embeddings\n'), ((3558, 3582), 'pytest.raises', 'pytest.raises', (['Exception'], {}), '(Exception)\n', (3571, 3582), False, 'import pytest\n'), ((2069, 2111), 'pyarrow.array', 'pa.array', (["['hello world', 'goodbye world']"], {}), "(['hello world', 'goodbye world'])\n", (2077, 2111), True, 'import pyarrow as pa\n'), ((2310, 2337), 'lancedb.conftest.MockTextEmbeddingFunction', 'MockTextEmbeddingFunction', ([], {}), '()\n', (2335, 2337), False, 'from lancedb.conftest import MockTextEmbeddingFunction\n'), ((1102, 1122), 'numpy.random.randn', 'np.random.randn', (['(128)'], {}), '(128)\n', (1117, 1122), True, 'import numpy as np\n'), ((1427, 1451), 'pyarrow.array', 'pa.array', (["['foo', 'bar']"], {}), "(['foo', 'bar'])\n", (1435, 1451), True, 'import pyarrow as pa\n'), ((1469, 1491), 'pyarrow.array', 'pa.array', (['[10.0, 20.0]'], {}), '([10.0, 20.0])\n', (1477, 1491), True, 'import pyarrow as pa\n'), ((2136, 2155), 'numpy.random.randn', 'np.random.randn', (['(10)'], {}), '(10)\n', (2151, 2155), True, 'import numpy as np\n'), ((2157, 2176), 'numpy.random.randn', 'np.random.randn', (['(10)'], {}), '(10)\n', (2172, 2176), True, 'import numpy as np\n')]