Commit 8280060e authored by Your Name's avatar Your Name

Fix OpenAIFormatter to not rely on litellm imports

The litellm library doesn't export Delta, Choices, etc. directly.
Rewrote the formatter to build response dictionaries directly.
parent 98a48640
import time
import uuid
from litellm import ModelResponse, ChatCompletionChunk, Choices, StreamingChoices, Delta, Message, Usage
class OpenAIFormatter:
......@@ -25,24 +24,37 @@ class OpenAIFormatter:
tool_calls: Optional list of tool calls to include
Returns:
Dictionary representation of ModelResponse
Dictionary representation of the response
"""
return ModelResponse(
id=self.id,
model=self.model_name,
object="chat.completion",
created=int(time.time()),
choices=[Choices(
finish_reason="tool_calls" if tool_calls else "stop",
index=0,
message=Message(content=text if not tool_calls else None, role="assistant", tool_calls=tool_calls)
)],
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens
)
).model_dump()
message = {
"role": "assistant",
"content": text if not tool_calls else None,
}
if tool_calls:
message["tool_calls"] = tool_calls
choice = {
"index": 0,
"message": message,
"finish_reason": "tool_calls" if tool_calls else "stop",
}
return {
"id": self.id,
"object": "chat.completion",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
"provider": {
"provider_name": "coderai",
"provider_id": "coderai",
},
}
def format_chunk(self, delta_text: str, is_final: bool = False, usage: dict = None) -> dict:
"""Format a streaming chunk response.
......@@ -53,20 +65,31 @@ class OpenAIFormatter:
usage: Optional usage information (typically only sent on final chunk)
Returns:
Dictionary representation of ChatCompletionChunk
Dictionary representation of the chunk
"""
return ChatCompletionChunk(
id=self.id,
model=self.model_name,
object="chat.completion.chunk",
created=int(time.time()),
choices=[StreamingChoices(
finish_reason="stop" if is_final else None,
index=0,
delta=Delta(content=delta_text, role="assistant")
)],
usage=usage # Only send usage on the final chunk
).model_dump()
delta = {
"content": delta_text,
"role": "assistant",
}
choice = {
"index": 0,
"delta": delta,
"finish_reason": "stop" if is_final else None,
}
chunk = {
"id": self.id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
}
if usage and is_final:
chunk["usage"] = usage
return chunk
def format_final_chunk(self, usage: dict = None) -> dict:
"""Format the final streaming chunk with usage information.
......@@ -75,17 +98,28 @@ class OpenAIFormatter:
usage: Usage statistics dictionary with prompt_tokens, completion_tokens, total_tokens
Returns:
Dictionary representation of the final ChatCompletionChunk
Dictionary representation of the final chunk
"""
return ChatCompletionChunk(
id=self.id,
model=self.model_name,
object="chat.completion.chunk",
created=int(time.time()),
choices=[StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(content=None, role="assistant")
)],
usage=usage
).model_dump()
delta = {
"content": None,
"role": "assistant",
}
choice = {
"index": 0,
"delta": delta,
"finish_reason": "stop",
}
chunk = {
"id": self.id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": self.model_name,
"choices": [choice],
}
if usage:
chunk["usage"] = usage
return chunk
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment