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

feat(vertex_ai.py): support gemini (vertex ai) function calling when streaming

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-18 11:47:27 -07:00
parent 8619499853
commit f49dc9a99f
2 changed files with 76 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

38
litellm/utils.py
View file

@ -9193,7 +9193,41 @@ class CustomStreamWrapper:
try:
if hasattr(chunk, "candidates") == True:
try:
completion_obj["content"] = chunk.text
try:
completion_obj["content"] = chunk.text
except Exception as e:
if "Part has no text." in str(e):
## check for function calling
function_call = (
chunk.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)
_delta_obj = litellm.utils.Delta(
content=None,
tool_calls=[
{
"id": f"call_{str(uuid.uuid4())}",
"function": {
"arguments": args_str,
"name": function_call.name,
},
"type": "function",
}
],
)
_streaming_response = StreamingChoices(
delta=_delta_obj
)
_model_response = ModelResponse(stream=True)
_model_response.choices = [_streaming_response]
response_obj = {"original_chunk": _model_response}
else:
raise e
if (
hasattr(chunk.candidates[0], "finish_reason")
and chunk.candidates[0].finish_reason.name
@ -9204,7 +9238,7 @@ class CustomStreamWrapper:
chunk.candidates[0].finish_reason.name
)
)
except:
except Exception as e:
if chunk.candidates[0].finish_reason.name == "SAFETY":
raise Exception(
f"The response was blocked by VertexAI. {str(chunk)}"