add tests for $schema vertex

This commit is contained in:
Ishaan Jaff 2024-11-01 10:37:42 +05:30
parent 134bd2cebb
commit c5f7182547
2 changed files with 90 additions and 10 deletions

View file

@ -774,6 +774,20 @@
"supports_vision": true,
"supports_prompt_caching": true
},
"azure/gpt-4o-mini-2024-07-18": {
"max_tokens": 16384,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"input_cost_per_token": 0.000000165,
"output_cost_per_token": 0.00000066,
"cache_read_input_token_cost": 0.000000075,
"litellm_provider": "azure",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_vision": true,
"supports_prompt_caching": true
},
"azure/gpt-4-turbo-2024-04-09": {
"max_tokens": 4096,
"max_input_tokens": 128000,

View file

@ -786,19 +786,85 @@ def test_unmapped_vertex_anthropic_model():
assert "max_retries" not in optional_params
@pytest.mark.parametrize(
"tools, key",
[
([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"),
([{"code_execution": {}}], "code_execution"),
],
)
def test_vertex_tool_params(tools, key):
def test_vertex_schema_field():
tools = [
{
"type": "function",
"function": {
"name": "json",
"description": "Respond with a JSON object.",
"parameters": {
"type": "object",
"properties": {
"thinking": {
"type": "string",
"description": "Your internal thoughts on different problem details given the guidance.",
},
"problems": {
"type": "array",
"items": {
"type": "object",
"properties": {
"icon": {
"type": "string",
"enum": [
"BarChart2",
"Bell",
],
"description": "The name of a Lucide icon to display",
},
"color": {
"type": "string",
"description": "A Tailwind color class for the icon, e.g., 'text-red-500'",
},
"problem": {
"type": "string",
"description": "The title of the problem being addressed, approximately 3-5 words.",
},
"description": {
"type": "string",
"description": "A brief explanation of the problem, approximately 20 words.",
},
"impacts": {
"type": "array",
"items": {"type": "string"},
"description": "A list of potential impacts or consequences of the problem, approximately 3 words each.",
},
"automations": {
"type": "array",
"items": {"type": "string"},
"description": "A list of potential automations to address the problem, approximately 3-5 words each.",
},
},
"required": [
"icon",
"color",
"problem",
"description",
"impacts",
"automations",
],
"additionalProperties": False,
},
"description": "Please generate problem cards that match this guidance.",
},
},
"required": ["thinking", "problems"],
"additionalProperties": False,
"$schema": "http://json-schema.org/draft-07/schema#",
},
},
}
]
optional_params = get_optional_params(
model="gemini-1.5-pro",
model="gemini-1.5-flash",
custom_llm_provider="vertex_ai",
tools=tools,
)
print(optional_params)
assert optional_params["tools"][0][key] == {}
print(optional_params["tools"][0]["function_declarations"][0])
assert (
"$schema"
not in optional_params["tools"][0]["function_declarations"][0]["parameters"]
)