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

fix: move to using pydantic obj for setting values

Browse source
This commit is contained in:
Krrish Dholakia 2024-07-11 13:18:36 -07:00
parent dd1048cb35
commit 6e9f048618
30 changed files with 1018 additions and 886 deletions
Download patch file Download diff file Expand all files Collapse all files

50
litellm/llms/vertex_ai.py
View file

@ -852,16 +852,14 @@ def completion(
## RESPONSE OBJECT
if isinstance(completion_response, litellm.Message):
model_response["choices"][0]["message"] = completion_response
model_response.choices[0].message = completion_response # type: ignore
elif len(str(completion_response)) > 0:
model_response["choices"][0]["message"]["content"] = str(
completion_response
)
model_response["created"] = int(time.time())
model_response["model"] = model
model_response.choices[0].message.content = str(completion_response) # type: ignore
model_response.created = int(time.time())
model_response.model = model
## CALCULATING USAGE
if model in litellm.vertex_language_models and response_obj is not None:
model_response["choices"][0].finish_reason = map_finish_reason(
model_response.choices[0].finish_reason = map_finish_reason(
response_obj.candidates[0].finish_reason.name
)
usage = Usage(
@ -912,7 +910,7 @@ async def async_completion(
request_str: str,
print_verbose: Callable,
logging_obj,
encoding=None,
encoding,
client_options=None,
instances=None,
vertex_project=None,
@ -1088,16 +1086,16 @@ async def async_completion(
## RESPONSE OBJECT
if isinstance(completion_response, litellm.Message):
model_response["choices"][0]["message"] = completion_response
model_response.choices[0].message = completion_response # type: ignore
elif len(str(completion_response)) > 0:
model_response["choices"][0]["message"]["content"] = str(
model_response.choices[0].message.content = str( # type: ignore
completion_response
)
model_response["created"] = int(time.time())
model_response["model"] = model
model_response.created = int(time.time())
model_response.model = model
## CALCULATING USAGE
if model in litellm.vertex_language_models and response_obj is not None:
model_response["choices"][0].finish_reason = map_finish_reason(
model_response.choices[0].finish_reason = map_finish_reason(
response_obj.candidates[0].finish_reason.name
)
usage = Usage(
@ -1377,16 +1375,16 @@ class VertexAITextEmbeddingConfig(BaseModel):
def embedding(
model: str,
input: Union[list, str],
print_verbose,
model_response: litellm.EmbeddingResponse,
optional_params: dict,
api_key: Optional[str] = None,
logging_obj=None,
model_response=None,
optional_params=None,
encoding=None,
vertex_project=None,
vertex_location=None,
vertex_credentials=None,
aembedding=False,
print_verbose=None,
):
# logic for parsing in - calling - parsing out model embedding calls
try:
@ -1484,15 +1482,15 @@ def embedding(
"embedding": embedding.values,
}
)
input_tokens += embedding.statistics.token_count
model_response["object"] = "list"
model_response["data"] = embedding_response
model_response["model"] = model
input_tokens += embedding.statistics.token_count # type: ignore
model_response.object = "list"
model_response.data = embedding_response
model_response.model = model
usage = Usage(
prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
)
model_response.usage = usage
setattr(model_response, "usage", usage)
return model_response
@ -1500,8 +1498,8 @@ def embedding(
async def async_embedding(
model: str,
input: Union[list, str],
model_response: litellm.EmbeddingResponse,
logging_obj=None,
model_response=None,
optional_params=None,
encoding=None,
client=None,
@ -1541,11 +1539,11 @@ async def async_embedding(
)
input_tokens += embedding.statistics.token_count
model_response["object"] = "list"
model_response["data"] = embedding_response
model_response["model"] = model
model_response.object = "list"
model_response.data = embedding_response
model_response.model = model
usage = Usage(
prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
)
model_response.usage = usage
setattr(model_response, "usage", usage)
return model_response