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 time
import uuid import uuid
from litellm import ModelResponse, ChatCompletionChunk, Choices, StreamingChoices, Delta, Message, Usage
class OpenAIFormatter: class OpenAIFormatter:
...@@ -25,24 +24,37 @@ class OpenAIFormatter: ...@@ -25,24 +24,37 @@ class OpenAIFormatter:
tool_calls: Optional list of tool calls to include tool_calls: Optional list of tool calls to include
Returns: Returns:
Dictionary representation of ModelResponse Dictionary representation of the response
""" """
return ModelResponse( message = {
id=self.id, "role": "assistant",
model=self.model_name, "content": text if not tool_calls else None,
object="chat.completion", }
created=int(time.time()), if tool_calls:
choices=[Choices( message["tool_calls"] = tool_calls
finish_reason="tool_calls" if tool_calls else "stop",
index=0, choice = {
message=Message(content=text if not tool_calls else None, role="assistant", tool_calls=tool_calls) "index": 0,
)], "message": message,
usage=Usage( "finish_reason": "tool_calls" if tool_calls else "stop",
prompt_tokens=prompt_tokens, }
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens return {
) "id": self.id,
).model_dump() "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: def format_chunk(self, delta_text: str, is_final: bool = False, usage: dict = None) -> dict:
"""Format a streaming chunk response. """Format a streaming chunk response.
...@@ -53,20 +65,31 @@ class OpenAIFormatter: ...@@ -53,20 +65,31 @@ class OpenAIFormatter:
usage: Optional usage information (typically only sent on final chunk) usage: Optional usage information (typically only sent on final chunk)
Returns: Returns:
Dictionary representation of ChatCompletionChunk Dictionary representation of the chunk
""" """
return ChatCompletionChunk( delta = {
id=self.id, "content": delta_text,
model=self.model_name, "role": "assistant",
object="chat.completion.chunk", }
created=int(time.time()),
choices=[StreamingChoices( choice = {
finish_reason="stop" if is_final else None, "index": 0,
index=0, "delta": delta,
delta=Delta(content=delta_text, role="assistant") "finish_reason": "stop" if is_final else None,
)], }
usage=usage # Only send usage on the final chunk
).model_dump() 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: def format_final_chunk(self, usage: dict = None) -> dict:
"""Format the final streaming chunk with usage information. """Format the final streaming chunk with usage information.
...@@ -75,17 +98,28 @@ class OpenAIFormatter: ...@@ -75,17 +98,28 @@ class OpenAIFormatter:
usage: Usage statistics dictionary with prompt_tokens, completion_tokens, total_tokens usage: Usage statistics dictionary with prompt_tokens, completion_tokens, total_tokens
Returns: Returns:
Dictionary representation of the final ChatCompletionChunk Dictionary representation of the final chunk
""" """
return ChatCompletionChunk( delta = {
id=self.id, "content": None,
model=self.model_name, "role": "assistant",
object="chat.completion.chunk", }
created=int(time.time()),
choices=[StreamingChoices( choice = {
finish_reason="stop", "index": 0,
index=0, "delta": delta,
delta=Delta(content=None, role="assistant") "finish_reason": "stop",
)], }
usage=usage
).model_dump() 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