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

fix: Fix chroma check exists bug

Browse Source
This commit is contained in:
Fangyin Cheng
2025-03-17 15:57:46 +08:00
parent 9f2e747698
commit 82392dc4eb
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py
View File

@@ -121,6 +121,7 @@ class ChromaStore(VectorStoreBase):
path=self.persist_dir, settings=chroma_settings path=self.persist_dir, settings=chroma_settings
) )
collection_metadata = collection_metadata or {"hnsw:space": "cosine"} collection_metadata = collection_metadata or {"hnsw:space": "cosine"}
self._collection_name = name
self._collection = self._chroma_client.get_or_create_collection( self._collection = self._chroma_client.get_or_create_collection(
name=name, name=name,
embedding_function=None, embedding_function=None,
@@ -196,13 +197,12 @@ class ChromaStore(VectorStoreBase):
def vector_name_exists(self) -> bool: def vector_name_exists(self) -> bool:
"""Whether vector name exists.""" """Whether vector name exists."""
logger.info(f"Check persist_dir: {self.persist_dir}") try:
if not os.path.exists(self.persist_dir): collection = self._chroma_client.get_collection(self._collection_name)
return collection.count() > 0
except Exception as _e:
logger.info(f"Collection {self._collection_name} does not exist")
return False return False
files = os.listdir(self.persist_dir)
# Skip default file: chroma.sqlite3
files = list(filter(lambda f: f != "chroma.sqlite3", files))
return len(files) > 0
def load_document(self, chunks: List[Chunk]) -> List[str]: def load_document(self, chunks: List[Chunk]) -> List[str]:
"""Load document to vector store.""" """Load document to vector store."""