Chroma

ChromaVectorStore

Bases: BasePydanticVectorStore

Chroma vector store.

In this vector store, embeddings are stored within a ChromaDB collection.

During query time, the index uses ChromaDB to query for the top k most similar nodes.

Parameters:

Name	Type	Description	Default
chroma_collection	Collection	ChromaDB collection instance	None

Examples:

pip install llama-index-vector-stores-chroma

import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore

# Create a Chroma client and collection
chroma_client = chromadb.EphemeralClient()
chroma_collection = chroma_client.create_collection("example_collection")

# Set up the ChromaVectorStore and StorageContext
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

Source code in llama-index-integrations/vector_stores/llama-index-vector-stores-chroma/llama_index/vector_stores/chroma/base.py

class ChromaVectorStore(BasePydanticVectorStore):
    """Chroma vector store.

    In this vector store, embeddings are stored within a ChromaDB collection.

    During query time, the index uses ChromaDB to query for the top
    k most similar nodes.

    Args:
        chroma_collection (chromadb.api.models.Collection.Collection):
            ChromaDB collection instance

    Examples:
        `pip install llama-index-vector-stores-chroma`

        ```python
        import chromadb
        from llama_index.vector_stores.chroma import ChromaVectorStore

        # Create a Chroma client and collection
        chroma_client = chromadb.EphemeralClient()
        chroma_collection = chroma_client.create_collection("example_collection")

        # Set up the ChromaVectorStore and StorageContext
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        ```

    """

    stores_text: bool = True
    flat_metadata: bool = True

    collection_name: Optional[str]
    host: Optional[str]
    port: Optional[str]
    ssl: bool
    headers: Optional[Dict[str, str]]
    persist_dir: Optional[str]
    collection_kwargs: Dict[str, Any] = Field(default_factory=dict)

    _collection: Any = PrivateAttr()

    def __init__(
        self,
        chroma_collection: Optional[Any] = None,
        collection_name: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[str] = None,
        ssl: bool = False,
        headers: Optional[Dict[str, str]] = None,
        persist_dir: Optional[str] = None,
        collection_kwargs: Optional[dict] = None,
        **kwargs: Any,
    ) -> None:
        """Init params."""
        collection_kwargs = collection_kwargs or {}
        if chroma_collection is None:
            client = chromadb.HttpClient(host=host, port=port, ssl=ssl, headers=headers)
            self._collection = client.get_or_create_collection(
                name=collection_name, **collection_kwargs
            )
        else:
            self._collection = cast(Collection, chroma_collection)

        super().__init__(
            host=host,
            port=port,
            ssl=ssl,
            headers=headers,
            collection_name=collection_name,
            persist_dir=persist_dir,
            collection_kwargs=collection_kwargs or {},
        )

    @classmethod
    def from_collection(cls, collection: Any) -> "ChromaVectorStore":
        try:
            from chromadb import Collection
        except ImportError:
            raise ImportError(import_err_msg)

        if not isinstance(collection, Collection):
            raise Exception("argument is not chromadb collection instance")

        return cls(chroma_collection=collection)

    @classmethod
    def from_params(
        cls,
        collection_name: str,
        host: Optional[str] = None,
        port: Optional[str] = None,
        ssl: bool = False,
        headers: Optional[Dict[str, str]] = None,
        persist_dir: Optional[str] = None,
        collection_kwargs: dict = {},
        **kwargs: Any,
    ) -> "ChromaVectorStore":
        if persist_dir:
            client = chromadb.PersistentClient(path=persist_dir)
            collection = client.get_or_create_collection(
                name=collection_name, **collection_kwargs
            )
        elif host and port:
            client = chromadb.HttpClient(host=host, port=port, ssl=ssl, headers=headers)
            collection = client.get_or_create_collection(
                name=collection_name, **collection_kwargs
            )
        else:
            raise ValueError(
                "Either `persist_dir` or (`host`,`port`) must be specified"
            )
        return cls(
            chroma_collection=collection,
            host=host,
            port=port,
            ssl=ssl,
            headers=headers,
            persist_dir=persist_dir,
            collection_kwargs=collection_kwargs,
            **kwargs,
        )

    @classmethod
    def class_name(cls) -> str:
        return "ChromaVectorStore"

    def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]:
        """Add nodes to index.

        Args:
            nodes: List[BaseNode]: list of nodes with embeddings

        """
        if not self._collection:
            raise ValueError("Collection not initialized")

        max_chunk_size = MAX_CHUNK_SIZE
        node_chunks = chunk_list(nodes, max_chunk_size)

        all_ids = []
        for node_chunk in node_chunks:
            embeddings = []
            metadatas = []
            ids = []
            documents = []
            for node in node_chunk:
                embeddings.append(node.get_embedding())
                metadata_dict = node_to_metadata_dict(
                    node, remove_text=True, flat_metadata=self.flat_metadata
                )
                for key in metadata_dict:
                    if metadata_dict[key] is None:
                        metadata_dict[key] = ""
                metadatas.append(metadata_dict)
                ids.append(node.node_id)
                documents.append(node.get_content(metadata_mode=MetadataMode.NONE))

            self._collection.add(
                embeddings=embeddings,
                ids=ids,
                metadatas=metadatas,
                documents=documents,
            )
            all_ids.extend(ids)

        return all_ids

    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
        """
        Delete nodes using with ref_doc_id.

        Args:
            ref_doc_id (str): The doc_id of the document to delete.

        """
        self._collection.delete(where={"document_id": ref_doc_id})

    @property
    def client(self) -> Any:
        """Return client."""
        return self._collection

    def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
        """Query index for top k most similar nodes.

        Args:
            query_embedding (List[float]): query embedding
            similarity_top_k (int): top k most similar nodes

        """
        if query.filters is not None:
            if "where" in kwargs:
                raise ValueError(
                    "Cannot specify metadata filters via both query and kwargs. "
                    "Use kwargs only for chroma specific items that are "
                    "not supported via the generic query interface."
                )
            where = _to_chroma_filter(query.filters)
        else:
            where = kwargs.pop("where", {})

        if not query.query_embedding:
            return self._get(limit=query.similarity_top_k, where=where, **kwargs)

        return self._query(
            query_embeddings=query.query_embedding,
            n_results=query.similarity_top_k,
            where=where,
            **kwargs,
        )

    def _query(
        self, query_embeddings: List["float"], n_results: int, where: dict, **kwargs
    ) -> VectorStoreQueryResult:
        results = self._collection.query(
            query_embeddings=query_embeddings,
            n_results=n_results,
            where=where,
            **kwargs,
        )

        logger.debug(f"> Top {len(results['documents'])} nodes:")
        nodes = []
        similarities = []
        ids = []
        for node_id, text, metadata, distance in zip(
            results["ids"][0],
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        ):
            try:
                node = metadata_dict_to_node(metadata)
                node.set_content(text)
            except Exception:
                # NOTE: deprecated legacy logic for backward compatibility
                metadata, node_info, relationships = legacy_metadata_dict_to_node(
                    metadata
                )

                node = TextNode(
                    text=text,
                    id_=node_id,
                    metadata=metadata,
                    start_char_idx=node_info.get("start", None),
                    end_char_idx=node_info.get("end", None),
                    relationships=relationships,
                )

            nodes.append(node)

            similarity_score = math.exp(-distance)
            similarities.append(similarity_score)

            logger.debug(
                f"> [Node {node_id}] [Similarity score: {similarity_score}] "
                f"{truncate_text(str(text), 100)}"
            )
            ids.append(node_id)

        return VectorStoreQueryResult(nodes=nodes, similarities=similarities, ids=ids)

    def _get(self, limit: int, where: dict, **kwargs) -> VectorStoreQueryResult:
        results = self._collection.get(
            limit=limit,
            where=where,
            **kwargs,
        )

        logger.debug(f"> Top {len(results['documents'])} nodes:")
        nodes = []
        ids = []

        if not results["ids"]:
            results["ids"] = [[]]

        for node_id, text, metadata in zip(
            results["ids"][0], results["documents"], results["metadatas"]
        ):
            try:
                node = metadata_dict_to_node(metadata)
                node.set_content(text)
            except Exception:
                # NOTE: deprecated legacy logic for backward compatibility
                metadata, node_info, relationships = legacy_metadata_dict_to_node(
                    metadata
                )

                node = TextNode(
                    text=text,
                    id_=node_id,
                    metadata=metadata,
                    start_char_idx=node_info.get("start", None),
                    end_char_idx=node_info.get("end", None),
                    relationships=relationships,
                )

            nodes.append(node)

            logger.debug(
                f"> [Node {node_id}] [Similarity score: N/A - using get()] "
                f"{truncate_text(str(text), 100)}"
            )
            ids.append(node_id)

        return VectorStoreQueryResult(nodes=nodes, ids=ids)

