from typing import Any, Iterable, List, Optional, Tuple
from uuid import uuid4
import numpy as np
import requests
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.utils import get_from_env
from langchain_core.vectorstores import VectorStore
from langchain_community.vectorstores.utils import DistanceStrategy
[docs]class SemaDB(VectorStore):
"""`SemaDB` vector store.
This vector store is a wrapper around the SemaDB database.
Example:
.. code-block:: python
from langchain_community.vectorstores import SemaDB
db = SemaDB('mycollection', 768, embeddings, DistanceStrategy.COSINE)
"""
HOST = "semadb.p.rapidapi.com"
BASE_URL = "https://" + HOST
[docs] def __init__(
self,
collection_name: str,
vector_size: int,
embedding: Embeddings,
distance_strategy: DistanceStrategy = DistanceStrategy.EUCLIDEAN_DISTANCE,
api_key: str = "",
):
"""initialize the SemaDB vector store."""
self.collection_name = collection_name
self.vector_size = vector_size
self.api_key = api_key or get_from_env("api_key", "SEMADB_API_KEY")
self._embedding = embedding
self.distance_strategy = distance_strategy
@property
def headers(self) -> dict:
"""Return the common headers."""
return {
"content-type": "application/json",
"X-RapidAPI-Key": self.api_key,
"X-RapidAPI-Host": SemaDB.HOST,
}
def _get_internal_distance_strategy(self) -> str:
"""Return the internal distance strategy."""
if self.distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
return "euclidean"
elif self.distance_strategy == DistanceStrategy.MAX_INNER_PRODUCT:
raise ValueError("Max inner product is not supported by SemaDB")
elif self.distance_strategy == DistanceStrategy.DOT_PRODUCT:
return "dot"
elif self.distance_strategy == DistanceStrategy.JACCARD:
raise ValueError("Max inner product is not supported by SemaDB")
elif self.distance_strategy == DistanceStrategy.COSINE:
return "cosine"
else:
raise ValueError(f"Unknown distance strategy {self.distance_strategy}")
[docs] def create_collection(self) -> bool:
"""Creates the corresponding collection in SemaDB."""
payload = {
"id": self.collection_name,
"vectorSize": self.vector_size,
"distanceMetric": self._get_internal_distance_strategy(),
}
response = requests.post(
SemaDB.BASE_URL + "/collections",
json=payload,
headers=self.headers,
)
return response.status_code == 200
[docs] def delete_collection(self) -> bool:
"""Deletes the corresponding collection in SemaDB."""
response = requests.delete(
SemaDB.BASE_URL + f"/collections/{self.collection_name}",
headers=self.headers,
)
return response.status_code == 200
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
batch_size: int = 1000,
**kwargs: Any,
) -> List[str]:
"""Add texts to the vector store."""
if not isinstance(texts, list):
texts = list(texts)
embeddings = self._embedding.embed_documents(texts)
# Check dimensions
if len(embeddings[0]) != self.vector_size:
raise ValueError(
f"Embedding size mismatch {len(embeddings[0])} != {self.vector_size}"
)
# Normalise if needed
if self.distance_strategy == DistanceStrategy.COSINE:
embed_matrix = np.array(embeddings)
embed_matrix = embed_matrix / np.linalg.norm(
embed_matrix, axis=1, keepdims=True
)
embeddings = embed_matrix.tolist()
# Create points
ids: List[str] = []
points = []
if metadatas is not None:
for text, embedding, metadata in zip(texts, embeddings, metadatas):
new_id = str(uuid4())
ids.append(new_id)
points.append(
{
"id": new_id,
"vector": embedding,
"metadata": {**metadata, **{"text": text}},
}
)
else:
for text, embedding in zip(texts, embeddings):
new_id = str(uuid4())
ids.append(new_id)
points.append(
{
"id": new_id,
"vector": embedding,
"metadata": {"text": text},
}
)
# Insert points in batches
for i in range(0, len(points), batch_size):
batch = points[i : i + batch_size]
response = requests.post(
SemaDB.BASE_URL + f"/collections/{self.collection_name}/points",
json={"points": batch},
headers=self.headers,
)
if response.status_code != 200:
print("HERE--", batch) # noqa: T201
raise ValueError(f"Error adding points: {response.text}")
failed_ranges = response.json()["failedRanges"]
if len(failed_ranges) > 0:
raise ValueError(f"Error adding points: {failed_ranges}")
# Return ids
return ids
@property
def embeddings(self) -> Embeddings:
"""Return the embeddings."""
return self._embedding
[docs] def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
"""Delete by vector ID or other criteria.
Args:
ids: List of ids to delete.
**kwargs: Other keyword arguments that subclasses might use.
Returns:
Optional[bool]: True if deletion is successful,
False otherwise, None if not implemented.
"""
payload = {
"ids": ids,
}
response = requests.delete(
SemaDB.BASE_URL + f"/collections/{self.collection_name}/points",
json=payload,
headers=self.headers,
)
return response.status_code == 200 and len(response.json()["failedPoints"]) == 0
def _search_points(self, embedding: List[float], k: int = 4) -> List[dict]:
"""Search points."""
# Normalise if needed
if self.distance_strategy == DistanceStrategy.COSINE:
vec = np.array(embedding)
vec = vec / np.linalg.norm(vec)
embedding = vec.tolist()
# Perform search request
payload = {
"vector": embedding,
"limit": k,
}
response = requests.post(
SemaDB.BASE_URL + f"/collections/{self.collection_name}/points/search",
json=payload,
headers=self.headers,
)
if response.status_code != 200:
raise ValueError(f"Error searching: {response.text}")
return response.json()["points"]
[docs] def similarity_search(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]:
"""Return docs most similar to query."""
query_embedding = self._embedding.embed_query(query)
return self.similarity_search_by_vector(query_embedding, k=k)
[docs] def similarity_search_with_score(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Tuple[Document, float]]:
"""Run similarity search with distance."""
query_embedding = self._embedding.embed_query(query)
points = self._search_points(query_embedding, k=k)
return [
(
Document(page_content=p["metadata"]["text"], metadata=p["metadata"]),
p["distance"],
)
for p in points
]
[docs] def similarity_search_by_vector(
self, embedding: List[float], k: int = 4, **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.
Returns:
List of Documents most similar to the query vector.
"""
points = self._search_points(embedding, k=k)
return [
Document(page_content=p["metadata"]["text"], metadata=p["metadata"])
for p in points
]
[docs] @classmethod
def from_texts(
cls,
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
collection_name: str = "",
vector_size: int = 0,
api_key: str = "",
distance_strategy: DistanceStrategy = DistanceStrategy.EUCLIDEAN_DISTANCE,
**kwargs: Any,
) -> "SemaDB":
"""Return VectorStore initialized from texts and embeddings."""
if not collection_name:
raise ValueError("Collection name must be provided")
if not vector_size:
raise ValueError("Vector size must be provided")
if not api_key:
raise ValueError("API key must be provided")
semadb = cls(
collection_name,
vector_size,
embedding,
distance_strategy=distance_strategy,
api_key=api_key,
)
if not semadb.create_collection():
raise ValueError("Error creating collection")
semadb.add_texts(texts, metadatas=metadatas)
return semadb