mseep-txtai 9.1.1__py3-none-any.whl

txtai/api/base.py

@@ -0,0 +1,159 @@
+"""
+API module
+"""
+
+import json
+
+from .cluster import Cluster
+
+from ..app import Application
+
+
+class API(Application):
+    """
+    Base API template. The API is an extended txtai application, adding the ability to cluster API instances together.
+
+    Downstream applications can extend this base template to add/modify functionality.
+    """
+
+    def __init__(self, config, loaddata=True):
+        super().__init__(config, loaddata)
+
+        # Embeddings cluster
+        self.cluster = None
+        if self.config.get("cluster"):
+            self.cluster = Cluster(self.config["cluster"])
+
+    # pylint: disable=W0221
+    def search(self, query, limit=None, weights=None, index=None, parameters=None, graph=False, request=None):
+        # When search is invoked via the API, limit is set from the request
+        # When search is invoked directly, limit is set using the method parameter
+        limit = self.limit(request.query_params.get("limit") if request and hasattr(request, "query_params") else limit)
+        weights = self.weights(request.query_params.get("weights") if request and hasattr(request, "query_params") else weights)
+        index = request.query_params.get("index") if request and hasattr(request, "query_params") else index
+        parameters = request.query_params.get("parameters") if request and hasattr(request, "query_params") else parameters
+        graph = request.query_params.get("graph") if request and hasattr(request, "query_params") else graph
+
+        # Decode parameters
+        parameters = json.loads(parameters) if parameters and isinstance(parameters, str) else parameters
+
+        if self.cluster:
+            return self.cluster.search(query, limit, weights, index, parameters, graph)
+
+        return super().search(query, limit, weights, index, parameters, graph)
+
+    def batchsearch(self, queries, limit=None, weights=None, index=None, parameters=None, graph=False):
+        if self.cluster:
+            return self.cluster.batchsearch(queries, self.limit(limit), weights, index, parameters, graph)
+
+        return super().batchsearch(queries, limit, weights, index, parameters, graph)
+
+    def add(self, documents):
+        """
+        Adds a batch of documents for indexing.
+
+        Downstream applications can override this method to also store full documents in an external system.
+
+        Args:
+            documents: list of {id: value, text: value}
+
+        Returns:
+            unmodified input documents
+        """
+
+        if self.cluster:
+            self.cluster.add(documents)
+        else:
+            super().add(documents)
+
+        return documents
+
+    def index(self):
+        """
+        Builds an embeddings index for previously batched documents.
+        """
+
+        if self.cluster:
+            self.cluster.index()
+        else:
+            super().index()
+
+    def upsert(self):
+        """
+        Runs an embeddings upsert operation for previously batched documents.
+        """
+
+        if self.cluster:
+            self.cluster.upsert()
+        else:
+            super().upsert()
+
+    def delete(self, ids):
+        """
+        Deletes from an embeddings index. Returns list of ids deleted.
+
+        Args:
+            ids: list of ids to delete
+
+        Returns:
+            ids deleted
+        """
+
+        if self.cluster:
+            return self.cluster.delete(ids)
+
+        return super().delete(ids)
+
+    def reindex(self, config, function=None):
+        """
+        Recreates this embeddings index using config. This method only works if document content storage is enabled.
+
+        Args:
+            config: new config
+            function: optional function to prepare content for indexing
+        """
+
+        if self.cluster:
+            self.cluster.reindex(config, function)
+        else:
+            super().reindex(config, function)
+
+    def count(self):
+        """
+        Total number of elements in this embeddings index.
+
+        Returns:
+            number of elements in embeddings index
+        """
+
+        if self.cluster:
+            return self.cluster.count()
+
+        return super().count()
+
+    def limit(self, limit):
+        """
+        Parses the number of results to return from the request. Allows range of 1-250, with a default of 10.
+
+        Args:
+            limit: limit parameter
+
+        Returns:
+            bounded limit
+        """
+
+        # Return between 1 and 250 results, defaults to 10
+        return max(1, min(250, int(limit) if limit else 10))
+
+    def weights(self, weights):
+        """
+        Parses the weights parameter from the request.
+
+        Args:
+            weights: weights parameter
+
+        Returns:
+            weights
+        """
+
+        return float(weights) if weights else weights

txtai/api/cluster.py

