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

Merge pull request #2577 from BerriAI/litellm_vertex_ai_streaming_func_call

Browse source 
feat(vertex_ai.py): support gemini (vertex ai) function calling when streaming
This commit is contained in:
Krish Dholakia 2024-03-18 20:10:00 -07:00 • committed by GitHub
commit c669943292
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 74 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

38
litellm/utils.py
View file

@ -9224,7 +9224,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
@ -9235,7 +9269,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)}"