Skip to content

Commit a943e98

Browse files
committed
fixup! fix LLMChoice with improved openai schema
1 parent 807f477 commit a943e98

File tree

5 files changed

+490
-33
lines changed

5 files changed

+490
-33
lines changed

runner/app/routes/llm.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,19 @@
1919
status.HTTP_500_INTERNAL_SERVER_ERROR: {"model": HTTPError},
2020
}
2121

22-
2322
@router.post(
2423
"/llm",
25-
response_model=LLMResponse,
24+
response_model=LLMResponse
25+
,
2626
responses=RESPONSES,
2727
operation_id="genLLM",
2828
description="Generate text using a language model.",
2929
summary="LLM",
3030
tags=["generate"],
3131
openapi_extra={"x-speakeasy-name-override": "llm"},
3232
)
33-
@router.post("/llm/", response_model=LLMResponse, responses=RESPONSES, include_in_schema=False)
33+
@router.post("/llm/", response_model=LLMResponse
34+
, responses=RESPONSES, include_in_schema=False)
3435
async def llm(
3536
request: LLMRequest,
3637
pipeline: Pipeline = Depends(get_pipeline),

runner/app/routes/utils.py

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,25 +77,39 @@ class LLMMessage(BaseModel):
7777
content: str
7878

7979

80-
class LLMChoice(BaseModel):
81-
delta: Optional[LLMMessage]
82-
message: Optional[LLMMessage]
80+
class LLMBaseChoice(BaseModel):
8381
index: int
84-
finish_reason: Optional[str]
82+
finish_reason: str = "" # Needs OpenAPI 3.1 support to make optional
8583

8684

8785
class LLMTokenUsage(BaseModel):
8886
prompt_tokens: int
8987
completion_tokens: int
9088
total_tokens: int
9189

90+
class LLMChoice(LLMBaseChoice):
91+
delta: LLMMessage = None
92+
message: LLMMessage = None
9293

9394
class LLMResponse(BaseModel):
94-
choices: List[LLMChoice]
95-
tokens_used: LLMTokenUsage
9695
id: str
9796
model: str
9897
created: int
98+
tokens_used: LLMTokenUsage
99+
choices: List[LLMChoice]
100+
101+
102+
# class LLMStreamChoice(LLMBaseChoice):
103+
# delta: LLMMessage
104+
105+
# class LLMNonStreamChoice(LLMBaseChoice):
106+
# message: LLMMessage
107+
108+
# class LLMStreamResponse(LLMBaseResponse):
109+
# choices: List[LLMStreamChoice]
110+
111+
# class LLMNonStreamResponse(LLMBaseResponse):
112+
# choices: List[LLMNonStreamChoice]
99113

100114

101115
class LLMRequest(BaseModel):

runner/gateway.openapi.yaml

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -868,18 +868,19 @@ components:
868868
description: Response model for text generation.
869869
LLMChoice:
870870
properties:
871-
delta:
872-
$ref: '#/components/schemas/LLMMessage'
873871
index:
874872
type: integer
875873
title: Index
876874
finish_reason:
877875
type: string
878876
title: Finish Reason
879877
default: ''
878+
delta:
879+
$ref: '#/components/schemas/LLMMessage'
880+
message:
881+
$ref: '#/components/schemas/LLMMessage'
880882
type: object
881883
required:
882-
- delta
883884
- index
884885
title: LLMChoice
885886
LLMMessage:
@@ -932,13 +933,6 @@ components:
932933
title: LLMRequest
933934
LLMResponse:
934935
properties:
935-
choices:
936-
items:
937-
$ref: '#/components/schemas/LLMChoice'
938-
type: array
939-
title: Choices
940-
tokens_used:
941-
$ref: '#/components/schemas/LLMTokenUsage'
942936
id:
943937
type: string
944938
title: Id
@@ -948,13 +942,20 @@ components:
948942
created:
949943
type: integer
950944
title: Created
945+
tokens_used:
946+
$ref: '#/components/schemas/LLMTokenUsage'
947+
choices:
948+
items:
949+
$ref: '#/components/schemas/LLMChoice'
950+
type: array
951+
title: Choices
951952
type: object
952953
required:
953-
- choices
954-
- tokens_used
955954
- id
956955
- model
957956
- created
957+
- tokens_used
958+
- choices
958959
title: LLMResponse
959960
LLMTokenUsage:
960961
properties:

runner/openapi.yaml

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1014,18 +1014,19 @@ components:
10141014
description: Response model for text generation.
10151015
LLMChoice:
10161016
properties:
1017-
delta:
1018-
$ref: '#/components/schemas/LLMMessage'
10191017
index:
10201018
type: integer
10211019
title: Index
10221020
finish_reason:
10231021
type: string
10241022
title: Finish Reason
10251023
default: ''
1024+
delta:
1025+
$ref: '#/components/schemas/LLMMessage'
1026+
message:
1027+
$ref: '#/components/schemas/LLMMessage'
10261028
type: object
10271029
required:
1028-
- delta
10291030
- index
10301031
title: LLMChoice
10311032
LLMMessage:
@@ -1078,13 +1079,6 @@ components:
10781079
title: LLMRequest
10791080
LLMResponse:
10801081
properties:
1081-
choices:
1082-
items:
1083-
$ref: '#/components/schemas/LLMChoice'
1084-
type: array
1085-
title: Choices
1086-
tokens_used:
1087-
$ref: '#/components/schemas/LLMTokenUsage'
10881082
id:
10891083
type: string
10901084
title: Id
@@ -1094,13 +1088,20 @@ components:
10941088
created:
10951089
type: integer
10961090
title: Created
1091+
tokens_used:
1092+
$ref: '#/components/schemas/LLMTokenUsage'
1093+
choices:
1094+
items:
1095+
$ref: '#/components/schemas/LLMChoice'
1096+
type: array
1097+
title: Choices
10971098
type: object
10981099
required:
1099-
- choices
1100-
- tokens_used
11011100
- id
11021101
- model
11031102
- created
1103+
- tokens_used
1104+
- choices
11041105
title: LLMResponse
11051106
LLMTokenUsage:
11061107
properties:

0 commit comments

Comments
 (0)