fix chroma init bug (#7639)

This commit is contained in:
Bagatur 2023-07-13 03:00:33 -04:00 committed by GitHub
parent f307ca094b
commit b08f903755
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 11 deletions

View File

@ -87,23 +87,26 @@ class Chroma(VectorStore):
) )
if client is not None: if client is not None:
self._client_settings = client_settings
self._client = client self._client = client
self._persist_directory = persist_directory
else: else:
if client_settings: if client_settings:
self._client_settings = client_settings _client_settings = client_settings
else: elif persist_directory:
self._client_settings = chromadb.config.Settings() _client_settings = chromadb.config.Settings(
if persist_directory is not None:
self._client_settings = chromadb.config.Settings(
chroma_db_impl="duckdb+parquet", chroma_db_impl="duckdb+parquet",
persist_directory=persist_directory, persist_directory=persist_directory,
) )
self._client = chromadb.Client(self._client_settings) else:
_client_settings = chromadb.config.Settings()
self._client_settings = _client_settings
self._client = chromadb.Client(_client_settings)
self._persist_directory = (
_client_settings.persist_directory or persist_directory
)
self._embedding_function = embedding_function self._embedding_function = embedding_function
self._persist_directory = (
self._client_settings.persist_directory or persist_directory
)
self._collection = self._client.get_or_create_collection( self._collection = self._client.get_or_create_collection(
name=collection_name, name=collection_name,
embedding_function=self._embedding_function.embed_documents embedding_function=self._embedding_function.embed_documents

View File

@ -267,3 +267,17 @@ def test_chroma_with_relevance_score_custom_normalization_fn() -> None:
(Document(page_content="bar", metadata={"page": "1"}), -0.0), (Document(page_content="bar", metadata={"page": "1"}), -0.0),
(Document(page_content="baz", metadata={"page": "2"}), -0.0), (Document(page_content="baz", metadata={"page": "2"}), -0.0),
] ]
def test_init_from_client() -> None:
import chromadb
client = chromadb.Client(chromadb.config.Settings())
Chroma(client=client)
def test_init_from_client_settings() -> None:
import chromadb
client_settings = chromadb.config.Settings()
Chroma(client_settings=client_settings)