mseep-txtai 9.1.1__py3-none-any.whl

txtai/agent/tool/function.py

@@ -0,0 +1,49 @@
+"""
+Function imports
+"""
+
+from smolagents import Tool
+
+
+class FunctionTool(Tool):
+    """
+    Creates a FunctionTool. A FunctionTool takes descriptive configuration and injects it along with a target function
+    into an LLM prompt.
+    """
+
+    # pylint: disable=W0231
+    def __init__(self, config):
+        """
+        Creates a FunctionTool.
+
+        Args:
+            config: `name`, `description`, `inputs`, `output` and `target` configuration
+        """
+
+        # Tool parameters
+        self.name = config["name"]
+        self.description = config["description"]
+        self.inputs = config["inputs"]
+        self.output_type = config.get("output", config.get("output_type", "any"))
+        self.target = config["target"]
+
+        # pylint: disable=C0103
+        # Skip forward signature validation
+        self.skip_forward_signature_validation = True
+
+        # Validate parameters and initialize tool
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        """
+        Runs target function.
+
+        Args:
+            args: positional args
+            kwargs: keyword args
+
+        Returns:
+            result
+        """
+
+        return self.target(*args, **kwargs)

txtai/ann/__init__.py

@@ -0,0 +1,7 @@
+"""
+ANN imports
+"""
+
+from .base import ANN
+from .dense import *
+from .sparse import *

txtai/ann/base.py

@@ -0,0 +1,153 @@
+"""
+ANN (Approximate Nearest Neighbor) module
+"""
+
+import datetime
+import platform
+
+from ..version import __version__
+
+
+class ANN:
+    """
+    Base class for ANN instances. This class builds vector indexes to support similarity search.
+    The built-in ANN backends store ids and vectors. Content storage is supported via database instances.
+    """
+
+    def __init__(self, config):
+        """
+        Creates a new ANN.
+
+        Args:
+            config: index configuration parameters
+        """
+
+        # ANN index
+        self.backend = None
+
+        # ANN configuration
+        self.config = config
+
+    def load(self, path):
+        """
+        Loads an ANN at path.
+
+        Args:
+            path: path to load ann index
+        """
+
+        raise NotImplementedError
+
+    def index(self, embeddings):
+        """
+        Builds an ANN index.
+
+        Args:
+            embeddings: embeddings array
+        """
+
+        raise NotImplementedError
+
+    def append(self, embeddings):
+        """
+        Append elements to an existing index.
+
+        Args:
+            embeddings: embeddings array
+        """
+
+        raise NotImplementedError
+
+    def delete(self, ids):
+        """
+        Deletes elements from existing index.
+
+        Args:
+            ids: ids to delete
+        """
+
+        raise NotImplementedError
+
+    def search(self, queries, limit):
+        """
+        Searches ANN index for query. Returns topn results.
+
+        Args:
+            queries: queries array
+            limit: maximum results
+
+        Returns:
+            query results
+        """
+
+        raise NotImplementedError
+
+    def count(self):
+        """
+        Number of elements in the ANN index.
+
+        Returns:
+            count
+        """
+
+        raise NotImplementedError
+
+    def save(self, path):
+        """
+        Saves an ANN index at path.
+
+        Args:
+            path: path to save ann index
+        """
+
+        raise NotImplementedError
+
+    def close(self):
+        """
+        Closes this ANN.
+        """
+
+        self.backend = None
+
+    def setting(self, name, default=None):
+        """
+        Looks up backend specific setting.
+
+        Args:
+            name: setting name
+            default: default value when setting not found
+
+        Returns:
+            setting value
+        """
+
+        # Get the backend-specific config object
+        backend = self.config.get(self.config["backend"])
+
+        # Get setting value, set default value if not found
+        setting = backend.get(name) if backend else None
+        return setting if setting else default
+
+    def metadata(self, settings=None):
+        """
+        Adds index build metadata.
+
+        Args:
+            settings: index build settings
+        """
+
+        # ISO 8601 timestamp
+        create = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+        # Set build metadata if this is not an update
+        if settings:
+            self.config["build"] = {
+                "create": create,
+                "python": platform.python_version(),
+                "settings": settings,
+                "system": f"{platform.system()} ({platform.machine()})",
+                "txtai": __version__,
+            }
+
+        # Set last update date
+        self.config["update"] = create

txtai/ann/dense/__init__.py

@@ -0,0 +1,11 @@
+"""
+Dense ANN imports
+"""
+
+from .annoy import Annoy
+from .factory import ANNFactory
+from .faiss import Faiss
+from .hnsw import HNSW
+from .numpy import NumPy
+from .pgvector import PGVector
+from .torch import Torch

