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

"""Module providing Infinispan as a VectorStore"""

from __future__ import annotations

import json
import logging
import uuid
from typing import Any, Iterable, List, Optional, Tuple, Type, cast

import requests
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore

logger = logging.getLogger(__name__)


class InfinispanVS(VectorStore):
    """`Infinispan` VectorStore interface.

        This class exposes the method to present Infinispan as a
        VectorStore. It relies on the Infinispan class (below) which takes care
        of the REST interface with the server.

    Example:
        ... code-block:: python
            from langchain_community.vectorstores import InfinispanVS
            from mymodels import RGBEmbeddings
            ...
            vectorDb = InfinispanVS.from_documents(docs,
                            embedding=RGBEmbeddings(),
                            output_fields=["texture", "color"],
                            lambda_key=lambda text,meta: str(meta["_key"]),
                            lambda_content=lambda item: item["color"])

        or an empty InfinispanVS instance can be created if preliminary setup
        is required before populating the store

        ... code-block:: python
            from langchain_community.vectorstores import InfinispanVS
            from mymodels import RGBEmbeddings
            ...
            ispnVS = InfinispanVS()
            # configure Infinispan here
            # i.e. create cache and schema

            # then populate the store
            vectorDb = InfinispanVS.from_documents(docs,
                            embedding=RGBEmbeddings(),
                            output_fields: ["texture", "color"],
                            lambda_key: lambda text,meta: str(meta["_key"]),
                            lambda_content: lambda item: item["color"]})
    """

    def __init__(
        self,
        embedding: Optional[Embeddings] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ):
        self.ispn = Infinispan(**kwargs)
        self._configuration = kwargs
        self._cache_name = str(self._configuration.get("cache_name", "vector"))
        self._entity_name = str(self._configuration.get("entity_name", "vector"))
        self._embedding = embedding
        self._textfield = self._configuration.get("textfield", "text")
        self._vectorfield = self._configuration.get("vectorfield", "vector")
        self._to_content = self._configuration.get(
            "lambda_content", lambda item: self._default_content(item)
        )
        self._to_metadata = self._configuration.get(
            "lambda_metadata", lambda item: self._default_metadata(item)
        )
        self._output_fields = self._configuration.get("output_fields")
        self._ids = ids

    def _default_metadata(self, item: dict) -> dict:
        meta = dict(item)
        meta.pop(self._vectorfield, None)
        meta.pop(self._textfield, None)
        meta.pop("_type", None)
        return meta

    def _default_content(self, item: dict[str, Any]) -> Any:
        return item.get(self._textfield)

    def schema_builder(self, templ: dict, dimension: int) -> str:
        metadata_proto_tpl = """
/**
* @Indexed
*/
message %s {
/**
* @Vector(dimension=%d)
*/
repeated float %s = 1;
"""
        metadata_proto = metadata_proto_tpl % (
            self._entity_name,
            dimension,
            self._vectorfield,
        )
        idx = 2
        for f, v in templ.items():
            if isinstance(v, str):
                metadata_proto += "optional string " + f + " = " + str(idx) + ";\n"
            elif isinstance(v, int):
                metadata_proto += "optional int64 " + f + " = " + str(idx) + ";\n"
            elif isinstance(v, float):
                metadata_proto += "optional double " + f + " = " + str(idx) + ";\n"
            elif isinstance(v, bytes):
                metadata_proto += "optional bytes " + f + " = " + str(idx) + ";\n"
            elif isinstance(v, bool):
                metadata_proto += "optional bool " + f + " = " + str(idx) + ";\n"
            else:
                raise Exception(
                    "Unable to build proto schema for metadata. "
                    "Unhandled type for field: " + f
                )
            idx += 1
        metadata_proto += "}\n"
        return metadata_proto

    def schema_create(self, proto: str) -> requests.Response:
        """Deploy the schema for the vector db
        Args:
            proto(str): protobuf schema
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.schema_post(self._entity_name + ".proto", proto)

    def schema_delete(self) -> requests.Response:
        """Delete the schema for the vector db
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.schema_delete(self._entity_name + ".proto")

    def cache_create(self, config: str = "") -> requests.Response:
        """Create the cache for the vector db
        Args:
            config(str): configuration of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        if config == "":
            config = (
                '''
            {
  "distributed-cache": {
    "owners": "2",
    "mode": "SYNC",
    "statistics": true,
    "encoding": {
      "media-type": "application/x-protostream"
    },
    "indexing": {
      "enabled": true,
      "storage": "filesystem",
      "startup-mode": "AUTO",
      "indexing-mode": "AUTO",
      "indexed-entities": [
        "'''
                + self._entity_name
                + """"
      ]
    }
  }
}
"""
            )
        return self.ispn.cache_post(self._cache_name, config)

    def cache_delete(self) -> requests.Response:
        """Delete the cache for the vector db
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.cache_delete(self._cache_name)

    def cache_clear(self) -> requests.Response:
        """Clear the cache for the vector db
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.cache_clear(self._cache_name)

    def cache_exists(self) -> bool:
        """Checks if the cache exists
        Returns:
            true if exists
        """
        return self.ispn.cache_exists(self._cache_name)

    def cache_index_clear(self) -> requests.Response:
        """Clear the index for the vector db
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.index_clear(self._cache_name)

    def cache_index_reindex(self) -> requests.Response:
        """Rebuild the for the vector db
        Returns:
            An http Response containing the result of the operation
        """
        return self.ispn.index_reindex(self._cache_name)

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        last_vector: Optional[List[float]] = None,
        **kwargs: Any,
    ) -> List[str]:
        result = []
        texts_l = list(texts)
        if last_vector:
            texts_l.pop()
        embeds = self._embedding.embed_documents(texts_l)  # type: ignore
        if last_vector:
            embeds.append(last_vector)
        if not metadatas:
            metadatas = [{} for _ in texts]
        ids = self._ids or [str(uuid.uuid4()) for _ in texts]
        data_input = list(zip(metadatas, embeds, ids))
        for metadata, embed, key in data_input:
            data = {"_type": self._entity_name, self._vectorfield: embed}
            data.update(metadata)
            data_str = json.dumps(data)
            self.ispn.put(key, data_str, self._cache_name)
            result.append(key)
        return result

    def similarity_search(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> List[Document]:
        """Return docs most similar to query."""
        documents = self.similarity_search_with_score(query=query, k=k)
        return [doc for doc, _ in documents]

    def similarity_search_with_score(
        self, query: str, k: int = 4, **kwargs: Any
    ) -> List[Tuple[Document, float]]:
        """Perform a search on a query string and return results with score.

        Args:
            query (str): The text being searched.
            k (int, optional): The amount of results to return. Defaults to 4.

        Returns:
            List[Tuple[Document, float]]
        """
        embed = self._embedding.embed_query(query)  # type: ignore
        documents = self.similarity_search_with_score_by_vector(embedding=embed, k=k)
        return documents

    def similarity_search_by_vector(
        self, embedding: List[float], k: int = 4, **kwargs: Any
    ) -> List[Document]:
        res = self.similarity_search_with_score_by_vector(embedding, k)
        return [doc for doc, _ in res]

    def similarity_search_with_score_by_vector(
        self, embedding: List[float], k: int = 4
    ) -> List[Tuple[Document, float]]:
        """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 pair (Documents, score) most similar to the query vector.
        """
        if self._output_fields is None:
            query_str = (
                "select v, score(v) from "
                + self._entity_name
                + " v where v."
                + self._vectorfield
                + " <-> "
                + json.dumps(embedding)
                + "~"
                + str(k)
            )
        else:
            query_proj = "select "
            for field in self._output_fields[:-1]:
                query_proj = query_proj + "v." + field + ","
            query_proj = query_proj + "v." + self._output_fields[-1]
            query_str = (
                query_proj
                + ", score(v) from "
                + self._entity_name
                + " v where v."
                + self._vectorfield
                + " <-> "
                + json.dumps(embedding)
                + "~"
                + str(k)
            )
        query_res = self.ispn.req_query(query_str, self._cache_name)
        result = json.loads(query_res.text)
        return self._query_result_to_docs(result)

    def _query_result_to_docs(
        self, result: dict[str, Any]
    ) -> List[Tuple[Document, float]]:
        documents = []
        for row in result["hits"]:
            hit = row["hit"] or {}
            if self._output_fields is None:
                entity = hit["*"]
            else:
                entity = {key: hit.get(key) for key in self._output_fields}
            doc = Document(
                page_content=self._to_content(entity),
                metadata=self._to_metadata(entity),
            )
            documents.append((doc, hit["score()"]))
        return documents

    def configure(self, metadata: dict, dimension: int) -> None:
        schema = self.schema_builder(metadata, dimension)
        output = self.schema_create(schema)
        assert output.ok, "Unable to create schema. Already exists? "
        "Consider using clear_old=True"
        assert json.loads(output.text)["error"] is None
        if not self.cache_exists():
            output = self.cache_create()
            assert output.ok, "Unable to create cache. Already exists? "
            "Consider using clear_old=True"
            # Ensure index is clean
            self.cache_index_clear()

    def config_clear(self) -> None:
        self.schema_delete()
        self.cache_delete()

    @classmethod
    def from_texts(
        cls: Type[InfinispanVS],
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        clear_old: Optional[bool] = True,
        auto_config: Optional[bool] = True,
        **kwargs: Any,
    ) -> InfinispanVS:
        """Return VectorStore initialized from texts and embeddings."""
        infinispanvs = cls(embedding=embedding, ids=ids, **kwargs)
        if auto_config and len(metadatas or []) > 0:
            if clear_old:
                infinispanvs.config_clear()
            vec = embedding.embed_query(texts[len(texts) - 1])
            metadatas = cast(List[dict], metadatas)
            infinispanvs.configure(metadatas[0], len(vec))
        else:
            if clear_old:
                infinispanvs.cache_clear()
            vec = embedding.embed_query(texts[len(texts) - 1])
        if texts:
            infinispanvs.add_texts(texts, metadatas, vector=vec)
        return infinispanvs


REST_TIMEOUT = 10


class Infinispan:
    """Helper class for `Infinispan` REST interface.

    This class exposes the Infinispan operations needed to
    create and set up a vector db.

    You need a running Infinispan (15+) server without authentication.
    You can easily start one, see:
    https://github.com/rigazilla/infinispan-vector#run-infinispan
    """

    def __init__(self, **kwargs: Any):
        self._configuration = kwargs
        self._schema = str(self._configuration.get("schema", "http"))
        self._host = str(self._configuration.get("hosts", ["127.0.0.1:11222"])[0])
        self._default_node = self._schema + "://" + self._host
        self._cache_url = str(self._configuration.get("cache_url", "/rest/v2/caches"))
        self._schema_url = str(self._configuration.get("cache_url", "/rest/v2/schemas"))
        self._use_post_for_query = str(
            self._configuration.get("use_post_for_query", True)
        )

    def req_query(
        self, query: str, cache_name: str, local: bool = False
    ) -> requests.Response:
        """Request a query
        Args:
            query(str): query requested
            cache_name(str): name of the target cache
            local(boolean): whether the query is local to clustered
        Returns:
            An http Response containing the result set or errors
        """
        if self._use_post_for_query:
            return self._query_post(query, cache_name, local)
        return self._query_get(query, cache_name, local)

    def _query_post(
        self, query_str: str, cache_name: str, local: bool = False
    ) -> requests.Response:
        api_url = (
            self._default_node
            + self._cache_url
            + "/"
            + cache_name
            + "?action=search&local="
            + str(local)
        )
        data = {"query": query_str}
        data_json = json.dumps(data)
        response = requests.post(
            api_url,
            data_json,
            headers={"Content-Type": "application/json"},
            timeout=REST_TIMEOUT,
        )
        return response

    def _query_get(
        self, query_str: str, cache_name: str, local: bool = False
    ) -> requests.Response:
        api_url = (
            self._default_node
            + self._cache_url
            + "/"
            + cache_name
            + "?action=search&query="
            + query_str
            + "&local="
            + str(local)
        )
        response = requests.get(api_url, timeout=REST_TIMEOUT)
        return response

    def post(self, key: str, data: str, cache_name: str) -> requests.Response:
        """Post an entry
        Args:
            key(str): key of the entry
            data(str): content of the entry in json format
            cache_name(str): target cache
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._cache_url + "/" + cache_name + "/" + key
        response = requests.post(
            api_url,
            data,
            headers={"Content-Type": "application/json"},
            timeout=REST_TIMEOUT,
        )
        return response

    def put(self, key: str, data: str, cache_name: str) -> requests.Response:
        """Put an entry
        Args:
            key(str): key of the entry
            data(str): content of the entry in json format
            cache_name(str): target cache
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._cache_url + "/" + cache_name + "/" + key
        response = requests.put(
            api_url,
            data,
            headers={"Content-Type": "application/json"},
            timeout=REST_TIMEOUT,
        )
        return response

    def get(self, key: str, cache_name: str) -> requests.Response:
        """Get an entry
        Args:
            key(str): key of the entry
            cache_name(str): target cache
        Returns:
            An http Response containing the entry or errors
        """
        api_url = self._default_node + self._cache_url + "/" + cache_name + "/" + key
        response = requests.get(
            api_url, headers={"Content-Type": "application/json"}, timeout=REST_TIMEOUT
        )
        return response

    def schema_post(self, name: str, proto: str) -> requests.Response:
        """Deploy a schema
        Args:
            name(str): name of the schema. Will be used as a key
            proto(str): protobuf schema
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._schema_url + "/" + name
        response = requests.post(api_url, proto, timeout=REST_TIMEOUT)
        return response

    def cache_post(self, name: str, config: str) -> requests.Response:
        """Create a cache
        Args:
            name(str): name of the cache.
            config(str): configuration of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._cache_url + "/" + name
        response = requests.post(
            api_url,
            config,
            headers={"Content-Type": "application/json"},
            timeout=REST_TIMEOUT,
        )
        return response

    def schema_delete(self, name: str) -> requests.Response:
        """Delete a schema
        Args:
            name(str): name of the schema.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._schema_url + "/" + name
        response = requests.delete(api_url, timeout=REST_TIMEOUT)
        return response

    def cache_delete(self, name: str) -> requests.Response:
        """Delete a cache
        Args:
            name(str): name of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = self._default_node + self._cache_url + "/" + name
        response = requests.delete(api_url, timeout=REST_TIMEOUT)
        return response

    def cache_clear(self, cache_name: str) -> requests.Response:
        """Clear a cache
        Args:
            cache_name(str): name of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = (
            self._default_node + self._cache_url + "/" + cache_name + "?action=clear"
        )
        response = requests.post(api_url, timeout=REST_TIMEOUT)
        return response

    def cache_exists(self, cache_name: str) -> bool:
        """Check if a cache exists
        Args:
            cache_name(str): name of the cache.
        Returns:
            True if cache exists
        """
        api_url = (
            self._default_node + self._cache_url + "/" + cache_name + "?action=clear"
        )
        return self.resource_exists(api_url)

    @staticmethod
    def resource_exists(api_url: str) -> bool:
        """Check if a resource exists
        Args:
            api_url(str): url of the resource.
        Returns:
            true if resource exists
        """
        response = requests.head(api_url, timeout=REST_TIMEOUT)
        return response.ok

    def index_clear(self, cache_name: str) -> requests.Response:
        """Clear an index on a cache
        Args:
            cache_name(str): name of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = (
            self._default_node
            + self._cache_url
            + "/"
            + cache_name
            + "/search/indexes?action=clear"
        )
        return requests.post(api_url, timeout=REST_TIMEOUT)

    def index_reindex(self, cache_name: str) -> requests.Response:
        """Rebuild index on a cache
        Args:
            cache_name(str): name of the cache.
        Returns:
            An http Response containing the result of the operation
        """
        api_url = (
            self._default_node
            + self._cache_url
            + "/"
            + cache_name
            + "/search/indexes?action=reindex"
        )
        return requests.post(api_url, timeout=REST_TIMEOUT)