@@ -0,0 +1,295 @@
+"""
+Cluster module
+"""
+
+import asyncio
+import json
+import random
+import urllib.parse
+import zlib
+
+import aiohttp
+
+from ..database.sql import Aggregate
+
+
+class Cluster:
+    """
+    Aggregates multiple embeddings shards into a single logical embeddings instance.
+    """
+
+    # pylint: disable = W0231
+    def __init__(self, config=None):
+        """
+        Creates a new Cluster.
+
+        Args:
+            config: cluster configuration
+        """
+
+        # Configuration
+        self.config = config
+
+        # Embeddings shard urls
+        self.shards = None
+        if "shards" in self.config:
+            self.shards = self.config["shards"]
+
+        # Query aggregator
+        self.aggregate = Aggregate()
+
+    def search(self, query, limit=None, weights=None, index=None, parameters=None, graph=False):
+        """
+        Finds documents most similar to the input query. This method will run either an index search
+        or an index + database search depending on if a database is available.
+
+        Args:
+            query: input query
+            limit: maximum results
+            weights: hybrid score weights, if applicable
+            index: index name, if applicable
+            parameters: dict of named parameters to bind to placeholders
+            graph: return graph results if True
+
+        Returns:
+            list of {id: value, score: value} for index search, list of dict for an index + database search
+        """
+
+        # Build URL
+        action = f"search?query={urllib.parse.quote_plus(query)}"
+        if limit:
+            action += f"&limit={limit}"
+        if weights:
+            action += f"&weights={weights}"
+        if index:
+            action += f"&index={index}"
+        if parameters:
+            action += f"&parameters={json.dumps(parameters) if isinstance(parameters, dict) else parameters}"
+        if graph is not None:
+            action += f"&graph={graph}"
+
+        # Run query and flatten results into single results list
+        results = []
+        for result in self.execute("get", action):
+            results.extend(result)
+
+        # Combine aggregate functions and sort
+        results = self.aggregate(query, results)
+
+        # Limit results
+        return results[: (limit if limit else 10)]
+
+    def batchsearch(self, queries, limit=None, weights=None, index=None, parameters=None, graph=False):
+        """
+        Finds documents most similar to the input queries. This method will run either an index search
+        or an index + database search depending on if a database is available.
+
+        Args:
+            queries: input queries
+            limit: maximum results
+            weights: hybrid score weights, if applicable
+            index: index name, if applicable
+            parameters: list of dicts of named parameters to bind to placeholders
+            graph: return graph results if True
+
+        Returns:
+            list of {id: value, score: value} per query for index search, list of dict per query for an index + database search
+        """
+
+        # POST parameters
+        params = {"queries": queries}
+        if limit:
+            params["limit"] = limit
+        if weights:
+            params["weights"] = weights
+        if index:
+            params["index"] = index
+        if parameters:
+            params["parameters"] = parameters
+        if graph is not None:
+            params["graph"] = graph
+
+        # Run query
+        batch = self.execute("post", "batchsearch", [params] * len(self.shards))
+
+        # Combine results per query
+        results = []
+        for x, query in enumerate(queries):
+            result = []
+            for section in batch:
+                result.extend(section[x])
+
+            # Aggregate, sort and limit results
+            results.append(self.aggregate(query, result)[: (limit if limit else 10)])
+
+        return results
+
+    def add(self, documents):
+        """
+        Adds a batch of documents for indexing.
+
+        Args:
+            documents: list of {id: value, text: value}
+        """
+
+        self.execute("post", "add", self.shard(documents))
+
+    def index(self):
+        """
+        Builds an embeddings index for previously batched documents.
+        """
+
+        self.execute("get", "index")
+
+    def upsert(self):
+        """
+        Runs an embeddings upsert operation for previously batched documents.
+        """
+
+        self.execute("get", "upsert")
+
+    def delete(self, ids):
+        """
+        Deletes from an embeddings cluster. Returns list of ids deleted.
+
+        Args:
+            ids: list of ids to delete
+
+        Returns:
+            ids deleted
+        """
+
+        return [uid for ids in self.execute("post", "delete", [ids] * len(self.shards)) for uid in ids]
+
+    def reindex(self, config, function=None):
+        """
+        Recreates this embeddings index using config. This method only works if document content storage is enabled.
+
+        Args:
+            config: new config
+            function: optional function to prepare content for indexing
+        """
+
+        self.execute("post", "reindex", [{"config": config, "function": function}] * len(self.shards))
+
+    def count(self):
+        """
+        Total number of elements in this embeddings cluster.
+
+        Returns:
+            number of elements in embeddings cluster
+        """
+
+        return sum(self.execute("get", "count"))
+
+    def shard(self, documents):
+        """
+        Splits documents into equal sized shards.
+
+        Args:
+            documents: input documents
+
+        Returns:
+            list of evenly sized shards with the last shard having the remaining elements
+        """
+
+        shards = [[] for _ in range(len(self.shards))]
+        for document in documents:
+            uid = document.get("id") if isinstance(document, dict) else document
+            if uid and isinstance(uid, str):
+                # Quick int hash of string to help derive shard id
+                uid = zlib.adler32(uid.encode("utf-8"))
+            elif uid is None:
+                # Get random shard id when uid isn't set
+                uid = random.randint(0, len(shards) - 1)
+
+            shards[uid % len(self.shards)].append(document)
+
+        return shards
+
+    def execute(self, method, action, data=None):
+        """
+        Executes a HTTP action asynchronously.
+
+        Args:
+            method: get or post
+            action: url action to perform
+            data: post parameters
+
+        Returns:
+            json results if any
+        """
+
+        # Get urls
+        urls = [f"{shard}/{action}" for shard in self.shards]
+        close = False
+
+        # Use existing loop if available, otherwise create one
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            close = True
+
+        try:
+            return loop.run_until_complete(self.run(urls, method, data))
+        finally:
+            # Close loop if it was created in this method
+            if close:
+                loop.close()
+
+    async def run(self, urls, method, data):
+        """
+        Runs an async action.
+
+        Args:
+            urls: run against this list of urls
+            method: get or post
+            data: list of data for each url or None
+
+        Returns:
+            json results if any
+        """
+
+        async with aiohttp.ClientSession(raise_for_status=True) as session:
+            tasks = []
+
+            for x, url in enumerate(urls):
+                if method == "post":
+                    if not data or data[x]:
+                        tasks.append(asyncio.ensure_future(self.post(session, url, data[x] if data else None)))
+                else:
+                    tasks.append(asyncio.ensure_future(self.get(session, url)))
+
+            return await asyncio.gather(*tasks)
+
+    async def get(self, session, url):
+        """
+        Runs an async HTTP GET request.
+
+        Args:
+            session: ClientSession
+            url: url
+
+        Returns:
+            json results if any
+        """
+
+        async with session.get(url) as resp:
+            return await resp.json()
+
+    async def post(self, session, url, data):
+        """
+        Runs an async HTTP POST request.
+
+        Args:
+            session: ClientSession
+            url: url
+            data: data to POST
+
+        Returns:
+            json results if any
+        """
+
+        async with session.post(url, json=data) as resp:
+            return await resp.json()

