fixup! fix LLMChoice with improved openai schema

livepeer · Jan 10, 2025 · a943e98 · a943e98
1 parent 807f477
commit a943e98
Show file tree

Hide file tree

Showing 5 changed files with 490 additions and 33 deletions.
diff --git a/runner/app/routes/llm.py b/runner/app/routes/llm.py
@@ -19,18 +19,19 @@
     status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
 }
 
-
 @router.post(
     "/llm",
-    response_model=LLMResponse,
+    response_model=LLMResponse
+,
     responses=RESPONSES,
     operation_id="genLLM",
     description="Generate text using a language model.",
     summary="LLM",
     tags=["generate"],
     openapi_extra={"x-speakeasy-name-override": "llm"},
 )
-@router.post("/llm/", response_model=LLMResponse, responses=RESPONSES, include_in_schema=False)
+@router.post("/llm/", response_model=LLMResponse
+, responses=RESPONSES, include_in_schema=False)
 async def llm(
     request: LLMRequest,
     pipeline: Pipeline = Depends(get_pipeline),

diff --git a/runner/app/routes/utils.py b/runner/app/routes/utils.py
@@ -77,25 +77,39 @@ class LLMMessage(BaseModel):
     content: str
 
 
-class LLMChoice(BaseModel):
-    delta: Optional[LLMMessage]
-    message: Optional[LLMMessage]
+class LLMBaseChoice(BaseModel):
     index: int
-    finish_reason: Optional[str]
+    finish_reason: str = "" # Needs OpenAPI 3.1 support to make optional
 
 
 class LLMTokenUsage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
 
+class LLMChoice(LLMBaseChoice):
+    delta: LLMMessage = None
+    message: LLMMessage = None
 
 class LLMResponse(BaseModel):
-    choices: List[LLMChoice]
-    tokens_used: LLMTokenUsage
     id: str
     model: str
     created: int
+    tokens_used: LLMTokenUsage
+    choices: List[LLMChoice]
+
+
+# class LLMStreamChoice(LLMBaseChoice):
+#     delta: LLMMessage
+
+# class LLMNonStreamChoice(LLMBaseChoice):
+#     message: LLMMessage
+
+# class LLMStreamResponse(LLMBaseResponse):
+#     choices: List[LLMStreamChoice]
+
+# class LLMNonStreamResponse(LLMBaseResponse):
+#     choices: List[LLMNonStreamChoice]
 
 
 class LLMRequest(BaseModel):

diff --git a/runner/gateway.openapi.yaml b/runner/gateway.openapi.yaml
@@ -868,18 +868,19 @@ components:
       description: Response model for text generation.
     LLMChoice:
       properties:
-        delta:
-          $ref: '#/components/schemas/LLMMessage'
         index:
           type: integer
           title: Index
         finish_reason:
           type: string
           title: Finish Reason
           default: ''
+        delta:
+          $ref: '#/components/schemas/LLMMessage'
+        message:
+          $ref: '#/components/schemas/LLMMessage'
       type: object
       required:
-      - delta
       - index
       title: LLMChoice
     LLMMessage:
@@ -932,13 +933,6 @@ components:
       title: LLMRequest
     LLMResponse:
       properties:
-        choices:
-          items:
-            $ref: '#/components/schemas/LLMChoice'
-          type: array
-          title: Choices
-        tokens_used:
-          $ref: '#/components/schemas/LLMTokenUsage'
         id:
           type: string
           title: Id
@@ -948,13 +942,20 @@ components:
         created:
           type: integer
           title: Created
+        tokens_used:
+          $ref: '#/components/schemas/LLMTokenUsage'
+        choices:
+          items:
+            $ref: '#/components/schemas/LLMChoice'
+          type: array
+          title: Choices
       type: object
       required:
-      - choices
-      - tokens_used
       - id
       - model
       - created
+      - tokens_used
+      - choices
       title: LLMResponse
     LLMTokenUsage:
       properties:

diff --git a/runner/openapi.yaml b/runner/openapi.yaml
@@ -1014,18 +1014,19 @@ components:
       description: Response model for text generation.
     LLMChoice:
       properties:
-        delta:
-          $ref: '#/components/schemas/LLMMessage'
         index:
           type: integer
           title: Index
         finish_reason:
           type: string
           title: Finish Reason
           default: ''
+        delta:
+          $ref: '#/components/schemas/LLMMessage'
+        message:
+          $ref: '#/components/schemas/LLMMessage'
       type: object
       required:
-      - delta
       - index
       title: LLMChoice
     LLMMessage:
@@ -1078,13 +1079,6 @@ components:
       title: LLMRequest
     LLMResponse:
       properties:
-        choices:
-          items:
-            $ref: '#/components/schemas/LLMChoice'
-          type: array
-          title: Choices
-        tokens_used:
-          $ref: '#/components/schemas/LLMTokenUsage'
         id:
           type: string
           title: Id
@@ -1094,13 +1088,20 @@ components:
         created:
           type: integer
           title: Created
+        tokens_used:
+          $ref: '#/components/schemas/LLMTokenUsage'
+        choices:
+          items:
+            $ref: '#/components/schemas/LLMChoice'
+          type: array
+          title: Choices
       type: object
       required:
-      - choices
-      - tokens_used
       - id
       - model
       - created
+      - tokens_used
+      - choices
       title: LLMResponse
     LLMTokenUsage:
       properties: