client.py

"""
Weaviate client module
"""

import weaviate

from txtai.pipeline import Pipeline

class Weaviate(Pipeline):
    """
    Weaviate pipeline client. Supports indexing and searching content with Weaviate.
    """

    def __init__(self, url="http://localhost:8080"):
        """
        Create a new client.

        Args:
            url: Weaviate service url
        """

        self.client = weaviate.Client(url)

        # Delete autogenerated schema
        self.client.schema.delete_all()

    def __call__(self, inputs, action="index"):
        """
        Executes an action with Weaviate.

        Args:
            inputs: data inputs
            action: action to perform - index or search

        Returns:
            results
        """

        if action == "index":
            return [self.index(data, vector) for data, vector in inputs]

        # Default to search action
        return [self.search(vector) for vector in inputs]

    def index(self, data, vector):
        """
        Indexes data-vector pair in Weaviate.

        Args:
            data: record metadata
            vector: record embeddings

        Returns:
            uuid from Weaviate
        """

        return self.client.data_object.create(
            {"content": data},
            "Post",
            vector = vector,
        )

    def search(self, vector):
        """
        Runs a search using input vector.

        Args:
            vector: input vector

        Returns:
            search results
        """

        nearvector = {"vector": vector}
        return self.client.query.get("Post", ["content", "_additional {certainty}"]).with_near_vector(nearvector).with_limit(1).do()