Rob's fork diff

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = "~=3.11"
 readme = "README.md"
 dependencies = [
     "fastapi>=0.115.9,<0.116",
+    "fastapi-utils>=0.8.0",
     "pandas>=2.2.3,<3",
     "prometheus-client>=0.21.1,<0.23",
     "pydantic>=2.4.2,<3",
@@ -29,11 +30,10 @@ dev = [
     "pytest-cov>=4.1.0,<7",
     "httpx>=0.25.0,<0.29",
 ]
-eval = ["lm-eval[api]==0.4.4", "fastapi-utils>=0.8.0", "typing-inspect==0.9.0"]
+eval = ["lm-eval[api]==0.4.4", "typing-inspect==0.9.0"]
 protobuf = ["numpy>=1.24.0,<3", "grpcio>=1.62.1,<2", "grpcio-tools>=1.62.1,<2"]
 mariadb = ["mariadb>=1.1.12", "javaobj-py3==0.4.4"]
 
-
 [tool.hatch.build.targets.sdist]
 include = ["src"]
 

src/endpoints/consumer/__init__.py

@@ -0,0 +1,74 @@
+from typing import Optional, Dict, List, Literal
+
+from pydantic import BaseModel
+
+
+PartialKind = Literal["request", "response"]
+
+class PartialPayloadId(BaseModel):
+    prediction_id: Optional[str] = None
+    kind: Optional[PartialKind] = None
+
+    def get_prediction_id(self) -> str:
+        return self.prediction_id
+
+    def set_prediction_id(self, id: str):
+        self.prediction_id = id
+
+    def get_kind(self) -> PartialKind:
+        return self.kind
+
+    def set_kind(self, kind: PartialKind):
+        self.kind = kind
+
+
+class InferencePartialPayload(BaseModel):
+    partialPayloadId: Optional[PartialPayloadId] = None
+    metadata: Optional[Dict[str, str]] = {}
+    data: Optional[str] = None
+    modelid: Optional[str] = None
+
+    def get_id(self) -> str:
+        return self.partialPayloadId.prediction_id if self.partialPayloadId else None
+
+    def set_id(self, id: str):
+        if not self.partialPayloadId:
+            self.partialPayloadId = PartialPayloadId()
+        self.partialPayloadId.prediction_id = id
+
+    def get_kind(self) -> PartialKind:
+        return self.partialPayloadId.kind if self.partialPayloadId else None
+
+    def set_kind(self, kind: PartialKind):
+        if not self.partialPayloadId:
+            self.partialPayloadId = PartialPayloadId()
+        self.partialPayloadId.kind = kind
+
+    def get_model_id(self) -> str:
+        return self.modelid
+
+    def set_model_id(self, model_id: str):
+        self.modelid = model_id
+
+
+class KServeData(BaseModel):
+    name: str
+    shape: List[int]
+    datatype: str
+    parameters: Optional[Dict[str, str]] = None
+    data: List
+
+
+class KServeInferenceRequest(BaseModel):
+    id: Optional[str] = None
+    parameters: Optional[Dict[str, str]] = None
+    inputs: List[KServeData]
+    outputs: Optional[List[KServeData]] = None
+
+
+class KServeInferenceResponse(BaseModel):
+    model_name: str = None
+    model_version: Optional[str] = None
+    id: Optional[str] = None
+    parameters: Optional[Dict[str, str]] = None
+    outputs: List[KServeData]