client property

client: Any

Return client.

add

add(nodes: List[BaseNode], **add_kwargs: Any) -> List[str]

Add nodes to index.

Parameters:

Name	Type	Description	Default
nodes	List[BaseNode]	List[BaseNode]: list of nodes with embeddings	required

Source code in llama-index-integrations/vector_stores/llama-index-vector-stores-chroma/llama_index/vector_stores/chroma/base.py

def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]:
    """Add nodes to index.

    Args:
        nodes: List[BaseNode]: list of nodes with embeddings

    """
    if not self._collection:
        raise ValueError("Collection not initialized")

    max_chunk_size = MAX_CHUNK_SIZE
    node_chunks = chunk_list(nodes, max_chunk_size)

    all_ids = []
    for node_chunk in node_chunks:
        embeddings = []
        metadatas = []
        ids = []
        documents = []
        for node in node_chunk:
            embeddings.append(node.get_embedding())
            metadata_dict = node_to_metadata_dict(
                node, remove_text=True, flat_metadata=self.flat_metadata
            )
            for key in metadata_dict:
                if metadata_dict[key] is None:
                    metadata_dict[key] = ""
            metadatas.append(metadata_dict)
            ids.append(node.node_id)
            documents.append(node.get_content(metadata_mode=MetadataMode.NONE))

        self._collection.add(
            embeddings=embeddings,
            ids=ids,
            metadatas=metadatas,
            documents=documents,
        )
        all_ids.extend(ids)

    return all_ids

