← back to Exo
feat: add Claude Messages API and OpenAI Responses API support (#1167)
c3537980bd1c8330dcee7455f76b98c03f8f69f1 · 2026-02-02 07:58:37 -0800 · Alex Cheema
## Motivation
Add support for Claude Messages API and OpenAI Responses API to allow
users to interact with exo using these popular API formats. This enables
broader compatibility with existing tooling and SDKs that expect these
API formats.
## Architecture
Adapter logic lives exclusively in the API layer
(`src/exo/master/adapters/`). On the way in, each adapter converts its
API-specific request type (`ChatCompletionRequest`,
`ClaudeMessagesRequest`, `ResponsesRequest`) into
`TextGenerationTaskParams`. On the way out, each adapter converts the
`TokenChunk` stream back into its API-specific response format.
Everything inside the application — commands, worker, runner, event
sourcing — only sees `TextGenerationTaskParams` and `TokenChunk`. No
API-specific types cross the boundary.
```
API layer │ Application internals
│
Chat Completions → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ChatCompletionResponse
Claude Messages → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ClaudeMessagesResponse
Responses API → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ResponsesResponse
```
## Changes
### New Files
- `src/exo/shared/types/claude_api.py` - Pydantic types for Claude
Messages API
- `src/exo/shared/types/openai_responses.py` - Pydantic types for OpenAI
Responses API
- `src/exo/shared/types/text_generation.py` - Shared
`TextGenerationTaskParams` internal type
- `src/exo/master/adapters/chat_completions.py` - Chat Completions
adapter (streaming/non-streaming)
- `src/exo/master/adapters/claude.py` - Claude Messages adapter
(streaming/non-streaming)
- `src/exo/master/adapters/responses.py` - OpenAI Responses adapter
(streaming/non-streaming)
### Modified Files
- `src/exo/master/api.py` - Refactored to use adapters uniformly for all
endpoints; extracted `_resolve_and_validate_text_model` helper to
deduplicate model validation across all text endpoints; removed ad-hoc
`try/except ValueError` blocks from non-streaming paths
### New Endpoints
- `POST /v1/messages` - Claude Messages API (streaming and
non-streaming)
- `POST /v1/responses` - OpenAI Responses API (streaming and
non-streaming)
## Why It Works
All APIs are implemented as pure conversion adapters at the edge of the
application:
1. Adapter functions in `src/exo/master/adapters/` convert incoming
requests to `TextGenerationTaskParams`
2. `api.py` wraps the params in a `TextGeneration` command and sends it
through the existing command/event flow
3. The worker, runner, and event sourcing layers only handle
`TextGenerationTaskParams` and `TokenChunk` — they have no awareness of
Chat Completions, Claude, or Responses API formats
4. On response, adapter functions convert the `TokenChunk` stream back
to the caller's expected format
5. Model validation is handled by a single shared helper
(`_resolve_and_validate_text_model`), mirroring the existing
`_validate_image_model` pattern for image endpoints
No changes to core inference logic were needed.
### Streaming Formats
- **Chat Completions**: Uses `data: {...}\n\n` with `[DONE]` terminator
- **Claude**: Uses event types `message_start`, `content_block_start`,
`content_block_delta`, `content_block_stop`, `message_delta`,
`message_stop`
- **OpenAI Responses**: Uses event types `response.created`,
`response.in_progress`, `response.output_item.added`,
`response.content_part.added`, `response.output_text.delta`,
`response.output_text.done`, `response.content_part.done`,
`response.output_item.done`, `response.completed`
## Test Plan
### Manual Testing
Hardware: MacBook Pro M3 Max
**Non-streaming tests:**
```bash
# Chat Completions API
curl -X POST http://localhost:52415/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 20}'
# Claude Messages API
curl -X POST http://localhost:52415/v1/messages \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "max_tokens": 50, "messages": [{"role": "user", "content": "Hello"}]}'
# OpenAI Responses API
curl -X POST http://localhost:52415/v1/responses \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "input": "Hello", "max_output_tokens": 20}'
```
**Streaming tests:**
```bash
# Chat Completions API (streaming)
curl -N -X POST http://localhost:52415/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "messages": [{"role": "user", "content": "Hello"}], "stream": true, "max_tokens": 20}'
# Claude Messages API (streaming)
curl -N -X POST http://localhost:52415/v1/messages \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "max_tokens": 50, "messages": [{"role": "user", "content": "Hello"}], "stream": true}'
# OpenAI Responses API (streaming)
curl -N -X POST http://localhost:52415/v1/responses \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "input": "Hello", "stream": true, "max_output_tokens": 20}'
```
All endpoints tested successfully with proper response formats and
streaming events.
### Automated Testing
- Tests in `src/exo/master/tests/` all pass (85 tests)
- Type checker (basedpyright) passes with 0 errors
- Linter (ruff) passes
---------
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Evan <evanev7@gmail.com>
Files touched
M app/EXO/EXO/Models/ClusterState.swiftA src/exo/master/adapters/__init__.pyA src/exo/master/adapters/chat_completions.pyA src/exo/master/adapters/claude.pyA src/exo/master/adapters/responses.pyM src/exo/master/api.pyM src/exo/master/main.pyA src/exo/master/tests/test_claude_api.pyA src/exo/master/tests/test_claude_tool_use.pyM src/exo/master/tests/test_master.pyA src/exo/master/tests/test_openai_responses_api.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/types/api.pyA src/exo/shared/types/claude_api.pyM src/exo/shared/types/commands.pyA src/exo/shared/types/openai_responses.pyM src/exo/shared/types/tasks.pyA src/exo/shared/types/text_generation.pyM src/exo/utils/info_gatherer/info_gatherer.pyM src/exo/worker/engines/image/generate.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.pyM src/exo/worker/tests/unittests/test_plan/test_task_forwarding.pyM src/exo/worker/tests/unittests/test_runner/test_event_ordering.pyA src/exo/worker/tests/unittests/test_runner/test_glm_tool_parsing.pyA src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.pyM tests/headless_runner.py
Diff
commit c3537980bd1c8330dcee7455f76b98c03f8f69f1
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Mon Feb 2 07:58:37 2026 -0800
feat: add Claude Messages API and OpenAI Responses API support (#1167)
## Motivation
Add support for Claude Messages API and OpenAI Responses API to allow
users to interact with exo using these popular API formats. This enables
broader compatibility with existing tooling and SDKs that expect these
API formats.
## Architecture
Adapter logic lives exclusively in the API layer
(`src/exo/master/adapters/`). On the way in, each adapter converts its
API-specific request type (`ChatCompletionRequest`,
`ClaudeMessagesRequest`, `ResponsesRequest`) into
`TextGenerationTaskParams`. On the way out, each adapter converts the
`TokenChunk` stream back into its API-specific response format.
Everything inside the application — commands, worker, runner, event
sourcing — only sees `TextGenerationTaskParams` and `TokenChunk`. No
API-specific types cross the boundary.
```
API layer │ Application internals
│
Chat Completions → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ChatCompletionResponse
Claude Messages → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ClaudeMessagesResponse
Responses API → [adapter] → TextGenerationTaskParams ──→ │ ──→ TextGeneration command → Runner → TokenChunk ──→ │ ──→ [adapter] → ResponsesResponse
```
## Changes
### New Files
- `src/exo/shared/types/claude_api.py` - Pydantic types for Claude
Messages API
- `src/exo/shared/types/openai_responses.py` - Pydantic types for OpenAI
Responses API
- `src/exo/shared/types/text_generation.py` - Shared
`TextGenerationTaskParams` internal type
- `src/exo/master/adapters/chat_completions.py` - Chat Completions
adapter (streaming/non-streaming)
- `src/exo/master/adapters/claude.py` - Claude Messages adapter
(streaming/non-streaming)
- `src/exo/master/adapters/responses.py` - OpenAI Responses adapter
(streaming/non-streaming)
### Modified Files
- `src/exo/master/api.py` - Refactored to use adapters uniformly for all
endpoints; extracted `_resolve_and_validate_text_model` helper to
deduplicate model validation across all text endpoints; removed ad-hoc
`try/except ValueError` blocks from non-streaming paths
### New Endpoints
- `POST /v1/messages` - Claude Messages API (streaming and
non-streaming)
- `POST /v1/responses` - OpenAI Responses API (streaming and
non-streaming)
## Why It Works
All APIs are implemented as pure conversion adapters at the edge of the
application:
1. Adapter functions in `src/exo/master/adapters/` convert incoming
requests to `TextGenerationTaskParams`
2. `api.py` wraps the params in a `TextGeneration` command and sends it
through the existing command/event flow
3. The worker, runner, and event sourcing layers only handle
`TextGenerationTaskParams` and `TokenChunk` — they have no awareness of
Chat Completions, Claude, or Responses API formats
4. On response, adapter functions convert the `TokenChunk` stream back
to the caller's expected format
5. Model validation is handled by a single shared helper
(`_resolve_and_validate_text_model`), mirroring the existing
`_validate_image_model` pattern for image endpoints
No changes to core inference logic were needed.
### Streaming Formats
- **Chat Completions**: Uses `data: {...}\n\n` with `[DONE]` terminator
- **Claude**: Uses event types `message_start`, `content_block_start`,
`content_block_delta`, `content_block_stop`, `message_delta`,
`message_stop`
- **OpenAI Responses**: Uses event types `response.created`,
`response.in_progress`, `response.output_item.added`,
`response.content_part.added`, `response.output_text.delta`,
`response.output_text.done`, `response.content_part.done`,
`response.output_item.done`, `response.completed`
## Test Plan
### Manual Testing
Hardware: MacBook Pro M3 Max
**Non-streaming tests:**
```bash
# Chat Completions API
curl -X POST http://localhost:52415/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 20}'
# Claude Messages API
curl -X POST http://localhost:52415/v1/messages \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "max_tokens": 50, "messages": [{"role": "user", "content": "Hello"}]}'
# OpenAI Responses API
curl -X POST http://localhost:52415/v1/responses \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "input": "Hello", "max_output_tokens": 20}'
```
**Streaming tests:**
```bash
# Chat Completions API (streaming)
curl -N -X POST http://localhost:52415/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "messages": [{"role": "user", "content": "Hello"}], "stream": true, "max_tokens": 20}'
# Claude Messages API (streaming)
curl -N -X POST http://localhost:52415/v1/messages \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "max_tokens": 50, "messages": [{"role": "user", "content": "Hello"}], "stream": true}'
# OpenAI Responses API (streaming)
curl -N -X POST http://localhost:52415/v1/responses \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.2-1b", "input": "Hello", "stream": true, "max_output_tokens": 20}'
```
All endpoints tested successfully with proper response formats and
streaming events.
### Automated Testing
- Tests in `src/exo/master/tests/` all pass (85 tests)
- Type checker (basedpyright) passes with 0 errors
- Linter (ruff) passes
---------
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Evan <evanev7@gmail.com>
---
app/EXO/EXO/Models/ClusterState.swift | 8 +-
src/exo/master/adapters/__init__.py | 1 +
src/exo/master/adapters/chat_completions.py | 212 ++++++++++++
src/exo/master/adapters/claude.py | 321 ++++++++++++++++++
src/exo/master/adapters/responses.py | 369 +++++++++++++++++++++
src/exo/master/api.py | 345 ++++++++++---------
src/exo/master/main.py | 30 +-
src/exo/master/tests/test_claude_api.py | 182 ++++++++++
src/exo/master/tests/test_claude_tool_use.py | 265 +++++++++++++++
src/exo/master/tests/test_master.py | 34 +-
src/exo/master/tests/test_openai_responses_api.py | 48 +++
src/exo/shared/models/model_cards.py | 2 +-
src/exo/shared/types/api.py | 35 +-
src/exo/shared/types/claude_api.py | 214 ++++++++++++
src/exo/shared/types/commands.py | 15 +-
src/exo/shared/types/openai_responses.py | 296 +++++++++++++++++
src/exo/shared/types/tasks.py | 13 +-
src/exo/shared/types/text_generation.py | 42 +++
src/exo/utils/info_gatherer/info_gatherer.py | 3 +-
src/exo/worker/engines/image/generate.py | 6 +-
src/exo/worker/engines/mlx/generator/generate.py | 77 +++--
src/exo/worker/engines/mlx/utils_mlx.py | 65 ++--
src/exo/worker/main.py | 4 +-
src/exo/worker/plan.py | 4 +-
src/exo/worker/runner/runner.py | 34 +-
.../worker/tests/unittests/test_mlx/conftest.py | 15 +-
.../unittests/test_mlx/test_kv_prefix_cache.py | 89 +++--
.../unittests/test_plan/test_task_forwarding.py | 28 +-
.../unittests/test_runner/test_event_ordering.py | 15 +-
.../unittests/test_runner/test_glm_tool_parsing.py | 110 ++++++
.../unittests/test_runner/test_parse_tool_calls.py | 87 +++++
tests/headless_runner.py | 18 +-
32 files changed, 2556 insertions(+), 431 deletions(-)
diff --git a/app/EXO/EXO/Models/ClusterState.swift b/app/EXO/EXO/Models/ClusterState.swift
index 33e6aee4..ee4c3816 100644
--- a/app/EXO/EXO/Models/ClusterState.swift
+++ b/app/EXO/EXO/Models/ClusterState.swift
@@ -293,7 +293,7 @@ struct ClusterTask {
let modelName: String?
let promptPreview: String?
let errorMessage: String?
- let parameters: ChatCompletionTaskParameters?
+ let parameters: TextGenerationTaskParameters?
var sortPriority: Int {
switch status {
@@ -330,12 +330,12 @@ struct ClusterTaskPayload: Decodable {
let taskStatus: TaskStatus?
let instanceId: String?
let commandId: String?
- let taskParams: ChatCompletionTaskParameters?
+ let taskParams: TextGenerationTaskParameters?
let errorType: String?
let errorMessage: String?
}
-struct ChatCompletionTaskParameters: Decodable, Equatable {
+struct TextGenerationTaskParameters: Decodable, Equatable {
let model: String?
let messages: [ChatCompletionMessage]?
let maxTokens: Int?
@@ -374,7 +374,7 @@ extension ClusterTask {
guard let id = payload.taskId else { return nil }
let status = payload.taskStatus ?? .unknown
switch kindKey {
- case "ChatCompletion":
+ case "TextGeneration":
self.init(
id: id,
status: status,
diff --git a/src/exo/master/adapters/__init__.py b/src/exo/master/adapters/__init__.py
new file mode 100644
index 00000000..60463871
--- /dev/null
+++ b/src/exo/master/adapters/__init__.py
@@ -0,0 +1 @@
+"""API adapters for different API formats (Claude, OpenAI Responses, etc.)."""
diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py
new file mode 100644
index 00000000..5a27664d
--- /dev/null
+++ b/src/exo/master/adapters/chat_completions.py
@@ -0,0 +1,212 @@
+"""OpenAI Chat Completions API adapter for converting requests/responses."""
+
+import time
+from collections.abc import AsyncGenerator
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.api import (
+ ChatCompletionChoice,
+ ChatCompletionMessage,
+ ChatCompletionMessageText,
+ ChatCompletionRequest,
+ ChatCompletionResponse,
+ ErrorInfo,
+ ErrorResponse,
+ FinishReason,
+ StreamingChoiceResponse,
+ ToolCall,
+)
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def chat_request_to_text_generation(
+ request: ChatCompletionRequest,
+) -> TextGenerationTaskParams:
+ instructions: str | None = None
+ input_messages: list[InputMessage] = []
+ chat_template_messages: list[dict[str, Any]] = []
+
+ for msg in request.messages:
+ # Normalize content to string
+ content: str
+ if msg.content is None:
+ content = ""
+ elif isinstance(msg.content, str):
+ content = msg.content
+ elif isinstance(msg.content, ChatCompletionMessageText):
+ content = msg.content.text
+ else:
+ # List of ChatCompletionMessageText
+ content = "\n".join(item.text for item in msg.content)
+
+ # Extract system message as instructions
+ if msg.role == "system":
+ if instructions is None:
+ instructions = content
+ else:
+ # Append additional system messages
+ instructions = f"{instructions}\n{content}"
+ chat_template_messages.append({"role": "system", "content": content})
+ else:
+ # Skip messages with no meaningful content
+ if msg.content is None and msg.thinking is None and msg.tool_calls is None:
+ continue
+
+ if msg.role in ("user", "assistant", "developer"):
+ input_messages.append(InputMessage(role=msg.role, content=content))
+
+ # Build full message dict for chat template (preserves tool_calls etc.)
+ # Normalize content for model_dump
+ msg_copy = msg.model_copy(update={"content": content})
+ dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
+ chat_template_messages.append(dumped)
+
+ return TextGenerationTaskParams(
+ model=request.model,
+ input=input_messages if input_messages else "",
+ instructions=instructions,
+ max_output_tokens=request.max_tokens,
+ temperature=request.temperature,
+ top_p=request.top_p,
+ top_k=request.top_k,
+ stop=request.stop,
+ seed=request.seed,
+ stream=request.stream,
+ tools=request.tools,
+ chat_template_messages=chat_template_messages
+ if chat_template_messages
+ else None,
+ )
+
+
+def chunk_to_response(
+ chunk: TokenChunk, command_id: CommandId
+) -> ChatCompletionResponse:
+ """Convert a TokenChunk to a streaming ChatCompletionResponse."""
+ return ChatCompletionResponse(
+ id=command_id,
+ created=int(time.time()),
+ model=chunk.model,
+ choices=[
+ StreamingChoiceResponse(
+ index=0,
+ delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+ finish_reason=chunk.finish_reason,
+ )
+ ],
+ )
+
+
+async def generate_chat_stream(
+ command_id: CommandId,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+ """Generate Chat Completions API streaming events from chunks."""
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ error_response = ErrorResponse(
+ error=ErrorInfo(
+ message=chunk.error_message or "Internal server error",
+ type="InternalServerError",
+ code=500,
+ )
+ )
+ yield f"data: {error_response.model_dump_json()}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+
+ if isinstance(chunk, ToolCallChunk):
+ tool_call_deltas = [
+ ToolCall(
+ id=str(uuid4()),
+ index=i,
+ function=tool,
+ )
+ for i, tool in enumerate(chunk.tool_calls)
+ ]
+ tool_response = ChatCompletionResponse(
+ id=command_id,
+ created=int(time.time()),
+ model=chunk.model,
+ choices=[
+ StreamingChoiceResponse(
+ index=0,
+ delta=ChatCompletionMessage(
+ role="assistant",
+ tool_calls=tool_call_deltas,
+ ),
+ finish_reason="tool_calls",
+ )
+ ],
+ )
+ yield f"data: {tool_response.model_dump_json()}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+
+ chunk_response = chunk_to_response(chunk, command_id)
+ yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+ if chunk.finish_reason is not None:
+ yield "data: [DONE]\n\n"
+
+
+async def collect_chat_response(
+ command_id: CommandId,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ChatCompletionResponse:
+ """Collect all token chunks and return a single ChatCompletionResponse."""
+ text_parts: list[str] = []
+ tool_calls: list[ToolCall] = []
+ model: str | None = None
+ finish_reason: FinishReason | None = None
+ error_message: str | None = None
+
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ error_message = chunk.error_message or "Internal server error"
+ break
+
+ if model is None:
+ model = chunk.model
+
+ if isinstance(chunk, TokenChunk):
+ text_parts.append(chunk.text)
+
+ if isinstance(chunk, ToolCallChunk):
+ tool_calls.extend(
+ ToolCall(
+ id=str(uuid4()),
+ index=i,
+ function=tool,
+ )
+ for i, tool in enumerate(chunk.tool_calls)
+ )
+
+ if chunk.finish_reason is not None:
+ finish_reason = chunk.finish_reason
+
+ if error_message is not None:
+ raise ValueError(error_message)
+
+ combined_text = "".join(text_parts)
+ assert model is not None
+
+ return ChatCompletionResponse(
+ id=command_id,
+ created=int(time.time()),
+ model=model,
+ choices=[
+ ChatCompletionChoice(
+ index=0,
+ message=ChatCompletionMessage(
+ role="assistant",
+ content=combined_text,
+ tool_calls=tool_calls if tool_calls else None,
+ ),
+ finish_reason=finish_reason,
+ )
+ ],
+ )
diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py
new file mode 100644
index 00000000..13398012
--- /dev/null
+++ b/src/exo/master/adapters/claude.py
@@ -0,0 +1,321 @@
+"""Claude Messages API adapter for converting requests/responses."""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.api import FinishReason
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.claude_api import (
+ ClaudeContentBlock,
+ ClaudeContentBlockDeltaEvent,
+ ClaudeContentBlockStartEvent,
+ ClaudeContentBlockStopEvent,
+ ClaudeInputJsonDelta,
+ ClaudeMessageDelta,
+ ClaudeMessageDeltaEvent,
+ ClaudeMessageDeltaUsage,
+ ClaudeMessagesRequest,
+ ClaudeMessagesResponse,
+ ClaudeMessageStart,
+ ClaudeMessageStartEvent,
+ ClaudeMessageStopEvent,
+ ClaudeStopReason,
+ ClaudeTextBlock,
+ ClaudeTextDelta,
+ ClaudeToolResultBlock,
+ ClaudeToolUseBlock,
+ ClaudeUsage,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def finish_reason_to_claude_stop_reason(
+ finish_reason: FinishReason | None,
+) -> ClaudeStopReason | None:
+ """Map OpenAI finish_reason to Claude stop_reason."""
+ if finish_reason is None:
+ return None
+ mapping: dict[FinishReason, ClaudeStopReason] = {
+ "stop": "end_turn",
+ "length": "max_tokens",
+ "tool_calls": "tool_use",
+ "content_filter": "end_turn",
+ "function_call": "tool_use",
+ }
+ return mapping.get(finish_reason, "end_turn")
+
+
+def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
+ """Extract plain text from a tool_result content field."""
+ if block.content is None:
+ return ""
+ if isinstance(block.content, str):
+ return block.content
+ return "".join(sub_block.text for sub_block in block.content)
+
+
+def claude_request_to_text_generation(
+ request: ClaudeMessagesRequest,
+) -> TextGenerationTaskParams:
+ # Handle system message
+ instructions: str | None = None
+ chat_template_messages: list[dict[str, Any]] = []
+
+ if request.system:
+ if isinstance(request.system, str):
+ instructions = request.system
+ else:
+ instructions = "".join(block.text for block in request.system)
+ chat_template_messages.append({"role": "system", "content": 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})
+ continue
+
+ # Process structured content blocks
+ text_parts: list[str] = []
+ tool_calls: list[dict[str, Any]] = []
+ tool_results: list[ClaudeToolResultBlock] = []
+
+ for block in msg.content:
+ if isinstance(block, ClaudeTextBlock):
+ text_parts.append(block.text)
+ elif isinstance(block, ClaudeToolUseBlock):
+ tool_calls.append(
+ {
+ "id": block.id,
+ "type": "function",
+ "function": {
+ "name": block.name,
+ "arguments": json.dumps(block.input),
+ },
+ }
+ )
+ elif isinstance(block, ClaudeToolResultBlock):
+ tool_results.append(block)
+
+ content = "".join(text_parts)
+
+ # Build InputMessage from text content
+ if msg.role in ("user", "assistant"):
+ input_messages.append(InputMessage(role=msg.role, content=content))
+
+ # Build chat_template_messages preserving tool structure
+ if tool_calls:
+ chat_template_messages.append(
+ {"role": "assistant", "content": content, "tool_calls": tool_calls}
+ )
+ elif tool_results:
+ for tr in tool_results:
+ chat_template_messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": tr.tool_use_id,
+ "content": _extract_tool_result_text(tr),
+ }
+ )
+ else:
+ chat_template_messages.append({"role": msg.role, "content": content})
+
+ # Convert Claude tool definitions to OpenAI-style function tools
+ tools: list[dict[str, Any]] | None = None
+ if request.tools:
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": tool.name,
+ "description": tool.description or "",
+ "parameters": tool.input_schema,
+ },
+ }
+ for tool in request.tools
+ ]
+
+ return TextGenerationTaskParams(
+ model=request.model,
+ input=input_messages if input_messages else "",
+ instructions=instructions,
+ max_output_tokens=request.max_tokens,
+ temperature=request.temperature,
+ top_p=request.top_p,
+ top_k=request.top_k,
+ stop=request.stop_sequences,
+ stream=request.stream,
+ tools=tools,
+ chat_template_messages=chat_template_messages
+ if chat_template_messages
+ else None,
+ )
+
+
+async def collect_claude_response(
+ command_id: CommandId,
+ model: str,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ClaudeMessagesResponse:
+ """Collect all token chunks and return a single ClaudeMessagesResponse."""
+ text_parts: list[str] = []
+ tool_use_blocks: list[ClaudeToolUseBlock] = []
+ stop_reason: ClaudeStopReason | None = None
+ last_stats = None
+ error_message: str | None = None
+
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ error_message = chunk.error_message or "Internal server error"
+ break
+
+ if isinstance(chunk, ToolCallChunk):
+ for tool in chunk.tool_calls:
+ tool_use_blocks.append(
+ ClaudeToolUseBlock(
+ id=f"toolu_{uuid4().hex[:24]}",
+ name=tool.name,
+ input=json.loads(tool.arguments), # pyright: ignore[reportAny]
+ )
+ )
+ last_stats = chunk.stats or last_stats
+ stop_reason = "tool_use"
+ continue
+
+ text_parts.append(chunk.text)
+ last_stats = chunk.stats or last_stats
+
+ if chunk.finish_reason is not None:
+ stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+ if error_message is not None:
+ raise ValueError(error_message)
+
+ combined_text = "".join(text_parts)
+
+ # Build content blocks
+ content: list[ClaudeContentBlock] = []
+ if combined_text:
+ content.append(ClaudeTextBlock(text=combined_text))
+ content.extend(tool_use_blocks)
+
+ # If no content at all, include empty text block
+ if not content:
+ content.append(ClaudeTextBlock(text=""))
+
+ # Use actual usage data from stats if available
+ input_tokens = last_stats.prompt_tokens if last_stats else 0
+ output_tokens = last_stats.generation_tokens if last_stats else 0
+
+ return ClaudeMessagesResponse(
+ id=f"msg_{command_id}",
+ model=model,
+ content=content,
+ stop_reason=stop_reason,
+ usage=ClaudeUsage(
+ input_tokens=input_tokens,
+ output_tokens=output_tokens,
+ ),
+ )
+
+
+async def generate_claude_stream(
+ command_id: CommandId,
+ model: str,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+ """Generate Claude Messages API streaming events from TokenChunks."""
+ # Initial message_start event
+ initial_message = ClaudeMessageStart(
+ id=f"msg_{command_id}",
+ model=model,
+ content=[],
+ stop_reason=None,
+ usage=ClaudeUsage(input_tokens=0, output_tokens=0),
+ )
+ start_event = ClaudeMessageStartEvent(message=initial_message)
+ yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
+
+ # content_block_start for text block at index 0
+ block_start = ClaudeContentBlockStartEvent(
+ index=0, content_block=ClaudeTextBlock(text="")
+ )
+ yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
+
+ output_tokens = 0
+ stop_reason: ClaudeStopReason | None = None
+ last_stats = None
+ next_block_index = 1 # text block is 0, tool blocks start at 1
+
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ # Close text block and bail
+ break
+
+ if isinstance(chunk, ToolCallChunk):
+ last_stats = chunk.stats or last_stats
+ stop_reason = "tool_use"
+
+ # Emit tool_use content blocks
+ for tool in chunk.tool_calls:
+ tool_id = f"toolu_{uuid4().hex[:24]}"
+ tool_input_json = tool.arguments
+
+ # content_block_start for tool_use
+ tool_block_start = ClaudeContentBlockStartEvent(
+ index=next_block_index,
+ content_block=ClaudeToolUseBlock(
+ id=tool_id, name=tool.name, input={}
+ ),
+ )
+ yield f"event: content_block_start\ndata: {tool_block_start.model_dump_json()}\n\n"
+
+ # content_block_delta with input_json_delta
+ tool_delta_event = ClaudeContentBlockDeltaEvent(
+ index=next_block_index,
+ delta=ClaudeInputJsonDelta(partial_json=tool_input_json),
+ )
+ yield f"event: content_block_delta\ndata: {tool_delta_event.model_dump_json()}\n\n"
+
+ # content_block_stop
+ tool_block_stop = ClaudeContentBlockStopEvent(index=next_block_index)
+ yield f"event: content_block_stop\ndata: {tool_block_stop.model_dump_json()}\n\n"
+
+ next_block_index += 1
+ continue
+
+ output_tokens += 1 # Count each chunk as one token
+ last_stats = chunk.stats or last_stats
+
+ # content_block_delta
+ delta_event = ClaudeContentBlockDeltaEvent(
+ index=0,
+ delta=ClaudeTextDelta(text=chunk.text),
+ )
+ yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+ if chunk.finish_reason is not None:
+ stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
+
+ # Use actual token count from stats if available
+ if last_stats is not None:
+ output_tokens = last_stats.generation_tokens
+
+ # content_block_stop for text block
+ block_stop = ClaudeContentBlockStopEvent(index=0)
+ yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
+
+ # message_delta
+ message_delta = ClaudeMessageDeltaEvent(
+ delta=ClaudeMessageDelta(stop_reason=stop_reason),
+ usage=ClaudeMessageDeltaUsage(output_tokens=output_tokens),
+ )
+ yield f"event: message_delta\ndata: {message_delta.model_dump_json()}\n\n"
+
+ # message_stop
+ message_stop = ClaudeMessageStopEvent()
+ yield f"event: message_stop\ndata: {message_stop.model_dump_json()}\n\n"
diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py
new file mode 100644
index 00000000..27d845cd
--- /dev/null
+++ b/src/exo/master/adapters/responses.py
@@ -0,0 +1,369 @@
+"""OpenAI Responses API adapter for converting requests/responses."""
+
+from collections.abc import AsyncGenerator
+from itertools import count
+from typing import Any
+from uuid import uuid4
+
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import (
+ FunctionCallInputItem,
+ ResponseCompletedEvent,
+ ResponseContentPart,
+ ResponseContentPartAddedEvent,
+ ResponseContentPartDoneEvent,
+ ResponseCreatedEvent,
+ ResponseFunctionCallArgumentsDeltaEvent,
+ ResponseFunctionCallArgumentsDoneEvent,
+ ResponseFunctionCallItem,
+ ResponseInProgressEvent,
+ ResponseInputMessage,
+ ResponseItem,
+ ResponseMessageItem,
+ ResponseOutputItemAddedEvent,
+ ResponseOutputItemDoneEvent,
+ ResponseOutputText,
+ ResponsesRequest,
+ ResponsesResponse,
+ ResponseTextDeltaEvent,
+ ResponseTextDoneEvent,
+ ResponseUsage,
+)
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def _extract_content(content: str | list[ResponseContentPart]) -> str:
+ """Extract plain text from a content field that may be a string or list of parts."""
+ if isinstance(content, str):
+ return content
+ return "".join(part.text for part in content)
+
+
+def responses_request_to_text_generation(
+ request: ResponsesRequest,
+) -> TextGenerationTaskParams:
+ input_value: str | list[InputMessage]
+ built_chat_template: list[dict[str, Any]] | None = None
+ if isinstance(request.input, str):
+ input_value = request.input
+ else:
+ input_messages: list[InputMessage] = []
+ chat_template_messages: list[dict[str, Any]] = []
+
+ if request.instructions is not None:
+ chat_template_messages.append(
+ {"role": "system", "content": request.instructions}
+ )
+
+ for item in request.input:
+ if isinstance(item, ResponseInputMessage):
+ content = _extract_content(item.content)
+ if item.role in ("user", "assistant", "developer"):
+ input_messages.append(InputMessage(role=item.role, content=content))
+ if item.role == "system":
+ chat_template_messages.append(
+ {"role": "system", "content": content}
+ )
+ else:
+ chat_template_messages.append(
+ {"role": item.role, "content": content}
+ )
+ elif isinstance(item, FunctionCallInputItem):
+ chat_template_messages.append(
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": item.call_id,
+ "type": "function",
+ "function": {
+ "name": item.name,
+ "arguments": item.arguments,
+ },
+ }
+ ],
+ }
+ )
+ else:
+ chat_template_messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": item.call_id,
+ "content": item.output,
+ }
+ )
+
+ input_value = input_messages if input_messages else ""
+ built_chat_template = chat_template_messages if chat_template_messages else None
+
+ return TextGenerationTaskParams(
+ model=request.model,
+ input=input_value,
+ instructions=request.instructions,
+ max_output_tokens=request.max_output_tokens,
+ temperature=request.temperature,
+ top_p=request.top_p,
+ stream=request.stream,
+ tools=request.tools,
+ top_k=request.top_k,
+ stop=request.stop,
+ seed=request.seed,
+ chat_template_messages=built_chat_template or request.chat_template_messages,
+ )
+
+
+async def collect_responses_response(
+ command_id: CommandId,
+ model: str,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> ResponsesResponse:
+ """Collect all token chunks and return a single ResponsesResponse."""
+ response_id = f"resp_{command_id}"
+ item_id = f"item_{command_id}"
+ accumulated_text = ""
+ function_call_items: list[ResponseFunctionCallItem] = []
+ last_stats = None
+ error_message: str | None = None
+
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ error_message = chunk.error_message or "Internal server error"
+ break
+
+ if isinstance(chunk, ToolCallChunk):
+ for tool in chunk.tool_calls:
+ function_call_items.append(
+ ResponseFunctionCallItem(
+ id=f"fc_{uuid4().hex[:24]}",
+ call_id=f"call_{uuid4().hex[:24]}",
+ name=tool.name,
+ arguments=tool.arguments,
+ )
+ )
+ last_stats = chunk.stats or last_stats
+ continue
+
+ accumulated_text += chunk.text
+ last_stats = chunk.stats or last_stats
+
+ if error_message is not None:
+ raise ValueError(error_message)
+
+ # Create usage from stats if available
+ usage = None
+ if last_stats is not None:
+ usage = ResponseUsage(
+ input_tokens=last_stats.prompt_tokens,
+ output_tokens=last_stats.generation_tokens,
+ total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+ )
+
+ output: list[ResponseItem] = [
+ ResponseMessageItem(
+ id=item_id,
+ content=[ResponseOutputText(text=accumulated_text)],
+ status="completed",
+ )
+ ]
+ output.extend(function_call_items)
+
+ return ResponsesResponse(
+ id=response_id,
+ model=model,
+ status="completed",
+ output=output,
+ output_text=accumulated_text,
+ usage=usage,
+ )
+
+
+async def generate_responses_stream(
+ command_id: CommandId,
+ model: str,
+ chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
+) -> AsyncGenerator[str, None]:
+ """Generate OpenAI Responses API streaming events from TokenChunks."""
+ response_id = f"resp_{command_id}"
+ item_id = f"item_{command_id}"
+ seq = count(1)
+
+ # response.created
+ initial_response = ResponsesResponse(
+ id=response_id,
+ model=model,
+ status="in_progress",
+ output=[],
+ output_text="",
+ )
+ created_event = ResponseCreatedEvent(
+ sequence_number=next(seq), response=initial_response
+ )
+ yield f"event: response.created\ndata: {created_event.model_dump_json()}\n\n"
+
+ # response.in_progress
+ in_progress_event = ResponseInProgressEvent(
+ sequence_number=next(seq), response=initial_response
+ )
+ yield f"event: response.in_progress\ndata: {in_progress_event.model_dump_json()}\n\n"
+
+ # response.output_item.added
+ initial_item = ResponseMessageItem(
+ id=item_id,
+ content=[ResponseOutputText(text="")],
+ status="in_progress",
+ )
+ item_added = ResponseOutputItemAddedEvent(
+ sequence_number=next(seq), output_index=0, item=initial_item
+ )
+ yield f"event: response.output_item.added\ndata: {item_added.model_dump_json()}\n\n"
+
+ # response.content_part.added
+ initial_part = ResponseOutputText(text="")
+ part_added = ResponseContentPartAddedEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=0,
+ content_index=0,
+ part=initial_part,
+ )
+ yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n"
+
+ accumulated_text = ""
+ function_call_items: list[ResponseFunctionCallItem] = []
+ last_stats = None
+ next_output_index = 1 # message item is at 0
+
+ async for chunk in chunk_stream:
+ if isinstance(chunk, ErrorChunk):
+ break
+
+ if isinstance(chunk, ToolCallChunk):
+ last_stats = chunk.stats or last_stats
+ for tool in chunk.tool_calls:
+ fc_id = f"fc_{uuid4().hex[:24]}"
+ call_id = f"call_{uuid4().hex[:24]}"
+
+ # response.output_item.added for function_call
+ fc_item = ResponseFunctionCallItem(
+ id=fc_id,
+ call_id=call_id,
+ name=tool.name,
+ arguments="",
+ status="in_progress",
+ )
+ fc_added = ResponseOutputItemAddedEvent(
+ sequence_number=next(seq),
+ output_index=next_output_index,
+ item=fc_item,
+ )
+ yield f"event: response.output_item.added\ndata: {fc_added.model_dump_json()}\n\n"
+
+ # response.function_call_arguments.delta
+ args_delta = ResponseFunctionCallArgumentsDeltaEvent(
+ sequence_number=next(seq),
+ item_id=fc_id,
+ output_index=next_output_index,
+ delta=tool.arguments,
+ )
+ yield f"event: response.function_call_arguments.delta\ndata: {args_delta.model_dump_json()}\n\n"
+
+ # response.function_call_arguments.done
+ args_done = ResponseFunctionCallArgumentsDoneEvent(
+ sequence_number=next(seq),
+ item_id=fc_id,
+ output_index=next_output_index,
+ name=tool.name,
+ arguments=tool.arguments,
+ )
+ yield f"event: response.function_call_arguments.done\ndata: {args_done.model_dump_json()}\n\n"
+
+ # response.output_item.done
+ fc_done_item = ResponseFunctionCallItem(
+ id=fc_id,
+ call_id=call_id,
+ name=tool.name,
+ arguments=tool.arguments,
+ status="completed",
+ )
+ fc_item_done = ResponseOutputItemDoneEvent(
+ sequence_number=next(seq),
+ output_index=next_output_index,
+ item=fc_done_item,
+ )
+ yield f"event: response.output_item.done\ndata: {fc_item_done.model_dump_json()}\n\n"
+
+ function_call_items.append(fc_done_item)
+ next_output_index += 1
+ continue
+
+ accumulated_text += chunk.text
+ last_stats = chunk.stats or last_stats
+
+ # response.output_text.delta
+ delta_event = ResponseTextDeltaEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=0,
+ content_index=0,
+ delta=chunk.text,
+ )
+ yield f"event: response.output_text.delta\ndata: {delta_event.model_dump_json()}\n\n"
+
+ # response.output_text.done
+ text_done = ResponseTextDoneEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=0,
+ content_index=0,
+ text=accumulated_text,
+ )
+ yield f"event: response.output_text.done\ndata: {text_done.model_dump_json()}\n\n"
+
+ # response.content_part.done
+ final_part = ResponseOutputText(text=accumulated_text)
+ part_done = ResponseContentPartDoneEvent(
+ sequence_number=next(seq),
+ item_id=item_id,
+ output_index=0,
+ content_index=0,
+ part=final_part,
+ )
+ yield f"event: response.content_part.done\ndata: {part_done.model_dump_json()}\n\n"
+
+ # response.output_item.done
+ final_message_item = ResponseMessageItem(
+ id=item_id,
+ content=[ResponseOutputText(text=accumulated_text)],
+ status="completed",
+ )
+ item_done = ResponseOutputItemDoneEvent(
+ sequence_number=next(seq), output_index=0, item=final_message_item
+ )
+ yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
+
+ # Create usage from stats if available
+ usage = None
+ if last_stats is not None:
+ usage = ResponseUsage(
+ input_tokens=last_stats.prompt_tokens,
+ output_tokens=last_stats.generation_tokens,
+ total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
+ )
+
+ # response.completed
+ output: list[ResponseItem] = [final_message_item]
+ output.extend(function_call_items)
+ final_response = ResponsesResponse(
+ id=response_id,
+ model=model,
+ status="completed",
+ output=output,
+ output_text=accumulated_text,
+ usage=usage,
+ )
+ completed_event = ResponseCompletedEvent(
+ sequence_number=next(seq), response=final_response
+ )
+ yield f"event: response.completed\ndata: {completed_event.model_dump_json()}\n\n"
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index c2323618..908716d8 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -2,7 +2,7 @@ import base64
import contextlib
import json
import time
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Awaitable, Callable
from http import HTTPStatus
from typing import Annotated, Literal, cast
from uuid import uuid4
@@ -19,6 +19,21 @@ from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
from loguru import logger
+from exo.master.adapters.chat_completions import (
+ chat_request_to_text_generation,
+ collect_chat_response,
+ generate_chat_stream,
+)
+from exo.master.adapters.claude import (
+ claude_request_to_text_generation,
+ collect_claude_response,
+ generate_claude_stream,
+)
+from exo.master.adapters.responses import (
+ collect_responses_response,
+ generate_responses_stream,
+ responses_request_to_text_generation,
+)
from exo.master.image_store import ImageStore
from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
@@ -35,12 +50,13 @@ from exo.shared.models.model_cards import (
)
from exo.shared.types.api import (
AdvancedImageParams,
+ BenchChatCompletionRequest,
BenchChatCompletionResponse,
- BenchChatCompletionTaskParams,
BenchImageGenerationResponse,
BenchImageGenerationTaskParams,
ChatCompletionChoice,
ChatCompletionMessage,
+ ChatCompletionRequest,
ChatCompletionResponse,
CreateInstanceParams,
CreateInstanceResponse,
@@ -51,7 +67,7 @@ from exo.shared.types.api import (
FinishReason,
GenerationStats,
ImageData,
- ImageEditsInternalParams,
+ ImageEditsTaskParams,
ImageGenerationResponse,
ImageGenerationStats,
ImageGenerationTaskParams,
@@ -64,10 +80,7 @@ from exo.shared.types.api import (
PlacementPreviewResponse,
StartDownloadParams,
StartDownloadResponse,
- StreamingChoiceResponse,
- StreamOptions,
ToolCall,
- Usage,
)
from exo.shared.types.chunks import (
ErrorChunk,
@@ -76,8 +89,11 @@ from exo.shared.types.chunks import (
TokenChunk,
ToolCallChunk,
)
+from exo.shared.types.claude_api import (
+ ClaudeMessagesRequest,
+ ClaudeMessagesResponse,
+)
from exo.shared.types.commands import (
- ChatCompletion,
Command,
CreateInstance,
DeleteDownload,
@@ -91,6 +107,7 @@ from exo.shared.types.commands import (
SendInputChunk,
StartDownload,
TaskFinished,
+ TextGeneration,
)
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.events import (
@@ -100,8 +117,11 @@ from exo.shared.types.events import (
IndexedEvent,
)
from exo.shared.types.memory import Memory
+from exo.shared.types.openai_responses import (
+ ResponsesRequest,
+ ResponsesResponse,
+)
from exo.shared.types.state import State
-from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.banner import print_startup_banner
@@ -114,36 +134,16 @@ def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None)
return f"image/{image_format or 'png'}"
-def chunk_to_response(
- chunk: TokenChunk | ToolCallChunk,
- command_id: CommandId,
- usage: Usage | None,
-) -> ChatCompletionResponse:
- return ChatCompletionResponse(
- id=command_id,
- created=int(time.time()),
- model=chunk.model,
- choices=[
- StreamingChoiceResponse(
- index=0,
- delta=ChatCompletionMessage(role="assistant", content=chunk.text)
- if isinstance(chunk, TokenChunk)
- else ChatCompletionMessage(
- role="assistant",
- tool_calls=[
- ToolCall(
- id=str(uuid4()),
- index=i,
- function=tool,
- )
- for i, tool in enumerate(chunk.tool_calls)
- ],
- ),
- finish_reason=chunk.finish_reason,
- )
- ],
- usage=usage,
- )
+async def resolve_model_card(model_id: ModelId) -> ModelCard:
+ if model_id in MODEL_CARDS:
+ model_card = MODEL_CARDS[model_id]
+ return model_card
+
+ for card in MODEL_CARDS.values():
+ if card.model_id == ModelId(model_id):
+ return card
+
+ return await ModelCard.from_hf(model_id)
class API:
@@ -176,6 +176,15 @@ class API:
self.paused_ev: anyio.Event = anyio.Event()
self.app = FastAPI()
+
+ @self.app.middleware("http")
+ async def _log_requests( # pyright: ignore[reportUnusedFunction]
+ request: Request,
+ call_next: Callable[[Request], Awaitable[StreamingResponse]],
+ ) -> StreamingResponse:
+ logger.debug(f"API request: {request.method} {request.url.path}")
+ return await call_next(request)
+
self._setup_exception_handlers()
self._setup_cors()
self._setup_routes()
@@ -189,7 +198,7 @@ class API:
name="dashboard",
)
- self._chat_completion_queues: dict[
+ self._text_generation_queues: dict[
CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk]
] = {}
self._image_generation_queues: dict[
@@ -203,7 +212,7 @@ class API:
self.state = State()
self.session_id = new_session_id
self.event_buffer = OrderedBuffer[Event]()
- self._chat_completion_queues = {}
+ self._text_generation_queues = {}
self._image_generation_queues = {}
self.unpause(result_clock)
@@ -260,6 +269,8 @@ class API:
self.app.post("/bench/images/edits")(self.bench_image_edits)
self.app.get("/images")(self.list_images)
self.app.get("/images/{image_id}")(self.get_image)
+ self.app.post("/v1/messages", response_model=None)(self.claude_messages)
+ self.app.post("/v1/responses", response_model=None)(self.openai_responses)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(lambda: self._event_log)
self.app.post("/download/start")(self.start_download)
@@ -484,13 +495,15 @@ class API:
instance_id=instance_id,
)
- async def _chat_chunk_stream(
+ async def _token_chunk_stream(
self, command_id: CommandId
) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
- """Yield `TokenChunk`s for a given command until completion."""
+ """Yield chunks for a given command until completion.
+ This is the internal low-level stream used by all API adapters.
+ """
try:
- self._chat_completion_queues[command_id], recv = channel[
+ self._text_generation_queues[command_id], recv = channel[
ErrorChunk | ToolCallChunk | TokenChunk
]()
@@ -511,103 +524,10 @@ class API:
finally:
command = TaskFinished(finished_command_id=command_id)
await self._send(command)
- if command_id in self._chat_completion_queues:
- del self._chat_completion_queues[command_id]
-
- async def _generate_chat_stream(
- self, command_id: CommandId, stream_options: StreamOptions | None = None
- ) -> AsyncGenerator[str, None]:
- """Generate chat completion stream as JSON strings."""
- include_usage = stream_options.include_usage if stream_options else False
-
- async for chunk in self._chat_chunk_stream(command_id):
- assert not isinstance(chunk, ImageChunk)
- if chunk.finish_reason == "error":
- error_response = ErrorResponse(
- error=ErrorInfo(
- message=chunk.error_message or "Internal server error",
- type="InternalServerError",
- code=500,
- )
- )
- yield f"data: {error_response.model_dump_json()}\n\n"
- yield "data: [DONE]\n\n"
- return
-
- usage = chunk.usage if include_usage else None
-
- chunk_response: ChatCompletionResponse = chunk_to_response(
- chunk, command_id, usage=usage
- )
- logger.debug(f"chunk_response: {chunk_response}")
-
- yield f"data: {chunk_response.model_dump_json()}\n\n"
-
- if chunk.finish_reason is not None:
- yield "data: [DONE]\n\n"
-
- async def _collect_chat_completion(
- self, command_id: CommandId
- ) -> ChatCompletionResponse:
- """Collect all token chunks for a chat completion and return a single response."""
-
- text_parts: list[str] = []
- tool_calls: list[ToolCall] = []
- model: ModelId | None = None
- finish_reason: FinishReason | None = None
- usage: Usage | None = None
-
- async for chunk in self._chat_chunk_stream(command_id):
- if isinstance(chunk, ErrorChunk):
- raise HTTPException(
- status_code=500,
- detail=chunk.error_message or "Internal server error",
- )
-
- if model is None:
- model = chunk.model
-
- if isinstance(chunk, TokenChunk):
- text_parts.append(chunk.text)
-
- if isinstance(chunk, ToolCallChunk):
- tool_calls.extend(
- ToolCall(
- id=str(uuid4()),
- index=i,
- function=tool,
- )
- for i, tool in enumerate(chunk.tool_calls)
- )
-
- if chunk.usage is not None:
- usage = chunk.usage
-
- if chunk.finish_reason is not None:
- finish_reason = chunk.finish_reason
-
- combined_text = "".join(text_parts)
- assert model is not None
-
- return ChatCompletionResponse(
- id=command_id,
- created=int(time.time()),
- model=model,
- choices=[
- ChatCompletionChoice(
- index=0,
- message=ChatCompletionMessage(
- role="assistant",
- content=combined_text,
- tool_calls=tool_calls,
- ),
- finish_reason=finish_reason,
- )
- ],
- usage=usage,
- )
+ if command_id in self._text_generation_queues:
+ del self._text_generation_queues[command_id]
- async def _collect_chat_completion_with_stats(
+ async def _collect_text_generation_with_stats(
self, command_id: CommandId
) -> BenchChatCompletionResponse:
text_parts: list[str] = []
@@ -617,7 +537,7 @@ class API:
stats: GenerationStats | None = None
- async for chunk in self._chat_chunk_stream(command_id):
+ async for chunk in self._token_chunk_stream(command_id):
if chunk.finish_reason == "error":
raise HTTPException(
status_code=500,
@@ -656,7 +576,9 @@ class API:
ChatCompletionChoice(
index=0,
message=ChatCompletionMessage(
- role="assistant", content=combined_text, tool_calls=tool_calls
+ role="assistant",
+ content=combined_text,
+ tool_calls=tool_calls if tool_calls else None,
),
finish_reason=finish_reason,
)
@@ -671,55 +593,66 @@ class API:
)
async def chat_completions(
- self, payload: ChatCompletionTaskParams
+ self, payload: ChatCompletionRequest
) -> ChatCompletionResponse | StreamingResponse:
- """Handle chat completions, supporting both streaming and non-streaming responses."""
- model_card = await ModelCard.load(ModelId(payload.model))
- payload.model = model_card.model_id
-
- if not any(
- instance.shard_assignments.model_id == payload.model
- for instance in self.state.instances.values()
- ):
- await self._trigger_notify_user_to_download_model(payload.model)
- raise HTTPException(
- status_code=404, detail=f"No instance found for model {payload.model}"
- )
-
- command = ChatCompletion(
- request_params=payload,
+ """OpenAI Chat Completions API - adapter."""
+ task_params = chat_request_to_text_generation(payload)
+ resolved_model = await self._resolve_and_validate_text_model(
+ ModelId(task_params.model)
)
+ task_params = task_params.model_copy(update={"model": resolved_model})
+
+ command = TextGeneration(task_params=task_params)
await self._send(command)
+
if payload.stream:
return StreamingResponse(
- self._generate_chat_stream(command.command_id, payload.stream_options),
+ generate_chat_stream(
+ command.command_id,
+ self._token_chunk_stream(command.command_id),
+ ),
media_type="text/event-stream",
)
- return await self._collect_chat_completion(command.command_id)
+ return await collect_chat_response(
+ command.command_id,
+ self._token_chunk_stream(command.command_id),
+ )
async def bench_chat_completions(
- self, payload: BenchChatCompletionTaskParams
+ self, payload: BenchChatCompletionRequest
) -> BenchChatCompletionResponse:
- model_card = await ModelCard.load(ModelId(payload.model))
- payload.model = model_card.model_id
+ task_params = chat_request_to_text_generation(payload)
+ resolved_model = await self._resolve_and_validate_text_model(
+ ModelId(task_params.model)
+ )
+ task_params = task_params.model_copy(update={"model": resolved_model})
+
+ task_params = task_params.model_copy(update={"stream": False, "bench": True})
+
+ command = TextGeneration(task_params=task_params)
+ await self._send(command)
+
+ response = await self._collect_text_generation_with_stats(command.command_id)
+ return response
+
+ async def _resolve_and_validate_text_model(self, model: ModelId) -> ModelId:
+ """Validate a text model exists and return the resolved model ID.
+ Raises HTTPException 404 if no instance is found for the model.
+ """
+ model_card = await resolve_model_card(model)
+ resolved = model_card.model_id
if not any(
- instance.shard_assignments.model_id == payload.model
+ instance.shard_assignments.model_id == resolved
for instance in self.state.instances.values()
):
- await self._trigger_notify_user_to_download_model(payload.model)
+ await self._trigger_notify_user_to_download_model(resolved)
raise HTTPException(
- status_code=404, detail=f"No instance found for model {payload.model}"
+ status_code=404,
+ detail=f"No instance found for model {resolved}",
)
-
- payload.stream = False
-
- command = ChatCompletion(request_params=payload)
- await self._send(command)
-
- response = await self._collect_chat_completion_with_stats(command.command_id)
- return response
+ return resolved
async def _validate_image_model(self, model: ModelId) -> ModelId:
"""Validate model exists and return resolved model ID.
@@ -775,7 +708,7 @@ class API:
payload.model = await self._validate_image_model(ModelId(payload.model))
command = ImageGeneration(
- request_params=payload,
+ task_params=payload,
)
await self._send(command)
@@ -1023,7 +956,7 @@ class API:
payload.partial_images = 0
command = ImageGeneration(
- request_params=payload,
+ task_params=payload,
)
await self._send(command)
@@ -1065,7 +998,7 @@ class API:
total_chunks = len(data_chunks)
command = ImageEdits(
- request_params=ImageEditsInternalParams(
+ task_params=ImageEditsTaskParams(
image_data="",
total_input_chunks=total_chunks,
prompt=prompt,
@@ -1209,6 +1142,62 @@ class API:
response_format=response_format,
)
+ async def claude_messages(
+ self, payload: ClaudeMessagesRequest
+ ) -> ClaudeMessagesResponse | StreamingResponse:
+ """Claude Messages API - adapter."""
+ task_params = claude_request_to_text_generation(payload)
+ resolved_model = await self._resolve_and_validate_text_model(
+ ModelId(task_params.model)
+ )
+ task_params = task_params.model_copy(update={"model": resolved_model})
+
+ command = TextGeneration(task_params=task_params)
+ await self._send(command)
+
+ if payload.stream:
+ return StreamingResponse(
+ generate_claude_stream(
+ command.command_id,
+ payload.model,
+ self._token_chunk_stream(command.command_id),
+ ),
+ media_type="text/event-stream",
+ )
+
+ return await collect_claude_response(
+ command.command_id,
+ payload.model,
+ self._token_chunk_stream(command.command_id),
+ )
+
+ async def openai_responses(
+ self, payload: ResponsesRequest
+ ) -> ResponsesResponse | StreamingResponse:
+ """OpenAI Responses API."""
+ task_params = responses_request_to_text_generation(payload)
+ resolved_model = await self._resolve_and_validate_text_model(task_params.model)
+ task_params = task_params.model_copy(update={"model": resolved_model})
+
+ command = TextGeneration(task_params=task_params)
+ await self._send(command)
+
+ if payload.stream:
+ return StreamingResponse(
+ generate_responses_stream(
+ command.command_id,
+ payload.model,
+ self._token_chunk_stream(command.command_id),
+ ),
+ media_type="text/event-stream",
+ )
+
+ return await collect_responses_response(
+ command.command_id,
+ payload.model,
+ self._token_chunk_stream(command.command_id),
+ )
+
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
total_available = Memory()
@@ -1281,14 +1270,14 @@ class API:
self._image_generation_queues.pop(
event.command_id, None
)
- if queue := self._chat_completion_queues.get(
+ if queue := self._text_generation_queues.get(
event.command_id, None
):
assert not isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except BrokenResourceError:
- self._chat_completion_queues.pop(event.command_id, None)
+ self._text_generation_queues.pop(event.command_id, None)
async def _pause_on_new_election(self):
with self.election_receiver as ems:
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 407f68d4..d99a96c9 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -12,7 +12,6 @@ from exo.master.placement import (
)
from exo.shared.apply import apply
from exo.shared.types.commands import (
- ChatCompletion,
CreateInstance,
DeleteInstance,
ForwarderCommand,
@@ -23,6 +22,7 @@ from exo.shared.types.commands import (
SendInputChunk,
TaskFinished,
TestCommand,
+ TextGeneration,
)
from exo.shared.types.common import CommandId, NodeId, SessionId
from exo.shared.types.events import (
@@ -37,9 +37,6 @@ from exo.shared.types.events import (
TaskDeleted,
)
from exo.shared.types.state import State
-from exo.shared.types.tasks import (
- ChatCompletion as ChatCompletionTask,
-)
from exo.shared.types.tasks import (
ImageEdits as ImageEditsTask,
)
@@ -50,6 +47,9 @@ from exo.shared.types.tasks import (
TaskId,
TaskStatus,
)
+from exo.shared.types.tasks import (
+ TextGeneration as TextGenerationTask,
+)
from exo.shared.types.worker.instances import InstanceId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import MultiSourceBuffer
@@ -117,11 +117,11 @@ class Master:
match command:
case TestCommand():
pass
- case ChatCompletion():
+ case TextGeneration():
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
- == command.request_params.model
+ == command.task_params.model
):
task_count = sum(
1
@@ -134,7 +134,7 @@ class Master:
if not instance_task_counts:
raise ValueError(
- f"No instance found for model {command.request_params.model}"
+ f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
@@ -148,12 +148,12 @@ class Master:
generated_events.append(
TaskCreated(
task_id=task_id,
- task=ChatCompletionTask(
+ task=TextGenerationTask(
task_id=task_id,
command_id=command.command_id,
instance_id=available_instance_ids[0],
task_status=TaskStatus.Pending,
- task_params=command.request_params,
+ task_params=command.task_params,
),
)
)
@@ -163,7 +163,7 @@ class Master:
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
- == command.request_params.model
+ == command.task_params.model
):
task_count = sum(
1
@@ -176,7 +176,7 @@ class Master:
if not instance_task_counts:
raise ValueError(
- f"No instance found for model {command.request_params.model}"
+ f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
@@ -195,7 +195,7 @@ class Master:
command_id=command.command_id,
instance_id=available_instance_ids[0],
task_status=TaskStatus.Pending,
- task_params=command.request_params,
+ task_params=command.task_params,
),
)
)
@@ -205,7 +205,7 @@ class Master:
for instance in self.state.instances.values():
if (
instance.shard_assignments.model_id
- == command.request_params.model
+ == command.task_params.model
):
task_count = sum(
1
@@ -218,7 +218,7 @@ class Master:
if not instance_task_counts:
raise ValueError(
- f"No instance found for model {command.request_params.model}"
+ f"No instance found for model {command.task_params.model}"
)
available_instance_ids = sorted(
@@ -237,7 +237,7 @@ class Master:
command_id=command.command_id,
instance_id=available_instance_ids[0],
task_status=TaskStatus.Pending,
- task_params=command.request_params,
+ task_params=command.task_params,
),
)
)
diff --git a/src/exo/master/tests/test_claude_api.py b/src/exo/master/tests/test_claude_api.py
new file mode 100644
index 00000000..08b2bef7
--- /dev/null
+++ b/src/exo/master/tests/test_claude_api.py
@@ -0,0 +1,182 @@
+"""Tests for Claude Messages API conversion functions and types."""
+
+import pydantic
+import pytest
+
+from exo.master.adapters.claude import (
+ claude_request_to_text_generation,
+ finish_reason_to_claude_stop_reason,
+)
+from exo.shared.types.claude_api import (
+ ClaudeMessage,
+ ClaudeMessagesRequest,
+ ClaudeTextBlock,
+)
+from exo.shared.types.common import ModelId
+
+
+class TestFinishReasonToClaudeStopReason:
+ """Tests for finish_reason to Claude stop_reason mapping."""
+
+ def test_stop_maps_to_end_turn(self):
+ assert finish_reason_to_claude_stop_reason("stop") == "end_turn"
+
+ def test_length_maps_to_max_tokens(self):
+ assert finish_reason_to_claude_stop_reason("length") == "max_tokens"
+
+ def test_tool_calls_maps_to_tool_use(self):
+ assert finish_reason_to_claude_stop_reason("tool_calls") == "tool_use"
+
+ def test_function_call_maps_to_tool_use(self):
+ assert finish_reason_to_claude_stop_reason("function_call") == "tool_use"
+
+ def test_content_filter_maps_to_end_turn(self):
+ assert finish_reason_to_claude_stop_reason("content_filter") == "end_turn"
+
+ def test_none_returns_none(self):
+ assert finish_reason_to_claude_stop_reason(None) is None
+
+
+class TestClaudeRequestToInternal:
+ """Tests for converting Claude Messages API requests to TextGenerationTaskParams."""
+
+ def test_basic_request_conversion(self):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ messages=[
+ ClaudeMessage(role="user", content="Hello"),
+ ],
+ )
+ params = claude_request_to_text_generation(request)
+
+ assert params.model == "claude-3-opus"
+ assert params.max_output_tokens == 100
+ assert isinstance(params.input, list)
+ assert len(params.input) == 1
+ assert params.input[0].role == "user"
+ assert params.input[0].content == "Hello"
+ assert params.instructions is None
+
+ def test_request_with_system_string(self):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ system="You are a helpful assistant.",
+ messages=[
+ ClaudeMessage(role="user", content="Hello"),
+ ],
+ )
+ params = claude_request_to_text_generation(request)
+
+ assert params.instructions == "You are a helpful assistant."
+ assert isinstance(params.input, list)
+ assert len(params.input) == 1
+ assert params.input[0].role == "user"
+ assert params.input[0].content == "Hello"
+
+ def test_request_with_system_text_blocks(self):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ system=[
+ ClaudeTextBlock(text="You are helpful. "),
+ ClaudeTextBlock(text="Be concise."),
+ ],
+ messages=[
+ ClaudeMessage(role="user", content="Hello"),
+ ],
+ )
+ params = 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):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ messages=[
+ ClaudeMessage(
+ role="user",
+ content=[
+ ClaudeTextBlock(text="First part. "),
+ ClaudeTextBlock(text="Second part."),
+ ],
+ ),
+ ],
+ )
+ params = 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):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ messages=[
+ ClaudeMessage(role="user", content="Hello"),
+ ClaudeMessage(role="assistant", content="Hi there!"),
+ ClaudeMessage(role="user", content="How are you?"),
+ ],
+ )
+ params = claude_request_to_text_generation(request)
+
+ assert isinstance(params.input, list)
+ assert len(params.input) == 3
+ assert params.input[0].role == "user"
+ assert params.input[1].role == "assistant"
+ assert params.input[2].role == "user"
+
+ def test_request_with_optional_parameters(self):
+ request = ClaudeMessagesRequest(
+ model=ModelId("claude-3-opus"),
+ max_tokens=100,
+ messages=[ClaudeMessage(role="user", content="Hello")],
+ temperature=0.7,
+ top_p=0.9,
+ top_k=40,
+ stop_sequences=["STOP", "END"],
+ stream=True,
+ )
+ params = claude_request_to_text_generation(request)
+
+ assert params.temperature == 0.7
+ assert params.top_p == 0.9
+ assert params.top_k == 40
+ assert params.stop == ["STOP", "END"]
+ assert params.stream is True
+
+
+class TestClaudeMessagesRequestValidation:
+ """Tests for Claude Messages API request validation."""
+
+ def test_request_requires_model(self):
+ with pytest.raises(pydantic.ValidationError):
+ ClaudeMessagesRequest.model_validate(
+ {
+ "max_tokens": 100,
+ "messages": [{"role": "user", "content": "Hello"}],
+ }
+ )
+
+ def test_request_requires_max_tokens(self):
+ with pytest.raises(pydantic.ValidationError):
+ ClaudeMessagesRequest.model_validate(
+ {
+ "model": "claude-3-opus",
+ "messages": [{"role": "user", "content": "Hello"}],
+ }
+ )
+
+ def test_request_requires_messages(self):
+ with pytest.raises(pydantic.ValidationError):
+ ClaudeMessagesRequest.model_validate(
+ {
+ "model": "claude-3-opus",
+ "max_tokens": 100,
+ }
+ )
diff --git a/src/exo/master/tests/test_claude_tool_use.py b/src/exo/master/tests/test_claude_tool_use.py
new file mode 100644
index 00000000..58300f10
--- /dev/null
+++ b/src/exo/master/tests/test_claude_tool_use.py
@@ -0,0 +1,265 @@
+"""Tests for Claude Messages API tool_use support in the adapter."""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
+from exo.master.adapters.claude import collect_claude_response, generate_claude_stream
+from exo.shared.types.api import ToolCallItem
+from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo.shared.types.common import CommandId, ModelId
+
+
+async def _chunks_to_stream(
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk],
+) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
+ for chunk in chunks:
+ yield chunk
+
+
+MODEL = ModelId("test-model")
+COMMAND_ID = CommandId("cmd_test123")
+
+
+def _parse_sse_events(events: list[str]) -> list[dict[str, Any]]:
+ """Parse SSE event strings into JSON dicts."""
+ parsed: list[dict[str, Any]] = []
+ for event_str in events:
+ for line in event_str.strip().split("\n"):
+ if line.startswith("data: "):
+ parsed.append(cast(dict[str, Any], json.loads(line[6:])))
+ return parsed
+
+
+class TestCollectClaudeResponseToolUse:
+ """Tests for non-streaming tool_use response collection."""
+
+ async def test_tool_call_chunk_produces_tool_use_blocks(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(
+ name="get_weather",
+ arguments='{"location": "San Francisco"}',
+ )
+ ],
+ ),
+ ]
+ response = await collect_claude_response(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ )
+
+ assert response.stop_reason == "tool_use"
+ tool_blocks = [b for b in response.content if b.type == "tool_use"]
+ assert len(tool_blocks) == 1
+ block = tool_blocks[0]
+ assert block.type == "tool_use"
+ assert block.name == "get_weather"
+ assert block.input == {"location": "San Francisco"}
+ assert block.id.startswith("toolu_")
+
+ async def test_multiple_tool_calls(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(
+ name="get_weather",
+ arguments='{"location": "SF"}',
+ ),
+ ToolCallItem(
+ name="get_time",
+ arguments='{"timezone": "PST"}',
+ ),
+ ],
+ ),
+ ]
+ response = await collect_claude_response(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ )
+
+ assert response.stop_reason == "tool_use"
+ tool_blocks = [b for b in response.content if b.type == "tool_use"]
+ assert len(tool_blocks) == 2
+ assert tool_blocks[0].name == "get_weather"
+ assert tool_blocks[1].name == "get_time"
+
+ async def test_mixed_text_and_tool_use(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ TokenChunk(model=MODEL, text="Let me check ", token_id=1, usage=None),
+ TokenChunk(model=MODEL, text="the weather.", token_id=2, usage=None),
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(
+ name="get_weather",
+ arguments='{"location": "NYC"}',
+ )
+ ],
+ ),
+ ]
+ response = await collect_claude_response(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ )
+
+ assert response.stop_reason == "tool_use"
+ text_blocks = [b for b in response.content if b.type == "text"]
+ tool_blocks = [b for b in response.content if b.type == "tool_use"]
+ assert len(text_blocks) == 1
+ assert text_blocks[0].text == "Let me check the weather."
+ assert len(tool_blocks) == 1
+ assert tool_blocks[0].name == "get_weather"
+
+ async def test_no_content_produces_empty_text_block(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = []
+ response = await collect_claude_response(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ )
+ assert len(response.content) == 1
+ assert response.content[0].type == "text"
+
+
+class TestGenerateClaudeStreamToolUse:
+ """Tests for streaming tool_use event generation."""
+
+ async def test_tool_call_emits_tool_use_events(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(
+ name="get_weather",
+ arguments='{"location": "SF"}',
+ )
+ ],
+ ),
+ ]
+ events: list[str] = []
+ async for event in generate_claude_stream(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ ):
+ events.append(event)
+
+ parsed = _parse_sse_events(events)
+
+ # Find tool_use content_block_start
+ tool_starts = [
+ e
+ for e in parsed
+ if e.get("type") == "content_block_start"
+ and cast(dict[str, Any], e.get("content_block", {})).get("type")
+ == "tool_use"
+ ]
+ assert len(tool_starts) == 1
+ content_block = cast(dict[str, Any], tool_starts[0]["content_block"])
+ assert content_block["name"] == "get_weather"
+ assert content_block["input"] == {}
+ assert cast(str, content_block["id"]).startswith("toolu_")
+
+ # Find input_json_delta
+ json_deltas = [
+ e
+ for e in parsed
+ if e.get("type") == "content_block_delta"
+ and cast(dict[str, Any], e.get("delta", {})).get("type")
+ == "input_json_delta"
+ ]
+ assert len(json_deltas) == 1
+ delta = cast(dict[str, Any], json_deltas[0]["delta"])
+ assert json.loads(cast(str, delta["partial_json"])) == {"location": "SF"}
+
+ # Find message_delta with tool_use stop reason
+ msg_deltas = [e for e in parsed if e.get("type") == "message_delta"]
+ assert len(msg_deltas) == 1
+ assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use"
+
+ async def test_streaming_mixed_text_and_tool_use(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ TokenChunk(model=MODEL, text="Hello ", token_id=1, usage=None),
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(
+ name="search",
+ arguments='{"query": "test"}',
+ )
+ ],
+ ),
+ ]
+ events: list[str] = []
+ async for event in generate_claude_stream(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ ):
+ events.append(event)
+
+ parsed = _parse_sse_events(events)
+
+ # Should have text delta at index 0
+ text_deltas = [
+ e
+ for e in parsed
+ if e.get("type") == "content_block_delta"
+ and cast(dict[str, Any], e.get("delta", {})).get("type") == "text_delta"
+ ]
+ assert len(text_deltas) == 1
+ assert text_deltas[0]["index"] == 0
+ assert cast(dict[str, Any], text_deltas[0]["delta"])["text"] == "Hello "
+
+ # Tool block at index 1
+ tool_starts = [
+ e
+ for e in parsed
+ if e.get("type") == "content_block_start"
+ and cast(dict[str, Any], e.get("content_block", {})).get("type")
+ == "tool_use"
+ ]
+ assert len(tool_starts) == 1
+ assert tool_starts[0]["index"] == 1
+
+ # Stop reason should be tool_use
+ msg_deltas = [e for e in parsed if e.get("type") == "message_delta"]
+ assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use"
+
+ async def test_streaming_tool_block_stop_events(self):
+ chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [
+ ToolCallChunk(
+ model=MODEL,
+ usage=None,
+ tool_calls=[
+ ToolCallItem(name="fn1", arguments="{}"),
+ ToolCallItem(name="fn2", arguments='{"a": 1}'),
+ ],
+ ),
+ ]
+ events: list[str] = []
+ async for event in generate_claude_stream(
+ COMMAND_ID, "test-model", _chunks_to_stream(chunks)
+ ):
+ events.append(event)
+
+ parsed = _parse_sse_events(events)
+
+ # Two tool block starts (at indices 1 and 2)
+ tool_starts = [
+ e
+ for e in parsed
+ if e.get("type") == "content_block_start"
+ and cast(dict[str, Any], e.get("content_block", {})).get("type")
+ == "tool_use"
+ ]
+ assert len(tool_starts) == 2
+ assert tool_starts[0]["index"] == 1
+ assert tool_starts[1]["index"] == 2
+
+ # Two tool block stops (at indices 1 and 2), plus text block stop at 0
+ block_stops = [e for e in parsed if e.get("type") == "content_block_stop"]
+ stop_indices = [e["index"] for e in block_stops]
+ assert 0 in stop_indices
+ assert 1 in stop_indices
+ assert 2 in stop_indices
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 4796c264..d9987727 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -7,15 +7,14 @@ from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
+from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.commands import (
- ChatCompletion,
CommandId,
ForwarderCommand,
PlaceInstance,
+ TextGeneration,
)
-from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.common import ModelId, NodeId, SessionId
from exo.shared.types.events import (
ForwarderEvent,
IndexedEvent,
@@ -27,8 +26,9 @@ from exo.shared.types.memory import Memory
from exo.shared.types.profiling import (
MemoryUsage,
)
-from exo.shared.types.tasks import ChatCompletion as ChatCompletionTask
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.worker.instances import (
InstanceMeta,
MlxRingInstance,
@@ -127,20 +127,16 @@ async def test_master():
logger.info("wait for an instance")
while len(master.state.instances.keys()) == 0:
await anyio.sleep(0.001)
- logger.info("inject a ChatCompletion Command")
+ logger.info("inject a TextGeneration Command")
await command_sender.send(
ForwarderCommand(
origin=node_id,
command=(
- ChatCompletion(
+ TextGeneration(
command_id=CommandId(),
- request_params=ChatCompletionTaskParams(
- model="llama-3.2-1b",
- messages=[
- ChatCompletionMessage(
- role="user", content="Hello, how are you?"
- )
- ],
+ task_params=TextGenerationTaskParams(
+ model=ModelId("llama-3.2-1b"),
+ input="Hello, how are you?",
),
)
),
@@ -190,12 +186,10 @@ async def test_master():
assert created_instance.ephemeral_port > 0
assert isinstance(events[2].event, TaskCreated)
assert events[2].event.task.task_status == TaskStatus.Pending
- assert isinstance(events[2].event.task, ChatCompletionTask)
- assert events[2].event.task.task_params == ChatCompletionTaskParams(
- model="llama-3.2-1b",
- messages=[
- ChatCompletionMessage(role="user", content="Hello, how are you?")
- ],
+ 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?",
)
await master.shutdown()
diff --git a/src/exo/master/tests/test_openai_responses_api.py b/src/exo/master/tests/test_openai_responses_api.py
new file mode 100644
index 00000000..78e3b741
--- /dev/null
+++ b/src/exo/master/tests/test_openai_responses_api.py
@@ -0,0 +1,48 @@
+"""Tests for OpenAI Responses API wire types.
+
+ResponsesRequest is the API wire type for the Responses endpoint.
+The responses adapter converts it to TextGenerationTaskParams for the pipeline.
+"""
+
+import pydantic
+import pytest
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.openai_responses import (
+ ResponseInputMessage,
+ ResponsesRequest,
+)
+
+
+class TestResponsesRequestValidation:
+ """Tests for OpenAI Responses API request validation."""
+
+ def test_request_requires_model(self):
+ with pytest.raises(pydantic.ValidationError):
+ ResponsesRequest.model_validate(
+ {
+ "input": "Hello",
+ }
+ )
+
+ def test_request_requires_input(self):
+ with pytest.raises(pydantic.ValidationError):
+ ResponsesRequest.model_validate(
+ {
+ "model": "gpt-4o",
+ }
+ )
+
+ def test_request_accepts_string_input(self):
+ request = ResponsesRequest(
+ model=ModelId("gpt-4o"),
+ input="Hello",
+ )
+ assert request.input == "Hello"
+
+ def test_request_accepts_message_array_input(self):
+ request = ResponsesRequest(
+ model=ModelId("gpt-4o"),
+ input=[ResponseInputMessage(role="user", content="Hello")],
+ )
+ assert len(request.input) == 1
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 0add111e..483d2b52 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -270,7 +270,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
),
"qwen3-80b-a3B-thinking-4bit": ModelCard(
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
- storage_size=Memory.from_mb(84700),
+ storage_size=Memory.from_mb(44900),
n_layers=48,
hidden_size=2048,
supports_tensor=True,
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 1dac848f..3f665179 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -2,7 +2,6 @@ import time
from collections.abc import Generator
from typing import Annotated, Any, Literal
-from fastapi import UploadFile
from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
@@ -11,7 +10,7 @@ from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, ConfigDict, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel
FinishReason = Literal[
"stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -174,10 +173,8 @@ class StreamOptions(BaseModel):
include_usage: bool = False
-class ChatCompletionTaskParams(TaggedModel):
- model_config = ConfigDict(extra="ignore")
-
- model: str
+class ChatCompletionRequest(BaseModel):
+ model: ModelId
frequency_penalty: float | None = None
messages: list[ChatCompletionMessage]
logit_bias: dict[str, int] | None = None
@@ -193,13 +190,14 @@ class ChatCompletionTaskParams(TaggedModel):
stream_options: StreamOptions | None = None
temperature: float | None = None
top_p: float | None = None
+ top_k: int | None = None
tools: list[dict[str, Any]] | None = None
tool_choice: str | dict[str, Any] | None = None
parallel_tool_calls: bool | None = None
user: str | None = None
-class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
+class BenchChatCompletionRequest(ChatCompletionRequest):
pass
@@ -283,28 +281,7 @@ class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
class ImageEditsTaskParams(BaseModel):
- image: UploadFile
- prompt: str
- background: str | None = None
- input_fidelity: float | None = None
- mask: UploadFile | None = None
- model: str
- n: int | None = 1
- output_compression: int | None = None
- output_format: Literal["png", "jpeg", "webp"] = "png"
- partial_images: int | None = 0
- quality: Literal["high", "medium", "low"] | None = "medium"
- response_format: Literal["url", "b64_json"] | None = "b64_json"
- size: str | None = "1024x1024"
- stream: bool | None = False
- user: str | None = None
- advanced_params: AdvancedImageParams | None = None
- # Internal flag for benchmark mode - set by API, preserved through serialization
- bench: bool = False
-
-
-class ImageEditsInternalParams(BaseModel):
- """Serializable version of ImageEditsTaskParams for distributed task execution."""
+ """Internal task params for image-editing requests."""
image_data: str = "" # Base64-encoded image (empty when using chunked transfer)
total_input_chunks: int = 0
diff --git a/src/exo/shared/types/claude_api.py b/src/exo/shared/types/claude_api.py
new file mode 100644
index 00000000..dcf8c515
--- /dev/null
+++ b/src/exo/shared/types/claude_api.py
@@ -0,0 +1,214 @@
+"""Claude Messages API types for request/response conversion."""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from exo.shared.types.common import ModelId
+
+# Tool definition types
+ClaudeToolInputSchema = dict[str, Any]
+
+
+class ClaudeToolDefinition(BaseModel, frozen=True):
+ """Tool definition in Claude Messages API request."""
+
+ name: str
+ description: str | None = None
+ input_schema: ClaudeToolInputSchema
+
+
+# Type aliases
+ClaudeRole = Literal["user", "assistant"]
+ClaudeStopReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
+
+
+# Content block types
+class ClaudeTextBlock(BaseModel, frozen=True):
+ """Text content block in Claude Messages API."""
+
+ type: Literal["text"] = "text"
+ text: str
+
+
+class ClaudeImageSource(BaseModel, frozen=True):
+ """Image source for Claude image blocks."""
+
+ type: Literal["base64", "url"]
+ media_type: str | None = None
+ data: str | None = None
+ url: str | None = None
+
+
+class ClaudeImageBlock(BaseModel, frozen=True):
+ """Image content block in Claude Messages API."""
+
+ type: Literal["image"] = "image"
+ source: ClaudeImageSource
+
+
+class ClaudeToolUseBlock(BaseModel, frozen=True):
+ """Tool use content block in Claude Messages API."""
+
+ type: Literal["tool_use"] = "tool_use"
+ id: str
+ name: str
+ input: dict[str, Any]
+
+
+class ClaudeToolResultBlock(BaseModel, frozen=True):
+ """Tool result content block in Claude Messages API request."""
+
+ type: Literal["tool_result"] = "tool_result"
+ tool_use_id: str
+ content: str | list[ClaudeTextBlock] | None = None
+ is_error: bool | None = None
+ cache_control: dict[str, str] | None = None
+
+
+ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock
+
+# Input content blocks can also include tool_result (sent by user after tool_use)
+ClaudeInputContentBlock = (
+ ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock
+)
+
+
+# Request types
+class ClaudeMessage(BaseModel, frozen=True):
+ """Message in Claude Messages API request."""
+
+ role: ClaudeRole
+ content: str | list[ClaudeInputContentBlock]
+
+
+class ClaudeMessagesRequest(BaseModel):
+ """Request body for Claude Messages API."""
+
+ model: ModelId
+ max_tokens: int
+ messages: list[ClaudeMessage]
+ system: str | list[ClaudeTextBlock] | None = None
+ stop_sequences: list[str] | None = None
+ stream: bool = False
+ temperature: float | None = None
+ top_p: float | None = None
+ top_k: int | None = None
+ tools: list[ClaudeToolDefinition] | None = None
+ metadata: dict[str, str] | None = None
+
+
+# Response types
+class ClaudeUsage(BaseModel, frozen=True):
+ """Token usage in Claude Messages API response."""
+
+ input_tokens: int
+ output_tokens: int
+
+
+class ClaudeMessagesResponse(BaseModel, frozen=True):
+ """Response body for Claude Messages API."""
+
+ id: str
+ type: Literal["message"] = "message"
+ role: Literal["assistant"] = "assistant"
+ content: list[ClaudeContentBlock]
+ model: str
+ stop_reason: ClaudeStopReason | None = None
+ stop_sequence: str | None = None
+ usage: ClaudeUsage
+
+
+# Streaming event types
+class ClaudeMessageStart(BaseModel, frozen=True):
+ """Partial message in message_start event."""
+
+ id: str
+ type: Literal["message"] = "message"
+ role: Literal["assistant"] = "assistant"
+ content: list[ClaudeTextBlock] = Field(default_factory=list)
+ model: str
+ stop_reason: ClaudeStopReason | None = None
+ stop_sequence: str | None = None
+ usage: ClaudeUsage
+
+
+class ClaudeMessageStartEvent(BaseModel, frozen=True):
+ """Event sent at start of message stream."""
+
+ type: Literal["message_start"] = "message_start"
+ message: ClaudeMessageStart
+
+
+class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
+ """Event sent at start of a content block."""
+
+ type: Literal["content_block_start"] = "content_block_start"
+ index: int
+ content_block: ClaudeTextBlock | ClaudeToolUseBlock
+
+
+class ClaudeTextDelta(BaseModel, frozen=True):
+ """Delta for text content block."""
+
+ type: Literal["text_delta"] = "text_delta"
+ text: str
+
+
+class ClaudeInputJsonDelta(BaseModel, frozen=True):
+ """Delta for tool use input JSON content block."""
+
+ type: Literal["input_json_delta"] = "input_json_delta"
+ partial_json: str
+
+
+class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
+ """Event sent for content block delta."""
+
+ type: Literal["content_block_delta"] = "content_block_delta"
+ index: int
+ delta: ClaudeTextDelta | ClaudeInputJsonDelta
+
+
+class ClaudeContentBlockStopEvent(BaseModel, frozen=True):
+ """Event sent at end of a content block."""
+
+ type: Literal["content_block_stop"] = "content_block_stop"
+ index: int
+
+
+class ClaudeMessageDeltaUsage(BaseModel, frozen=True):
+ """Usage in message_delta event."""
+
+ output_tokens: int
+
+
+class ClaudeMessageDelta(BaseModel, frozen=True):
+ """Delta in message_delta event."""
+
+ stop_reason: ClaudeStopReason | None = None
+ stop_sequence: str | None = None
+
+
+class ClaudeMessageDeltaEvent(BaseModel, frozen=True):
+ """Event sent with final message delta."""
+
+ type: Literal["message_delta"] = "message_delta"
+ delta: ClaudeMessageDelta
+ usage: ClaudeMessageDeltaUsage
+
+
+class ClaudeMessageStopEvent(BaseModel, frozen=True):
+ """Event sent at end of message stream."""
+
+ type: Literal["message_stop"] = "message_stop"
+
+
+ClaudeStreamEvent = (
+ ClaudeMessageStartEvent
+ | ClaudeContentBlockStartEvent
+ | ClaudeContentBlockDeltaEvent
+ | ClaudeContentBlockStopEvent
+ | ClaudeMessageDeltaEvent
+ | ClaudeMessageStopEvent
+)
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index a6d4d5c1..115df719 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -2,13 +2,12 @@ from pydantic import Field
from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.api import (
- BenchChatCompletionTaskParams,
- ChatCompletionTaskParams,
- ImageEditsInternalParams,
+ ImageEditsTaskParams,
ImageGenerationTaskParams,
)
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -22,16 +21,16 @@ class TestCommand(BaseCommand):
__test__ = False
-class ChatCompletion(BaseCommand):
- request_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
+class TextGeneration(BaseCommand):
+ task_params: TextGenerationTaskParams
class ImageGeneration(BaseCommand):
- request_params: ImageGenerationTaskParams
+ task_params: ImageGenerationTaskParams
class ImageEdits(BaseCommand):
- request_params: ImageEditsInternalParams
+ task_params: ImageEditsTaskParams
class PlaceInstance(BaseCommand):
@@ -79,7 +78,7 @@ DownloadCommand = StartDownload | DeleteDownload
Command = (
TestCommand
| RequestEventLog
- | ChatCompletion
+ | TextGeneration
| ImageGeneration
| ImageEdits
| PlaceInstance
diff --git a/src/exo/shared/types/openai_responses.py b/src/exo/shared/types/openai_responses.py
new file mode 100644
index 00000000..e29fb253
--- /dev/null
+++ b/src/exo/shared/types/openai_responses.py
@@ -0,0 +1,296 @@
+"""OpenAI Responses API wire types.
+
+These types model the OpenAI Responses API request/response format.
+ResponsesRequest is the API-level wire type; for the canonical internal
+task params type used by the inference pipeline, see
+``exo.shared.types.text_generation.TextGenerationTaskParams``.
+"""
+
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+from exo.shared.types.common import ModelId
+
+# Type aliases
+ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
+ResponseRole = Literal["user", "assistant", "system", "developer"]
+
+
+# Request input content part types
+class ResponseInputTextPart(BaseModel, frozen=True):
+ """Text content part in a Responses API input message."""
+
+ type: Literal["input_text"] = "input_text"
+ text: str
+
+
+class ResponseOutputTextPart(BaseModel, frozen=True):
+ """Output text content part (used when replaying assistant messages in input)."""
+
+ type: Literal["output_text"] = "output_text"
+ text: str
+
+
+ResponseContentPart = ResponseInputTextPart | ResponseOutputTextPart
+
+
+# Request input item types
+class ResponseInputMessage(BaseModel, frozen=True):
+ """Input message for Responses API."""
+
+ role: ResponseRole
+ content: str | list[ResponseContentPart]
+ type: Literal["message"] = "message"
+
+
+class FunctionCallInputItem(BaseModel, frozen=True):
+ """Function call item replayed in input (from a previous assistant response)."""
+
+ type: Literal["function_call"] = "function_call"
+ id: str | None = None
+ call_id: str
+ name: str
+ arguments: str
+ status: ResponseStatus | None = None
+
+
+class FunctionCallOutputInputItem(BaseModel, frozen=True):
+ """Function call output item in input (user providing tool results)."""
+
+ type: Literal["function_call_output"] = "function_call_output"
+ call_id: str
+ output: str
+ id: str | None = None
+ status: ResponseStatus | None = None
+
+
+ResponseInputItem = (
+ ResponseInputMessage | FunctionCallInputItem | FunctionCallOutputInputItem
+)
+
+
+class ResponsesRequest(BaseModel, frozen=True):
+ """Request body for OpenAI Responses API.
+
+ This is the API wire type for the Responses endpoint. The canonical
+ internal task params type is ``TextGenerationTaskParams``; see the
+ ``responses_request_to_text_generation`` adapter for conversion.
+ """
+
+ # --- OpenAI Responses API standard fields ---
+ model: ModelId
+ input: str | list[ResponseInputItem]
+ instructions: str | None = None
+ max_output_tokens: int | None = None
+ temperature: float | None = None
+ top_p: float | None = None
+ stream: bool = False
+ tools: list[dict[str, Any]] | None = None
+ metadata: dict[str, str] | None = None
+
+ # --- exo extensions (not in OpenAI Responses API spec) ---
+ top_k: int | None = Field(
+ default=None,
+ description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",
+ json_schema_extra={"x-exo-extension": True},
+ )
+ stop: str | list[str] | None = Field(
+ default=None,
+ description="[exo extension] Stop sequence(s). Not part of the OpenAI Responses API.",
+ json_schema_extra={"x-exo-extension": True},
+ )
+ seed: int | None = Field(
+ default=None,
+ description="[exo extension] Seed for deterministic sampling. Not part of the OpenAI Responses API.",
+ json_schema_extra={"x-exo-extension": True},
+ )
+
+ # --- Internal fields (preserved during serialization, hidden from OpenAPI schema) ---
+ chat_template_messages: list[dict[str, Any]] | None = Field(
+ default=None,
+ description="Internal: pre-formatted messages for tokenizer chat template. Not part of the OpenAI Responses API.",
+ json_schema_extra={"x-exo-internal": True},
+ )
+
+
+# Response types
+class ResponseOutputText(BaseModel, frozen=True):
+ """Text content in response output."""
+
+ type: Literal["output_text"] = "output_text"
+ text: str
+ annotations: list[dict[str, str]] = Field(default_factory=list)
+
+
+class ResponseMessageItem(BaseModel, frozen=True):
+ """Message item in response output array."""
+
+ type: Literal["message"] = "message"
+ id: str
+ role: Literal["assistant"] = "assistant"
+ content: list[ResponseOutputText]
+ status: ResponseStatus = "completed"
+
+
+class ResponseFunctionCallItem(BaseModel, frozen=True):
+ """Function call item in response output array."""
+
+ type: Literal["function_call"] = "function_call"
+ id: str
+ call_id: str
+ name: str
+ arguments: str
+ status: ResponseStatus = "completed"
+
+
+ResponseItem = ResponseMessageItem | ResponseFunctionCallItem
+
+
+class ResponseUsage(BaseModel, frozen=True):
+ """Token usage in Responses API response."""
+
+ input_tokens: int
+ output_tokens: int
+ total_tokens: int
+
+
+class ResponsesResponse(BaseModel, frozen=True):
+ """Response body for OpenAI Responses API."""
+
+ id: str
+ object: Literal["response"] = "response"
+ created_at: int = Field(default_factory=lambda: int(time.time()))
+ status: ResponseStatus = "completed"
+ model: str
+ output: list[ResponseItem]
+ output_text: str
+ usage: ResponseUsage | None = None
+
+
+# Streaming event types
+class ResponseCreatedEvent(BaseModel, frozen=True):
+ """Event sent when response is created."""
+
+ type: Literal["response.created"] = "response.created"
+ sequence_number: int
+ response: ResponsesResponse
+
+
+class ResponseInProgressEvent(BaseModel, frozen=True):
+ """Event sent when response starts processing."""
+
+ type: Literal["response.in_progress"] = "response.in_progress"
+ sequence_number: int
+ response: ResponsesResponse
+
+
+class ResponseOutputItemAddedEvent(BaseModel, frozen=True):
+ """Event sent when an output item is added."""
+
+ type: Literal["response.output_item.added"] = "response.output_item.added"
+ sequence_number: int
+ output_index: int
+ item: ResponseItem
+
+
+class ResponseContentPartAddedEvent(BaseModel, frozen=True):
+ """Event sent when a content part is added."""
+
+ type: Literal["response.content_part.added"] = "response.content_part.added"
+ sequence_number: int
+ item_id: str
+ output_index: int
+ content_index: int
+ part: ResponseOutputText
+
+
+class ResponseTextDeltaEvent(BaseModel, frozen=True):
+ """Event sent for text delta during streaming."""
+
+ type: Literal["response.output_text.delta"] = "response.output_text.delta"
+ sequence_number: int
+ item_id: str
+ output_index: int
+ content_index: int
+ delta: str
+
+
+class ResponseTextDoneEvent(BaseModel, frozen=True):
+ """Event sent when text content is done."""
+
+ type: Literal["response.output_text.done"] = "response.output_text.done"
+ sequence_number: int
+ item_id: str
+ output_index: int
+ content_index: int
+ text: str
+
+
+class ResponseContentPartDoneEvent(BaseModel, frozen=True):
+ """Event sent when a content part is done."""
+
+ type: Literal["response.content_part.done"] = "response.content_part.done"
+ sequence_number: int
+ item_id: str
+ output_index: int
+ content_index: int
+ part: ResponseOutputText
+
+
+class ResponseOutputItemDoneEvent(BaseModel, frozen=True):
+ """Event sent when an output item is done."""
+
+ type: Literal["response.output_item.done"] = "response.output_item.done"
+ sequence_number: int
+ output_index: int
+ item: ResponseItem
+
+
+class ResponseFunctionCallArgumentsDeltaEvent(BaseModel, frozen=True):
+ """Event sent for function call arguments delta during streaming."""
+
+ type: Literal["response.function_call_arguments.delta"] = (
+ "response.function_call_arguments.delta"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ delta: str
+
+
+class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True):
+ """Event sent when function call arguments are complete."""
+
+ type: Literal["response.function_call_arguments.done"] = (
+ "response.function_call_arguments.done"
+ )
+ sequence_number: int
+ item_id: str
+ output_index: int
+ name: str
+ arguments: str
+
+
+class ResponseCompletedEvent(BaseModel, frozen=True):
+ """Event sent when response is completed."""
+
+ type: Literal["response.completed"] = "response.completed"
+ sequence_number: int
+ response: ResponsesResponse
+
+
+ResponsesStreamEvent = (
+ ResponseCreatedEvent
+ | ResponseInProgressEvent
+ | ResponseOutputItemAddedEvent
+ | ResponseContentPartAddedEvent
+ | ResponseTextDeltaEvent
+ | ResponseTextDoneEvent
+ | ResponseContentPartDoneEvent
+ | ResponseOutputItemDoneEvent
+ | ResponseFunctionCallArgumentsDeltaEvent
+ | ResponseFunctionCallArgumentsDoneEvent
+ | ResponseCompletedEvent
+)
diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py
index 4f61a59a..91fb44b1 100644
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -3,12 +3,11 @@ from enum import Enum
from pydantic import Field
from exo.shared.types.api import (
- BenchChatCompletionTaskParams,
- ChatCompletionTaskParams,
- ImageEditsInternalParams,
+ ImageEditsTaskParams,
ImageGenerationTaskParams,
)
from exo.shared.types.common import CommandId, Id
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import RunnerId
from exo.shared.types.worker.shards import ShardMetadata
@@ -53,9 +52,9 @@ class StartWarmup(BaseTask): # emitted by Worker
pass
-class ChatCompletion(BaseTask): # emitted by Master
+class TextGeneration(BaseTask): # emitted by Master
command_id: CommandId
- task_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
+ task_params: TextGenerationTaskParams
error_type: str | None = Field(default=None)
error_message: str | None = Field(default=None)
@@ -71,7 +70,7 @@ class ImageGeneration(BaseTask): # emitted by Master
class ImageEdits(BaseTask): # emitted by Master
command_id: CommandId
- task_params: ImageEditsInternalParams
+ task_params: ImageEditsTaskParams
error_type: str | None = Field(default=None)
error_message: str | None = Field(default=None)
@@ -87,7 +86,7 @@ Task = (
| ConnectToGroup
| LoadModel
| StartWarmup
- | ChatCompletion
+ | TextGeneration
| ImageGeneration
| ImageEdits
| Shutdown
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
new file mode 100644
index 00000000..b9c5565c
--- /dev/null
+++ b/src/exo/shared/types/text_generation.py
@@ -0,0 +1,42 @@
+"""Canonical internal type for text generation task parameters.
+
+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 pydantic import BaseModel
+
+from exo.shared.types.common import ModelId
+
+MessageRole = Literal["user", "assistant", "system", "developer"]
+
+
+class InputMessage(BaseModel, frozen=True):
+ """Internal message for text generation pipelines."""
+
+ role: MessageRole
+ content: str
+
+
+class TextGenerationTaskParams(BaseModel, frozen=True):
+ """Canonical internal task params for text generation.
+
+ Every API adapter converts its wire type into this before handing
+ off to the master/worker pipeline.
+ """
+
+ model: ModelId
+ input: str | list[InputMessage]
+ instructions: str | None = None
+ max_output_tokens: int | None = None
+ temperature: float | None = None
+ top_p: float | None = None
+ stream: bool = False
+ tools: list[dict[str, Any]] | None = None
+ bench: bool = False
+ top_k: int | None = None
+ stop: str | list[str] | None = None
+ seed: int | None = None
+ chat_template_messages: list[dict[str, Any]] | None = None
diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py
index 0b931eda..203a8da3 100644
--- a/src/exo/utils/info_gatherer/info_gatherer.py
+++ b/src/exo/utils/info_gatherer/info_gatherer.py
@@ -13,6 +13,7 @@ from anyio.abc import TaskGroup
from anyio.streams.buffered import BufferedByteReceiveStream
from anyio.streams.text import TextReceiveStream
from loguru import logger
+from pydantic import ValidationError
from exo.shared.constants import EXO_CONFIG_FILE
from exo.shared.types.memory import Memory
@@ -267,7 +268,7 @@ class NodeConfig(TaggedModel):
contents = (await f.read()).decode("utf-8")
data = tomllib.loads(contents)
return cls.model_validate(data)
- except (tomllib.TOMLDecodeError, UnicodeDecodeError):
+ except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError):
logger.warning("Invalid config file, skipping...")
return None
diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py
index 014b3b2f..a59e4eed 100644
--- a/src/exo/worker/engines/image/generate.py
+++ b/src/exo/worker/engines/image/generate.py
@@ -11,7 +11,7 @@ from PIL import Image
from exo.shared.types.api import (
AdvancedImageParams,
- ImageEditsInternalParams,
+ ImageEditsTaskParams,
ImageGenerationStats,
ImageGenerationTaskParams,
)
@@ -67,7 +67,7 @@ def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
def generate_image(
model: DistributedImageModel,
- task: ImageGenerationTaskParams | ImageEditsInternalParams,
+ task: ImageGenerationTaskParams | ImageEditsTaskParams,
) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
"""Generate image(s), optionally yielding partial results.
@@ -105,7 +105,7 @@ def generate_image(
image_path: Path | None = None
with tempfile.TemporaryDirectory() as tmpdir:
- if isinstance(task, ImageEditsInternalParams):
+ if isinstance(task, ImageEditsTaskParams):
# Decode base64 image data and save to temp file
image_path = Path(tmpdir) / "input.png"
image_path.write_bytes(base64.b64decode(task.image_data))
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index c392f6a8..c9585f57 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -8,17 +8,16 @@ from mlx_lm.sample_utils import make_sampler
from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.types.api import (
- BenchChatCompletionTaskParams,
- ChatCompletionMessage,
CompletionTokensDetails,
FinishReason,
GenerationStats,
PromptTokensDetails,
Usage,
)
+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.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.runner_response import (
GenerationResponse,
)
@@ -99,14 +98,9 @@ def warmup_inference(
warmup_prompt = apply_chat_template(
tokenizer=tokenizer,
- chat_task_data=ChatCompletionTaskParams(
- model="",
- messages=[
- ChatCompletionMessage(
- role="user",
- content=content,
- )
- ],
+ task_params=TextGenerationTaskParams(
+ model=ModelId(""),
+ input=content,
),
)
@@ -164,23 +158,17 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
def mlx_generate(
model: Model,
tokenizer: TokenizerWrapper,
- task: ChatCompletionTaskParams,
+ task: TextGenerationTaskParams,
prompt: str,
kv_prefix_cache: KVPrefixCache | None = None,
) -> Generator[GenerationResponse]:
# Ensure that generation stats only contains peak memory for this generation
mx.reset_peak_memory()
- is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
-
- logger.info(f"{is_bench=}")
-
- # Currently we support chat-completion tasks only.
- logger.debug(f"task_params: {task}")
-
if task.seed is not None:
mx.random.seed(task.seed)
# Do not use the prefix cache if we are trying to do benchmarks.
+ is_bench = task.bench
if is_bench:
kv_prefix_cache = None
@@ -206,7 +194,16 @@ def mlx_generate(
sampler = make_sampler(
temp=task.temperature if task.temperature is not None else 0.7,
top_p=task.top_p if task.top_p is not None else 1.0,
+ top_k=task.top_k if task.top_k is not None else 0,
+ )
+
+ # Normalize stop sequences to a list
+ stop_sequences: list[str] = (
+ ([task.stop] if isinstance(task.stop, str) else task.stop)
+ if task.stop is not None
+ else []
)
+ max_stop_len = max((len(s) for s in stop_sequences), default=0)
# Prefill cache with all tokens except the last one
prefill_tps, prefill_tokens = prefill(
@@ -216,7 +213,8 @@ def mlx_generate(
# stream_generate starts from the last token
last_token = prompt_tokens[-1:]
- max_tokens = task.max_tokens or MAX_TOKENS
+ max_tokens = task.max_output_tokens or MAX_TOKENS
+ accumulated_text = ""
generated_text_parts: list[str] = []
generation_start_time = time.perf_counter()
usage: Usage | None = None
@@ -242,6 +240,7 @@ def mlx_generate(
):
generated_text_parts.append(out.text)
logger.info(out.text)
+ accumulated_text += out.text
if think_start is not None and out.text == think_start:
in_thinking = True
@@ -250,8 +249,29 @@ def mlx_generate(
if in_thinking:
reasoning_tokens += 1
+ # Check for stop sequences
+ text = out.text
+ finish_reason: FinishReason | None = cast(
+ FinishReason | None, out.finish_reason
+ )
+ stop_matched = False
+
+ if stop_sequences:
+ for stop_seq in stop_sequences:
+ if stop_seq in accumulated_text:
+ # Trim text to just before the stop sequence
+ stop_index = accumulated_text.find(stop_seq)
+ text_before_stop = accumulated_text[:stop_index]
+ chunk_start = len(accumulated_text) - len(out.text)
+ text = text_before_stop[chunk_start:]
+ finish_reason = "stop"
+ stop_matched = True
+ break
+
+ is_done = finish_reason is not None
+
stats: GenerationStats | None = None
- if out.finish_reason is not None:
+ if is_done:
stats = GenerationStats(
prompt_tps=float(prefill_tps or out.prompt_tps),
generation_tps=float(out.generation_tps),
@@ -259,10 +279,7 @@ def mlx_generate(
generation_tokens=int(out.generation_tokens),
peak_memory_usage=Memory.from_gb(out.peak_memory),
)
-
- if out.finish_reason not in get_args(FinishReason):
- # We don't throw here as this failure case is really not all that bad
- # Just log the error and move on
+ if not stop_matched and out.finish_reason not in get_args(FinishReason):
logger.warning(
f"Model generated unexpected finish_reason: {out.finish_reason}"
)
@@ -280,14 +297,14 @@ def mlx_generate(
)
yield GenerationResponse(
- text=out.text,
+ text=text,
token=out.token,
- finish_reason=cast(FinishReason | None, out.finish_reason),
+ finish_reason=finish_reason,
stats=stats,
usage=usage,
)
- if out.finish_reason is not None:
+ if is_done:
# Log generation stats
generation_elapsed = time.perf_counter() - generation_start_time
generated_tokens = len(generated_text_parts)
@@ -310,4 +327,8 @@ def mlx_generate(
kv_prefix_cache.add_kv_cache(full_prompt, caches)
break
+ # Limit accumulated_text to what's needed for stop sequence detection
+ if max_stop_len > 0 and len(accumulated_text) > max_stop_len:
+ accumulated_text = accumulated_text[-max_stop_len:]
+
# TODO: Do we want an mx_barrier?
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 883925b6..de5cb190 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -39,10 +39,9 @@ from mlx_lm.utils import load_model
from pydantic import RootModel
from exo.download.download_utils import build_model_path
-from exo.shared.types.api import ChatCompletionMessageText
from exo.shared.types.common import Host
from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import (
BoundInstance,
MlxJacclInstance,
@@ -389,8 +388,7 @@ def load_tokenizer_for_model_id(
def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
- """
- Normalize tool_calls in a message dict.
+ """Normalize tool_calls in a message dict.
OpenAI format has tool_calls[].function.arguments as a JSON string,
but some chat templates (e.g., GLM) expect it as a dict.
@@ -413,42 +411,47 @@ def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
def apply_chat_template(
tokenizer: TokenizerWrapper,
- chat_task_data: ChatCompletionTaskParams,
+ task_params: TextGenerationTaskParams,
) -> str:
- messages = chat_task_data.messages
- tools = chat_task_data.tools
+ """Convert TextGenerationTaskParams to a chat template prompt.
- formatted_messages: list[dict[str, Any]] = []
- for message in messages:
- if isinstance(message.content, ChatCompletionMessageText):
- message.content = message.content.text
- if isinstance(message.content, list):
- if len(message.content) == 0:
- logger.warning("Received prompt with no content, skipping")
- continue
-
- message.content = "\n".join(c.text for c in message.content).strip()
- if (
- message.content is None
- and message.thinking is None
- and message.tool_calls is None
- ):
- continue
+ Converts the internal format (input + instructions) to a messages list
+ that can be processed by the tokenizer's chat template.
- # Null values are not valid when applying templates in tokenizer
- dumped: dict[str, Any] = message.model_dump()
- msg_dict: dict[str, Any] = {k: v for k, v in dumped.items() if v is not None} # pyright: ignore[reportAny]
-
- # Parse tool_calls arguments from JSON string to dict for templates that expect dicts
- _normalize_tool_calls(msg_dict)
+ When chat_template_messages is available (from Chat Completions API),
+ uses those directly to preserve tool_calls, thinking, and other fields.
+ Otherwise builds messages from the task params input/instructions.
+ """
+ formatted_messages: list[dict[str, Any]] = []
+ if task_params.chat_template_messages is not None:
+ # Use pre-formatted messages that preserve tool_calls, thinking, etc.
+ formatted_messages = list(task_params.chat_template_messages)
+ for msg in formatted_messages:
+ _normalize_tool_calls(msg)
+ else:
+ # Add system message (instructions) if present
+ if task_params.instructions:
+ formatted_messages.append(
+ {"role": "system", "content": task_params.instructions}
+ )
- formatted_messages.append(msg_dict)
+ # 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})
prompt: str = tokenizer.apply_chat_template(
formatted_messages,
tokenize=False,
add_generation_prompt=True,
- tools=tools,
+ tools=task_params.tools,
)
logger.info(prompt)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 3473be41..2aea72be 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -10,7 +10,7 @@ from loguru import logger
from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import ImageEditsInternalParams
+from exo.shared.types.api import ImageEditsTaskParams
from exo.shared.types.commands import (
ForwarderCommand,
ForwarderDownloadCommand,
@@ -244,7 +244,7 @@ class Worker:
command_id=task.command_id,
instance_id=task.instance_id,
task_status=task.task_status,
- task_params=ImageEditsInternalParams(
+ task_params=ImageEditsTaskParams(
image_data=assembled,
total_input_chunks=task.task_params.total_input_chunks,
prompt=task.task_params.prompt,
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index cb6fbbf3..c173c143 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -4,7 +4,6 @@ from collections.abc import Mapping, Sequence
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.tasks import (
- ChatCompletion,
ConnectToGroup,
CreateRunner,
DownloadModel,
@@ -16,6 +15,7 @@ from exo.shared.types.tasks import (
Task,
TaskId,
TaskStatus,
+ TextGeneration,
)
from exo.shared.types.worker.downloads import (
DownloadCompleted,
@@ -275,7 +275,7 @@ def _pending_tasks(
for task in tasks.values():
# for now, just forward chat completions
# TODO(ciaran): do this better!
- if not isinstance(task, (ChatCompletion, ImageGeneration, ImageEdits)):
+ if not isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
continue
if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
continue
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 61205f38..0bc2e846 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -18,7 +18,7 @@ from pydantic import ValidationError
from exo.shared.constants import EXO_MAX_CHUNK_SIZE
from exo.shared.models.model_cards import ModelId, ModelTask
-from exo.shared.types.api import ChatCompletionMessageText, ImageGenerationStats
+from exo.shared.types.api import ImageGenerationStats
from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk
from exo.shared.types.common import CommandId
from exo.shared.types.events import (
@@ -29,7 +29,6 @@ from exo.shared.types.events import (
TaskStatusUpdated,
)
from exo.shared.types.tasks import (
- ChatCompletion,
ConnectToGroup,
ImageEdits,
ImageGeneration,
@@ -39,7 +38,9 @@ from exo.shared.types.tasks import (
Task,
TaskId,
TaskStatus,
+ TextGeneration,
)
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_response import (
GenerationResponse,
@@ -219,7 +220,7 @@ def main(
current_status = RunnerReady()
logger.info("runner ready")
- case ChatCompletion(task_params=task_params, command_id=command_id) if (
+ case TextGeneration(task_params=task_params, command_id=command_id) if (
isinstance(current_status, RunnerReady)
):
logger.info(f"received chat request: {task}")
@@ -232,10 +233,9 @@ def main(
)
assert model and not isinstance(model, DistributedImageModel)
assert tokenizer
- assert task_params.messages[0].content is not None
try:
- _check_for_debug_prompts(task_params.messages[0].content)
+ _check_for_debug_prompts(task_params)
# Build prompt once - used for both generation and thinking detection
prompt = apply_chat_template(tokenizer, task_params)
@@ -862,17 +862,23 @@ EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
-def _check_for_debug_prompts(
- prompt: str | ChatCompletionMessageText | list[ChatCompletionMessageText],
-):
- if isinstance(prompt, list):
- if len(prompt) == 0:
- logger.debug("Empty message prompt received in debug prompt")
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+ """Check for debug prompt triggers in the input.
+
+ 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 = prompt[0]
+ prompt = task_params.input[0].content
- if isinstance(prompt, ChatCompletionMessageText):
- prompt = prompt.text
+ if not prompt:
+ return
if EXO_RUNNER_MUST_FAIL in prompt:
logger.info("raising exception")
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index 77b3fc11..7015d70b 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -12,10 +12,9 @@ import mlx.nn as nn
from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.api import ChatCompletionMessage
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import 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
@@ -113,10 +112,10 @@ def run_gpt_oss_pipeline_device(
tokens = tokens[:prompt_tokens]
prompt_text = tokenizer.decode(tokens)
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content=prompt_text)],
- max_tokens=max_tokens,
+ input=prompt_text,
+ max_output_tokens=max_tokens,
)
prompt = apply_chat_template(tokenizer, task)
@@ -181,10 +180,10 @@ def run_gpt_oss_tensor_parallel_device(
tokens = tokens[:prompt_tokens]
prompt_text = tokenizer.decode(tokens)
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content=prompt_text)],
- max_tokens=max_tokens,
+ input=prompt_text,
+ max_output_tokens=max_tokens,
)
prompt = apply_chat_template(tokenizer, task)
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
index d0d9e406..bf0a95fb 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
@@ -8,9 +8,8 @@ import pytest
from mlx_lm.models.cache import KVCache
from mlx_lm.sample_utils import make_sampler
-from exo.shared.types.api import ChatCompletionMessage
from exo.shared.types.common import ModelId
-from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.cache import (
KVPrefixCache,
@@ -134,10 +133,10 @@ class TestKVPrefixCacheWithModel:
def test_prefill_populates_cache(self, model_and_tokenizer):
model, tokenizer = model_and_tokenizer
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Hello!!")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Hello!!")],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
@@ -151,10 +150,10 @@ class TestKVPrefixCacheWithModel:
def test_add_and_get_exact_match(self, model_and_tokenizer):
model, tokenizer = model_and_tokenizer
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Test exact")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Test exact")],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
@@ -183,10 +182,10 @@ class TestKVPrefixCacheWithModel:
"""get_kv_cache with a longer prompt sharing prefix should return partial match."""
model, tokenizer = model_and_tokenizer
- short_task = ChatCompletionTaskParams(
+ short_task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Hi")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Hi")],
+ max_output_tokens=1,
)
short_prompt = apply_chat_template(tokenizer, short_task)
short_tokens = encode_prompt(tokenizer, short_prompt)
@@ -198,12 +197,10 @@ class TestKVPrefixCacheWithModel:
kv_prefix_cache.add_kv_cache(short_prompt, cache)
# Query with longer prompt that shares the chat template prefix
- long_task = ChatCompletionTaskParams(
+ long_task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[
- ChatCompletionMessage(role="user", content="Hi there, how are you?")
- ],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Hi there, how are you?")],
+ max_output_tokens=1,
)
long_prompt = apply_chat_template(tokenizer, long_task)
long_tokens = encode_prompt(tokenizer, long_prompt)
@@ -229,10 +226,10 @@ class TestKVPrefixCacheWithModel:
"""Getting a cache and then mutating it (as generation does) must not corrupt stored cache."""
model, tokenizer = model_and_tokenizer
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Mutation test")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Mutation test")],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
@@ -267,10 +264,10 @@ class TestKVPrefixCacheWithModel:
"""Multiple get+mutate cycles (like repeated user requests) must not corrupt cache."""
model, tokenizer = model_and_tokenizer
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Repeat test")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="Repeat test")],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
@@ -302,10 +299,10 @@ class TestKVPrefixCacheWithModel:
model, tokenizer = model_and_tokenizer
kv_prefix_cache = KVPrefixCache(tokenizer)
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Hello")],
- max_tokens=5,
+ input=[InputMessage(role="user", content="Hello")],
+ max_output_tokens=5,
)
prompt = apply_chat_template(tokenizer, task)
prompt_tokens = encode_prompt(tokenizer, prompt)
@@ -332,10 +329,10 @@ class TestKVPrefixCacheWithModel:
model, tokenizer = model_and_tokenizer
kv_prefix_cache = KVPrefixCache(tokenizer)
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Reuse test")],
- max_tokens=5,
+ input=[InputMessage(role="user", content="Reuse test")],
+ max_output_tokens=5,
)
prompt = apply_chat_template(tokenizer, task)
prompt_tokens = encode_prompt(tokenizer, prompt)
@@ -376,10 +373,10 @@ class TestKVPrefixCacheWithModel:
repeats = (1200 // len(base_tokens)) + 2
long_content = base_text * repeats
- task1 = ChatCompletionTaskParams(
+ task1 = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content=long_content)],
- max_tokens=5,
+ input=[InputMessage(role="user", content=long_content)],
+ max_output_tokens=5,
)
prompt1 = apply_chat_template(tokenizer, task1)
prompt1_tokens = encode_prompt(tokenizer, prompt1)
@@ -403,14 +400,14 @@ class TestKVPrefixCacheWithModel:
first_cache_length = cache_length(kv_prefix_cache.caches[0])
# Second generation: same long prompt + extra content (simulating multi-turn)
- task2 = ChatCompletionTaskParams(
+ task2 = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[
- ChatCompletionMessage(role="user", content=long_content),
- ChatCompletionMessage(role="assistant", content="Sure, I can help."),
- ChatCompletionMessage(role="user", content="Tell me more."),
+ input=[
+ InputMessage(role="user", content=long_content),
+ InputMessage(role="assistant", content="Sure, I can help."),
+ InputMessage(role="user", content="Tell me more."),
],
- max_tokens=5,
+ max_output_tokens=5,
)
prompt2 = apply_chat_template(tokenizer, task2)
prompt2_tokens = encode_prompt(tokenizer, prompt2)
@@ -448,10 +445,10 @@ class TestKVPrefixCacheWithModel:
model, tokenizer = model_and_tokenizer
kv_prefix_cache = KVPrefixCache(tokenizer)
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="Immutable test")],
- max_tokens=5,
+ input=[InputMessage(role="user", content="Immutable test")],
+ max_output_tokens=5,
)
prompt = apply_chat_template(tokenizer, task)
@@ -489,10 +486,10 @@ class TestKVPrefixCacheWithModel:
# Add three cache entries with different prompts
prompts = ["First entry", "Second entry", "Third entry"]
for i, content in enumerate(prompts):
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content=content)],
- max_tokens=1,
+ input=[InputMessage(role="user", content=content)],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
@@ -523,10 +520,10 @@ class TestKVPrefixCacheWithModel:
),
):
# Trigger eviction by adding a new entry
- task = ChatCompletionTaskParams(
+ task = TextGenerationTaskParams(
model=DEFAULT_GPT_OSS_MODEL_ID,
- messages=[ChatCompletionMessage(role="user", content="New entry")],
- max_tokens=1,
+ input=[InputMessage(role="user", content="New entry")],
+ max_output_tokens=1,
)
prompt = apply_chat_template(tokenizer, task)
tokens = encode_prompt(tokenizer, prompt)
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 6c8fea8d..07376787 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
@@ -1,8 +1,8 @@
from typing import cast
import exo.worker.plan as plan_mod
-from exo.shared.types.api import ChatCompletionTaskParams
-from exo.shared.types.tasks import ChatCompletion, Task, TaskId, TaskStatus
+from exo.shared.types.tasks import Task, TaskId, TaskStatus, TextGeneration
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import BoundInstance, InstanceId
from exo.shared.types.worker.runners import (
RunnerIdle,
@@ -29,7 +29,7 @@ from exo.worker.tests.unittests.conftest import (
def test_plan_forwards_pending_chat_completion_when_runner_ready():
"""
- When there is a pending ChatCompletion for the local instance and all
+ When there is a pending TextGeneration for the local instance and all
runners are Ready/Running, plan() should forward that task.
"""
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
@@ -54,12 +54,12 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
RUNNER_2_ID: RunnerReady(),
}
- task = ChatCompletion(
+ task = TextGeneration(
task_id=TASK_1_ID,
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+ task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
)
result = plan_mod.plan(
@@ -76,7 +76,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
"""
- Even with a pending ChatCompletion, plan() should not forward it unless
+ Even with a pending TextGeneration, plan() should not forward it unless
all runners for the instance are Ready/Running.
"""
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
@@ -101,12 +101,12 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
RUNNER_2_ID: RunnerIdle(),
}
- task = ChatCompletion(
+ task = TextGeneration(
task_id=TASK_1_ID,
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+ task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
)
result = plan_mod.plan(
@@ -123,7 +123,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
def test_plan_does_not_forward_tasks_for_other_instances():
"""
- plan() should ignore pending ChatCompletion tasks whose instance_id does
+ plan() should ignore pending TextGeneration tasks whose instance_id does
not match the local instance.
"""
shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0)
@@ -145,12 +145,12 @@ def test_plan_does_not_forward_tasks_for_other_instances():
all_runners = {RUNNER_1_ID: RunnerReady()}
other_instance_id = InstanceId("instance-2")
- foreign_task = ChatCompletion(
+ foreign_task = TextGeneration(
task_id=TaskId("other-task"),
instance_id=other_instance_id,
task_status=TaskStatus.Pending,
command_id=COMMAND_1_ID,
- task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+ task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
)
result = plan_mod.plan(
@@ -167,7 +167,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
def test_plan_ignores_non_pending_or_non_chat_tasks():
"""
- _pending_tasks should not forward tasks that are either not ChatCompletion
+ _pending_tasks should not forward tasks that are either not TextGeneration
or not in Pending/Running states.
"""
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
@@ -193,12 +193,12 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
RUNNER_2_ID: RunnerReady(),
}
- completed_task = ChatCompletion(
+ completed_task = TextGeneration(
task_id=TASK_1_ID,
instance_id=INSTANCE_1_ID,
task_status=TaskStatus.Complete,
command_id=COMMAND_1_ID,
- task_params=ChatCompletionTaskParams(model=MODEL_A_ID, messages=[]),
+ task_params=TextGenerationTaskParams(model=MODEL_A_ID, input=""),
)
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 3dfbc6f5..9d7703d8 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
@@ -5,7 +5,6 @@ from typing import Callable
import pytest
import exo.worker.runner.runner as mlx_runner
-from exo.shared.types.api import ChatCompletionMessage
from exo.shared.types.chunks import TokenChunk
from exo.shared.types.events import (
ChunkGenerated,
@@ -15,15 +14,15 @@ from exo.shared.types.events import (
TaskStatusUpdated,
)
from exo.shared.types.tasks import (
- ChatCompletion,
- ChatCompletionTaskParams,
ConnectToGroup,
LoadModel,
Shutdown,
StartWarmup,
Task,
TaskStatus,
+ TextGeneration,
)
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.shared.types.worker.runners import (
RunnerConnected,
@@ -85,15 +84,15 @@ SHUTDOWN_TASK = Shutdown(
runner_id=RUNNER_1_ID,
)
-CHAT_PARAMS = ChatCompletionTaskParams(
- model=str(MODEL_A_ID),
- messages=[ChatCompletionMessage(role="user", content="hello")],
+CHAT_PARAMS = TextGenerationTaskParams(
+ model=MODEL_A_ID,
+ input="hello",
stream=True,
- max_tokens=4,
+ max_output_tokens=4,
temperature=0.0,
)
-CHAT_TASK = ChatCompletion(
+CHAT_TASK = TextGeneration(
task_id=CHAT_COMPLETION_TASK_ID,
command_id=COMMAND_1_ID,
task_params=CHAT_PARAMS,
diff --git a/src/exo/worker/tests/unittests/test_runner/test_glm_tool_parsing.py b/src/exo/worker/tests/unittests/test_runner/test_glm_tool_parsing.py
new file mode 100644
index 00000000..9b8eb092
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_runner/test_glm_tool_parsing.py
@@ -0,0 +1,110 @@
+"""Tests for GLM tool call argument parsing regex."""
+
+import regex as re
+
+# Replicate the regex patterns from runner.py to test them in isolation
+_func_name_regex = re.compile(r"^(.*?)<arg_key>", re.DOTALL)
+_func_arg_regex = re.compile(
+ r"<arg_key>(.*?)</arg_key>(?:\n|\s)*<arg_value>(.*?)(?:</arg_value>|(?=<arg_key>)|$)",
+ re.DOTALL,
+)
+
+
+def _parse_args(text: str) -> list[tuple[str, str]]:
+ """Extract (key, value) pairs from GLM tool call text."""
+ pairs = _func_arg_regex.findall(text)
+ return [(k.strip(), v.strip()) for k, v in pairs] # pyright: ignore[reportAny]
+
+
+def _parse_func_name(text: str) -> str:
+ """Extract function name from GLM tool call text."""
+ match = _func_name_regex.search(text)
+ if match is None:
+ raise ValueError(f"Could not parse function name: {text!r}")
+ return match.group(1).strip()
+
+
+class TestGlmToolParsingWithClosingTags:
+ """Tests for normal format with closing tags present."""
+
+ def test_single_argument(self):
+ text = (
+ "get_weather<arg_key>location</arg_key><arg_value>San Francisco</arg_value>"
+ )
+ assert _parse_func_name(text) == "get_weather"
+ pairs = _parse_args(text)
+ assert pairs == [("location", "San Francisco")]
+
+ def test_multiple_arguments(self):
+ text = (
+ "search<arg_key>query</arg_key><arg_value>python</arg_value>"
+ "<arg_key>limit</arg_key><arg_value>10</arg_value>"
+ )
+ assert _parse_func_name(text) == "search"
+ pairs = _parse_args(text)
+ assert pairs == [("query", "python"), ("limit", "10")]
+
+ def test_arguments_with_whitespace_between(self):
+ text = (
+ "fn<arg_key>a</arg_key>\n<arg_value>1</arg_value>\n"
+ "<arg_key>b</arg_key> <arg_value>2</arg_value>"
+ )
+ pairs = _parse_args(text)
+ assert pairs == [("a", "1"), ("b", "2")]
+
+
+class TestGlmToolParsingMissingClosingTags:
+ """Tests for format where </arg_value> closing tags are missing."""
+
+ def test_single_argument_no_closing(self):
+ text = "get_weather<arg_key>location</arg_key><arg_value>San Francisco"
+ assert _parse_func_name(text) == "get_weather"
+ pairs = _parse_args(text)
+ assert pairs == [("location", "San Francisco")]
+
+ def test_multiple_arguments_no_closing(self):
+ text = (
+ "search<arg_key>query</arg_key><arg_value>python"
+ "<arg_key>limit</arg_key><arg_value>10"
+ )
+ assert _parse_func_name(text) == "search"
+ pairs = _parse_args(text)
+ assert pairs == [("query", "python"), ("limit", "10")]
+
+ def test_mixed_closing_tags(self):
+ """First arg has closing tag, second does not."""
+ text = (
+ "fn<arg_key>a</arg_key><arg_value>1</arg_value>"
+ "<arg_key>b</arg_key><arg_value>2"
+ )
+ pairs = _parse_args(text)
+ assert pairs == [("a", "1"), ("b", "2")]
+
+ def test_value_with_trailing_whitespace(self):
+ text = "fn<arg_key>x</arg_key><arg_value>hello world \n"
+ pairs = _parse_args(text)
+ assert pairs == [("x", "hello world")]
+
+ def test_value_with_newlines_no_closing(self):
+ text = "fn<arg_key>data</arg_key><arg_value>line1\nline2"
+ pairs = _parse_args(text)
+ assert pairs == [("data", "line1\nline2")]
+
+
+class TestGlmToolParsingEdgeCases:
+ """Edge case tests for GLM tool call parsing."""
+
+ def test_empty_value_with_closing(self):
+ text = "fn<arg_key>empty</arg_key><arg_value></arg_value>"
+ pairs = _parse_args(text)
+ assert pairs == [("empty", "")]
+
+ def test_value_with_json_content(self):
+ text = 'fn<arg_key>data</arg_key><arg_value>{"key": "value"}</arg_value>'
+ pairs = _parse_args(text)
+ assert pairs == [("data", '{"key": "value"}')]
+
+ def test_value_with_json_no_closing(self):
+ text = 'fn<arg_key>data</arg_key><arg_value>{"key": "value"}'
+ pairs = _parse_args(text)
+ assert pairs == [("data", '{"key": "value"}')]
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
new file mode 100644
index 00000000..31f4822e
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py
@@ -0,0 +1,87 @@
+"""Tests for parse_tool_calls generator, especially unclosed tool call handling."""
+
+from collections.abc import Generator
+from typing import Any
+
+from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.runner.runner import parse_tool_calls
+
+
+def _make_responses(
+ texts: list[str],
+ finish_on_last: bool = True,
+) -> Generator[GenerationResponse | ToolCallResponse]:
+ """Create a sequence of GenerationResponses from text strings."""
+ for i, text in enumerate(texts):
+ is_last = i == len(texts) - 1
+ yield GenerationResponse(
+ text=text,
+ token=i,
+ finish_reason="stop" if (is_last and finish_on_last) else None,
+ usage=None,
+ )
+
+
+def _dummy_parser(text: str) -> dict[str, Any]:
+ return {"name": "test_fn", "arguments": {"arg": text}}
+
+
+class TestParseToolCalls:
+ """Tests for parse_tool_calls generator."""
+
+ def test_closed_tool_call_works_normally(self):
+ """Normal tool call flow should not be affected."""
+ texts = ["<tool_call>", "test_fn", "</tool_call>"]
+ results = list(
+ parse_tool_calls(
+ _make_responses(texts, finish_on_last=False),
+ "<tool_call>",
+ "</tool_call>",
+ _dummy_parser,
+ )
+ )
+
+ assert len(results) == 1
+ assert isinstance(results[0], ToolCallResponse)
+
+ def test_no_tool_call_passes_through(self):
+ """Responses without tool calls should pass through unchanged."""
+ texts = ["Hello", " world"]
+ results = list(
+ parse_tool_calls(
+ _make_responses(texts),
+ "<tool_call>",
+ "</tool_call>",
+ _dummy_parser,
+ )
+ )
+
+ assert len(results) == 2
+ assert all(isinstance(r, GenerationResponse) for r in results)
+ r0 = results[0]
+ r1 = results[1]
+ assert isinstance(r0, GenerationResponse)
+ assert isinstance(r1, GenerationResponse)
+ assert r0.text == "Hello"
+ assert r1.text == " world"
+ assert r1.finish_reason == "stop"
+
+ def test_failed_parse_yields_text(self):
+ """When tool call parsing fails, the text should be yielded as-is."""
+
+ def _failing_parser(text: str) -> dict[str, Any]:
+ raise ValueError("parse failed")
+
+ texts = ["<tool_call>", "bad content", "</tool_call>"]
+ results = list(
+ parse_tool_calls(
+ _make_responses(texts, finish_on_last=False),
+ "<tool_call>",
+ "</tool_call>",
+ _failing_parser,
+ )
+ )
+
+ assert len(results) == 1
+ assert isinstance(results[0], GenerationResponse)
+ assert results[0].text == "<tool_call>bad content</tool_call>"
diff --git a/tests/headless_runner.py b/tests/headless_runner.py
index 30fd19d4..4c67656a 100644
--- a/tests/headless_runner.py
+++ b/tests/headless_runner.py
@@ -17,18 +17,18 @@ from exo.download.impl_shard_downloader import (
)
from exo.shared.logging import InterceptLogger, logger_setup
from exo.shared.models.model_cards import MODEL_CARDS, ModelId
-from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
from exo.shared.types.commands import CommandId
from exo.shared.types.common import Host, NodeId
from exo.shared.types.events import Event
from exo.shared.types.tasks import (
- ChatCompletion,
ConnectToGroup,
LoadModel,
Shutdown,
StartWarmup,
Task,
+ TextGeneration,
)
+from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import (
BoundInstance,
Instance,
@@ -179,17 +179,11 @@ async def execute_test(test: Tests, instance: Instance, hn: str):
case "inference":
send.send(StartWarmup(instance_id=iid))
send.send(
- ChatCompletion(
- task_params=ChatCompletionTaskParams(
+ TextGeneration(
+ task_params=TextGenerationTaskParams(
model=test.model_id,
- messages=[
- ChatCompletionMessage(
- role="system", content="You are a helpful assistant"
- ),
- ChatCompletionMessage(
- role="user", content="What is the capital of France?"
- ),
- ],
+ instructions="You are a helpful assistant",
+ input="What is the capital of France?",
),
command_id=CommandId("yo"),
instance_id=iid,
← 21d477f1 Update exo bench (#1357)
·
back to Exo
·
chore: gitignore hosts_*.json files (#1343) d826d309 →