Skip to content

Commit d02701c

Browse files
guillaqguillaq
authored
Merge pull request #101 from WorkflowAI/guillaume/add-use-fallback
Add use fallback
2 parents 693339a + e97cd03 commit d02701c

File tree

7 files changed

+58
-12
lines changed

7 files changed

+58
-12
lines changed

‎tests/integration/conftest.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,8 @@ def check_request(
121121
assert request.headers["x-workflowai-source"] == "sdk"
122122
assert request.headers["x-workflowai-language"] == "python"
123123

124+
return body
125+
124126

125127
@pytest.fixture
126128
def test_client(httpx_mock: HTTPXMock) -> IntTestClient:

‎tests/integration/run_test.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,22 @@ async def city_to_capital(task_input: CityToCapitalTaskInput) -> Run[CityToCapit
3838
test_client.check_request()
3939

4040

41+
async def test_run_task_run_with_fallback(test_client: IntTestClient) -> None:
42+
@workflowai.agent(schema_id=1)
43+
async def city_to_capital(task_input: CityToCapitalTaskInput) -> Run[CityToCapitalTaskOutput]: ...
44+
45+
test_client.mock_response()
46+
47+
task_input = CityToCapitalTaskInput(city="Hello")
48+
with_run = await city_to_capital(task_input, use_fallback="never")
49+
50+
assert with_run.id == "123"
51+
assert with_run.output.capital == "Tokyo"
52+
53+
body = test_client.check_request()
54+
assert body["use_fallback"] == "never"
55+
56+
4157
async def test_run_task_run_version(test_client: IntTestClient) -> None:
4258
@workflowai.agent(schema_id=1, version="staging")
4359
async def city_to_capital(task_input: CityToCapitalTaskInput) -> Run[CityToCapitalTaskOutput]: ...

‎workflowai/core/_common_types.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@
22
Annotated,
33
Any,
44
Generic,
5+
Literal,
56
Optional,
67
Protocol,
78
TypeVar,
9+
Union,
810
)
911

1012
from pydantic import BaseModel
@@ -42,6 +44,7 @@ class VersionRunParams(TypedDict):
4244

4345
class OtherRunParams(TypedDict):
4446
use_cache: NotRequired["CacheUsage"]
47+
use_fallback: NotRequired[Union[Literal["auto", "never"], list[str]]]
4548

4649
max_retry_delay: NotRequired[float]
4750
max_retry_count: NotRequired[float]

‎workflowai/core/client/_models.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ class RunRequest(BaseModel):
2828

2929
use_cache: Optional[CacheUsage] = None
3030

31+
use_fallback: Optional[Union[Literal["auto", "never"], list[str]]] = None
32+
3133
metadata: Optional[dict[str, Any]] = None
3234

3335
labels: Optional[set[str]] = None # deprecated, to be included in metadata

‎workflowai/core/client/agent.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ async def _prepare_run(self, agent_input: AgentInput, stream: bool, **kwargs: Un
209209
version=version,
210210
stream=stream,
211211
use_cache=self._get_run_param("use_cache", kwargs),
212+
use_fallback=self._get_run_param("use_fallback", kwargs),
212213
metadata=kwargs.get("metadata"),
213214
)
214215

‎workflowai/core/domain/model.py

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,22 +12,13 @@
1212
# higher, comment out the line where it should be in "natural" order, and add another one wherever
1313
# needed for the order
1414
class Model(str, Enum):
15-
# --------------------------------------------------------------------------
16-
# OpenAI Models
17-
# --------------------------------------------------------------------------
18-
19-
GPT_41_LATEST = "gpt-4.1-latest"
20-
GPT_41_2025_04_14 = "gpt-4.1-2025-04-14"
21-
GPT_41_MINI_LATEST = "gpt-4.1-mini-latest"
22-
GPT_41_MINI_2025_04_14 = "gpt-4.1-mini-2025-04-14"
23-
GPT_41_NANO_LATEST = "gpt-4.1-nano-latest"
24-
GPT_41_NANO_2025_04_14 = "gpt-4.1-nano-2025-04-14"
2515
GPT_4O_LATEST = "gpt-4o-latest"
2616
GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
2717
GPT_4O_2024_08_06 = "gpt-4o-2024-08-06"
2818
GPT_4O_2024_05_13 = "gpt-4o-2024-05-13"
2919
GPT_4O_MINI_LATEST = "gpt-4o-mini-latest"
3020
GPT_4O_MINI_2024_07_18 = "gpt-4o-mini-2024-07-18"
21+
GPT_IMAGE_1 = "gpt-image-1"
3122
O3_LATEST_HIGH_REASONING_EFFORT = "o3-latest-high"
3223
O3_LATEST_MEDIUM_REASONING_EFFORT = "o3-latest-medium"
3324
O3_LATEST_LOW_REASONING_EFFORT = "o3-latest-low"
@@ -52,7 +43,14 @@ class Model(str, Enum):
5243
O1_PREVIEW_2024_09_12 = "o1-preview-2024-09-12"
5344
O1_MINI_LATEST = "o1-mini-latest"
5445
O1_MINI_2024_09_12 = "o1-mini-2024-09-12"
46+
GPT_41_LATEST = "gpt-4.1-latest"
47+
GPT_41_2025_04_14 = "gpt-4.1-2025-04-14"
48+
GPT_41_MINI_LATEST = "gpt-4.1-mini-latest"
49+
GPT_41_MINI_2025_04_14 = "gpt-4.1-mini-2025-04-14"
50+
GPT_41_NANO_LATEST = "gpt-4.1-nano-latest"
51+
GPT_41_NANO_2025_04_14 = "gpt-4.1-nano-2025-04-14"
5552
GPT_45_PREVIEW_2025_02_27 = "gpt-4.5-preview-2025-02-27"
53+
GPT_4O_AUDIO_PREVIEW_2025_06_03 = "gpt-4o-audio-preview-2025-06-03"
5654
GPT_4O_AUDIO_PREVIEW_2024_12_17 = "gpt-4o-audio-preview-2024-12-17"
5755
GPT_40_AUDIO_PREVIEW_2024_10_01 = "gpt-4o-audio-preview-2024-10-01"
5856
GPT_4_TURBO_2024_04_09 = "gpt-4-turbo-2024-04-09"
@@ -66,6 +64,10 @@ class Model(str, Enum):
6664
# Gemini Models
6765
# --------------------------------------------------------------------------
6866
GEMINI_2_0_FLASH_LATEST = "gemini-2.0-flash-latest"
67+
GEMINI_2_5_PRO_PREVIEW_0605 = "gemini-2.5-pro-preview-06-05"
68+
GEMINI_2_5_FLASH_PREVIEW_0520 = "gemini-2.5-flash-preview-05-20"
69+
GEMINI_2_5_FLASH_THINKING_PREVIEW_0520 = "gemini-2.5-flash-thinking-preview-05-20"
70+
GEMINI_2_5_PRO_PREVIEW_0506 = "gemini-2.5-pro-preview-05-06"
6971
GEMINI_2_5_FLASH_PREVIEW_0417 = "gemini-2.5-flash-preview-04-17"
7072
GEMINI_2_5_FLASH_THINKING_PREVIEW_0417 = "gemini-2.5-flash-thinking-preview-04-17"
7173
GEMINI_2_5_PRO_PREVIEW_0325 = "gemini-2.5-pro-preview-03-25"
@@ -93,9 +95,18 @@ class Model(str, Enum):
9395
GEMINI_1_0_PRO_001 = "gemini-1.0-pro-001"
9496
GEMINI_1_0_PRO_VISION_001 = "gemini-1.0-pro-vision-001"
9597

98+
IMAGEN_3_0_LATEST = "imagen-3.0-generate-latest"
99+
IMAGEN_3_0_002 = "imagen-3.0-generate-002"
100+
IMAGEN_3_0_001 = "imagen-3.0-generate-001"
101+
IMAGEN_3_0_FAST_001 = "imagen-3.0-fast-generate-001"
102+
96103
# --------------------------------------------------------------------------
97104
# Claude Models
98105
# --------------------------------------------------------------------------
106+
CLAUDE_4_SONNET_LATEST = "claude-sonnet-4-latest"
107+
CLAUDE_4_SONNET_20250514 = "claude-sonnet-4-20250514"
108+
CLAUDE_4_OPUS_LATEST = "claude-opus-4-latest"
109+
CLAUDE_4_OPUS_20250514 = "claude-opus-4-20250514"
99110
CLAUDE_3_7_SONNET_LATEST = "claude-3-7-sonnet-latest"
100111
CLAUDE_3_7_SONNET_20250219 = "claude-3-7-sonnet-20250219"
101112
CLAUDE_3_5_SONNET_LATEST = "claude-3-5-sonnet-latest"
@@ -139,6 +150,7 @@ class Model(str, Enum):
139150
MISTRAL_LARGE_2_2407 = "mistral-large-2-2407"
140151
MISTRAL_LARGE_LATEST = "mistral-large-latest"
141152
MISTRAL_LARGE_2411 = "mistral-large-2411"
153+
MISTRAL_MEDIUM_2505 = "mistral-medium-2505"
142154
PIXTRAL_LARGE_LATEST = "pixtral-large-latest"
143155
PIXTRAL_LARGE_2411 = "pixtral-large-2411"
144156
PIXTRAL_12B_2409 = "pixtral-12b-2409"
@@ -149,6 +161,8 @@ class Model(str, Enum):
149161
MISTRAL_SMALL_2501 = "mistral-small-2501"
150162
MISTRAL_SMALL_2409 = "mistral-small-2409"
151163
MISTRAL_SABA_2502 = "mistral-saba-2502"
164+
MAGISTRAL_SMALL_2506 = "magistral-small-2506"
165+
MAGISTRAL_MEDIUM_2506 = "magistral-medium-2506"
152166
CODESTRAL_2501 = "codestral-2501"
153167
CODESTRAL_MAMBA_2407 = "codestral-mamba-2407"
154168

@@ -157,6 +171,8 @@ class Model(str, Enum):
157171
# --------------------------------------------------------------------------
158172
QWEN_QWQ_32B = "qwen-qwq-32b"
159173
QWEN_QWQ_32B_PREVIEW = "qwen-v3p2-32b-instruct"
174+
QWEN3_235B_A22B = "qwen3-235b-a22b"
175+
QWEN3_30B_A3B = "qwen3-30b-a3b"
160176

161177
# --------------------------------------------------------------------------
162178
# DeepSeek Models
@@ -166,6 +182,7 @@ class Model(str, Enum):
166182
DEEPSEEK_V3_LATEST = "deepseek-v3-latest"
167183
DEEPSEEK_R1_2501 = "deepseek-r1-2501"
168184
DEEPSEEK_R1_2501_BASIC = "deepseek-r1-2501-basic"
185+
DEEPSEEK_R1_0528 = "deepseek-r1-0528"
169186

170187
# --------------------------------------------------------------------------
171188
# XAI Models

‎workflowai/core/domain/model_test.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,15 @@
66
async def test_model_exhaustive():
77
"""Make sure the list of models is synchronized with the prod API"""
88
async with httpx.AsyncClient() as client:
9-
response = await client.get("https://run.workflowai.com/v1/models")
9+
response = await client.get("https://run.workflowai.com/v1/models?raw=true")
1010
response.raise_for_status()
1111
models: list[str] = response.json()
1212

1313
# Converting to a set of strings should not be needed
1414
# but it makes pytest errors prettier
15-
assert set(models) == {m.value for m in Model}
15+
expected_models = set(models)
16+
actual_models = {m.value for m in Model}
17+
missing_models = expected_models - actual_models
18+
extra_models = actual_models - expected_models
19+
assert not extra_models, f"Extra models: {extra_models}"
20+
assert not missing_models, f"Missing models: {missing_models}"

0 commit comments

Comments
 (0)