yangdx
commited on
Commit
·
1ee263c
1
Parent(s):
f5df192
Refactor Faiss index access with helper method to improve code organization
Browse files- 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 |
-
|
146 |
-
|
147 |
-
|
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.
|
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 |
|