← back to Exo
Normalize TextGenerationTaskParams.input to list[InputMessage] (#1360)
acb97127bf1326d55f07d2301da6e6f19da19df5 · 2026-02-03 06:01:56 -0800 · Alex Cheema
## Motivation
With the addition of the Responses API, we introduced `str |
list[InputMessage]` as the type for `TextGenerationTaskParams.input`
since the Responses API supports sending input as a plain string. But
there was no reason to leak that flexibility past the API adapter
boundary — it just meant every downstream consumer had to do `if
isinstance(messages, str):` checks, adding complexity for no benefit.
## Changes
- Changed `TextGenerationTaskParams.input` from `str |
list[InputMessage]` to `list[InputMessage]`
- Each API adapter (Chat Completions, Claude Messages, Responses) now
normalizes to `list[InputMessage]` at the boundary
- Removed `isinstance(task_params.input, str)` branches in
`utils_mlx.py` and `runner.py`
- Wrapped string inputs in `[InputMessage(role="user", content=...)]` in
the warmup path and all test files
## Why It Works
The API adapters are the only place where we deal with raw user input
formats. By normalizing there, all downstream code (worker, runner, MLX
engine) can just assume `list[InputMessage]` and skip the type-checking
branches. The type system (`basedpyright`) catches any missed call sites
at compile time.
## Test Plan
### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes
- `nix fmt` — applied
- `uv run pytest` — 174 passed, 1 skipped
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Files touched
M src/exo/master/adapters/chat_completions.pyM src/exo/master/adapters/claude.pyM src/exo/master/adapters/responses.pyM src/exo/master/tests/test_master.pyM src/exo/shared/types/text_generation.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_plan/test_task_forwarding.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.pyM tests/headless_runner.py
Diff
commit acb97127bf1326d55f07d2301da6e6f19da19df5
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Tue Feb 3 06:01:56 2026 -0800
Normalize TextGenerationTaskParams.input to list[InputMessage] (#1360)
## Motivation
With the addition of the Responses API, we introduced `str |
list[InputMessage]` as the type for `TextGenerationTaskParams.input`
since the Responses API supports sending input as a plain string. But
there was no reason to leak that flexibility past the API adapter
boundary — it just meant every downstream consumer had to do `if
isinstance(messages, str):` checks, adding complexity for no benefit.
## Changes
- Changed `TextGenerationTaskParams.input` from `str |
list[InputMessage]` to `list[InputMessage]`
- Each API adapter (Chat Completions, Claude Messages, Responses) now
normalizes to `list[InputMessage]` at the boundary
- Removed `isinstance(task_params.input, str)` branches in
`utils_mlx.py` and `runner.py`
- Wrapped string inputs in `[InputMessage(role="user", content=...)]` in
the warmup path and all test files
## Why It Works
The API adapters are the only place where we deal with raw user input
formats. By normalizing there, all downstream code (worker, runner, MLX
engine) can just assume `list[InputMessage]` and skip the type-checking
branches. The type system (`basedpyright`) catches any missed call sites
at compile time.
## Test Plan
### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes
- `nix fmt` — applied
- `uv run pytest` — 174 passed, 1 skipped
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
---
src/exo/master/adapters/chat_completions.py | 4 +++-
src/exo/master/adapters/claude.py | 4 +++-
src/exo/master/adapters/responses.py | 10 +++++++---
src/exo/master/tests/test_master.py | 8 +++++---
src/exo/shared/types/text_generation.py | 2 +-
src/exo/worker/engines/mlx/generator/generate.py | 4 ++--
src/exo/worker/engines/mlx/utils_mlx.py | 15 +++++----------
src/exo/worker/runner/runner.py | 13 ++++---------
src/exo/worker/tests/unittests/test_mlx/conftest.py | 6 +++---
.../tests/unittests/test_plan/test_task_forwarding.py | 18 +++++++++++++-----
.../tests/unittests/test_runner/test_event_ordering.py | 4 ++--
tests/headless_runner.py | 8 ++++++--
12 files changed, 54 insertions(+), 42 deletions(-)
diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py
index 5a27664d..e144696b 100644
--- a/src/exo/master/adapters/chat_completions.py
+++ b/src/exo/master/adapters/chat_completions.py
@@ -66,7 +66,9 @@ def chat_request_to_text_generation(
return TextGenerationTaskParams(
model=request.model,
- input=input_messages if input_messages else "",
+ input=input_messages
+ if input_messages
+ else [InputMessage(role="user", content="")],
instructions=instructions,
max_output_tokens=request.max_tokens,
temperature=request.temperature,
diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py
index 13398012..6c17b49c 100644
--- a/src/exo/master/adapters/claude.py
+++ b/src/exo/master/adapters/claude.py
@@ -141,7 +141,9 @@ def claude_request_to_text_generation(
return TextGenerationTaskParams(
model=request.model,
- input=input_messages if input_messages else "",
+ input=input_messages
+ if input_messages
+ else [InputMessage(role="user", content="")],
instructions=instructions,
max_output_tokens=request.max_tokens,
temperature=request.temperature,
diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py
index 27d845cd..c2a416ac 100644
--- a/src/exo/master/adapters/responses.py
+++ b/src/exo/master/adapters/responses.py
@@ -43,10 +43,10 @@ def _extract_content(content: str | list[ResponseContentPart]) -> str:
def responses_request_to_text_generation(
request: ResponsesRequest,
) -> TextGenerationTaskParams:
- input_value: str | list[InputMessage]
+ input_value: list[InputMessage]
built_chat_template: list[dict[str, Any]] | None = None
if isinstance(request.input, str):
- input_value = request.input
+ input_value = [InputMessage(role="user", content=request.input)]
else:
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
@@ -95,7 +95,11 @@ def responses_request_to_text_generation(
}
)
- input_value = input_messages if input_messages else ""
+ input_value = (
+ input_messages
+ if input_messages
+ else [InputMessage(role="user", content="")]
+ )
built_chat_template = chat_template_messages if chat_template_messages else None
return TextGenerationTaskParams(
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index d9987727..ddf9aec8 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -28,7 +28,7 @@ from exo.shared.types.profiling import (
)
from exo.shared.types.tasks import TaskStatus
from exo.shared.types.tasks import TextGeneration as TextGenerationTask
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.instances import (
InstanceMeta,
MlxRingInstance,
@@ -136,7 +136,9 @@ async def test_master():
command_id=CommandId(),
task_params=TextGenerationTaskParams(
model=ModelId("llama-3.2-1b"),
- input="Hello, how are you?",
+ input=[
+ InputMessage(role="user", content="Hello, how are you?")
+ ],
),
)
),
@@ -189,7 +191,7 @@ async def test_master():
assert isinstance(events[2].event.task, TextGenerationTask)
assert events[2].event.task.task_params == TextGenerationTaskParams(
model=ModelId("llama-3.2-1b"),
- input="Hello, how are you?",
+ input=[InputMessage(role="user", content="Hello, how are you?")],
)
await master.shutdown()
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index b9c5565c..31f97a70 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -28,7 +28,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
"""
model: ModelId
- input: str | list[InputMessage]
+ input: list[InputMessage]
instructions: str | None = None
max_output_tokens: int | None = None
temperature: float | None = None
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index c9585f57..a38d70c5 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -17,7 +17,7 @@ from exo.shared.types.api import (
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.runner_response import (
GenerationResponse,
)
@@ -100,7 +100,7 @@ def warmup_inference(
tokenizer=tokenizer,
task_params=TextGenerationTaskParams(
model=ModelId(""),
- input=content,
+ input=[InputMessage(role="user", content=content)],
),
)
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index de5cb190..d7fb9958 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -436,16 +436,11 @@ def apply_chat_template(
)
# Convert input to messages
- if isinstance(task_params.input, str):
- # Simple string input becomes a single user message
- formatted_messages.append({"role": "user", "content": task_params.input})
- else:
- # List of InputMessage
- for msg in task_params.input:
- if not msg.content:
- logger.warning("Received message with empty content, skipping")
- continue
- formatted_messages.append({"role": msg.role, "content": msg.content})
+ for msg in task_params.input:
+ if not msg.content:
+ logger.warning("Received message with empty content, skipping")
+ continue
+ formatted_messages.append({"role": msg.role, "content": msg.content})
prompt: str = tokenizer.apply_chat_template(
formatted_messages,
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 5439b72b..0b527318 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -918,15 +918,10 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
Extracts the first user input text and checks for debug triggers.
"""
- prompt: str
- if isinstance(task_params.input, str):
- prompt = task_params.input
- else:
- # List of InputMessage - get first message content
- if len(task_params.input) == 0:
- logger.debug("Empty message list in debug prompt check")
- return
- prompt = task_params.input[0].content
+ if len(task_params.input) == 0:
+ logger.debug("Empty message list in debug prompt check")
+ return
+ prompt = task_params.input[0].content
if not prompt:
return
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index 7015d70b..9e897141 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -14,7 +14,7 @@ from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.generator.generate import mlx_generate
@@ -114,7 +114,7 @@ def run_gpt_oss_pipeline_device(
task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- input=prompt_text,
+ input=[InputMessage(role="user", content=prompt_text)],
max_output_tokens=max_tokens,
)
@@ -182,7 +182,7 @@ def run_gpt_oss_tensor_parallel_device(
task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- input=prompt_text,
+ input=[InputMessage(role="user", content=prompt_text)],
max_output_tokens=max_tokens,
)
diff --git a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py b/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
index 07376787..64be74e6 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
@@ -2,7 +2,7 @@ from typing import cast
import exo.worker.plan as plan_mod
from exo.shared.types.tasks import Task, TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import (
RunnerIdle,
@@ -59,7 +59,9 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
+ task_params=TextGenerationTaskParams(
+ model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ ),
)
result = plan_mod.plan(
@@ -106,7 +108,9 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
+ task_params=TextGenerationTaskParams(
+ model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ ),
)
result = plan_mod.plan(
@@ -150,7 +154,9 @@ def test_plan_does_not_forward_tasks_for_other_instances():
instance_id=other_instance_id,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
+ task_params=TextGenerationTaskParams(
+ model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ ),
)
result = plan_mod.plan(
@@ -198,7 +204,9 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Complete,
command_id=COMMAND_1_ID,
- task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
+ task_params=TextGenerationTaskParams(
+ model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ ),
)
other_task_id = TaskId("other-task")
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index 9d7703d8..16a43f2b 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -22,7 +22,7 @@ from exo.shared.types.tasks import (
TaskStatus,
TextGeneration,
)
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.shared.types.worker.runners import (
RunnerConnected,
@@ -86,7 +86,7 @@ SHUTDOWN_TASK = Shutdown(
CHAT_PARAMS = TextGenerationTaskParams(
model=MODEL_A_ID,
- input="hello",
+ input=[InputMessage(role="user", content="hello")],
stream=True,
max_output_tokens=4,
temperature=0.0,
diff --git a/tests/headless_runner.py b/tests/headless_runner.py
index ed57823b..56fb2632 100644
--- a/tests/headless_runner.py
+++ b/tests/headless_runner.py
@@ -23,7 +23,7 @@ from exo.shared.types.tasks import (
Task,
TextGeneration,
)
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.instances import (
BoundInstance,
Instance,
@@ -196,7 +196,11 @@ async def execute_test(test: Tests, instance: Instance, hn: str) -> list[Event]:
task_params=TextGenerationTaskParams(
model=test.model_id,
instructions="You are a helpful assistant",
- input="What is the capital of France?",
+ input=[
+ InputMessage(
+ role="user", content="What is the capital of France?"
+ )
+ ],
),
command_id=CommandId("yo"),
instance_id=iid,
← d90605f1 migrate model cards to .toml files (#1354)
·
back to Exo
·
Reduce reliance on internet (#1363) a0f4f363 →