txtai/ann/dense/annoy.py

@@ -0,0 +1,72 @@
+"""
+Annoy module
+"""
+
+# Conditional import
+try:
+    from annoy import AnnoyIndex
+
+    ANNOY = True
+except ImportError:
+    ANNOY = False
+
+from ..base import ANN
+
+
+# pylint: disable=W0223
+class Annoy(ANN):
+    """
+    Builds an ANN index using the Annoy library.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not ANNOY:
+            raise ImportError('Annoy is not available - install "ann" extra to enable')
+
+    def load(self, path):
+        # Load index
+        self.backend = AnnoyIndex(self.config["dimensions"], self.config["metric"])
+        self.backend.load(path)
+
+    def index(self, embeddings):
+        # Inner product is equal to cosine similarity on normalized vectors
+        self.config["metric"] = "dot"
+
+        # Create index
+        self.backend = AnnoyIndex(self.config["dimensions"], self.config["metric"])
+
+        # Add items - position in embeddings is used as the id
+        for x in range(embeddings.shape[0]):
+            self.backend.add_item(x, embeddings[x])
+
+        # Build index
+        ntrees = self.setting("ntrees", 10)
+        self.backend.build(ntrees)
+
+        # Add index build metadata
+        self.metadata({"ntrees": ntrees})
+
+    def search(self, queries, limit):
+        # Lookup search k setting
+        searchk = self.setting("searchk", -1)
+
+        # Annoy doesn't have a built in batch query method
+        results = []
+        for query in queries:
+            # Run the query
+            ids, scores = self.backend.get_nns_by_vector(query, n=limit, search_k=searchk, include_distances=True)
+
+            # Map results to [(id, score)]
+            results.append(list(zip(ids, scores)))
+
+        return results
+
+    def count(self):
+        # Number of items in index
+        return self.backend.get_n_items()
+
+    def save(self, path):
+        # Write index
+        self.backend.save(path)

txtai/ann/dense/factory.py

@@ -0,0 +1,76 @@
+"""
+Factory module
+"""
+
+from ...util import Resolver
+
+from .annoy import Annoy
+from .faiss import Faiss
+from .hnsw import HNSW
+from .numpy import NumPy
+from .pgvector import PGVector
+from .sqlite import SQLite
+from .torch import Torch
+
+
+class ANNFactory:
+    """
+    Methods to create ANN indexes.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Create an ANN.
+
+        Args:
+            config: index configuration parameters
+
+        Returns:
+            ANN
+        """
+
+        # ANN instance
+        ann = None
+        backend = config.get("backend", "faiss")
+
+        # Create ANN instance
+        if backend == "annoy":
+            ann = Annoy(config)
+        elif backend == "faiss":
+            ann = Faiss(config)
+        elif backend == "hnsw":
+            ann = HNSW(config)
+        elif backend == "numpy":
+            ann = NumPy(config)
+        elif backend == "pgvector":
+            ann = PGVector(config)
+        elif backend == "sqlite":
+            ann = SQLite(config)
+        elif backend == "torch":
+            ann = Torch(config)
+        else:
+            ann = ANNFactory.resolve(backend, config)
+
+        # Store config back
+        config["backend"] = backend
+
+        return ann
+
+    @staticmethod
+    def resolve(backend, config):
+        """
+        Attempt to resolve a custom backend.
+
+        Args:
+            backend: backend class
+            config: index configuration parameters
+
+        Returns:
+            ANN
+        """
+
+        try:
+            return Resolver()(backend)(config)
+        except Exception as e:
+            raise ImportError(f"Unable to resolve ann backend: '{backend}'") from e

txtai/ann/dense/faiss.py

@@ -0,0 +1,233 @@
+"""
+Faiss module
+"""
+
+import math
+import platform
+
+import numpy as np
+
+from faiss import omp_set_num_threads
+from faiss import index_factory, IO_FLAG_MMAP, METRIC_INNER_PRODUCT, read_index, write_index
+from faiss import index_binary_factory, read_index_binary, write_index_binary, IndexBinaryIDMap
+
+from ..base import ANN
+
+if platform.system() == "Darwin":
+    # Workaround for a Faiss issue causing segmentation faults on macOS. See txtai FAQ for more.
+    omp_set_num_threads(1)
+
+
+class Faiss(ANN):
+    """
+    Builds an ANN index using the Faiss library.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Scalar quantization
+        quantize = self.config.get("quantize")
+        self.qbits = quantize if quantize and isinstance(quantize, int) and not isinstance(quantize, bool) else None
+
+    def load(self, path):
+        # Get read function
+        readindex = read_index_binary if self.qbits else read_index
+
+        # Load index
+        self.backend = readindex(path, IO_FLAG_MMAP if self.setting("mmap") is True else 0)
+
+    def index(self, embeddings):
+        # Compute model training size
+        train, sample = embeddings, self.setting("sample")
+        if sample:
+            # Get sample for training
+            rng = np.random.default_rng(0)
+            indices = sorted(rng.choice(train.shape[0], int(sample * train.shape[0]), replace=False, shuffle=False))
+            train = train[indices]
+
+        # Configure embeddings index. Inner product is equal to cosine similarity on normalized vectors.
+        params = self.configure(embeddings.shape[0], train.shape[0])
+
+        # Create index
+        self.backend = self.create(embeddings, params)
+
+        # Train model
+        self.backend.train(train)
+
+        # Add embeddings - position in embeddings is used as the id
+        self.backend.add_with_ids(embeddings, np.arange(embeddings.shape[0], dtype=np.int64))
+
+        # Add id offset and index build metadata
+        self.config["offset"] = embeddings.shape[0]
+        self.metadata({"components": params})
+
+    def append(self, embeddings):
+        new = embeddings.shape[0]
+
+        # Append new ids - position in embeddings + existing offset is used as the id
+        self.backend.add_with_ids(embeddings, np.arange(self.config["offset"], self.config["offset"] + new, dtype=np.int64))
+
+        # Update id offset and index metadata
+        self.config["offset"] += new
+        self.metadata()
+
+    def delete(self, ids):
+        # Remove specified ids
+        self.backend.remove_ids(np.array(ids, dtype=np.int64))
+
+    def search(self, queries, limit):
+        # Set nprobe and nflip search parameters
+        self.backend.nprobe = self.nprobe()
+        self.backend.nflip = self.setting("nflip", self.backend.nprobe)
+
+        # Run the query
+        scores, ids = self.backend.search(queries, limit)
+
+        # Map results to [(id, score)]
+        results = []
+        for x, score in enumerate(scores):
+            # Transform scores and add results
+            results.append(list(zip(ids[x].tolist(), self.scores(score))))
+
+        return results
+
+    def count(self):
+        return self.backend.ntotal
+
+    def save(self, path):
+        # Get write function
+        writeindex = write_index_binary if self.qbits else write_index
+
+        # Write index
+        writeindex(self.backend, path)
+
+    def configure(self, count, train):
+        """
+        Configures settings for a new index.
+
+        Args:
+            count: initial number of embeddings rows
+            train: number of rows selected for model training
+
+        Returns:
+            user-specified or generated components setting
+        """
+
+        # Lookup components setting
+        components = self.setting("components")
+
+        if components:
+            # Format and return components string
+            return self.components(components, train)
+
+        # Derive quantization. Prefer backend-specific setting. Fallback to root-level parameter.
+        quantize = self.setting("quantize", self.config.get("quantize"))
+        quantize = 8 if isinstance(quantize, bool) else quantize
+
+        # Get storage setting
+        storage = f"SQ{quantize}" if quantize else "Flat"
+
+        # Small index, use storage directly with IDMap
+        if count <= 5000:
+            return "BFlat" if self.qbits else f"IDMap,{storage}"
+
+        x = self.cells(train)
+        components = f"BIVF{x}" if self.qbits else f"IVF{x},{storage}"
+
+        return components
+
+    def create(self, embeddings, params):
+        """
+        Creates a new index.
+
+        Args:
+            embeddings: embeddings to index
+            params: index parameters
+
+        Returns:
+            new index
+        """
+
+        # Create binary index
+        if self.qbits:
+            index = index_binary_factory(embeddings.shape[1] * 8, params)
+
+            # Wrap with BinaryIDMap, if necessary
+            if any(x in params for x in ["BFlat", "BHNSW"]):
+                index = IndexBinaryIDMap(index)
+
+            return index
+
+        # Create standard float index
+        return index_factory(embeddings.shape[1], params, METRIC_INNER_PRODUCT)
+
+    def cells(self, count):
+        """
+        Calculates the number of IVF cells for an IVF index.
+
+        Args:
+            count: number of embeddings rows
+
+        Returns:
+            number of IVF cells
+        """
+
+        # Calculate number of IVF cells where x = min(4 * sqrt(embeddings count), embeddings count / 39)
+        # Faiss requires at least 39 points per cluster
+        return max(min(round(4 * math.sqrt(count)), int(count / 39)), 1)
+
+    def components(self, components, train):
+        """
+        Formats a components string. This method automatically calculates the optimal number of IVF cells, if omitted.
+
+        Args:
+            components: input components string
+            train: number of rows selected for model training
+
+        Returns:
+            formatted components string
+        """
+
+        # Optimal number of IVF cells
+        x = self.cells(train)
+
+        # Add number of IVF cells, if missing
+        components = [f"IVF{x}" if component == "IVF" else component for component in components.split(",")]
+
+        # Return components string
+        return ",".join(components)
+
+    def nprobe(self):
+        """
+        Gets or derives the nprobe search parameter.
+
+        Returns:
+            nprobe setting
+        """
+
+        # Get size of embeddings index
+        count = self.count()
+
+        default = 6 if count <= 5000 else round(self.cells(count) / 16)
+        return self.setting("nprobe", default)
+
+    def scores(self, scores):
+        """
+        Calculates the index score from the input score. This method returns the hamming score
+        (1.0 - (hamming distance / total number of bits)) for binary indexes and the input
+        scores otherwise.
+
+        Args:
+            scores: input scores
+
+        Returns:
+            index scores
+        """
+
+        # Calculate hamming score, bound between 0.0 - 1.0
+        if self.qbits:
+            return np.clip(1.0 - (scores / (self.config["dimensions"] * 8)), 0.0, 1.0).tolist()
+
+        # Standard scoring
+        return scores.tolist()