txtai/api/extension.py

@@ -0,0 +1,19 @@
+"""
+Extension module
+"""
+
+
+class Extension:
+    """
+    Defines an API extension. API extensions can expose custom pipelines or other custom logic.
+    """
+
+    def __call__(self, app):
+        """
+        Hook to register custom routing logic and/or modify the FastAPI instance.
+
+        Args:
+            app: FastAPI application instance
+        """
+
+        return

txtai/api/factory.py

@@ -0,0 +1,40 @@
+"""
+API factory module
+"""
+
+from ..util import Resolver
+
+
+class APIFactory:
+    """
+    API factory. Creates new API instances.
+    """
+
+    @staticmethod
+    def get(api):
+        """
+        Gets a new instance of api class.
+
+        Args:
+            api: API instance class
+
+        Returns:
+            API
+        """
+
+        return Resolver()(api)
+
+    @staticmethod
+    def create(config, api):
+        """
+        Creates a new API instance.
+
+        Args:
+            config: API configuration
+            api: API instance class
+
+        Returns:
+            API instance
+        """
+
+        return APIFactory.get(api)(config)

txtai/api/responses/__init__.py

@@ -0,0 +1,7 @@
+"""
+Responses imports
+"""
+
+from .factory import ResponseFactory
+from .json import JSONEncoder, JSONResponse
+from .messagepack import MessagePackResponse

txtai/api/responses/factory.py

@@ -0,0 +1,30 @@
+"""
+Factory module
+"""
+
+from .json import JSONResponse
+from .messagepack import MessagePackResponse
+
+
+class ResponseFactory:
+    """
+    Methods to create Response classes.
+    """
+
+    @staticmethod
+    def create(request):
+        """
+        Gets a response class for request using the Accept header.
+
+        Args:
+            request: request
+
+        Returns:
+            response class
+        """
+
+        # Get Accept header
+        accept = request.headers.get("Accept")
+
+        # Get response class
+        return MessagePackResponse if accept == MessagePackResponse.media_type else JSONResponse

txtai/api/responses/json.py

@@ -0,0 +1,56 @@
+"""
+JSON module
+"""
+
+import base64
+import json
+
+from io import BytesIO
+from typing import Any
+
+import fastapi.responses
+
+from PIL.Image import Image
+
+
+class JSONEncoder(json.JSONEncoder):
+    """
+    Extended JSONEncoder that serializes images and byte streams as base64 encoded text.
+    """
+
+    def default(self, o):
+        # Convert Image to BytesIO
+        if isinstance(o, Image):
+            buffered = BytesIO()
+            o.save(buffered, format=o.format, quality="keep")
+            o = buffered
+
+        # Unpack bytes from BytesIO
+        if isinstance(o, BytesIO):
+            o = o.getvalue()
+
+        # Base64 encode bytes instances
+        if isinstance(o, bytes):
+            return base64.b64encode(o).decode("utf-8")
+
+        # Default handler
+        return super().default(o)
+
+
+class JSONResponse(fastapi.responses.JSONResponse):
+    """
+    Extended JSONResponse that serializes images and byte streams as base64 encoded text.
+    """
+
+    def render(self, content: Any) -> bytes:
+        """
+        Renders content to the response.
+
+        Args:
+            content: input content
+
+        Returns:
+            rendered content as bytes
+        """
+
+        return json.dumps(content, ensure_ascii=False, allow_nan=False, indent=None, separators=(",", ":"), cls=JSONEncoder).encode("utf-8")

txtai/api/responses/messagepack.py

@@ -0,0 +1,51 @@
+"""
+MessagePack module
+"""
+
+from io import BytesIO
+from typing import Any
+
+import msgpack
+
+from fastapi import Response
+from PIL.Image import Image
+
+
+class MessagePackResponse(Response):
+    """
+    Encodes responses with MessagePack.
+    """
+
+    media_type = "application/msgpack"
+
+    def render(self, content: Any) -> bytes:
+        """
+        Renders content to the response.
+
+        Args:
+            content: input content
+
+        Returns:
+            rendered content as bytes
+        """
+
+        return msgpack.packb(content, default=MessagePackEncoder())
+
+
+class MessagePackEncoder:
+    """
+    Extended MessagePack encoder that converts images to bytes.
+    """
+
+    def __call__(self, o):
+        # Convert Image to bytes
+        if isinstance(o, Image):
+            buffered = BytesIO()
+            o.save(buffered, format=o.format, quality="keep")
+            o = buffered
+
+        # Get bytes from BytesIO
+        if isinstance(o, BytesIO):
+            o = o.getvalue()
+
+        return o

txtai/api/route.py

@@ -0,0 +1,41 @@
+"""
+Route module
+"""
+
+from fastapi.routing import APIRoute, get_request_handler
+
+from .responses import ResponseFactory
+
+
+class EncodingAPIRoute(APIRoute):
+    """
+    Extended APIRoute that encodes responses based on HTTP Accept header.
+    """
+
+    def get_route_handler(self):
+        """
+        Resolves a response class based on the HTTP Accept header.
+
+        Returns:
+            route handler function
+        """
+
+        async def handler(request):
+            route = get_request_handler(
+                dependant=self.dependant,
+                body_field=self.body_field,
+                status_code=self.status_code,
+                response_class=ResponseFactory.create(request),
+                response_field=self.secure_cloned_response_field,
+                response_model_include=self.response_model_include,
+                response_model_exclude=self.response_model_exclude,
+                response_model_by_alias=self.response_model_by_alias,
+                response_model_exclude_unset=self.response_model_exclude_unset,
+                response_model_exclude_defaults=self.response_model_exclude_defaults,
+                response_model_exclude_none=self.response_model_exclude_none,
+                dependency_overrides_provider=self.dependency_overrides_provider,
+            )
+
+            return await route(request)
+
+        return handler