← back to Exo
Truncate long logs with repr (#1854)
abd75ae06c150a0c5b26fcd7311fd03b891a950e · 2026-04-10 15:53:48 +0100 · rltakashige
## Motivation
<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->
## Changes
<!-- Describe what you changed in detail -->
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
---------
Co-authored-by: Evan <evanev7@gmail.com>
Files touched
M src/exo/api/adapters/chat_completions.pyM src/exo/api/adapters/claude.pyM src/exo/api/adapters/ollama.pyM src/exo/api/adapters/responses.pyM src/exo/api/main.pyM src/exo/api/tests/test_claude_api.pyM src/exo/api/tests/test_instance_deleted_stream_cleanup.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/shared/types/common.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/engines/mlx/vision.pyM src/exo/worker/main.pyM src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.pyM src/exo/worker/tests/unittests/test_plan/test_task_forwarding.pyM src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.pyM src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
Diff
commit abd75ae06c150a0c5b26fcd7311fd03b891a950e
Author: rltakashige <rl.takashige@gmail.com>
Date: Fri Apr 10 15:53:48 2026 +0100
Truncate long logs with repr (#1854)
## Motivation
<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->
## Changes
<!-- Describe what you changed in detail -->
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
---------
Co-authored-by: Evan <evanev7@gmail.com>
---
src/exo/api/adapters/chat_completions.py | 22 ++++---
src/exo/api/adapters/claude.py | 70 +++++++++++++++-------
src/exo/api/adapters/ollama.py | 31 ++++++----
src/exo/api/adapters/responses.py | 18 ++++--
src/exo/api/main.py | 19 ++----
src/exo/api/tests/test_claude_api.py | 24 ++++----
.../tests/test_instance_deleted_stream_cleanup.py | 8 ++-
src/exo/master/tests/test_master.py | 17 +++++-
src/exo/master/tests/test_placement.py | 8 ++-
src/exo/shared/types/common.py | 30 ++++++++--
src/exo/shared/types/text_generation.py | 56 ++++++++++++++---
src/exo/worker/engines/mlx/generator/generate.py | 10 +++-
src/exo/worker/engines/mlx/utils_mlx.py | 8 +--
src/exo/worker/engines/mlx/vision.py | 12 ++--
src/exo/worker/main.py | 11 ++--
.../test_mlx/test_prefix_cache_architectures.py | 10 +++-
.../unittests/test_plan/test_task_forwarding.py | 18 ++++--
.../tests/unittests/test_runner/test_dsml_e2e.py | 5 +-
.../unittests/test_runner/test_event_ordering.py | 8 ++-
.../test_runner/test_runner_supervisor.py | 8 ++-
20 files changed, 270 insertions(+), 123 deletions(-)
diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py
index 17ebbf6b..6dd5d7c1 100644
--- a/src/exo/api/adapters/chat_completions.py
+++ b/src/exo/api/adapters/chat_completions.py
@@ -31,20 +31,22 @@ from exo.shared.types.chunks import (
)
from exo.shared.types.common import CommandId
from exo.shared.types.text_generation import (
+ Base64Image,
InputMessage,
+ InputMessageContent,
TextGenerationTaskParams,
resolve_reasoning_params,
)
-def extract_base64_from_data_url(data_url: str) -> str:
+def extract_base64_from_data_url(data_url: str) -> Base64Image:
match = re.match(r"data:[^;]+;base64,(.+)", data_url)
if match:
- return match.group(1)
- return data_url
+ return Base64Image(match.group(1))
+ return Base64Image(data_url)
-async def fetch_image_url(url: str) -> str:
+async def fetch_image_url(url: str) -> Base64Image:
headers = {"User-Agent": "exo/1.0"}
async with (
create_http_session(timeout_profile="short") as session,
@@ -52,7 +54,7 @@ async def fetch_image_url(url: str) -> str:
):
resp.raise_for_status()
data = await resp.read()
- return base64.b64encode(data).decode("ascii")
+ return Base64Image(base64.b64encode(data).decode("ascii"))
async def chat_request_to_text_generation(
@@ -61,7 +63,7 @@ async def chat_request_to_text_generation(
instructions: str | None = None
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
- images: list[str] = []
+ images: list[Base64Image] = []
for msg in request.messages:
# Normalize content to string
@@ -115,7 +117,9 @@ async def chat_request_to_text_generation(
continue
if msg.role in ("user", "assistant", "developer"):
- input_messages.append(InputMessage(role=msg.role, content=content))
+ input_messages.append(
+ InputMessage(role=msg.role, content=InputMessageContent(content))
+ )
# Build full message dict for chat template (preserves tool_calls etc.)
# Normalize content for model_dump
@@ -144,8 +148,8 @@ async def chat_request_to_text_generation(
model=request.model,
input=input_messages
if input_messages
- else [InputMessage(role="user", content="")],
- instructions=instructions,
+ else [InputMessage(role="user", content=InputMessageContent(""))],
+ instructions=InputMessageContent(instructions) if instructions else None,
max_output_tokens=request.max_tokens,
temperature=request.temperature,
top_p=request.top_p,
diff --git a/src/exo/api/adapters/claude.py b/src/exo/api/adapters/claude.py
index 0b53e330..a54de764 100644
--- a/src/exo/api/adapters/claude.py
+++ b/src/exo/api/adapters/claude.py
@@ -5,6 +5,7 @@ import re
from collections.abc import AsyncGenerator
from typing import Any
+from exo.api.adapters.chat_completions import fetch_image_url
from exo.api.types import FinishReason, Usage
from exo.api.types.claude_api import (
ClaudeContentBlock,
@@ -30,6 +31,7 @@ from exo.api.types.claude_api import (
ClaudeToolUseBlock,
ClaudeUsage,
)
+from exo.shared.logging import logger
from exo.shared.types.chunks import (
ErrorChunk,
PrefillProgressChunk,
@@ -37,7 +39,13 @@ from exo.shared.types.chunks import (
ToolCallChunk,
)
from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ Base64Image,
+ ChatTemplateValue,
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
def finish_reason_to_claude_stop_reason(
@@ -83,13 +91,27 @@ def _strip_volatile_headers(text: str) -> str:
return _VOLATILE_HEADER_RE.sub("", text)
-def claude_request_to_text_generation(
+async def handle_image_block(block: ClaudeImageBlock) -> Base64Image | None:
+ if block.source.type == "base64" and block.source.data:
+ return Base64Image(block.source.data)
+ elif block.source.type == "url" and block.source.url:
+ try:
+ return await fetch_image_url(block.source.url)
+ except Exception:
+ logger.opt(exception=True).warning(
+ f"Failed to fetch image at {block.source.url}"
+ )
+
+ return None
+
+
+async def claude_request_to_text_generation(
request: ClaudeMessagesRequest,
) -> TextGenerationTaskParams:
# Handle system message
instructions: str | None = None
- chat_template_messages: list[dict[str, Any]] = []
- images: list[str] = []
+ chat_template_messages: list[dict[str, ChatTemplateValue]] = []
+ images: list[Base64Image] = []
if request.system:
if isinstance(request.system, str):
@@ -98,14 +120,20 @@ def claude_request_to_text_generation(
instructions = "".join(block.text for block in request.system)
instructions = _strip_volatile_headers(instructions)
- chat_template_messages.append({"role": "system", "content": instructions})
+ chat_template_messages.append(
+ {"role": "system", "content": InputMessageContent(instructions)}
+ )
# Convert messages to input
input_messages: list[InputMessage] = []
for msg in request.messages:
if isinstance(msg.content, str):
- input_messages.append(InputMessage(role=msg.role, content=msg.content))
- chat_template_messages.append({"role": msg.role, "content": msg.content})
+ input_messages.append(
+ InputMessage(role=msg.role, content=InputMessageContent(msg.content))
+ )
+ chat_template_messages.append(
+ {"role": msg.role, "content": InputMessageContent(msg.content)}
+ )
continue
# Process structured content blocks
@@ -119,12 +147,9 @@ def claude_request_to_text_generation(
if isinstance(block, ClaudeTextBlock):
text_parts.append(block.text)
elif isinstance(block, ClaudeImageBlock):
- if block.source.type == "base64" and block.source.data:
- images.append(block.source.data)
- has_images = True
- elif block.source.type == "url" and block.source.url:
- images.append(block.source.url)
+ if (img := await handle_image_block(block)) is not None:
has_images = True
+ images.append(img)
elif isinstance(block, ClaudeThinkingBlock):
thinking_parts.append(block.thinking)
elif isinstance(block, ClaudeToolUseBlock):
@@ -142,20 +167,21 @@ def claude_request_to_text_generation(
tool_results.append(block)
if isinstance(block.content, list):
for sub in block.content:
- if isinstance(sub, ClaudeImageBlock):
- if sub.source.type == "base64" and sub.source.data:
- images.append(sub.source.data)
- has_images = True
- elif sub.source.type == "url" and sub.source.url:
- images.append(sub.source.url)
- has_images = True
+ if (
+ isinstance(sub, ClaudeImageBlock)
+ and (img := await handle_image_block(sub)) is not None
+ ):
+ has_images = True
+ images.append(img)
content = "".join(text_parts)
reasoning_content = "".join(thinking_parts) if thinking_parts else None
# Build InputMessage from text content
if msg.role in ("user", "assistant"):
- input_messages.append(InputMessage(role=msg.role, content=content))
+ input_messages.append(
+ InputMessage(role=msg.role, content=InputMessageContent(content))
+ )
# Build chat_template_messages preserving tool structure
if tool_calls:
@@ -216,8 +242,8 @@ def claude_request_to_text_generation(
model=request.model,
input=input_messages
if input_messages
- else [InputMessage(role="user", content="")],
- instructions=instructions,
+ else [InputMessage(role="user", content=InputMessageContent(""))],
+ instructions=InputMessageContent(instructions) if instructions else None,
max_output_tokens=request.max_tokens,
temperature=request.temperature,
top_p=request.top_p,
diff --git a/src/exo/api/adapters/ollama.py b/src/exo/api/adapters/ollama.py
index 76436370..1f027616 100644
--- a/src/exo/api/adapters/ollama.py
+++ b/src/exo/api/adapters/ollama.py
@@ -21,7 +21,12 @@ from exo.shared.types.chunks import (
ToolCallChunk,
)
from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ Base64Image,
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
def _map_done_reason(
@@ -82,14 +87,14 @@ def ollama_request_to_text_generation(
instructions: str | None = None
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
- images: list[str] = []
+ images: list[Base64Image] = []
tool_message_index = 0
for msg in request.messages:
content = msg.content or ""
has_images = False
if msg.images:
- images.extend(msg.images)
+ images.extend(map(Base64Image, msg.images))
has_images = True
if msg.role == "system":
@@ -103,7 +108,9 @@ def ollama_request_to_text_generation(
if msg.role in ("user", "assistant") and (
msg.content is not None or msg.thinking is not None or msg.tool_calls
):
- input_messages.append(InputMessage(role=msg.role, content=content))
+ input_messages.append(
+ InputMessage(role=msg.role, content=InputMessageContent(content))
+ )
if has_images:
multimodal: list[dict[str, Any]] = [
@@ -113,7 +120,9 @@ def ollama_request_to_text_generation(
multimodal.append({"type": "text", "text": content})
chat_template_messages.append({"role": msg.role, "content": multimodal})
if msg.role in ("user", "assistant"):
- input_messages.append(InputMessage(role=msg.role, content=content))
+ input_messages.append(
+ InputMessage(role=msg.role, content=InputMessageContent(content))
+ )
continue
dumped: dict[str, Any] = {"role": msg.role, "content": content}
if msg.thinking is not None:
@@ -153,8 +162,8 @@ def ollama_request_to_text_generation(
model=request.model,
input=input_messages
if input_messages
- else [InputMessage(role="user", content="")],
- instructions=instructions,
+ else [InputMessage(role="user", content=InputMessageContent(""))],
+ instructions=InputMessageContent(instructions) if instructions else None,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
@@ -327,11 +336,11 @@ def ollama_generate_request_to_text_generation(
) -> TextGenerationTaskParams:
"""Convert Ollama generate request to exo's internal text generation format."""
chat_template_messages: list[dict[str, Any]] = []
- images: list[str] = []
+ images: list[Base64Image] = []
if request.system:
chat_template_messages.append({"role": "system", "content": request.system})
if request.images:
- images.extend(request.images)
+ images.extend(map(Base64Image, request.images))
multimodal: list[dict[str, Any]] = [{"type": "image"} for _ in request.images]
if request.prompt:
multimodal.append({"type": "text", "text": request.prompt})
@@ -342,8 +351,8 @@ def ollama_generate_request_to_text_generation(
options = request.options
return TextGenerationTaskParams(
model=request.model,
- input=[InputMessage(role="user", content=request.prompt)],
- instructions=request.system,
+ input=[InputMessage(role="user", content=InputMessageContent(request.prompt))],
+ instructions=InputMessageContent(request.system) if request.system else None,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
diff --git a/src/exo/api/adapters/responses.py b/src/exo/api/adapters/responses.py
index 2e8baf14..a3c248c0 100644
--- a/src/exo/api/adapters/responses.py
+++ b/src/exo/api/adapters/responses.py
@@ -74,7 +74,9 @@ from exo.shared.types.chunks import (
)
from exo.shared.types.common import CommandId
from exo.shared.types.text_generation import (
+ Base64Image,
InputMessage,
+ InputMessageContent,
TextGenerationTaskParams,
resolve_reasoning_params,
)
@@ -99,9 +101,11 @@ async def responses_request_to_text_generation(
) -> TextGenerationTaskParams:
input_value: list[InputMessage]
built_chat_template: list[dict[str, Any]] | None = None
- images: list[str] = []
+ images: list[Base64Image] = []
if isinstance(request.input, str):
- input_value = [InputMessage(role="user", content=request.input)]
+ input_value = [
+ InputMessage(role="user", content=InputMessageContent(request.input))
+ ]
else:
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
@@ -130,7 +134,9 @@ async def responses_request_to_text_generation(
has_images = True
if item.role in ("user", "assistant", "developer"):
input_messages.append(
- InputMessage(role=item.role, content=content)
+ InputMessage(
+ role=item.role, content=InputMessageContent(content)
+ )
)
if item.role == "system":
chat_template_messages.append(
@@ -327,7 +333,7 @@ async def responses_request_to_text_generation(
input_value = (
input_messages
if input_messages
- else [InputMessage(role="user", content="")]
+ else [InputMessage(role="user", content=InputMessageContent(""))]
)
built_chat_template = chat_template_messages if chat_template_messages else None
@@ -361,7 +367,9 @@ async def responses_request_to_text_generation(
return TextGenerationTaskParams(
model=request.model,
input=input_value,
- instructions=request.instructions,
+ instructions=InputMessageContent(request.instructions)
+ if request.instructions
+ else None,
max_output_tokens=request.max_output_tokens,
temperature=request.temperature,
top_p=request.top_p,
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 73488ff0..7bc0a462 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -25,7 +25,6 @@ from loguru import logger
from exo.api.adapters.chat_completions import (
chat_request_to_text_generation,
collect_chat_response,
- fetch_image_url,
generate_chat_stream,
)
from exo.api.adapters.claude import (
@@ -186,7 +185,7 @@ from exo.shared.types.tasks import (
from exo.shared.types.tasks import (
TextGeneration as TextGenerationTask,
)
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
@@ -760,9 +759,11 @@ class API:
self._sent_image_hashes.add(h)
new_images.append((idx, img))
+ wrapped_hashes = {idx: Base64Image(h) for idx, h in cached_hashes.items()}
+
if not new_images:
task_params = task_params.model_copy(
- update={"images": [], "image_hashes": cached_hashes}
+ update={"images": [], "image_hashes": wrapped_hashes}
)
command = TextGeneration(task_params=task_params)
await self._send(command)
@@ -776,7 +777,7 @@ class API:
task_params = task_params.model_copy(
update={
"images": [],
- "image_hashes": cached_hashes,
+ "image_hashes": wrapped_hashes,
"total_input_chunks": len(all_chunks),
"image_count": len(new_images),
}
@@ -1403,15 +1404,7 @@ class API:
self, payload: ClaudeMessagesRequest
) -> ClaudeMessagesResponse | StreamingResponse:
"""Claude Messages API - adapter."""
- task_params = claude_request_to_text_generation(payload)
- if task_params.images:
- resolved_images: list[str] = []
- for img in task_params.images:
- if img.startswith(("http://", "https://")):
- resolved_images.append(await fetch_image_url(img))
- else:
- resolved_images.append(img)
- task_params = task_params.model_copy(update={"images": resolved_images})
+ task_params = await claude_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
diff --git a/src/exo/api/tests/test_claude_api.py b/src/exo/api/tests/test_claude_api.py
index 27e8cc4e..f31e057a 100644
--- a/src/exo/api/tests/test_claude_api.py
+++ b/src/exo/api/tests/test_claude_api.py
@@ -40,7 +40,7 @@ class TestFinishReasonToClaudeStopReason:
class TestClaudeRequestToInternal:
"""Tests for converting Claude Messages API requests to TextGenerationTaskParams."""
- def test_basic_request_conversion(self):
+ async def test_basic_request_conversion(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -48,7 +48,7 @@ class TestClaudeRequestToInternal:
ClaudeMessage(role="user", content="Hello"),
],
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert params.model == "claude-3-opus"
assert params.max_output_tokens == 100
@@ -58,7 +58,7 @@ class TestClaudeRequestToInternal:
assert params.input[0].content == "Hello"
assert params.instructions is None
- def test_request_with_system_string(self):
+ async def test_request_with_system_string(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -67,7 +67,7 @@ class TestClaudeRequestToInternal:
ClaudeMessage(role="user", content="Hello"),
],
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert params.instructions == "You are a helpful assistant."
assert isinstance(params.input, list)
@@ -75,7 +75,7 @@ class TestClaudeRequestToInternal:
assert params.input[0].role == "user"
assert params.input[0].content == "Hello"
- def test_request_with_system_text_blocks(self):
+ async def test_request_with_system_text_blocks(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -87,13 +87,13 @@ class TestClaudeRequestToInternal:
ClaudeMessage(role="user", content="Hello"),
],
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert params.instructions == "You are helpful. Be concise."
assert isinstance(params.input, list)
assert len(params.input) == 1
- def test_request_with_content_blocks(self):
+ async def test_request_with_content_blocks(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -107,13 +107,13 @@ class TestClaudeRequestToInternal:
),
],
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert isinstance(params.input, list)
assert len(params.input) == 1
assert params.input[0].content == "First part. Second part."
- def test_request_with_multi_turn_conversation(self):
+ async def test_request_with_multi_turn_conversation(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -123,7 +123,7 @@ class TestClaudeRequestToInternal:
ClaudeMessage(role="user", content="How are you?"),
],
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert isinstance(params.input, list)
assert len(params.input) == 3
@@ -131,7 +131,7 @@ class TestClaudeRequestToInternal:
assert params.input[1].role == "assistant"
assert params.input[2].role == "user"
- def test_request_with_optional_parameters(self):
+ async def test_request_with_optional_parameters(self):
request = ClaudeMessagesRequest(
model=ModelId("claude-3-opus"),
max_tokens=100,
@@ -142,7 +142,7 @@ class TestClaudeRequestToInternal:
stop_sequences=["STOP", "END"],
stream=True,
)
- params = claude_request_to_text_generation(request)
+ params = await claude_request_to_text_generation(request)
assert params.temperature == 0.7
assert params.top_p == 0.9
diff --git a/src/exo/api/tests/test_instance_deleted_stream_cleanup.py b/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
index 843554b9..fb03c5da 100644
--- a/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
+++ b/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
@@ -8,7 +8,11 @@ from exo.api.types import ImageGenerationTaskParams
from exo.shared.types.common import CommandId, ModelId
from exo.shared.types.state import State
from exo.shared.types.tasks import ImageGeneration, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.instances import InstanceId
@@ -29,7 +33,7 @@ def _make_text_gen_task(
command_id=command_id,
task_params=TextGenerationTaskParams(
model=ModelId("test-model"),
- input=[InputMessage(role="user", content="hello")],
+ input=[InputMessage(role="user", content=InputMessageContent("hello"))],
),
)
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 9d4e5793..c4a1cff0 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -31,7 +31,11 @@ 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 InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.instances import (
InstanceMeta,
MlxRingInstance,
@@ -159,7 +163,10 @@ async def test_master():
task_params=TextGenerationTaskParams(
model=ModelId("llama-3.2-1b"),
input=[
- InputMessage(role="user", content="Hello, how are you?")
+ InputMessage(
+ role="user",
+ content=InputMessageContent("Hello, how are you?"),
+ )
],
),
)
@@ -213,7 +220,11 @@ 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=[InputMessage(role="user", content="Hello, how are you?")],
+ input=[
+ InputMessage(
+ role="user", content=InputMessageContent("Hello, how are you?")
+ )
+ ],
)
ev_send.close()
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 3a8a36c0..40530ad2 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -23,7 +23,11 @@ from exo.shared.types.memory import Memory
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.downloads import (
DownloadCompleted,
@@ -481,7 +485,7 @@ def _make_task(
command_id=CommandId(),
task_params=TextGenerationTaskParams(
model=ModelId("test-model"),
- input=[InputMessage(role="user", content="hello")],
+ input=[InputMessage(role="user", content=InputMessageContent("hello"))],
),
)
diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py
index b5674b4b..e539386f 100644
--- a/src/exo/shared/types/common.py
+++ b/src/exo/shared/types/common.py
@@ -1,8 +1,8 @@
-from typing import Self
+from typing import Any, Self
from uuid import uuid4
from pydantic import GetCoreSchemaHandler, field_validator
-from pydantic_core import core_schema
+from pydantic_core import CoreSchema, core_schema
from exo.utils.pydantic_ext import CamelCaseModel
@@ -37,15 +37,33 @@ class ModelId(Id):
return self.split("/")[-1]
+class CommandId(Id):
+ pass
+
+
+class TruncatingString(str):
+ truncate_length: int = -1
+
+ @classmethod
+ def __get_pydantic_core_schema__(
+ cls,
+ source_type: Any, # pyright: ignore[reportAny]
+ handler: GetCoreSchemaHandler,
+ ) -> CoreSchema:
+ return core_schema.no_info_after_validator_function(cls, handler(str))
+
+ def __repr__(self):
+ tl = type(self).truncate_length
+ return (
+ f"<{type(self).__name__}: {self[:tl] + '...' if len(self) > tl else self}>"
+ )
+
+
class SessionId(CamelCaseModel):
master_node_id: NodeId
election_clock: int
-class CommandId(Id):
- pass
-
-
class Host(CamelCaseModel):
ip: str
port: int
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index 29affa67..4cc0a1ec 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -4,13 +4,13 @@ All external API formats (Chat Completions, Claude Messages, OpenAI Responses)
are converted to TextGenerationTaskParams at the API boundary via adapters.
"""
-from typing import Any, Literal
+from typing import Annotated, Any, Literal
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, WrapValidator
-from exo.shared.types.common import ModelId
+from exo.shared.types.common import ModelId, TruncatingString
-MessageRole = Literal["user", "assistant", "system", "developer"]
+MessageRole = Literal["user", "assistant", "system", "developer", "tool"]
ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
@@ -36,11 +36,49 @@ def resolve_reasoning_params(
return resolved_effort, resolved_thinking
+class InputMessageContent(TruncatingString):
+ truncate_length = 100
+
+
class InputMessage(BaseModel, frozen=True):
"""Internal message for text generation pipelines."""
role: MessageRole
- content: str
+ content: InputMessageContent
+
+
+class Base64Image(TruncatingString):
+ truncate_length = 10
+
+
+class Base64ImageHash(TruncatingString):
+ truncate_length = 10
+
+
+def _wrap_chat_value(x: Any) -> Any: # pyright: ignore[reportAny]
+ if isinstance(x, (InputMessageContent, Base64Image)):
+ return x
+ if isinstance(x, str):
+ return InputMessageContent(x)
+ if isinstance(x, dict):
+ return {k: _wrap_chat_value(v) for k, v in x.items()} # pyright: ignore[reportUnknownVariableType]
+ if isinstance(x, list):
+ return [_wrap_chat_value(i) for i in x] # pyright: ignore[reportUnknownVariableType]
+ return x # pyright: ignore[reportAny]
+
+
+type ChatTemplateValue = Annotated[
+ InputMessageContent
+ | Base64Image
+ | dict[str, ChatTemplateValue]
+ | list[ChatTemplateValue]
+ | str
+ | int
+ | float
+ | MessageRole
+ | bool,
+ WrapValidator(lambda a, b: b(_wrap_chat_value(a))), # pyright: ignore[reportAny]
+]
class TextGenerationTaskParams(BaseModel, frozen=True):
@@ -52,7 +90,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
model: ModelId
input: list[InputMessage]
- instructions: str | None = None
+ instructions: InputMessageContent | None = None
max_output_tokens: int | None = None
temperature: float | None = None
top_p: float | None = None
@@ -62,7 +100,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
top_k: int | None = None
stop: str | list[str] | None = None
seed: int | None = None
- chat_template_messages: list[dict[str, Any]] | None = None
+ chat_template_messages: list[dict[str, ChatTemplateValue]] | None = None
reasoning_effort: ReasoningEffort | None = None
enable_thinking: bool | None = None
logprobs: bool = False
@@ -70,7 +108,7 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
min_p: float | None = None
repetition_penalty: float | None = None
repetition_context_size: int | None = None
- images: list[str] = Field(default_factory=list)
- image_hashes: dict[int, str] = Field(default_factory=dict)
+ images: list[Base64Image] = Field(default_factory=list)
+ image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
total_input_chunks: int = 0
image_count: int = 0
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 43a38bca..3c0cb2ac 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -25,7 +25,11 @@ from exo.api.types import (
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType, Model
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.runner_response import (
GenerationResponse,
)
@@ -355,7 +359,9 @@ def warmup_inference(
) -> int:
logger.info(f"warming up inference for instance: {model_id}")
- content = "Prompt to warm up the inference engine. Repeat this."
+ content = InputMessageContent(
+ "Prompt to warm up the inference engine. Repeat this."
+ )
warmup_task_params = TextGenerationTaskParams(
model=model_id,
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 790dcd8d..b4eca6aa 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -45,7 +45,7 @@ from exo.shared.types.common import Host
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import Model
from exo.shared.types.tasks import TaskId, TextGeneration
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import ChatTemplateValue, TextGenerationTaskParams
from exo.shared.types.worker.instances import (
BoundInstance,
MlxJacclInstance,
@@ -611,10 +611,10 @@ def apply_chat_template(
tokenizer: TokenizerWrapper,
task_params: TextGenerationTaskParams,
) -> str:
- messages: list[dict[str, Any]] = []
+ messages: list[dict[str, ChatTemplateValue]] = []
if task_params.chat_template_messages is not None:
# Use pre-formatted messages that preserve tool_calls, thinking, etc.
- messages = list(task_params.chat_template_messages)
+ messages = task_params.chat_template_messages
else:
# Add system message (instructions) if present
if task_params.instructions:
@@ -642,7 +642,7 @@ def system_prompt_token_count(
if task_params.chat_template_messages is not None:
for msg in task_params.chat_template_messages:
if msg.get("role") in ("system", "developer"):
- content = msg.get("content", "") # type: ignore
+ content = msg.get("content", "")
if isinstance(content, str):
parts.append(content)
else:
diff --git a/src/exo/worker/engines/mlx/vision.py b/src/exo/worker/engines/mlx/vision.py
index 49cb024d..b2cc8ea4 100644
--- a/src/exo/worker/engines/mlx/vision.py
+++ b/src/exo/worker/engines/mlx/vision.py
@@ -25,7 +25,7 @@ from exo.download.download_utils import build_model_path
from exo.shared.models.model_cards import VisionCardConfig
from exo.shared.types.common import ModelId
from exo.shared.types.mlx import Model
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
from exo.worker.engines.mlx.cache import encode_prompt
from exo.worker.engines.mlx.utils_mlx import (
fix_unmatched_think_end_tokens,
@@ -324,7 +324,7 @@ class VisionEncoder:
n_vision = sum(v.size for _, v in vision_weights.items()) # type: ignore
logger.info(f"Vision encoder loaded: {n_vision / 1e6:.1f}M params")
- def encode_images(self, images: list[str]) -> tuple[mx.array, list[int]]:
+ def encode_images(self, images: list[Base64Image]) -> tuple[mx.array, list[int]]:
self.ensure_loaded()
assert self._vision_tower is not None
assert self._processor is not None
@@ -432,7 +432,7 @@ def create_vision_embeddings(
def _find_media_regions(
prompt_tokens: mx.array,
- images: list[str],
+ images: list[Base64Image],
image_token_id: int,
) -> list[MediaRegion]:
tokens_np = np.array(prompt_tokens)
@@ -483,7 +483,7 @@ class VisionProcessor:
def load(self) -> None:
self._encoder.ensure_loaded()
- def _image_cache_key(self, images: list[str]) -> str:
+ def _image_cache_key(self, images: list[Base64Image]) -> str:
h = hashlib.sha256()
for img in images:
pil = decode_base64_image(img)
@@ -492,7 +492,7 @@ class VisionProcessor:
def process(
self,
- images: list[str],
+ images: list[Base64Image],
chat_template_messages: list[dict[str, Any]],
tokenizer: TokenizerWrapper,
model: Model,
@@ -567,7 +567,7 @@ class VisionProcessor:
def prepare_vision(
- images: list[str] | None,
+ images: list[Base64Image] | None,
chat_template_messages: list[dict[str, Any]] | None,
vision_processor: VisionProcessor,
tokenizer: TokenizerWrapper,
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index e3da3bdb..8cdc3ffc 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -46,6 +46,7 @@ from exo.shared.types.tasks import (
TaskStatus,
TextGeneration,
)
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import InstanceId
@@ -86,7 +87,7 @@ class Worker:
# Buffer for input image chunks (for image editing)
self.input_chunk_buffer: dict[CommandId, dict[int, InputImageChunk]] = {}
self.input_chunk_counts: dict[CommandId, int] = {}
- self.image_cache: dict[str, str] = {}
+ self.image_cache: dict[Base64ImageHash, Base64Image] = {}
self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
self._instance_backoff: KeyedBackoff[InstanceId] = KeyedBackoff(
@@ -309,7 +310,7 @@ class Worker:
or task.task_params.total_input_chunks > 0
):
cmd_id = task.command_id
- by_index: dict[int, str] = {}
+ by_index: dict[int, Base64Image] = {}
for idx, h in task.task_params.image_hashes.items():
assert h in self.image_cache
@@ -326,9 +327,11 @@ class Worker:
sorted_chunks = sorted(
per_image[img_idx], key=lambda c: c.chunk_index
)
- img = "".join(c.data for c in sorted_chunks)
+ img = Base64Image("".join(c.data for c in sorted_chunks))
self.image_cache[
- hashlib.sha256(img.encode("ascii")).hexdigest()
+ Base64ImageHash(
+ hashlib.sha256(img.encode("ascii")).hexdigest()
+ )
] = img
by_index[img_idx] = img
logger.info(
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
index 609ea867..fec5082a 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
@@ -16,7 +16,11 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.types.common import ModelId
from exo.shared.types.mlx import Model
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.generator.generate import mlx_generate
from exo.worker.engines.mlx.utils_mlx import (
@@ -203,7 +207,9 @@ def _make_task() -> TextGenerationTaskParams:
input=[
InputMessage(
role="user",
- content="Use the calculator to compute 1847 * 263 + 5921",
+ content=InputMessageContent(
+ "Use the calculator to compute 1847 * 263 + 5921"
+ ),
)
],
max_output_tokens=20,
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 0aa901ce..bb87268f 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,11 @@ 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 InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import (
RunnerIdle,
@@ -61,7 +65,8 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
task_params=TextGenerationTaskParams(
- model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ model=MODEL_A_ID,
+ input=[InputMessage(role="user", content=InputMessageContent(""))],
),
)
@@ -113,7 +118,8 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
task_params=TextGenerationTaskParams(
- model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ model=MODEL_A_ID,
+ input=[InputMessage(role="user", content=InputMessageContent(""))],
),
)
@@ -162,7 +168,8 @@ def test_plan_does_not_forward_tasks_for_other_instances():
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
task_params=TextGenerationTaskParams(
- model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ model=MODEL_A_ID,
+ input=[InputMessage(role="user", content=InputMessageContent(""))],
),
)
@@ -215,7 +222,8 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
task_status=TaskStatus.Complete,
command_id=COMMAND_1_ID,
task_params=TextGenerationTaskParams(
- model=MODEL_A_ID, input=[InputMessage(role="user", content="")]
+ model=MODEL_A_ID,
+ input=[InputMessage(role="user", content=InputMessageContent(""))],
),
)
diff --git a/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py
index 26c11fe5..74efbef1 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py
@@ -988,6 +988,7 @@ class TestApplyChatTemplateWithToolCalls:
def test_dsml_encoding_with_tool_calls_in_history(self):
from exo.shared.types.text_generation import (
InputMessage,
+ InputMessageContent,
TextGenerationTaskParams,
)
from exo.worker.engines.mlx.utils_mlx import apply_chat_template
@@ -1022,8 +1023,8 @@ class TestApplyChatTemplateWithToolCalls:
params = TextGenerationTaskParams(
model=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
- input=[InputMessage(role="user", content="Thanks!")],
- instructions="You are a helpful assistant.",
+ input=[InputMessage(role="user", content=InputMessageContent("Thanks!"))],
+ instructions=InputMessageContent("You are a helpful assistant."),
enable_thinking=True,
chat_template_messages=chat_template_messages,
tools=_WEATHER_TOOLS,
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 cd4c18d7..ffd8fbfd 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
@@ -27,7 +27,11 @@ from exo.shared.types.tasks import (
TaskStatus,
TextGeneration,
)
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.shared.types.worker.runners import (
RunnerConnected,
@@ -91,7 +95,7 @@ SHUTDOWN_TASK = Shutdown(
CHAT_PARAMS = TextGenerationTaskParams(
model=MODEL_A_ID,
- input=[InputMessage(role="user", content="hello")],
+ input=[InputMessage(role="user", content=InputMessageContent("hello"))],
stream=True,
max_output_tokens=4,
temperature=0.0,
diff --git a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
index ecb07f94..82612d6d 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
@@ -9,7 +9,11 @@ from exo.shared.types.chunks import ErrorChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
from exo.shared.types.tasks import Task, TaskId, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ InputMessage,
+ InputMessageContent,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import RunnerFailed, RunnerId
from exo.utils.channels import channel, mp_channel
@@ -68,7 +72,7 @@ async def test_check_runner_emits_error_chunk_for_inflight_text_generation() ->
command_id=command_id,
task_params=TextGenerationTaskParams(
model=bound_instance.bound_shard.model_card.model_id,
- input=[InputMessage(role="user", content="hi")],
+ input=[InputMessage(role="user", content=InputMessageContent("hi"))],
stream=True,
),
)
← ee2e505b fix: handle BrokenResourceError in download progress callbac
·
back to Exo
·
Fix pdf inputs on Safari (#1865) 2962ebee →