Add support for binary content and MessagePack, closes #630, closes #658

Makefile

@@ -33,7 +33,8 @@ test:
 # Run tests while calculating code coverage
 coverage:
 	coverage run -m unittest discover -v -k testann -s ${TEST_DIR}
-	coverage run -m unittest discover -v -k testapi.testauthorization -k testapi.testcluster -k testapi.testextension -s ${TEST_DIR}
+	coverage run -m unittest discover -v -k testapi.testauthorization -k testapi.testcluster -s ${TEST_DIR}
+	coverage run -m unittest discover -v -k testapi.testencoding -k testapi.testextension -s ${TEST_DIR}
 	coverage run -m unittest discover -v -k testapp -s ${TEST_DIR}
 	coverage run -m unittest discover -v -k testarchive -s ${TEST_DIR}
 	coverage run -m unittest discover -v -k testcloud -s ${TEST_DIR}

setup.py

@@ -30,6 +30,8 @@ extras["dev"] = [
 extras["api"] = [
     "aiohttp>=3.8.1",
     "fastapi>=0.94.0",
+    "msgpack>=1.0.7",
+    "python-multipart>=0.0.7",
     "uvicorn>=0.12.1",
 ]
 

src/python/txtai/api/__init__.py

@@ -9,8 +9,10 @@ try:
     from .base import API
     from .cluster import Cluster
     from .extension import Extension
-    from .factory import Factory
+    from .factory import APIFactory
+    from .responses import *
     from .routers import *
+    from .route import EncodingAPIRoute
 except ImportError as missing:
     # pylint: disable=W0707
     raise ImportError('API is not available - install "api" extra to enable') from missing

src/python/txtai/api/application.py

@@ -10,7 +10,7 @@ from fastapi import APIRouter, Depends, FastAPI
 
 from .authorization import Authorization
 from .base import API
-from .factory import Factory
+from .factory import APIFactory
 
 from ..app import Application
 
@@ -44,7 +44,7 @@ def create():
     if deps:
         for dep in deps.split(","):
             # Create and add dependency
-            dep = Factory.get(dep.strip())()
+            dep = APIFactory.get(dep.strip())()
             dependencies.append(Depends(dep))
 
     # Create FastAPI application
@@ -86,7 +86,7 @@ def lifespan(application):
 
     # Instantiate API instance
     api = os.environ.get("API_CLASS")
-    INSTANCE = Factory.create(config, api) if api else API(config)
+    INSTANCE = APIFactory.create(config, api) if api else API(config)
 
     # Get all known routers
     routers = apirouters()
@@ -109,7 +109,7 @@ def lifespan(application):
     if extensions:
         for extension in extensions.split(","):
             # Create instance and execute extension
-            extension = Factory.get(extension.strip())()
+            extension = APIFactory.get(extension.strip())()
             extension(application)
 
     yield

src/python/txtai/api/factory.py

@@ -5,7 +5,7 @@ API factory module
 from ..util import Resolver
 
 
-class Factory:
+class APIFactory:
     """
     API factory. Creates new API instances.
     """
@@ -37,4 +37,4 @@ class Factory:
             API instance
         """
 
-        return Factory.get(api)(config)
+        return APIFactory.get(api)(config)

src/python/txtai/api/responses/__init__.py

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

src/python/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

src/python/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")

src/python/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

src/python/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

src/python/txtai/api/routers/caption.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/caption")

src/python/txtai/api/routers/embeddings.py

@@ -2,14 +2,21 @@
 Defines API paths for embeddings endpoints.
 """
 
+from io import BytesIO
 from typing import List
 
-from fastapi import APIRouter, Body, HTTPException, Request
+import PIL
+
+from fastapi import APIRouter, Body, File, Form, HTTPException, Request, UploadFile
+from fastapi.encoders import jsonable_encoder
 
 from .. import application
+from ..responses import ResponseFactory
+from ..route import EncodingAPIRoute
+
 from ...app import ReadOnlyError
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/search")
@@ -26,13 +33,25 @@ def search(query: str, request: Request):
         list of {id: value, score: value} for index search, list of dict for an index + database search
     """
 
-    return application.get().search(query, request=request)
+    # Execute search
+    results = application.get().search(query, request=request)
+
+    # Encode using standard FastAPI encoder but skip certain classes
+    results = jsonable_encoder(results, custom_encoder={bytes: lambda x: x, BytesIO: lambda x: x, PIL.Image.Image: lambda x: x})
+
+    # Return raw response to prevent duplicate encoding
+    response = ResponseFactory.create(request)
+    return response(results)
 
 
 # pylint: disable=W0621
 @router.post("/batchsearch")
 def batchsearch(
-    queries: List[str] = Body(...), limit: int = Body(default=None), weights: float = Body(default=None), index: str = Body(default=None)
+    request: Request,
+    queries: List[str] = Body(...),
+    limit: int = Body(default=None),
+    weights: float = Body(default=None),
+    index: str = Body(default=None),
 ):
     """
     Finds documents most similar to the input queries. This method will run either an index search
@@ -48,7 +67,15 @@ def batchsearch(
         list of {id: value, score: value} per query for index search, list of dict per query for an index + database search
     """
 
-    return application.get().batchsearch(queries, limit, weights, index)
+    # Execute search
+    results = application.get().batchsearch(queries, limit, weights, index)
+
+    # Encode using standard FastAPI encoder but skip certain classes
+    results = jsonable_encoder(results, custom_encoder={bytes: lambda x: x, BytesIO: lambda x: x, PIL.Image.Image: lambda x: x})
+
+    # Return raw response to prevent duplicate encoding
+    response = ResponseFactory.create(request)
+    return response(results)
 
 
 @router.post("/add")
@@ -66,6 +93,48 @@ def add(documents: List[dict] = Body(...)):
         raise HTTPException(status_code=403, detail=e.args[0]) from e
 
 
+@router.post("/addobject")
+def addobject(data: List[bytes] = File(), uid: List[str] = Form(default=None), field: str = Form(default=None)):
+    """
+    Adds a batch of binary documents for indexing.
+
+    Args:
+        data: list of binary objects
+        uid: list of corresponding ids
+        field: optional object field name
+    """
+
+    if uid and len(data) != len(uid):
+        raise HTTPException(status_code=422, detail="Length of data and document lists must match")
+
+    try:
+        # Add objects
+        application.get().addobject(data, uid, field)
+    except ReadOnlyError as e:
+        raise HTTPException(status_code=403, detail=e.args[0]) from e
+
+
+@router.post("/addimage")
+def addimage(data: List[UploadFile] = File(), uid: List[str] = Form(), field: str = Form(default=None)):
+    """
+    Adds a batch of images for indexing.
+
+    Args:
+        data: list of images
+        uid: list of corresponding ids
+        field: optional object field name
+    """
+
+    if uid and len(data) != len(uid):
+        raise HTTPException(status_code=422, detail="Length of data and uid lists must match")
+
+    try:
+        # Add images
+        application.get().addobject([PIL.Image.open(content.file) for content in data], uid, field)
+    except ReadOnlyError as e:
+        raise HTTPException(status_code=403, detail=e.args[0]) from e
+
+
 @router.get("/index")
 def index():
     """

src/python/txtai/api/routers/entity.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/entity")

src/python/txtai/api/routers/extractor.py

@@ -7,8 +7,9 @@ from typing import List, Optional
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.post("/extract")

src/python/txtai/api/routers/labels.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.post("/label")

src/python/txtai/api/routers/objects.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/objects")

src/python/txtai/api/routers/segmentation.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/segment")

src/python/txtai/api/routers/similarity.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.post("/similarity")

src/python/txtai/api/routers/summary.py

@@ -7,8 +7,9 @@ from typing import List, Optional
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/summary")

src/python/txtai/api/routers/tabular.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/tabular")

src/python/txtai/api/routers/textractor.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/textract")

src/python/txtai/api/routers/transcription.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/transcribe")

src/python/txtai/api/routers/translation.py

@@ -7,8 +7,9 @@ from typing import List, Optional
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.get("/translate")

src/python/txtai/api/routers/workflow.py

@@ -7,8 +7,9 @@ from typing import List
 from fastapi import APIRouter, Body
 
 from .. import application
+from ..route import EncodingAPIRoute
 
-router = APIRouter()
+router = APIRouter(route_class=EncodingAPIRoute)
 
 
 @router.post("/workflow")

src/python/txtai/app/base.py

@@ -387,6 +387,36 @@ class Application:
         # Return unmodified input documents
         return documents
 
+    def addobject(self, data, uid, field):
+        """
+        Helper method that builds a batch of object documents.
+
+        Args:
+            data: object content
+            uid: optional list of corresponding uids
+            field: optional field to set
+
+        Returns:
+            documents
+        """
+
+        # Raise error if index is not writable
+        if not self.config.get("writable"):
+            raise ReadOnlyError("Attempting to add documents to a read-only index (writable != True)")
+
+        documents = []
+        for x, content in enumerate(data):
+            if field:
+                row = {"id": uid[x], field: content} if uid else {field: content}
+            elif uid:
+                row = (uid[x], content)
+            else:
+                row = content
+
+            documents.append(row)
+
+        return self.add(documents)
+
     def index(self):
         """
         Builds an embeddings index for previously batched documents.

test/python/testapi/testencoding.py

@@ -0,0 +1,163 @@
+"""
+Encoding module tests
+"""
+
+import base64
+import os
+import tempfile
+import unittest
+import urllib.parse
+
+from io import BytesIO
+from unittest.mock import patch
+
+import msgpack
+import numpy as np
+import PIL
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+from txtai.api.responses import JSONEncoder
+
+# pylint: disable = C0411
+from utils import Utils
+
+# Configuration for image storage
+INDEX = """
+# Allow indexing of documents
+writable: %s
+
+embeddings:
+    defaults: False
+    content: True
+    objects: %s
+"""
+
+
+class TestEncoding(unittest.TestCase):
+    """
+    API tests for response encoding
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start(yaml):
+        """
+        Starts a mock FastAPI client.
+
+        Args:
+            yaml: input configuration
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(yaml)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestEncoding.start(INDEX % ("True", "image"))
+
+    def testImages(self):
+        """
+        Test image encoding
+        """
+
+        with open(Utils.PATH + "/books.jpg", "rb") as f:
+            self.client.post("addimage", data={"uid": 0}, files={"data": f})
+            self.client.get("index")
+
+        query = urllib.parse.quote_plus("select id, object from txtai limit 1")
+        results = self.client.get(f"search?query={query}").json()
+
+        # Test reading image
+        self.assertIsInstance(PIL.Image.open(BytesIO(base64.b64decode(results[0]["object"]))), PIL.Image.Image)
+
+    def testInvalidInputs(self):
+        """
+        Tests invalid parameter inputs
+        """
+
+        response = self.client.post("addimage", data={"uid": [0, 1]}, files={"data": b"123"})
+        self.assertEqual(response.status_code, 422)
+
+        response = self.client.post("addobject", data={"uid": [0, 1]}, files={"data": b"123"})
+        self.assertEqual(response.status_code, 422)
+
+    def testInvalidJSON(self):
+        """
+        Tests that invalid JSON raises an exception
+        """
+
+        with self.assertRaises(TypeError):
+            JSONEncoder().encode(np.random.rand(1, 1))
+
+    def testMessagePack(self):
+        """
+        Test message pack encoding
+        """
+
+        # Validate binary encoding
+        results = self.client.get("count", headers={"Accept": "application/msgpack"}).content
+        self.assertEqual(results, b"\x01")
+
+        # Validate query result
+        query = urllib.parse.quote_plus("select id, object from txtai limit 1")
+        results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
+        results = msgpack.unpackb(results)
+
+        # Test reading image
+        self.assertIsInstance(PIL.Image.open(BytesIO(results[0]["object"])), PIL.Image.Image)
+
+    def testObjects(self):
+        """
+        Test object encoding
+        """
+
+        # Recreate model with standard object encoding
+        self.client = TestEncoding.start(INDEX % ("True", "True"))
+
+        # Test various formats
+        self.client.post("addobject", data={"uid": "id0"}, files={"data": b"1234"})
+        self.client.post("addobject", files={"data": b"ABC"})
+        self.client.post("addobject", data={"uid": "id1", "field": "object"}, files={"data": b"A1234"})
+        self.client.get("index")
+
+        query = urllib.parse.quote_plus("select id, object from txtai where id = 'id0' limit 1")
+        results = self.client.get(f"search?query={query}").json()
+        self.assertEqual(base64.b64decode(results[0]["object"]), b"1234")
+
+        # Test with messagepack encoding
+        results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
+        results = msgpack.unpackb(results)
+        self.assertEqual(results[0]["object"], b"1234")
+
+        count = self.client.get("count").json()
+        self.assertEqual(count, 3)
+
+    def testReadOnly(self):
+        """
+        Tests read only indexes
+        """
+
+        # Recreate model with standard object encoding
+        self.client = TestEncoding.start(INDEX % ("False", "True"))
+
+        # Test errors raised for write operations
+        with open(Utils.PATH + "/books.jpg", "rb") as f:
+            response = self.client.post("addimage", data={"uid": 0}, files={"data": f})
+            self.assertEqual(response.status_code, 403)
+
+        self.assertEqual(self.client.post("addobject", data={"uid": 0}, files={"data": b"1234"}).status_code, 403)