File size: 4,416 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | from aistudio_api.infrastructure.gateway.model_defaults import resolve_model_defaults
from aistudio_api.infrastructure.gateway.wire_types import AistudioImageOutputMode
def test_resolve_model_defaults_uses_repo_config_for_gemma():
defaults = resolve_model_defaults("models/gemma-4-31b-it")
assert defaults.default_tools == ("google_search",)
assert defaults.is_image_model is False
assert defaults.safety_settings == (
(None, None, 7, 5),
(None, None, 8, 5),
(None, None, 9, 5),
(None, None, 10, 5),
)
def test_resolve_model_defaults_merges_yaml_profiles(tmp_path):
config_path = tmp_path / "config.yaml"
config_path.write_text(
"""
model_defaults:
profiles:
- name: gemini_models
match:
prefixes: [gemini-]
generation_config_defaults:
response_mime_type: application/json
- name: image_models
match:
contains: [image]
is_image_model: true
generation_config_defaults:
image_output_mode: text_and_image
thinking_config:
level: HIGH
mode: 9
media_resolution: HIGH
clear_generation_config_indexes: [7, 17]
disable_safety_settings: true
models:
gemini-3.1-flash-image-preview:
generation_config_defaults:
response_mime_type: null
"""
)
defaults = resolve_model_defaults("models/gemini-3.1-flash-image-preview", config_path=config_path)
assert defaults.is_image_model is True
assert defaults.disable_safety_settings is True
assert defaults.clear_generation_config_indexes == (7, 17)
assert defaults.generation_config_defaults["response_mime_type"] is None
assert defaults.generation_config_defaults["image_output_mode"] == AistudioImageOutputMode.text_and_image()
assert defaults.generation_config_defaults["thinking_config"] == [9, None, None, 3]
assert defaults.generation_config_defaults["media_resolution"] == 3
def test_resolve_model_defaults_coerces_safety_settings(tmp_path):
config_path = tmp_path / "config.yaml"
config_path.write_text(
"""
model_defaults:
profiles:
- name: gemini_models
match:
prefixes: [gemini-]
safety_settings:
Harassment: 1
Hate: 2
Sexually Explicit: 3
Dangerous Content: 5
"""
)
defaults = resolve_model_defaults("models/gemini-2.5-flash", config_path=config_path)
assert defaults.safety_settings == (
(None, None, 7, 1),
(None, None, 8, 2),
(None, None, 9, 3),
(None, None, 10, 5),
)
def test_resolve_model_defaults_supports_image_default_tool_names(tmp_path):
config_path = tmp_path / "config.yaml"
config_path.write_text(
"""
model_defaults:
profiles:
- name: image_models
match:
contains: [image]
is_image_model: true
default_tools:
- google_search_and_image_search
"""
)
defaults = resolve_model_defaults("models/gemini-3.1-flash-image-preview", config_path=config_path)
assert defaults.is_image_model is True
assert defaults.default_tools == ("google_search_and_image_search",)
def test_resolve_model_defaults_keeps_gemini_four_tool_config(tmp_path):
config_path = tmp_path / "config.yaml"
config_path.write_text(
"""
model_defaults:
profiles:
- name: gemini_models
match:
prefixes: [gemini-]
default_tools:
- google_search
- code_execution
- google_maps
- url_context
"""
)
defaults = resolve_model_defaults("models/gemini-3.5-flash", config_path=config_path)
assert defaults.default_tools == ("google_search", "code_execution", "google_maps", "url_context")
def test_resolve_model_defaults_model_override_replaces_profile_tools_and_flags(tmp_path):
config_path = tmp_path / "config.yaml"
config_path.write_text(
"""
model_defaults:
profiles:
- name: gemini_models
match:
prefixes: [gemini-]
default_tools:
- google_search
disable_safety_settings: true
models:
gemini-3.5-flash:
default_tools:
- code_execution
disable_safety_settings: false
"""
)
defaults = resolve_model_defaults("models/gemini-3.5-flash", config_path=config_path)
assert defaults.default_tools == ("code_execution",)
assert defaults.disable_safety_settings is False
|