phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

have response type inherit from openaiobject

Browse source
This commit is contained in:
Krrish Dholakia 2023-08-18 05:13:30 -07:00
parent 3d1934a829
commit ee20e6a1e6
6 changed files with 60 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

76
litellm/utils.py
View file

@ -51,67 +51,35 @@ local_cache = {}
# 'usage': {'prompt_tokens': 18, 'completion_tokens': 23, 'total_tokens': 41}
# }
class Message:
def __init__(self):
self.content: str = "default"
self.role: str = "assistant"
class Message(OpenAIObject):
def __init__(self, content="default", role="assistant", **params):
super(Message, self).__init__(**params)
self.content = content
self.role = role
def __getitem__(self, key):
return getattr(self, key)
class Choices(OpenAIObject):
def __init__(self, finish_reason="stop", index=0, message=Message(), **params):
super(Choices, self).__init__(**params)
self.finish_reason = finish_reason
self.index = index
self.message = message
def __setitem__(self, key, value):
setattr(self, key, value)
def __iter__(self):
return iter(vars(self))
def __str__(self):
result = f"{{\n 'role': '{self.role}',\n 'content': \"{self.content}\"\n}}"
return result
class Choices:
def __init__(self):
self.finish_reason: str = "stop"
self.index: int = 0
self.message: Message = Message()
def __getitem__(self, key):
return getattr(self, key)
def __setitem__(self, key, value):
setattr(self, key, value)
def __iter__(self):
return iter(vars(self))
def __str__(self):
result = f"{{\n 'finish_reason': '{self.finish_reason}',\n 'index': {self.index},\n 'message': {self.message}\n}}"
return result
class ModelResponse(dict):
def __init__(self):
self.choices: List[Choices] = [Choices()]
self.created: str = None
self.model: str = None
self.usage: Dict[str, Union[int, None]] = {
class ModelResponse(OpenAIObject):
def __init__(self, choices=None, created=None, model=None, usage=None, **params):
super(ModelResponse, self).__init__(**params)
self.choices = choices if choices else [Choices()]
self.created = created
self.model = model
self.usage = usage if usage else {
"prompt_tokens": None,
"completion_tokens": None,
"total_tokens": None
}
def __getitem__(self, key):
return getattr(self, key)
def __setitem__(self, key, value):
setattr(self, key, value)
def __iter__(self):
return iter(vars(self))
def __str__(self):
choices_str = ",\n".join(str(choice) for choice in self.choices)
result = f"{{\n 'choices': [\n{choices_str}\n ],\n 'created': {self.created},\n 'model': '{self.model}',\n 'usage': {self.usage}\n}}"
return result
def to_dict_recursive(self):
d = super().to_dict_recursive()
d['choices'] = [choice.to_dict_recursive() for choice in self.choices]
return d
############################################################
def print_verbose(print_statement):
if litellm.set_verbose: