github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-05-22 07:27:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added client_settings support for chromadb vecstore (#1528)

Browse Source 
# Problem

The ChromaDB vecstore only supported local connection. There was no way
to use a chromadb server.

# Fix
Added `client_settings` as Chroma attribute. 

# Usage

```
from chromadb.config import Settings
from langchain.vectorstores import Chroma

chroma_settings = Settings(chroma_api_impl="rest",
                            chroma_server_host="localhost",
                            chroma_server_http_port="80")

docsearch = Chroma.from_documents(chunks, embeddings, metadatas=metadatas, client_settings=chroma_settings, collection_name=COLLECTION_NAME)
```
This commit is contained in:
Florian Leuerer 2023-03-09 02:42:09 +01:00 committed by GitHub
parent 7eba828e1b
commit 64c6435545
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
langchain/vectorstores/chroma.py
View File

@ -3,12 +3,16 @@ from __future__ import annotations
import logging import logging
import uuid import uuid
from typing import Any, Dict, Iterable, List, Optional, Tuple from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple
from langchain.docstore.document import Document from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings from langchain.embeddings.base import Embeddings
from langchain.vectorstores.base import VectorStore from langchain.vectorstores.base import VectorStore
if TYPE_CHECKING:
import chromadb
import chromadb.config
logger = logging.getLogger() logger = logging.getLogger()
@ -34,6 +38,7 @@ class Chroma(VectorStore):
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
embedding_function: Optional[Embeddings] = None, embedding_function: Optional[Embeddings] = None,
persist_directory: Optional[str] = None, persist_directory: Optional[str] = None,
client_settings: Optional[chromadb.config.Settings] = None,
) -> None: ) -> None:
"""Initialize with Chroma client.""" """Initialize with Chroma client."""
try: try:
@ -45,12 +50,14 @@ class Chroma(VectorStore):
"Please install it with `pip install chromadb`." "Please install it with `pip install chromadb`."
) )
# TODO: Add support for custom client. For now this is in-memory only. if client_settings:
self._client_settings = chromadb.config.Settings() self._client_settings = client_settings
if persist_directory is not None: else:
self._client_settings = chromadb.config.Settings( self._client_settings = chromadb.config.Settings()
chroma_db_impl="duckdb+parquet", persist_directory=persist_directory if persist_directory is not None:
) self._client_settings = chromadb.config.Settings(
chroma_db_impl="duckdb+parquet", persist_directory=persist_directory
)
self._client = chromadb.Client(self._client_settings) self._client = chromadb.Client(self._client_settings)
self._embedding_function = embedding_function self._embedding_function = embedding_function
self._persist_directory = persist_directory self._persist_directory = persist_directory
@ -185,6 +192,7 @@ class Chroma(VectorStore):
ids: Optional[List[str]] = None, ids: Optional[List[str]] = None,
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
persist_directory: Optional[str] = None, persist_directory: Optional[str] = None,
client_settings: Optional[chromadb.config.Settings] = None,
**kwargs: Any, **kwargs: Any,
) -> Chroma: ) -> Chroma:
"""Create a Chroma vectorstore from a raw documents. """Create a Chroma vectorstore from a raw documents.
@ -207,6 +215,7 @@ class Chroma(VectorStore):
collection_name=collection_name, collection_name=collection_name,
embedding_function=embedding, embedding_function=embedding,
persist_directory=persist_directory, persist_directory=persist_directory,
client_settings=client_settings,
) )
chroma_collection.add_texts(texts=texts, metadatas=metadatas, ids=ids) chroma_collection.add_texts(texts=texts, metadatas=metadatas, ids=ids)
return chroma_collection return chroma_collection
@ -219,6 +228,7 @@ class Chroma(VectorStore):
ids: Optional[List[str]] = None, ids: Optional[List[str]] = None,
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
persist_directory: Optional[str] = None, persist_directory: Optional[str] = None,
client_settings: Optional[chromadb.config.Settings] = None,
**kwargs: Any, **kwargs: Any,
) -> Chroma: ) -> Chroma:
"""Create a Chroma vectorstore from a list of documents. """Create a Chroma vectorstore from a list of documents.
@ -244,4 +254,5 @@ class Chroma(VectorStore):
ids=ids, ids=ids,
collection_name=collection_name, collection_name=collection_name,
persist_directory=persist_directory, persist_directory=persist_directory,
client_settings=client_settings,
) )