langchain/libs/community/langchain_community/vectorstores/semadb.py

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


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: str = "semadb.p.rapidapi.com"
    BASE_URL = "https://" + HOST

    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}")

    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

    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

    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

    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"]

    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)

    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
        ]

    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
        ]

    @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