phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

refactor: fixing linting issues

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-11 18:52:28 -08:00
parent ae35c13015
commit 45b6f8b853
25 changed files with 223 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

20
litellm/llms/vllm.py
View file

@ -4,7 +4,7 @@ from enum import Enum
import requests
import time
from typing import Callable, Any
from litellm.utils import ModelResponse
from litellm.utils import ModelResponse, Usage
from .prompt_templates.factory import prompt_factory, custom_prompt
llm = None
class VLLMError(Exception):
@ -90,9 +90,12 @@ def completion(
model_response["created"] = time.time()
model_response["model"] = model
model_response.usage.completion_tokens = completion_tokens
model_response.usage.prompt_tokens = prompt_tokens
model_response.usage.total_tokens = prompt_tokens + completion_tokens
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens
)
model_response.usage = usage
return model_response
def batch_completions(
@ -170,9 +173,12 @@ def batch_completions(
model_response["created"] = time.time()
model_response["model"] = model
model_response.usage.completion_tokens = completion_tokens
model_response.usage.prompt_tokens = prompt_tokens
model_response.usage.total_tokens = prompt_tokens + completion_tokens
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens
)
model_response.usage = usage
final_outputs.append(model_response)
return final_outputs