diff --git a/packages/dbgpt-core/src/dbgpt/storage/base.py b/packages/dbgpt-core/src/dbgpt/storage/base.py index c52f6ec24..b14e0661b 100644 --- a/packages/dbgpt-core/src/dbgpt/storage/base.py +++ b/packages/dbgpt-core/src/dbgpt/storage/base.py @@ -1,4 +1,5 @@ """Index store base class.""" + import asyncio import logging import time @@ -192,13 +193,15 @@ class IndexStoreBase(ABC): logger.info(f"Loaded {loaded_cnt} chunks, total {len(chunks)} chunks.") return ids + async def _run_tasks_with_concurrency(self, tasks, max_concurrent): results = [] for i in range(0, len(tasks), max_concurrent): - batch = tasks[i:i + max_concurrent] + batch = tasks[i : i + max_concurrent] batch_results = await asyncio.gather(*batch, return_exceptions=True) results.extend(batch_results) return results + def similar_search( self, text: str, topk: int, filters: Optional[MetadataFilters] = None ) -> List[Chunk]: diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py index c2a1f01ed..ef91ced42 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py @@ -291,7 +291,7 @@ class ChromaStore(VectorStoreBase): logger.info(f"Total IDs to delete: {total}") for i in range(0, total, batch_size): - batch_ids = id_list[i:i + batch_size] + batch_ids = id_list[i : i + batch_size] logger.info(f"Deleting batch {i // batch_size + 1}: {len(batch_ids)} IDs") try: self._collection.delete(ids=batch_ids)