fix(vertex_ai.py): fix list tool call responses

Closes https://github.com/BerriAI/litellm/issues/3147
This commit is contained in:
Krrish Dholakia 2024-05-10 20:05:38 -07:00
parent 363cdb1a0c
commit 94f3d361b0
3 changed files with 55 additions and 51 deletions

View file

@ -417,8 +417,10 @@ def completion(
from google.cloud import aiplatform # type: ignore
from google.protobuf import json_format # type: ignore
from google.protobuf.struct_pb2 import Value # type: ignore
from google.protobuf.json_format import MessageToDict
from google.cloud.aiplatform_v1beta1.types import content as gapic_content_types # type: ignore
import google.auth # type: ignore
import proto
## Load credentials with the correct quota project ref: https://github.com/googleapis/python-aiplatform/issues/2557#issuecomment-1709284744
print_verbose(
@ -605,9 +607,21 @@ def completion(
):
function_call = response.candidates[0].content.parts[0].function_call
args_dict = {}
for k, v in function_call.args.items():
args_dict[k] = v
args_str = json.dumps(args_dict)
# Check if it's a RepeatedComposite instance
for key, val in function_call.args.items():
if isinstance(
val, proto.marshal.collections.repeated.RepeatedComposite
):
# If so, convert to list
args_dict[key] = [v for v in val]
else:
args_dict[key] = val
try:
args_str = json.dumps(args_dict)
except Exception as e:
raise VertexAIError(status_code=422, message=str(e))
message = litellm.Message(
content=None,
tool_calls=[
@ -810,6 +824,8 @@ def completion(
setattr(model_response, "usage", usage)
return model_response
except Exception as e:
if isinstance(e, VertexAIError):
raise e
raise VertexAIError(status_code=500, message=str(e))