llama-index

1
import abc
2
import json
3
from typing import TYPE_CHECKING, List
4

5
if TYPE_CHECKING:
6
    from botocore.response import StreamingBody
7

8
from llama_index.legacy.bridge.pydantic import Field
9

10

11
class BaseIOHandler(metaclass=abc.ABCMeta):
12
    content_type: str = Field(
13
        description="The MIME type of the input data in the request body.",
14
    )
15
    accept: str = Field(
16
        description="The desired MIME type of the inference response from the model container.",
17
    )
18

19
    @classmethod
20
    def __subclasshook__(cls, subclass: type) -> bool:
21
        return (
22
            hasattr(subclass, "content_type")
23
            and hasattr(subclass, "accept")
24
            and hasattr(subclass, "serialize_input")
25
            and callable(subclass.serialize_input)
26
            and hasattr(subclass, "deserialize_output")
27
            and callable(subclass.deserialize_output)
28
            or NotImplemented
29
        )
30

31
    @abc.abstractmethod
32
    def serialize_input(self, request: List[str], model_kwargs: dict) -> bytes:
33
        raise NotImplementedError
34

35
    @abc.abstractmethod
36
    def deserialize_output(self, response: "StreamingBody") -> List[List[float]]:
37
        raise NotImplementedError
38

39

40
class IOHandler(BaseIOHandler):
41
    content_type: str = "application/json"
42
    accept: str = "application/json"
43

44
    def serialize_input(self, request: List[str], model_kwargs: dict) -> bytes:
45
        request_str = json.dumps({"inputs": request, **model_kwargs})
46
        return request_str.encode("utf-8")
47

48
    def deserialize_output(self, response: "StreamingBody") -> List[List[float]]:
49
        response_json = json.loads(response.read().decode("utf-8"))
50
        return response_json["vectors"]
51