返回提交历史
Modified
etc/unittest/__main__.py
+1
-0
Added
etc/unittest/test_gemini.py
+201
-0
Modified
g4f/Provider/needs_auth/Gemini.py
+699
-251
Added
g4f/Provider/needs_auth/gemini_utils.py
+262
-0
Modified
g4f/models.py
+25
-1
XFEstudio/gpt4free
Refactor Gemini provider: extract utils and add tests
Extract Gemini helper functions into a new gemini_utils module for better code organization and maintainability. Add comprehensive unit tests covering model parsing, account discovery, error handling, and streaming. Key changes: - New gemini_utils.py module with utility functions for parsing, validation, and header building - New test_gemini.py with 171 lines of unit tests - Improved model validation and account model discovery with caching - Better error handling and retry logic - More structured request building with explicit field indexing - Support for authenticated and unauthenticated account states - Enhanced reasoning extraction from responses
479f7272
代码差异
5 个文件
+1188
-252
@@ -18,5 +18,6 @@ from .models import *
18
18
from .mcp import *
19
19
from .tool_support_provider import *
20
20
from .config_provider import *
21
from .test_gemini import *
21
22
22
23
unittest.main()
@@ -0,0 +1,201 @@
1
from __future__ import annotations
2
3
import asyncio
4
import json
5
import unittest
6
7
from g4f.Provider.needs_auth.Gemini import (
8
ACCOUNT_STATUS_AVAILABLE,
9
ACCOUNT_STATUS_UNAUTHENTICATED,
10
MODEL_HEADER_KEY,
11
Gemini,
12
_build_model_headers,
13
_extract_gemini_error_code,
14
_extract_reasoning,
15
_iter_response_lines,
16
_parse_account_models,
17
_parse_google_frames,
18
_resolve_model,
19
)
20
from g4f.errors import MissingAuthError, ResponseError
21
from g4f.models import ModelRegistry
22
23
24
def build_account_response(status: int) -> tuple[str, dict[str, dict]]:
25
body = [None] * 18
26
body[14] = status
27
body[15] = [
28
["fbb127bbb056c959", "Flash", "All-around help"],
29
["5bf011840784117a", "Thinking", "Solves complex problems"],
30
["9d8ca3786ebdfbea", "Pro", "Advanced math & code"],
31
]
32
body[16] = []
33
body[17] = []
34
response = json.dumps([["wrb.fr", "otAQ7b", json.dumps(body)]])
35
parsed_status, registry = _parse_account_models(response)
36
assert parsed_status == status
37
return response, registry
38
39
40
class GeminiHelpersTest(unittest.TestCase):
41
def test_parse_account_models_and_availability(self):
42
_, registry = build_account_response(ACCOUNT_STATUS_UNAUTHENTICATED)
43
44
self.assertTrue(registry["fbb127bbb056c959"]["available"])
45
self.assertFalse(registry["5bf011840784117a"]["available"])
46
self.assertFalse(registry["9d8ca3786ebdfbea"]["available"])
47
48
def test_missing_account_status_means_available(self):
49
body = [None] * 18
50
body[15] = [["e6fa609c3fa255c0", "Pro", "Advanced model"]]
51
response = json.dumps([["wrb.fr", "otAQ7b", json.dumps(body)]])
52
53
status, registry = _parse_account_models(response)
54
55
self.assertEqual(status, ACCOUNT_STATUS_AVAILABLE)
56
self.assertTrue(registry["e6fa609c3fa255c0"]["available"])
57
58
def test_parse_utf16_length_prefixed_frame(self):
59
raw = json.dumps([["wrb.fr", "rpc", "emoji: 😀"]], ensure_ascii=False)
60
payload = f"\n{raw}\n"
61
utf16_length = len(payload.encode("utf-16-le")) // 2
62
63
frames, remaining = _parse_google_frames(f"{utf16_length}{payload}")
64
65
self.assertEqual(frames[0][0], "wrb.fr")
66
self.assertEqual(remaining, "")
67
68
def test_model_header_capacity_fields(self):
69
field_12 = json.loads(_build_model_headers("model", 4, 12)[MODEL_HEADER_KEY])
70
field_13 = json.loads(_build_model_headers("model", 2, 13)[MODEL_HEADER_KEY])
71
72
self.assertEqual(field_12[11], 4)
73
self.assertIsNone(field_13[11])
74
self.assertEqual(field_13[12], 2)
75
76
def test_mode_categories_and_default_thinking_depths(self):
77
expected = {
78
"gemini-3.5-flash": (1, 4),
79
"gemini-3.5-flash-thinking": (2, 0),
80
"gemini-3.1-pro": (3, 4),
81
"gemini-auto": (4, 4),
82
"gemini-3.5-flash-thinking-lite": (5, 0),
83
"gemini-flash-lite": (6, 4),
84
}
85
86
for requested, (mode, default_think) in expected.items():
87
with self.subTest(model=requested):
88
model, think = _resolve_model(requested)
89
request = Gemini.build_request(
90
"test", "en", model, think, request_uuid="test-request"
91
)
92
self.assertEqual(request[79], mode)
93
self.assertEqual(request[17], [[default_think]])
94
95
def test_explicit_thinking_depths_are_preserved(self):
96
for requested_depth in range(5):
97
with self.subTest(depth=requested_depth):
98
model, think = _resolve_model(
99
f"gemini-3.5-flash-thinking@think={requested_depth}"
100
)
101
request = Gemini.build_request(
102
"test", "en", model, think, request_uuid="test-request"
103
)
104
self.assertEqual(request[17], [[requested_depth]])
105
106
def test_legacy_model_names_resolve_to_current_modes(self):
107
expected = {
108
"gemini-2.0": "gemini-3.5-flash",
109
"gemini-2.0-flash": "gemini-3.5-flash",
110
"gemini-2.0-flash-thinking": "gemini-3.5-flash-thinking",
111
"gemini-2.0-flash-thinking-with-apps": "gemini-3.5-flash-thinking",
112
"gemini-2.5-flash": "gemini-3.5-flash",
113
"gemini-2.5-pro": "gemini-3.1-pro",
114
"gemini-3.1-flash-lite": "gemini-flash-lite",
115
}
116
117
for legacy_name, current_name in expected.items():
118
with self.subTest(model=legacy_name):
119
resolved, _ = _resolve_model(legacy_name)
120
self.assertEqual(resolved, current_name)
121
122
def test_public_model_registry_exposes_current_models(self):
123
self.assertEqual(ModelRegistry.get("gemini").name, "gemini-3.5-flash")
124
for model in (
125
"gemini-3.5-flash-thinking",
126
"gemini-auto",
127
"gemini-3.5-flash-thinking-lite",
128
"gemini-flash-lite",
129
):
130
with self.subTest(model=model):
131
self.assertEqual(ModelRegistry.get(model).name, model)
132
133
def test_extract_reasoning_from_dedicated_field(self):
134
candidate = [None] * 38
135
candidate[37] = [["private reasoning"]]
136
response = [None] * 5
137
response[4] = [candidate]
138
139
self.assertEqual(_extract_reasoning(response), "private reasoning")
140
141
def test_extract_structured_error_code(self):
142
frame = ["wrb.fr", None, None, None, None, [None, None, [[None, [1037]]]]]
143
144
self.assertEqual(_extract_gemini_error_code(frame), 1037)
145
146
def test_reject_silent_model_fallback(self):
147
_, registry = build_account_response(ACCOUNT_STATUS_UNAUTHENTICATED)
148
149
class ProbeGemini(Gemini):
150
_account_status = ACCOUNT_STATUS_UNAUTHENTICATED
151
_account_models = registry
152
153
with self.assertRaises(MissingAuthError):
154
ProbeGemini.validate_model_access("gemini-3.1-pro")
155
ProbeGemini.validate_model_access("gemini-3.5-flash")
156
ProbeGemini.validate_model_access(
157
"gemini-3.1-pro", allow_model_fallback=True
158
)
159
160
def test_dynamic_headers_only_for_available_pro(self):
161
_, registry = build_account_response(ACCOUNT_STATUS_AVAILABLE)
162
163
class ProbeGemini(Gemini):
164
_account_status = ACCOUNT_STATUS_AVAILABLE
165
_account_models = registry
166
167
pro_header = json.loads(
168
ProbeGemini.get_model_headers("gemini-3.1-pro")[MODEL_HEADER_KEY]
169
)
170
self.assertEqual(pro_header[4], "9d8ca3786ebdfbea")
171
self.assertEqual(ProbeGemini.get_model_headers("gemini-3.5-flash"), {})
172
self.assertEqual(
173
ProbeGemini.get_model_headers("gemini-3.5-flash-thinking"), {}
174
)
175
176
177
class GeminiStreamTest(unittest.IsolatedAsyncioTestCase):
178
async def test_stream_idle_timeout(self):
179
class SlowContent:
180
async def iter_any(self):
181
await asyncio.sleep(0.05)
182
yield b"late\n"
183
184
with self.assertRaises(ResponseError):
185
async for _ in _iter_response_lines(SlowContent(), idle_timeout=0.01):
186
pass
187
188
async def test_stream_reassembles_split_unicode_line(self):
189
data = '[["wrb.fr","rpc","😀"]]\n'.encode()
190
191
class SplitContent:
192
async def iter_any(self):
193
for chunk in (data[:20], data[20:23], data[23:]):
194
yield chunk
195
196
lines = [line async for line in _iter_response_lines(SplitContent(), 1)]
197
self.assertEqual(json.loads(lines[0])[0][2], "😀")
198
199
200
if __name__ == "__main__":
201
unittest.main()
@@ -0,0 +1,262 @@
1
from __future__ import annotations
2
3
import json
4
import re
5
6
from ...errors import RateLimitError, ResponseError
7
8
9
BARD_ERROR_PATTERN = re.compile(r"BardErrorInfo\s*\[(\d+)\]")
10
LENGTH_MARKER_PATTERN = re.compile(r"(\d+)\n")
11
ACCOUNT_STATUS_AVAILABLE = 1000
12
ACCOUNT_STATUS_UNAUTHENTICATED = 1016
13
MODEL_HEADER_KEY = "x-goog-ext-525001261-jspb"
14
MODEL_HEADER_AUXILIARY = {
15
"x-goog-ext-73010989-jspb": "[0]",
16
"x-goog-ext-73010990-jspb": "[0]",
17
}
18
MODEL_FAMILIES = {
19
"gemini-3.5-flash": "flash",
20
"gemini-3.5-flash-thinking": "thinking",
21
"gemini-3.1-pro": "pro",
22
}
23
ANONYMOUS_MODELS = {"gemini-3.5-flash", "gemini-auto"}
24
KNOWN_MODEL_IDS = {
25
"fbb127bbb056c959": "flash",
26
"5bf011840784117a": "thinking",
27
"9d8ca3786ebdfbea": "pro",
28
}
29
GEMINI_ERROR_MESSAGES = {
30
1013: "Gemini encountered a temporary generation error",
31
1037: "Gemini usage limit exceeded for the requested model",
32
1050: "The requested Gemini model is inconsistent with the conversation",
33
1052: "The requested Gemini model header is invalid or unavailable",
34
1060: "Gemini temporarily blocked this IP address",
35
}
36
37
38
def iter_wrb_payloads(value):
39
if isinstance(value, list):
40
if len(value) >= 3 and value[0] == "wrb.fr" and isinstance(value[2], str):
41
yield value[2]
42
return
43
for item in value:
44
yield from iter_wrb_payloads(item)
45
46
47
def get_nested_value(value, path, default=None):
48
current = value
49
for key in path:
50
if isinstance(key, int):
51
if not isinstance(current, list) or not -len(current) <= key < len(current):
52
return default
53
elif not isinstance(current, dict) or key not in current:
54
return default
55
current = current[key]
56
return default if current is None else current
57
58
59
def _utf16_char_count(value: str, start: int, units: int) -> tuple[int, int]:
60
count = found = 0
61
while found < units and start + count < len(value):
62
size = 2 if ord(value[start + count]) > 0xFFFF else 1
63
if found + size > units:
64
break
65
found += size
66
count += 1
67
return count, found
68
69
70
def parse_google_frames(content: str) -> tuple[list, str]:
71
"""Parse Google's length-prefixed frames, including UTF-16 lengths."""
72
frames = []
73
position = 0
74
while position < len(content):
75
while position < len(content) and content[position].isspace():
76
position += 1
77
if position >= len(content):
78
break
79
match = LENGTH_MARKER_PATTERN.match(content, position)
80
if match is None:
81
break
82
length = int(match.group(1))
83
start = match.start() + len(match.group(1))
84
char_count, units_found = _utf16_char_count(content, start, length)
85
if units_found < length:
86
break
87
end = start + char_count
88
chunk = content[start:end].strip()
89
position = end
90
if not chunk:
91
continue
92
try:
93
parsed = json.loads(chunk)
94
except ValueError:
95
continue
96
if isinstance(parsed, list):
97
frames.extend(parsed)
98
else:
99
frames.append(parsed)
100
return frames, content[position:]
101
102
103
def _iter_google_json(content: str):
104
content = content.lstrip()
105
if content.startswith(")]}'"):
106
content = content[4:].lstrip()
107
frames, _ = parse_google_frames(content)
108
if frames:
109
yield from frames
110
return
111
try:
112
parsed = json.loads(content)
113
except ValueError:
114
parsed = None
115
if parsed is not None:
116
if isinstance(parsed, list):
117
yield from parsed
118
else:
119
yield parsed
120
return
121
for line in content.splitlines():
122
line = line.strip()
123
if not line or line.isdigit():
124
continue
125
try:
126
parsed = json.loads(line)
127
except ValueError:
128
continue
129
if isinstance(parsed, list):
130
yield from parsed
131
else:
132
yield parsed
133
134
135
def _compute_model_capacity(tier_flags: list, capability_flags: list) -> tuple[int, int]:
136
if 21 in tier_flags:
137
return 1, 13
138
if 22 in tier_flags:
139
return 2, 13
140
if 115 in capability_flags:
141
return 4, 12
142
if 16 in tier_flags or 106 in capability_flags:
143
return 3, 12
144
if 8 in tier_flags or (106 not in capability_flags and 19 in capability_flags):
145
return 2, 12
146
return 1, 12
147
148
149
def _model_family(model_id: str, display_name: str, description: str) -> str | None:
150
if model_id in KNOWN_MODEL_IDS:
151
return KNOWN_MODEL_IDS[model_id]
152
label = f"{display_name} {description}".lower()
153
if "thinking" in label:
154
return "thinking"
155
if "pro" in label:
156
return "pro"
157
if "flash" in label:
158
return "flash"
159
return None
160
161
162
def build_model_headers(model_id: str, capacity: int, capacity_field: int) -> dict[str, str]:
163
header = [1, None, None, None, model_id, None, None, 0, [4], None, None]
164
if capacity_field == 13:
165
header.extend([None, capacity])
166
else:
167
header.append(capacity)
168
return {
169
MODEL_HEADER_KEY: json.dumps(header, separators=(",", ":")),
170
**MODEL_HEADER_AUXILIARY,
171
}
172
173
174
def parse_account_models(content: str) -> tuple[int | None, dict[str, dict]]:
175
status_code = None
176
registry = {}
177
for frame in _iter_google_json(content):
178
for payload in iter_wrb_payloads(frame):
179
try:
180
body = json.loads(payload)
181
except (TypeError, ValueError):
182
continue
183
current_status = get_nested_value(body, [14])
184
if isinstance(current_status, int):
185
status_code = current_status
186
models_list = get_nested_value(body, [15], [])
187
if not isinstance(models_list, list):
188
continue
189
tier_flags = get_nested_value(body, [16], [])
190
capability_flags = get_nested_value(body, [17], [])
191
tier_flags = tier_flags if isinstance(tier_flags, list) else []
192
capability_flags = capability_flags if isinstance(capability_flags, list) else []
193
capacity, capacity_field = _compute_model_capacity(
194
tier_flags, capability_flags
195
)
196
for model_data in models_list:
197
if not isinstance(model_data, list):
198
continue
199
model_id = get_nested_value(model_data, [0], "")
200
display_name = get_nested_value(model_data, [1], "")
201
description = get_nested_value(model_data, [2], "")
202
if not isinstance(model_id, str) or not model_id:
203
continue
204
family = _model_family(model_id, display_name, description)
205
# The authenticated response commonly omits field 14. Gemini's
206
# web client treats an absent status as the normal/available
207
# state; explicit non-1000 values describe restricted states.
208
available = status_code in (None, ACCOUNT_STATUS_AVAILABLE)
209
if status_code == ACCOUNT_STATUS_UNAUTHENTICATED:
210
available = family == "flash"
211
registry[model_id] = {
212
"model_id": model_id,
213
"family": family,
214
"display_name": display_name,
215
"description": description,
216
"capacity": capacity,
217
"capacity_field": capacity_field,
218
"available": available,
219
"headers": build_model_headers(
220
model_id, capacity, capacity_field
221
),
222
}
223
if registry and status_code is None:
224
status_code = ACCOUNT_STATUS_AVAILABLE
225
return status_code, registry
226
227
228
def extract_reasoning(response_part: list) -> str | None:
229
candidates = get_nested_value(response_part, [4], [])
230
if not isinstance(candidates, list):
231
return None
232
for candidate in candidates:
233
reasoning = get_nested_value(candidate, [37, 0, 0])
234
if isinstance(reasoning, str) and reasoning:
235
return reasoning
236
return None
237
238
239
def extract_gemini_error_code(value) -> int | None:
240
if isinstance(value, str):
241
match = BARD_ERROR_PATTERN.search(value)
242
return int(match.group(1)) if match else None
243
if not isinstance(value, list):
244
return None
245
if value and value[0] == "wrb.fr":
246
code = get_nested_value(value, [5, 2, 0, 1, 0])
247
if isinstance(code, int):
248
return code
249
for item in value:
250
code = extract_gemini_error_code(item)
251
if code is not None:
252
return code
253
return None
254
255
256
def raise_gemini_error(code: int, model: str) -> None:
257
message = GEMINI_ERROR_MESSAGES.get(
258
code, f"Gemini rejected the request with error code {code}"
259
)
260
if code in (1037, 1060):
261
raise RateLimitError(f"{message}: {model}")
262
raise ResponseError(f"{message}: {model}")
@@ -475,6 +475,30 @@ gemini_3_5_flash = Model(
475
475
best_provider = "Gemini"
476
476
)
477
477
478
gemini_3_5_flash_thinking = Model(
479
name = 'gemini-3.5-flash-thinking',
480
base_provider = 'Google',
481
best_provider = "Gemini"
482
)
483
484
gemini_auto = Model(
485
name = 'gemini-auto',
486
base_provider = 'Google',
487
best_provider = "Gemini"
488
)
489
490
gemini_3_5_flash_thinking_lite = Model(
491
name = 'gemini-3.5-flash-thinking-lite',
492
base_provider = 'Google',
493
best_provider = "Gemini"
494
)
495
496
gemini_flash_lite = Model(
497
name = 'gemini-flash-lite',
498
base_provider = 'Google',
499
best_provider = "Gemini"
500
)
501
478
502
479
503
480
504
@@ -817,7 +841,7 @@ class ModelUtils:
817
841
ModelRegistry._aliases[alias] = model_name
818
842
819
843
# Register special aliases after all models are created
820
ModelRegistry._aliases["gemini"] = "gemini-2.0"
844
ModelRegistry._aliases["gemini"] = "gemini-3.5-flash"
821
845
822
846
# Fill the convert dictionary
823
847
ModelUtils.convert = ModelRegistry.all_models()