delete

delete(ref_doc_id: str, **delete_kwargs: Any) -> None

Delete nodes using with ref_doc_id.

Parameters:

Name	Type	Description	Default
ref_doc_id	str	The doc_id of the document to delete.	required

Source code in llama-index-integrations/vector_stores/llama-index-vector-stores-chroma/llama_index/vector_stores/chroma/base.py

def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
    """
    Delete nodes using with ref_doc_id.

    Args:
        ref_doc_id (str): The doc_id of the document to delete.

    """
    self._collection.delete(where={"document_id": ref_doc_id})

query

query(query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult

Query index for top k most similar nodes.

Parameters:

Name	Type	Description	Default
query_embedding	List[float]	query embedding	required
similarity_top_k	int	top k most similar nodes	required

Source code in llama-index-integrations/vector_stores/llama-index-vector-stores-chroma/llama_index/vector_stores/chroma/base.py

def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
    """Query index for top k most similar nodes.

    Args:
        query_embedding (List[float]): query embedding
        similarity_top_k (int): top k most similar nodes

    """
    if query.filters is not None:
        if "where" in kwargs:
            raise ValueError(
                "Cannot specify metadata filters via both query and kwargs. "
                "Use kwargs only for chroma specific items that are "
                "not supported via the generic query interface."
            )
        where = _to_chroma_filter(query.filters)
    else:
        where = kwargs.pop("where", {})

    if not query.query_embedding:
        return self._get(limit=query.similarity_top_k, where=where, **kwargs)

    return self._query(
        query_embeddings=query.query_embedding,
        n_results=query.similarity_top_k,
        where=where,
        **kwargs,
    )