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.
452 lines
16 KiB
Python
452 lines
16 KiB
Python
"""Wrapper around Annoy vector database."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import pickle
|
|
import uuid
|
|
from configparser import ConfigParser
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from langchain.docstore.base import Docstore
|
|
from langchain.docstore.document import Document
|
|
from langchain.docstore.in_memory import InMemoryDocstore
|
|
from langchain.embeddings.base import Embeddings
|
|
from langchain.vectorstores.base import VectorStore
|
|
from langchain.vectorstores.utils import maximal_marginal_relevance
|
|
|
|
INDEX_METRICS = frozenset(["angular", "euclidean", "manhattan", "hamming", "dot"])
|
|
DEFAULT_METRIC = "angular"
|
|
|
|
|
|
def dependable_annoy_import() -> Any:
|
|
"""Import annoy if available, otherwise raise error."""
|
|
try:
|
|
import annoy
|
|
except ImportError:
|
|
raise ValueError(
|
|
"Could not import annoy python package. "
|
|
"Please install it with `pip install --user annoy` "
|
|
)
|
|
return annoy
|
|
|
|
|
|
class Annoy(VectorStore):
|
|
"""Wrapper around Annoy vector database.
|
|
|
|
To use, you should have the ``annoy`` python package installed.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
from langchain import Annoy
|
|
db = Annoy(embedding_function, index, docstore, index_to_docstore_id)
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
embedding_function: Callable,
|
|
index: Any,
|
|
metric: str,
|
|
docstore: Docstore,
|
|
index_to_docstore_id: Dict[int, str],
|
|
):
|
|
"""Initialize with necessary components."""
|
|
self.embedding_function = embedding_function
|
|
self.index = index
|
|
self.metric = metric
|
|
self.docstore = docstore
|
|
self.index_to_docstore_id = index_to_docstore_id
|
|
|
|
def add_texts(
|
|
self,
|
|
texts: Iterable[str],
|
|
metadatas: Optional[List[dict]] = None,
|
|
**kwargs: Any,
|
|
) -> List[str]:
|
|
raise NotImplementedError(
|
|
"Annoy does not allow to add new data once the index is build."
|
|
)
|
|
|
|
def process_index_results(
|
|
self, idxs: List[int], dists: List[float]
|
|
) -> List[Tuple[Document, float]]:
|
|
"""Turns annoy results into a list of documents and scores.
|
|
|
|
Args:
|
|
idxs: List of indices of the documents in the index.
|
|
dists: List of distances of the documents in the index.
|
|
Returns:
|
|
List of Documents and scores.
|
|
"""
|
|
docs = []
|
|
for idx, dist in zip(idxs, dists):
|
|
_id = self.index_to_docstore_id[idx]
|
|
doc = self.docstore.search(_id)
|
|
if not isinstance(doc, Document):
|
|
raise ValueError(f"Could not find document for id {_id}, got {doc}")
|
|
docs.append((doc, dist))
|
|
return docs
|
|
|
|
def similarity_search_with_score_by_vector(
|
|
self, embedding: List[float], k: int = 4, search_k: int = -1
|
|
) -> List[Tuple[Document, float]]:
|
|
"""Return docs most similar to query.
|
|
|
|
Args:
|
|
query: Text to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
Returns:
|
|
List of Documents most similar to the query and score for each
|
|
"""
|
|
idxs, dists = self.index.get_nns_by_vector(
|
|
embedding, k, search_k=search_k, include_distances=True
|
|
)
|
|
return self.process_index_results(idxs, dists)
|
|
|
|
def similarity_search_with_score_by_index(
|
|
self, docstore_index: int, k: int = 4, search_k: int = -1
|
|
) -> List[Tuple[Document, float]]:
|
|
"""Return docs most similar to query.
|
|
|
|
Args:
|
|
query: Text to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
Returns:
|
|
List of Documents most similar to the query and score for each
|
|
"""
|
|
idxs, dists = self.index.get_nns_by_item(
|
|
docstore_index, k, search_k=search_k, include_distances=True
|
|
)
|
|
return self.process_index_results(idxs, dists)
|
|
|
|
def similarity_search_with_score(
|
|
self, query: str, k: int = 4, search_k: int = -1
|
|
) -> List[Tuple[Document, float]]:
|
|
"""Return docs most similar to query.
|
|
|
|
Args:
|
|
query: Text to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
|
|
Returns:
|
|
List of Documents most similar to the query and score for each
|
|
"""
|
|
embedding = self.embedding_function(query)
|
|
docs = self.similarity_search_with_score_by_vector(embedding, k, search_k)
|
|
return docs
|
|
|
|
def similarity_search_by_vector(
|
|
self, embedding: List[float], k: int = 4, search_k: int = -1, **kwargs: Any
|
|
) -> List[Document]:
|
|
"""Return docs most similar to embedding vector.
|
|
|
|
Args:
|
|
embedding: Embedding to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
|
|
Returns:
|
|
List of Documents most similar to the embedding.
|
|
"""
|
|
docs_and_scores = self.similarity_search_with_score_by_vector(
|
|
embedding, k, search_k
|
|
)
|
|
return [doc for doc, _ in docs_and_scores]
|
|
|
|
def similarity_search_by_index(
|
|
self, docstore_index: int, k: int = 4, search_k: int = -1, **kwargs: Any
|
|
) -> List[Document]:
|
|
"""Return docs most similar to docstore_index.
|
|
|
|
Args:
|
|
docstore_index: Index of document in docstore
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
|
|
Returns:
|
|
List of Documents most similar to the embedding.
|
|
"""
|
|
docs_and_scores = self.similarity_search_with_score_by_index(
|
|
docstore_index, k, search_k
|
|
)
|
|
return [doc for doc, _ in docs_and_scores]
|
|
|
|
def similarity_search(
|
|
self, query: str, k: int = 4, search_k: int = -1, **kwargs: Any
|
|
) -> List[Document]:
|
|
"""Return docs most similar to query.
|
|
|
|
Args:
|
|
query: Text to look up documents similar to.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
search_k: inspect up to search_k nodes which defaults
|
|
to n_trees * n if not provided
|
|
|
|
Returns:
|
|
List of Documents most similar to the query.
|
|
"""
|
|
docs_and_scores = self.similarity_search_with_score(query, k, search_k)
|
|
return [doc for doc, _ in docs_and_scores]
|
|
|
|
def max_marginal_relevance_search_by_vector(
|
|
self,
|
|
embedding: List[float],
|
|
k: int = 4,
|
|
fetch_k: int = 20,
|
|
lambda_mult: float = 0.5,
|
|
**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.
|
|
fetch_k: Number of Documents to fetch to pass to MMR algorithm.
|
|
k: Number of Documents to return. Defaults to 4.
|
|
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.
|
|
|
|
Returns:
|
|
List of Documents selected by maximal marginal relevance.
|
|
"""
|
|
idxs = self.index.get_nns_by_vector(
|
|
embedding, fetch_k, search_k=-1, include_distances=False
|
|
)
|
|
embeddings = [self.index.get_item_vector(i) for i in idxs]
|
|
mmr_selected = maximal_marginal_relevance(
|
|
np.array([embedding], dtype=np.float32),
|
|
embeddings,
|
|
k=k,
|
|
lambda_mult=lambda_mult,
|
|
)
|
|
# ignore the -1's if not enough docs are returned/indexed
|
|
selected_indices = [idxs[i] for i in mmr_selected if i != -1]
|
|
|
|
docs = []
|
|
for i in selected_indices:
|
|
_id = self.index_to_docstore_id[i]
|
|
doc = self.docstore.search(_id)
|
|
if not isinstance(doc, Document):
|
|
raise ValueError(f"Could not find document for id {_id}, got {doc}")
|
|
docs.append(doc)
|
|
return docs
|
|
|
|
def max_marginal_relevance_search(
|
|
self,
|
|
query: str,
|
|
k: int = 4,
|
|
fetch_k: int = 20,
|
|
lambda_mult: float = 0.5,
|
|
**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.
|
|
Returns:
|
|
List of Documents selected by maximal marginal relevance.
|
|
"""
|
|
embedding = self.embedding_function(query)
|
|
docs = self.max_marginal_relevance_search_by_vector(
|
|
embedding, k, fetch_k, lambda_mult=lambda_mult
|
|
)
|
|
return docs
|
|
|
|
@classmethod
|
|
def __from(
|
|
cls,
|
|
texts: List[str],
|
|
embeddings: List[List[float]],
|
|
embedding: Embeddings,
|
|
metadatas: Optional[List[dict]] = None,
|
|
metric: str = DEFAULT_METRIC,
|
|
trees: int = 100,
|
|
n_jobs: int = -1,
|
|
**kwargs: Any,
|
|
) -> Annoy:
|
|
if metric not in INDEX_METRICS:
|
|
raise ValueError(
|
|
(
|
|
f"Unsupported distance metric: {metric}. "
|
|
f"Expected one of {list(INDEX_METRICS)}"
|
|
)
|
|
)
|
|
annoy = dependable_annoy_import()
|
|
if not embeddings:
|
|
raise ValueError("embeddings must be provided to build AnnoyIndex")
|
|
f = len(embeddings[0])
|
|
index = annoy.AnnoyIndex(f, metric=metric)
|
|
for i, emb in enumerate(embeddings):
|
|
index.add_item(i, emb)
|
|
index.build(trees, n_jobs=n_jobs)
|
|
|
|
documents = []
|
|
for i, text in enumerate(texts):
|
|
metadata = metadatas[i] if metadatas else {}
|
|
documents.append(Document(page_content=text, metadata=metadata))
|
|
index_to_id = {i: str(uuid.uuid4()) for i in range(len(documents))}
|
|
docstore = InMemoryDocstore(
|
|
{index_to_id[i]: doc for i, doc in enumerate(documents)}
|
|
)
|
|
return cls(embedding.embed_query, index, metric, docstore, index_to_id)
|
|
|
|
@classmethod
|
|
def from_texts(
|
|
cls,
|
|
texts: List[str],
|
|
embedding: Embeddings,
|
|
metadatas: Optional[List[dict]] = None,
|
|
metric: str = DEFAULT_METRIC,
|
|
trees: int = 100,
|
|
n_jobs: int = -1,
|
|
**kwargs: Any,
|
|
) -> Annoy:
|
|
"""Construct Annoy wrapper from raw documents.
|
|
|
|
Args:
|
|
texts: List of documents to index.
|
|
embedding: Embedding function to use.
|
|
metadatas: List of metadata dictionaries to associate with documents.
|
|
metric: Metric to use for indexing. Defaults to "angular".
|
|
trees: Number of trees to use for indexing. Defaults to 100.
|
|
n_jobs: Number of jobs to use for indexing. Defaults to -1.
|
|
|
|
This is a user friendly interface that:
|
|
1. Embeds documents.
|
|
2. Creates an in memory docstore
|
|
3. Initializes the Annoy database
|
|
|
|
This is intended to be a quick way to get started.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
from langchain import Annoy
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
embeddings = OpenAIEmbeddings()
|
|
index = Annoy.from_texts(texts, embeddings)
|
|
"""
|
|
embeddings = embedding.embed_documents(texts)
|
|
return cls.__from(
|
|
texts, embeddings, embedding, metadatas, metric, trees, n_jobs, **kwargs
|
|
)
|
|
|
|
@classmethod
|
|
def from_embeddings(
|
|
cls,
|
|
text_embeddings: List[Tuple[str, List[float]]],
|
|
embedding: Embeddings,
|
|
metadatas: Optional[List[dict]] = None,
|
|
metric: str = DEFAULT_METRIC,
|
|
trees: int = 100,
|
|
n_jobs: int = -1,
|
|
**kwargs: Any,
|
|
) -> Annoy:
|
|
"""Construct Annoy wrapper from embeddings.
|
|
|
|
Args:
|
|
text_embeddings: List of tuples of (text, embedding)
|
|
embedding: Embedding function to use.
|
|
metadatas: List of metadata dictionaries to associate with documents.
|
|
metric: Metric to use for indexing. Defaults to "angular".
|
|
trees: Number of trees to use for indexing. Defaults to 100.
|
|
n_jobs: Number of jobs to use for indexing. Defaults to -1
|
|
|
|
This is a user friendly interface that:
|
|
1. Creates an in memory docstore with provided embeddings
|
|
2. Initializes the Annoy database
|
|
|
|
This is intended to be a quick way to get started.
|
|
|
|
Example:
|
|
.. code-block:: python
|
|
|
|
from langchain import Annoy
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
embeddings = OpenAIEmbeddings()
|
|
text_embeddings = embeddings.embed_documents(texts)
|
|
text_embedding_pairs = list(zip(texts, text_embeddings))
|
|
db = Annoy.from_embeddings(text_embedding_pairs, embeddings)
|
|
"""
|
|
texts = [t[0] for t in text_embeddings]
|
|
embeddings = [t[1] for t in text_embeddings]
|
|
|
|
return cls.__from(
|
|
texts, embeddings, embedding, metadatas, metric, trees, n_jobs, **kwargs
|
|
)
|
|
|
|
def save_local(self, folder_path: str, prefault: bool = False) -> None:
|
|
"""Save Annoy index, docstore, and index_to_docstore_id to disk.
|
|
|
|
Args:
|
|
folder_path: folder path to save index, docstore,
|
|
and index_to_docstore_id to.
|
|
prefault: Whether to pre-load the index into memory.
|
|
"""
|
|
path = Path(folder_path)
|
|
os.makedirs(path, exist_ok=True)
|
|
# save index, index config, docstore and index_to_docstore_id
|
|
config_object = ConfigParser()
|
|
config_object["ANNOY"] = {
|
|
"f": self.index.f,
|
|
"metric": self.metric,
|
|
}
|
|
self.index.save(str(path / "index.annoy"), prefault=prefault)
|
|
with open(path / "index.pkl", "wb") as file:
|
|
pickle.dump((self.docstore, self.index_to_docstore_id, config_object), file)
|
|
|
|
@classmethod
|
|
def load_local(
|
|
cls,
|
|
folder_path: str,
|
|
embeddings: Embeddings,
|
|
) -> Annoy:
|
|
"""Load Annoy index, docstore, and index_to_docstore_id to disk.
|
|
|
|
Args:
|
|
folder_path: folder path to load index, docstore,
|
|
and index_to_docstore_id from.
|
|
embeddings: Embeddings to use when generating queries.
|
|
"""
|
|
path = Path(folder_path)
|
|
# load index separately since it is not picklable
|
|
annoy = dependable_annoy_import()
|
|
# load docstore and index_to_docstore_id
|
|
with open(path / "index.pkl", "rb") as file:
|
|
docstore, index_to_docstore_id, config_object = pickle.load(file)
|
|
|
|
f = int(config_object["ANNOY"]["f"])
|
|
metric = config_object["ANNOY"]["metric"]
|
|
|
|
index = annoy.AnnoyIndex(f, metric=metric)
|
|
index.load(str(path / "index.annoy"))
|
|
|
|
return cls(
|
|
embeddings.embed_query, index, metric, docstore, index_to_docstore_id
|
|
)
|