txtai/ann/dense/hnsw.py

@@ -0,0 +1,104 @@
+"""
+HNSW module
+"""
+
+import numpy as np
+
+# Conditional import
+try:
+    # pylint: disable=E0611
+    from hnswlib import Index
+
+    HNSWLIB = True
+except ImportError:
+    HNSWLIB = False
+
+from ..base import ANN
+
+
+class HNSW(ANN):
+    """
+    Builds an ANN index using the hnswlib library.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not HNSWLIB:
+            raise ImportError('HNSW is not available - install "ann" extra to enable')
+
+    def load(self, path):
+        # Load index
+        self.backend = Index(dim=self.config["dimensions"], space=self.config["metric"])
+        self.backend.load_index(path)
+
+    def index(self, embeddings):
+        # Inner product is equal to cosine similarity on normalized vectors
+        self.config["metric"] = "ip"
+
+        # Lookup index settings
+        efconstruction = self.setting("efconstruction", 200)
+        m = self.setting("m", 16)
+        seed = self.setting("randomseed", 100)
+
+        # Create index
+        self.backend = Index(dim=self.config["dimensions"], space=self.config["metric"])
+        self.backend.init_index(max_elements=embeddings.shape[0], ef_construction=efconstruction, M=m, random_seed=seed)
+
+        # Add items - position in embeddings is used as the id
+        self.backend.add_items(embeddings, np.arange(embeddings.shape[0], dtype=np.int64))
+
+        # Add id offset, delete counter and index build metadata
+        self.config["offset"] = embeddings.shape[0]
+        self.config["deletes"] = 0
+        self.metadata({"efconstruction": efconstruction, "m": m, "seed": seed})
+
+    def append(self, embeddings):
+        new = embeddings.shape[0]
+
+        # Resize index
+        self.backend.resize_index(self.config["offset"] + new)
+
+        # Append new ids - position in embeddings + existing offset is used as the id
+        self.backend.add_items(embeddings, np.arange(self.config["offset"], self.config["offset"] + new, dtype=np.int64))
+
+        # Update id offset and index metadata
+        self.config["offset"] += new
+        self.metadata()
+
+    def delete(self, ids):
+        # Mark elements as deleted to omit from search results
+        for uid in ids:
+            try:
+                self.backend.mark_deleted(uid)
+                self.config["deletes"] += 1
+            except RuntimeError:
+                # Ignore label not found error
+                continue
+
+    def search(self, queries, limit):
+        # Set ef query param
+        ef = self.setting("efsearch")
+        if ef:
+            self.backend.set_ef(ef)
+
+        # Run the query
+        ids, distances = self.backend.knn_query(queries, k=limit)
+
+        # Map results to [(id, score)]
+        results = []
+        for x, distance in enumerate(distances):
+            # Convert distances to similarity scores
+            scores = [1 - d for d in distance.tolist()]
+
+            # Build (id, score) tuples, convert np.int64 to python int
+            results.append(list(zip(ids[x].tolist(), scores)))
+
+        return results
+
+    def count(self):
+        return self.backend.get_current_count() - self.config["deletes"]
+
+    def save(self, path):
+        # Write index
+        self.backend.save_index(path)