Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add overwrite to load index #897

Merged
merged 1 commit into from
Feb 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions presets/ragengine/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,10 +310,11 @@ async def persist_index(
)
async def load_index(
index_name: str,
path: str = Query(DEFAULT_VECTOR_DB_PERSIST_DIR, description="Path to load the index from")
path: str = Query(DEFAULT_VECTOR_DB_PERSIST_DIR, description="Path to load the index from"),
overwrite: bool = Query(False, description="Overwrite the existing index if it already exists")
):
try:
await rag_ops.load(index_name, path)
await rag_ops.load(index_name, path, overwrite)
return {"message": f"Successfully loaded index {index_name} from {path}."}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
Expand Down
10 changes: 5 additions & 5 deletions presets/ragengine/vector_store/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,21 +290,21 @@
logger.error(f"Failed to persist index {index_name}. Error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Persistence failed: {str(e)}")

async def load(self, index_name: str, path: str):
async def load(self, index_name: str, path: str, overwrite: bool):
"""Common logic for loading an index."""
# Check path existence before acquiring any lock
if not os.path.exists(path):
raise HTTPException(status_code=404, detail=f"Path does not exist: {path}")
if self.use_rwlock:
async with self.rwlock.writer_lock:
await self._load_internal(index_name, path)
await self._load_internal(index_name, path, overwrite)
else:
await self._load_internal(index_name, path)
await self._load_internal(index_name, path, overwrite)

Check warning on line 302 in presets/ragengine/vector_store/base.py

View check run for this annotation

Codecov / codecov/patch

presets/ragengine/vector_store/base.py#L302

Added line #L302 was not covered by tests

async def _load_internal(self, index_name: str, path: str):
async def _load_internal(self, index_name: str, path: str, overwrite: bool):
"""Common logic for loading an index."""
try:
if index_name in self.index_map:
if index_name in self.index_map and not overwrite:
raise HTTPException(
status_code=409,
detail=f"Index '{index_name}' already exists. Use a different name or delete the existing index first."
Expand Down
4 changes: 2 additions & 2 deletions presets/ragengine/vector_store_manager/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,9 @@ async def persist(self, index_name: str, path: str) -> None:
"""Persist existing index."""
return await self.vector_store.persist(index_name, path)

async def load(self, index_name: str, path: str) -> None:
async def load(self, index_name: str, path: str, overwrite: bool) -> None:
"""Load existing index."""
return await self.vector_store.load(index_name, path)
return await self.vector_store.load(index_name, path, overwrite)

async def shutdown(self):
"""Shutdown the manager."""
Expand Down