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

fix(main.py): fix translation to text_completions format for async text completion calls

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-30 09:02:51 -07:00
parent 89471ba4c5
commit c0204310ee
3 changed files with 41 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/main.py
View file

@ -2952,7 +2952,26 @@ async def atext_completion(*args, **kwargs):
model=model,
)
else:
return response
transformed_logprobs = None
# only supported for TGI models
try:
raw_response = response._hidden_params.get("original_response", None)
transformed_logprobs = litellm.utils.transform_logprobs(raw_response)
except Exception as e:
print_verbose(f"LiteLLM non blocking exception: {e}")
text_completion_response = TextCompletionResponse()
text_completion_response["id"] = response.get("id", None)
text_completion_response["object"] = "text_completion"
text_completion_response["created"] = response.get("created", None)
text_completion_response["model"] = response.get("model", None)
text_choices = TextChoices()
text_choices["text"] = response["choices"][0]["message"]["content"]
text_choices["index"] = response["choices"][0]["index"]
text_choices["logprobs"] = transformed_logprobs
text_choices["finish_reason"] = response["choices"][0]["finish_reason"]
text_completion_response["choices"] = [text_choices]
text_completion_response["usage"] = response.get("usage", None)
return text_completion_response
except Exception as e:
custom_llm_provider = custom_llm_provider or "openai"
raise exception_type(
@ -3165,6 +3184,7 @@ def text_completion(
transformed_logprobs = litellm.utils.transform_logprobs(raw_response)
except Exception as e:
print_verbose(f"LiteLLM non blocking exception: {e}")
text_completion_response["id"] = response.get("id", None)
text_completion_response["object"] = "text_completion"
text_completion_response["created"] = response.get("created", None)
@ -3176,6 +3196,7 @@ def text_completion(
text_choices["finish_reason"] = response["choices"][0]["finish_reason"]
text_completion_response["choices"] = [text_choices]
text_completion_response["usage"] = response.get("usage", None)
return text_completion_response