yangdx commited on
Commit
1ee263c
·
1 Parent(s): f5df192

Refactor Faiss index access with helper method to improve code organization

Browse files
Files changed (1) hide show
  1. lightrag/kg/faiss_impl.py +12 -9
lightrag/kg/faiss_impl.py CHANGED
@@ -74,6 +74,13 @@ class FaissVectorDBStorage(BaseVectorStorage):
74
  self._id_to_meta.update({})
75
  self._load_faiss_index()
76
 
 
 
 
 
 
 
 
77
  async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
78
  """
79
  Insert or update vectors in the Faiss index.
@@ -142,11 +149,9 @@ class FaissVectorDBStorage(BaseVectorStorage):
142
  self._remove_faiss_ids(existing_ids_to_remove)
143
 
144
  # Step 2: Add new vectors
145
- start_idx = (self._index.value if is_multiprocess else self._index).ntotal
146
- if is_multiprocess:
147
- self._index.value.add(embeddings)
148
- else:
149
- self._index.add(embeddings)
150
 
151
  # Step 3: Store metadata + vector for each new ID
152
  for i, meta in enumerate(list_data):
@@ -173,9 +178,7 @@ class FaissVectorDBStorage(BaseVectorStorage):
173
 
174
  # Perform the similarity search
175
  with self._storage_lock:
176
- distances, indices = (
177
- self._index.value if is_multiprocess else self._index
178
- ).search(embedding, top_k)
179
 
180
  distances = distances[0]
181
  indices = indices[0]
@@ -303,7 +306,7 @@ class FaissVectorDBStorage(BaseVectorStorage):
303
  """
304
  with self._storage_lock:
305
  faiss.write_index(
306
- self._index.value if is_multiprocess else self._index,
307
  self._faiss_index_file,
308
  )
309
 
 
74
  self._id_to_meta.update({})
75
  self._load_faiss_index()
76
 
77
+ def _get_index(self):
78
+ """
79
+ Helper method to get the correct index object based on multiprocess mode.
80
+ Returns the actual index object that can be used for operations.
81
+ """
82
+ return self._index.value if is_multiprocess else self._index
83
+
84
  async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
85
  """
86
  Insert or update vectors in the Faiss index.
 
149
  self._remove_faiss_ids(existing_ids_to_remove)
150
 
151
  # Step 2: Add new vectors
152
+ index = self._get_index()
153
+ start_idx = index.ntotal
154
+ index.add(embeddings)
 
 
155
 
156
  # Step 3: Store metadata + vector for each new ID
157
  for i, meta in enumerate(list_data):
 
178
 
179
  # Perform the similarity search
180
  with self._storage_lock:
181
+ distances, indices = self._get_index().search(embedding, top_k)
 
 
182
 
183
  distances = distances[0]
184
  indices = indices[0]
 
306
  """
307
  with self._storage_lock:
308
  faiss.write_index(
309
+ self._get_index(),
310
  self._faiss_index_file,
311
  )
312