|
| 1 | +from typing import List, Optional |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +import pyarrow as pa |
| 5 | +import lancedb |
| 6 | +from gptcache.manager.vector_data.base import VectorBase, VectorData |
| 7 | +from gptcache.utils import import_lancedb, import_torch |
| 8 | + |
| 9 | +import_torch() |
| 10 | +import_lancedb() |
| 11 | + |
| 12 | + |
| 13 | +class LanceDB(VectorBase): |
| 14 | + """Vector store: LanceDB |
| 15 | + :param persist_directory: The directory to persist, defaults to '/tmp/lancedb'. |
| 16 | + :type persist_directory: str |
| 17 | + :param table_name: The name of the table in LanceDB, defaults to 'gptcache'. |
| 18 | + :type table_name: str |
| 19 | + :param top_k: The number of the vectors results to return, defaults to 1. |
| 20 | + :type top_k: int |
| 21 | + """ |
| 22 | + |
| 23 | + def __init__( |
| 24 | + self, |
| 25 | + persist_directory: Optional[str] = "/tmp/lancedb", |
| 26 | + table_name: str = "gptcache", |
| 27 | + top_k: int = 1, |
| 28 | + ): |
| 29 | + self._persist_directory = persist_directory |
| 30 | + self._table_name = table_name |
| 31 | + self._top_k = top_k |
| 32 | + |
| 33 | + # Initialize LanceDB database |
| 34 | + self._db = lancedb.connect(self._persist_directory) |
| 35 | + |
| 36 | + # Initialize or open table |
| 37 | + if self._table_name not in self._db.table_names(): |
| 38 | + self._table = None # Table will be created with the first insertion |
| 39 | + else: |
| 40 | + self._table = self._db.open_table(self._table_name) |
| 41 | + |
| 42 | + def mul_add(self, datas: List[VectorData]): |
| 43 | + """Add multiple vectors to the LanceDB table""" |
| 44 | + vectors, vector_ids = map(list, zip(*((data.data.tolist(), str(data.id)) for data in datas))) |
| 45 | + # Infer the dimension of the vectors |
| 46 | + vector_dim = len(vectors[0]) if vectors else 0 |
| 47 | + |
| 48 | + # Create table with the inferred schema if it doesn't exist |
| 49 | + if self._table is None: |
| 50 | + schema = pa.schema([ |
| 51 | + pa.field("id", pa.string()), |
| 52 | + pa.field("vector", pa.list_(pa.float32(), list_size=vector_dim)) |
| 53 | + ]) |
| 54 | + self._table = self._db.create_table(self._table_name, schema=schema) |
| 55 | + |
| 56 | + # Prepare and add data to the table |
| 57 | + self._table.add(({"id": vector_id, "vector": vector} for vector_id, vector in zip(vector_ids, vectors))) |
| 58 | + |
| 59 | + def search(self, data: np.ndarray, top_k: int = -1): |
| 60 | + """Search for the most similar vectors in the LanceDB table""" |
| 61 | + if len(self._table) == 0: |
| 62 | + return [] |
| 63 | + |
| 64 | + if top_k == -1: |
| 65 | + top_k = self._top_k |
| 66 | + |
| 67 | + results = self._table.search(data.tolist()).limit(top_k).to_list() |
| 68 | + return [(result["_distance"], int(result["id"])) for result in results] |
| 69 | + |
| 70 | + def delete(self, ids: List[int]): |
| 71 | + """Delete vectors from the LanceDB table based on IDs""" |
| 72 | + for vector_id in ids: |
| 73 | + self._table.delete(f"id = '{vector_id}'") |
| 74 | + |
| 75 | + def rebuild(self, ids: Optional[List[int]] = None): |
| 76 | + """Rebuild the index, if applicable""" |
| 77 | + return True |
| 78 | + |
| 79 | + def count(self): |
| 80 | + """Return the total number of vectors in the table""" |
| 81 | + return len(self._table) |
0 commit comments