mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(vertex_ai.py): fix list tool call responses
Closes https://github.com/BerriAI/litellm/issues/3147
This commit is contained in:
Krrish Dholakia
parent
363cdb1a0c
commit
94f3d361b0
3 changed files with 55 additions and 51 deletions
|
|
@ -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))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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: <your 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
|
||||
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"
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue