feat(sdks): implement SDK-specific API customizations (livepeer#191)

* feat(sdks): implement SDK-specific API customizations This commit introduces several SDK-specific OpenAPI configurations to the runner API configuration. These customizations will enhance the SDKs we are planning to release. Co-authored-by: Victor Elias <[email protected]>
JJassonn69 · Sep 18, 2024 · b770c82 · b770c82
1 parent b816ef2
commit b770c82
Show file tree

Hide file tree

Showing 18 changed files with 418 additions and 1,314 deletions.
diff --git a/cmd/examples/image-to-image/main.go b/cmd/examples/image-to-image/main.go
@@ -69,7 +69,7 @@ func main() {
 	imageFile := types.File{}
 	imageFile.InitFromBytes(imageBytes, imagePath)
 
-	req := worker.ImageToImageMultipartRequestBody{
+	req := worker.GenImageToImageMultipartRequestBody{
 		Image:   imageFile,
 		ModelId: &modelID,
 		Prompt:  prompt,

diff --git a/cmd/examples/image-to-video/main.go b/cmd/examples/image-to-video/main.go
@@ -68,7 +68,7 @@ func main() {
 	imageFile := types.File{}
 	imageFile.InitFromBytes(imageBytes, imagePath)
 
-	req := worker.ImageToVideoMultipartRequestBody{
+	req := worker.GenImageToVideoMultipartRequestBody{
 		Image:   imageFile,
 		ModelId: &modelID,
 	}

diff --git a/cmd/examples/text-to-image/main.go b/cmd/examples/text-to-image/main.go
@@ -59,7 +59,7 @@ func main() {
 
 	prompt := args[1]
 
-	req := worker.TextToImageJSONRequestBody{
+	req := worker.GenTextToImageJSONRequestBody{
 		ModelId: &modelID,
 		Prompt:  prompt,
 	}

diff --git a/runner/app/main.py b/runner/app/main.py
@@ -23,8 +23,6 @@ async def lifespan(app: FastAPI):
     app.pipeline = load_pipeline(pipeline, model_id)
     app.include_router(load_route(pipeline))
 
-    use_route_names_as_operation_ids(app)
-
     logger.info(f"Started up with pipeline {app.pipeline}")
     yield
     logger.info("Shutting down")

diff --git a/runner/app/routes/audio_to_text.py b/runner/app/routes/audio_to_text.py
@@ -15,6 +15,15 @@
 logger = logging.getLogger(__name__)
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_413_REQUEST_ENTITY_TOO_LARGE: {"model": HTTPError},
@@ -52,6 +61,10 @@ def handle_pipeline_error(e: Exception) -> JSONResponse:
     response_model=TextResponse,
     responses=RESPONSES,
     description="Transcribe audio files to text.",
+    operation_id="genAudioToText",
+    summary="Audio To Text",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "audioToText"},
 )
 @router.post(
     "/audio-to-text/",

diff --git a/runner/app/routes/health.py b/runner/app/routes/health.py
@@ -8,7 +8,7 @@ class HealthCheck(BaseModel):
     status: str = "OK"
 
 
-@router.get("/health", response_model=HealthCheck)
+@router.get("/health", operation_id="health", response_model=HealthCheck)
 @router.get("/health/", response_model=HealthCheck, include_in_schema=False)
 def health() -> HealthCheck:
     return HealthCheck(status="OK")
diff --git a/runner/app/routes/image_to_image.py b/runner/app/routes/image_to_image.py
@@ -19,6 +19,15 @@
 
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
@@ -32,6 +41,10 @@
     response_model=ImageResponse,
     responses=RESPONSES,
     description="Apply image transformations to a provided image.",
+    operation_id="genImageToImage",
+    summary="Image To Image",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "imageToImage"},
 )
 @router.post(
     "/image-to-image/",

diff --git a/runner/app/routes/image_to_video.py b/runner/app/routes/image_to_video.py
@@ -18,6 +18,15 @@
 logger = logging.getLogger(__name__)
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
@@ -31,6 +40,10 @@
     response_model=VideoResponse,
     responses=RESPONSES,
     description="Generate a video from a provided image.",
+    operation_id="genImageToVideo",
+    summary="Image To Video",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "imageToVideo"},
 )
 @router.post(
     "/image-to-video/",

diff --git a/runner/app/routes/segment_anything_2.py b/runner/app/routes/segment_anything_2.py
@@ -24,6 +24,15 @@
 logger = logging.getLogger(__name__)
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
@@ -37,6 +46,10 @@
     response_model=MasksResponse,
     responses=RESPONSES,
     description="Segment objects in an image.",
+    operation_id="genSegmentAnything2",
+    summary="Segment Anything 2",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "segmentAnything2"},
 )
 @router.post(
     "/segment-anything-2/",

diff --git a/runner/app/routes/text_to_image.py b/runner/app/routes/text_to_image.py
@@ -92,6 +92,15 @@ class TextToImageParams(BaseModel):
 
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
@@ -103,6 +112,10 @@ class TextToImageParams(BaseModel):
     response_model=ImageResponse,
     responses=RESPONSES,
     description="Generate images from text prompts.",
+    operation_id="genTextToImage",
+    summary="Text To Image",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "textToImage"},
 )
 @router.post(
     "/text-to-image/",

diff --git a/runner/app/routes/upscale.py b/runner/app/routes/upscale.py
@@ -19,6 +19,15 @@
 
 
 RESPONSES = {
+    status.HTTP_200_OK: {
+        "content": {
+            "application/json": {
+                "schema": {
+                    "x-speakeasy-name-override": "data",
+                }
+            }
+        },
+    },
     status.HTTP_400_BAD_REQUEST: {"model": HTTPError},
     status.HTTP_401_UNAUTHORIZED: {"model": HTTPError},
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
@@ -32,6 +41,10 @@
     response_model=ImageResponse,
     responses=RESPONSES,
     description="Upscale an image by increasing its resolution.",
+    operation_id="genUpscale",
+    summary="Upscale",
+    tags=["generate"],
+    openapi_extra={"x-speakeasy-name-override": "upscale"},
 )
 @router.post(
     "/upscale/",