Test fixes in openai_compat

This commit is contained in:
Swapna Lekkala 2025-09-17 16:50:46 -07:00
parent e56a3f266c
commit d60514b57b
7 changed files with 221 additions and 7 deletions

View file

@ -451,6 +451,20 @@ class ChatCompletionResponseStreamChunk(MetricResponseMixin):
event: ChatCompletionResponseEvent
@json_schema_type
class UsageInfo(BaseModel):
"""Usage information for a model.
:param completion_tokens: Number of tokens generated
:param prompt_tokens: Number of tokens in the prompt
:param total_tokens: Total number of tokens processed
"""
completion_tokens: int
prompt_tokens: int
total_tokens: int
@json_schema_type
class ChatCompletionResponse(MetricResponseMixin):
"""Response from a chat completion request.
@ -461,6 +475,7 @@ class ChatCompletionResponse(MetricResponseMixin):
completion_message: CompletionMessage
logprobs: list[TokenLogProbs] | None = None
usage: UsageInfo | None = None
@json_schema_type
@ -818,7 +833,21 @@ class OpenAIChoice(BaseModel):
@json_schema_type
class OpenAIChatCompletion(BaseModel):
class OpenAIChatCompletionUsage(BaseModel):
"""Usage information for an OpenAI-compatible chat completion response.
:param prompt_tokens: The number of tokens in the prompt
:param completion_tokens: The number of tokens in the completion
:param total_tokens: The total number of tokens used
"""
prompt_tokens: int
completion_tokens: int
total_tokens: int
@json_schema_type
class OpenAIChatCompletion(MetricResponseMixin):
"""Response from an OpenAI-compatible chat completion request.
:param id: The ID of the chat completion
@ -833,6 +862,7 @@ class OpenAIChatCompletion(BaseModel):
object: Literal["chat.completion"] = "chat.completion"
created: int
model: str
usage: OpenAIChatCompletionUsage | None = None
@json_schema_type