From 94f3d361b0359439c033b07ab7a00d6c41da5c77 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 10 May 2024 20:05:38 -0700 Subject: [PATCH] fix(vertex_ai.py): fix list tool call responses Closes https://github.com/BerriAI/litellm/issues/3147 --- litellm/llms/vertex_ai.py | 22 ++++++- litellm/proxy/_super_secret_config.yaml | 24 ++++---- .../tests/test_amazing_vertex_completion.py | 60 ++++++++----------- 3 files changed, 55 insertions(+), 51 deletions(-) diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index cab7ae19f25..cc4db1f9d30 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -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)) diff --git a/litellm/proxy/_super_secret_config.yaml b/litellm/proxy/_super_secret_config.yaml index 4ea98461146..752cd281da6 100644 --- a/litellm/proxy/_super_secret_config.yaml +++ b/litellm/proxy/_super_secret_config.yaml @@ -20,22 +20,20 @@ model_list: - litellm_params: model: together_ai/codellama/CodeLlama-13b-Instruct-hf model_name: CodeLlama-13b-Instruct -router_settings: - num_retries: 0 - enable_pre_call_checks: true - redis_host: os.environ/REDIS_HOST - redis_password: os.environ/REDIS_PASSWORD - redis_port: os.environ/REDIS_PORT router_settings: - routing_strategy: "latency-based-routing" + redis_host: redis + # redis_password: + redis_port: 6379 litellm_settings: - success_callback: ["langfuse"] + set_verbose: True + # service_callback: ["prometheus_system"] + # success_callback: ["prometheus"] + # failure_callback: ["prometheus"] general_settings: - alerting: ["slack"] - alert_types: ["llm_exceptions", "daily_reports"] - alerting_args: - daily_report_frequency: 60 # every minute - report_check_interval: 5 # every 5s \ No newline at end of file + enable_jwt_auth: True + disable_reset_budget: True + proxy_batch_write_at: 60 # 👈 Frequency of batch writing logs to server (in seconds) + routing_strategy: simple-shuffle # Literal["simple-shuffle", "least-busy", "usage-based-routing","latency-based-routing"], default="simple-shuffle" \ No newline at end of file diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 91fd444742b..a56d7fe5a9a 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -590,47 +590,37 @@ def test_gemini_pro_vision_base64(): pytest.fail(f"An exception occurred - {str(e)}") +@pytest.mark.asyncio def test_gemini_pro_function_calling(): try: load_vertex_ai_credentials() - tools = [ - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - }, - "unit": { - "type": "string", - "enum": ["celsius", "fahrenheit"], + response = litellm.completion( + model="vertex_ai/gemini-pro", + messages=[ + { + "role": "user", + "content": "Call the submit_cities function with San Francisco and New York", + } + ], + tools=[ + { + "type": "function", + "function": { + "name": "submit_cities", + "description": "Submits a list of cities", + "parameters": { + "type": "object", + "properties": { + "cities": {"type": "array", "items": {"type": "string"}} }, + "required": ["cities"], }, - "required": ["location"], }, - }, - } - ] - messages = [ - { - "role": "user", - "content": "What's the weather like in Boston today in fahrenheit?", - } - ] - completion = litellm.completion( - model="gemini-pro", messages=messages, tools=tools, tool_choice="auto" + } + ], ) - print(f"completion: {completion}") - # assert completion.choices[0].message.content is None ## GEMINI PRO is very chatty. - if hasattr(completion.choices[0].message, "tool_calls") and isinstance( - completion.choices[0].message.tool_calls, list - ): - assert len(completion.choices[0].message.tool_calls) == 1 + + print(f"response: {response}") except litellm.APIError as e: pass except litellm.RateLimitError as e: @@ -639,7 +629,7 @@ def test_gemini_pro_function_calling(): if "429 Quota exceeded" in str(e): pass else: - return + pytest.fail("An unexpected exception occurred - {}".format(str(e))) # gemini_pro_function_calling()