forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
445 lines
16 KiB
Python
445 lines
16 KiB
Python
"""Wrapper around ChromaDB embeddings platform."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Type
|
|
|
|
import numpy as np
|
|
|
|
from langchain.docstore.document import Document
|
|
from langchain.embeddings.base import Embeddings
|
|
from langchain.utils import xor_args
|
|
from langchain.vectorstores.base import VectorStore
|
|
from langchain.vectorstores.utils import maximal_marginal_relevance
|
|
|
|
if TYPE_CHECKING:
|
|
import chromadb
|
|
import chromadb.config
|
|
|
|
logger = logging.getLogger()
|
|
DEFAULT_K = 4 # Number of Documents to return.
|
|
|
|
|
|
def _results_to_docs(results: Any) -> List[Document]:
|
|
return [doc for doc, _ in _results_to_docs_and_scores(results)]
|
|
|
|
|
|
def _results_to_docs_and_scores(results: Any) -> List[Tuple[Document, float]]:
|
|
return [
|
|
# TODO: Chroma can do batch querying,
|
|
# we shouldn't hard code to the 1st result
|
|
(Document(page_content=result[0], metadata=result[1] or {}), result[2])
|
|
for result in zip(
|
|
results["documents"][0],
|
|
results["metadatas"][0],
|
|
results["distances"][0],
|
|
)
|
|
]
|
|
|
|
|
|
class Chroma(VectorStore):
|
|
"""Wrapper around ChromaDB embeddings platform.
|
|
|
|
To use, you should have the ``chromadb`` python package installed.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
from langchain.vectorstores import Chroma
|
|
from langchain.embeddings.openai import OpenAIEmbeddings
|
|
|
|
embeddings = OpenAIEmbeddings()
|
|
vectorstore = Chroma("langchain_store", embeddings)
|
|
"""
|
|
|
|
_LANGCHAIN_DEFAULT_COLLECTION_NAME = "langchain"
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
|
|
embedding_function: Optional[Embeddings] = None,
|
|
persist_directory: Optional[str] = None,
|
|
client_settings: Optional[chromadb.config.Settings] = None,
|
|
collection_metadata: Optional[Dict] = None,
|
|
client: Optional[chromadb.Client] = None,
|
|
) -> None:
|
|
"""Initialize with Chroma client."""
|
|
try:
|
|
import chromadb
|
|
import chromadb.config
|
|
except ImportError:
|
|
raise ValueError(
|
|
"Could not import chromadb python package. "
|
|
"Please install it with `pip install chromadb`."
|
|
)
|
|
|
|
if client is not None:
|
|
self._client = client
|
|
else:
|
|
if client_settings:
|
|
self._client_settings = client_settings
|
|
else:
|
|
self._client_settings = chromadb.config.Settings()
|
|
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._embedding_function = embedding_function
|
|
self._persist_directory = persist_directory
|
|
self._collection = self._client.get_or_create_collection(
|
|
name=collection_name,
|
|
embedding_function=self._embedding_function.embed_documents
|
|
if self._embedding_function is not None
|
|
else None,
|
|
metadata=collection_metadata,
|
|
)
|
|
|
|
@xor_args(("query_texts", "query_embeddings"))
|
|
def __query_collection(
|
|
self,
|
|
query_texts: Optional[List[str]] = None,
|
|
query_embeddings: Optional[List[List[float]]] = None,
|
|
n_results: int = 4,
|
|
where: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Query the chroma collection."""
|
|
try:
|
|
import chromadb # noqa: F401
|
|
except ImportError:
|
|
raise ValueError(
|
|
"Could not import chromadb python package. "
|
|
"Please install it with `pip install chromadb`."
|
|
)
|
|
return self._collection.query(
|
|
query_texts=query_texts,
|
|
query_embeddings=query_embeddings,
|
|
n_results=n_results,
|
|
where=where,
|
|
**kwargs,
|
|
)
|
|
|
|
def add_texts(
|
|
self,
|
|
texts: Iterable[str],
|
|
metadatas: Optional[List[dict]] = None,
|
|
ids: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[str]:
|
|
"""Run more texts through the embeddings and add to the vectorstore.
|
|
|
|
Args:
|
|
texts (Iterable[str]): Texts to add to the vectorstore.
|
|
metadatas (Optional[List[dict]], optional): Optional list of metadatas.
|
|
ids (Optional[List[str]], optional): Optional list of IDs.
|
|
|
|
Returns:
|
|
List[str]: List of IDs of the added texts.
|
|
"""
|
|
# TODO: Handle the case where the user doesn't provide ids on the Collection
|
|
if ids is None:
|
|
ids = [str(uuid.uuid1()) for _ in texts]
|
|
embeddings = None
|
|
if self._embedding_function is not None:
|
|
embeddings = self._embedding_function.embed_documents(list(texts))
|
|
self._collection.add(
|
|
metadatas=metadatas, embeddings=embeddings, documents=texts, ids=ids
|
|
)
|
|
return ids
|
|
|
|
def similarity_search(
|
|
self,
|
|
query: str,
|
|
k: int = DEFAULT_K,
|
|
filter: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Run similarity search with Chroma.
|
|
|
|
Args:
|
|
query (str): Query text to search for.
|
|
k (int): Number of results to return. Defaults to 4.
|
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
|
|
|
Returns:
|
|
List[Document]: List of documents most similar to the query text.
|
|
"""
|
|
docs_and_scores = self.similarity_search_with_score(query, k, filter=filter)
|
|
return [doc for doc, _ in docs_and_scores]
|
|
|
|
def similarity_search_by_vector(
|
|
self,
|
|
embedding: List[float],
|
|
k: int = DEFAULT_K,
|
|
filter: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Return docs most similar to embedding vector.
|
|
Args:
|
|
embedding (str): Embedding to look up documents similar to.
|
|
k (int): Number of Documents to return. Defaults to 4.
|
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
|
Returns:
|
|
List of Documents most similar to the query vector.
|
|
"""
|
|
results = self.__query_collection(
|
|
query_embeddings=embedding, n_results=k, where=filter
|
|
)
|
|
return _results_to_docs(results)
|
|
|
|
def similarity_search_with_score(
|
|
self,
|
|
query: str,
|
|
k: int = DEFAULT_K,
|
|
filter: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Tuple[Document, float]]:
|
|
"""Run similarity search with Chroma with distance.
|
|
|
|
Args:
|
|
query (str): Query text to search for.
|
|
k (int): Number of results to return. Defaults to 4.
|
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
|
|
|
Returns:
|
|
List[Tuple[Document, float]]: List of documents most similar to
|
|
the query text and cosine distance in float for each.
|
|
Lower score represents more similarity.
|
|
"""
|
|
if self._embedding_function is None:
|
|
results = self.__query_collection(
|
|
query_texts=[query], n_results=k, where=filter
|
|
)
|
|
else:
|
|
query_embedding = self._embedding_function.embed_query(query)
|
|
results = self.__query_collection(
|
|
query_embeddings=[query_embedding], n_results=k, where=filter
|
|
)
|
|
|
|
return _results_to_docs_and_scores(results)
|
|
|
|
def _similarity_search_with_relevance_scores(
|
|
self,
|
|
query: str,
|
|
k: int = 4,
|
|
**kwargs: Any,
|
|
) -> List[Tuple[Document, float]]:
|
|
return self.similarity_search_with_score(query, k)
|
|
|
|
def max_marginal_relevance_search_by_vector(
|
|
self,
|
|
embedding: List[float],
|
|
k: int = DEFAULT_K,
|
|
fetch_k: int = 20,
|
|
lambda_mult: float = 0.5,
|
|
filter: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Return docs selected using the maximal marginal relevance.
|
|
Maximal marginal relevance optimizes for similarity to query AND diversity
|
|
among selected documents.
|
|
|
|
Args:
|
|
embedding: Embedding to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
fetch_k: Number of Documents to fetch to pass to MMR algorithm.
|
|
lambda_mult: Number between 0 and 1 that determines the degree
|
|
of diversity among the results with 0 corresponding
|
|
to maximum diversity and 1 to minimum diversity.
|
|
Defaults to 0.5.
|
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
|
|
|
Returns:
|
|
List of Documents selected by maximal marginal relevance.
|
|
"""
|
|
|
|
results = self.__query_collection(
|
|
query_embeddings=embedding,
|
|
n_results=fetch_k,
|
|
where=filter,
|
|
include=["metadatas", "documents", "distances", "embeddings"],
|
|
)
|
|
mmr_selected = maximal_marginal_relevance(
|
|
np.array(embedding, dtype=np.float32),
|
|
results["embeddings"][0],
|
|
k=k,
|
|
lambda_mult=lambda_mult,
|
|
)
|
|
|
|
candidates = _results_to_docs(results)
|
|
|
|
selected_results = [r for i, r in enumerate(candidates) if i in mmr_selected]
|
|
return selected_results
|
|
|
|
def max_marginal_relevance_search(
|
|
self,
|
|
query: str,
|
|
k: int = DEFAULT_K,
|
|
fetch_k: int = 20,
|
|
lambda_mult: float = 0.5,
|
|
filter: Optional[Dict[str, str]] = None,
|
|
**kwargs: Any,
|
|
) -> List[Document]:
|
|
"""Return docs selected using the maximal marginal relevance.
|
|
Maximal marginal relevance optimizes for similarity to query AND diversity
|
|
among selected documents.
|
|
|
|
Args:
|
|
query: Text to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
fetch_k: Number of Documents to fetch to pass to MMR algorithm.
|
|
lambda_mult: Number between 0 and 1 that determines the degree
|
|
of diversity among the results with 0 corresponding
|
|
to maximum diversity and 1 to minimum diversity.
|
|
Defaults to 0.5.
|
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
|
|
|
Returns:
|
|
List of Documents selected by maximal marginal relevance.
|
|
"""
|
|
if self._embedding_function is None:
|
|
raise ValueError(
|
|
"For MMR search, you must specify an embedding function on" "creation."
|
|
)
|
|
|
|
embedding = self._embedding_function.embed_query(query)
|
|
docs = self.max_marginal_relevance_search_by_vector(
|
|
embedding, k, fetch_k, lambda_mul=lambda_mult, filter=filter
|
|
)
|
|
return docs
|
|
|
|
def delete_collection(self) -> None:
|
|
"""Delete the collection."""
|
|
self._client.delete_collection(self._collection.name)
|
|
|
|
def get(self, include: Optional[List[str]] = None) -> Dict[str, Any]:
|
|
"""Gets the collection.
|
|
|
|
Args:
|
|
include (Optional[List[str]]): List of fields to include from db.
|
|
Defaults to None.
|
|
"""
|
|
if include is not None:
|
|
return self._collection.get(include=include)
|
|
else:
|
|
return self._collection.get()
|
|
|
|
def persist(self) -> None:
|
|
"""Persist the collection.
|
|
|
|
This can be used to explicitly persist the data to disk.
|
|
It will also be called automatically when the object is destroyed.
|
|
"""
|
|
if self._persist_directory is None:
|
|
raise ValueError(
|
|
"You must specify a persist_directory on"
|
|
"creation to persist the collection."
|
|
)
|
|
self._client.persist()
|
|
|
|
def update_document(self, document_id: str, document: Document) -> None:
|
|
"""Update a document in the collection.
|
|
|
|
Args:
|
|
document_id (str): ID of the document to update.
|
|
document (Document): Document to update.
|
|
"""
|
|
text = document.page_content
|
|
metadata = document.metadata
|
|
if self._embedding_function is None:
|
|
raise ValueError(
|
|
"For update, you must specify an embedding function on creation."
|
|
)
|
|
embeddings = self._embedding_function.embed_documents([text])
|
|
|
|
self._collection.update(
|
|
ids=[document_id],
|
|
embeddings=embeddings,
|
|
documents=[text],
|
|
metadatas=[metadata],
|
|
)
|
|
|
|
@classmethod
|
|
def from_texts(
|
|
cls: Type[Chroma],
|
|
texts: List[str],
|
|
embedding: Optional[Embeddings] = None,
|
|
metadatas: Optional[List[dict]] = None,
|
|
ids: Optional[List[str]] = None,
|
|
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
|
|
persist_directory: Optional[str] = None,
|
|
client_settings: Optional[chromadb.config.Settings] = None,
|
|
client: Optional[chromadb.Client] = None,
|
|
**kwargs: Any,
|
|
) -> Chroma:
|
|
"""Create a Chroma vectorstore from a raw documents.
|
|
|
|
If a persist_directory is specified, the collection will be persisted there.
|
|
Otherwise, the data will be ephemeral in-memory.
|
|
|
|
Args:
|
|
texts (List[str]): List of texts to add to the collection.
|
|
collection_name (str): Name of the collection to create.
|
|
persist_directory (Optional[str]): Directory to persist the collection.
|
|
embedding (Optional[Embeddings]): Embedding function. Defaults to None.
|
|
metadatas (Optional[List[dict]]): List of metadatas. Defaults to None.
|
|
ids (Optional[List[str]]): List of document IDs. Defaults to None.
|
|
client_settings (Optional[chromadb.config.Settings]): Chroma client settings
|
|
|
|
Returns:
|
|
Chroma: Chroma vectorstore.
|
|
"""
|
|
chroma_collection = cls(
|
|
collection_name=collection_name,
|
|
embedding_function=embedding,
|
|
persist_directory=persist_directory,
|
|
client_settings=client_settings,
|
|
client=client,
|
|
)
|
|
chroma_collection.add_texts(texts=texts, metadatas=metadatas, ids=ids)
|
|
return chroma_collection
|
|
|
|
@classmethod
|
|
def from_documents(
|
|
cls: Type[Chroma],
|
|
documents: List[Document],
|
|
embedding: Optional[Embeddings] = None,
|
|
ids: Optional[List[str]] = None,
|
|
collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
|
|
persist_directory: Optional[str] = None,
|
|
client_settings: Optional[chromadb.config.Settings] = None,
|
|
client: Optional[chromadb.Client] = None, # Add this line
|
|
**kwargs: Any,
|
|
) -> Chroma:
|
|
"""Create a Chroma vectorstore from a list of documents.
|
|
|
|
If a persist_directory is specified, the collection will be persisted there.
|
|
Otherwise, the data will be ephemeral in-memory.
|
|
|
|
Args:
|
|
collection_name (str): Name of the collection to create.
|
|
persist_directory (Optional[str]): Directory to persist the collection.
|
|
ids (Optional[List[str]]): List of document IDs. Defaults to None.
|
|
documents (List[Document]): List of documents to add to the vectorstore.
|
|
embedding (Optional[Embeddings]): Embedding function. Defaults to None.
|
|
client_settings (Optional[chromadb.config.Settings]): Chroma client settings
|
|
Returns:
|
|
Chroma: Chroma vectorstore.
|
|
"""
|
|
texts = [doc.page_content for doc in documents]
|
|
metadatas = [doc.metadata for doc in documents]
|
|
return cls.from_texts(
|
|
texts=texts,
|
|
embedding=embedding,
|
|
metadatas=metadatas,
|
|
ids=ids,
|
|
collection_name=collection_name,
|
|
persist_directory=persist_directory,
|
|
client_settings=client_settings,
|
|
client=client,
|
|
)
|