github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-07 12:00:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(storage): Fix load big documents error

Browse Source
This commit is contained in:
Fangyin Cheng
2025-04-11 20:46:48 +08:00
parent 12170e2504
commit c04e3c7cb0
11 changed files with 88 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
packages/dbgpt-core/src/dbgpt/rag/retriever/time_weighted.py
View File

@@ -149,7 +149,7 @@ class TimeWeightedEmbeddingRetriever(EmbeddingRetriever):
self._save_memory_stream()
# Add to vector store
return self._index_store.load_document(dup_docs)
return self._index_store.load_document_with_limit(dup_docs)
def _retrieve(
self, query: str, filters: Optional[MetadataFilters] = None

23
packages/dbgpt-core/src/dbgpt/storage/base.py
View File

@@ -27,9 +27,16 @@ class IndexStoreConfig(BaseParameters):
class IndexStoreBase(ABC):
"""Index store base class."""
def __init__(self, executor: Optional[Executor] = None):
def __init__(
self,
executor: Optional[Executor] = None,
max_chunks_once_load: Optional[int] = None,
max_threads: Optional[int] = None,
):
"""Init index store."""
self._executor = executor or ThreadPoolExecutor()
self._max_chunks_once_load = max_chunks_once_load or 10
self._max_threads = max_threads or 1
@abstractmethod
def get_config(self) -> IndexStoreConfig:
@@ -102,7 +109,10 @@ class IndexStoreBase(ABC):
return True
def load_document_with_limit(
self, chunks: List[Chunk], max_chunks_once_load: int = 10, max_threads: int = 1
self,
chunks: List[Chunk],
max_chunks_once_load: Optional[int] = None,
max_threads: Optional[int] = None,
) -> List[str]:
"""Load document in index database with specified limit.
@@ -114,6 +124,8 @@ class IndexStoreBase(ABC):
Return:
List[str]: Chunk ids.
"""
max_chunks_once_load = max_chunks_once_load or self._max_chunks_once_load
max_threads = max_threads or self._max_threads
# Group the chunks into chunks of size max_chunks
chunk_groups = [
chunks[i : i + max_chunks_once_load]
@@ -141,7 +153,10 @@ class IndexStoreBase(ABC):
return ids
async def aload_document_with_limit(
self, chunks: List[Chunk], max_chunks_once_load: int = 10, max_threads: int = 1
self,
chunks: List[Chunk],
max_chunks_once_load: Optional[int] = None,
max_threads: Optional[int] = None,
) -> List[str]:
"""Load document in index database with specified limit.
@@ -153,6 +168,8 @@ class IndexStoreBase(ABC):
Return:
List[str]: Chunk ids.
"""
max_chunks_once_load = max_chunks_once_load or self._max_chunks_once_load
max_threads = max_threads or self._max_threads
chunk_groups = [
chunks[i : i + max_chunks_once_load]
for i in range(0, len(chunks), max_chunks_once_load)

29
packages/dbgpt-core/src/dbgpt/storage/vector_store/base.py
View File

@@ -88,6 +88,24 @@ class VectorStoreConfig(IndexStoreConfig, RegisterParameters):
),
},
)
max_chunks_once_load: Optional[int] = field(
default=None,
metadata={
"help": _(
"The max chunks once load in vector store, "
"if not set, will use the default value 10."
),
},
)
max_threads: Optional[int] = field(
default=None,
metadata={
"help": _(
"The max threads in vector store, "
"if not set, will use the default value 1."
),
},
)
def create_store(self, **kwargs) -> "VectorStoreBase":
"""Create a new index store from the config."""
@@ -97,9 +115,16 @@ class VectorStoreConfig(IndexStoreConfig, RegisterParameters):
class VectorStoreBase(IndexStoreBase, ABC):
"""Vector store base class."""
def __init__(self, executor: Optional[ThreadPoolExecutor] = None):
def __init__(
self,
executor: Optional[ThreadPoolExecutor] = None,
max_chunks_once_load: Optional[int] = None,
max_threads: Optional[int] = None,
):
"""Initialize vector store."""
super().__init__(executor)
super().__init__(
executor, max_chunks_once_load=max_chunks_once_load, max_threads=max_threads
)
@abstractmethod
def get_config(self) -> VectorStoreConfig: