From 01b86a9e814da97c0ab9b70c251c9e6caf215352 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Thu, 5 Feb 2026 05:21:26 -0800 Subject: [PATCH 1/2] feat: add uncertainty visualization with token-level logprobs (#1180) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation Adds uncertainty visualization to the chat interface, allowing users to see token-level confidence scores and regenerate responses from any point in the generation. This enables users to: - Understand model confidence at each token - Explore alternative completions by regenerating from uncertain tokens - Debug and analyze model behavior ## Changes ### Uncertainty Visualization - Add `TokenHeatmap` component showing token-level probability coloring - Toggle uncertainty view per message with bar chart icon - Display tooltip with probability, logprob, and top alternative tokens on hover ### Regenerate from Token - Add "Regenerate from here" button in token tooltip - Use `continue_final_message` in chat template to continue within same turn (no EOS tokens) - Add `continue_from_prefix` flag to `ChatCompletionTaskParams` ### Request Cancellation - Add `AbortController` to cancel in-flight requests when regenerating mid-generation - Handle `BrokenResourceError` server-side when client disconnects gracefully ### Additional APIs - Add Claude Messages API support (`/v1/messages`) - Add OpenAI Responses API support (`/v1/responses`) ## Why It Works - **Proper continuation**: Using `continue_final_message=True` instead of `add_generation_prompt=True` keeps the assistant turn open, allowing the model to continue naturally from the prefix without end-of-turn markers - **Clean cancellation**: AbortController aborts the HTTP request, and server catches `BrokenResourceError` to avoid crashes - **Stable hover during generation**: TokenHeatmap tracks hover by index (stable across re-renders) with longer hide delay during generation ## Test Plan ### Manual Testing - Send a message and verify logprobs are collected - Enable uncertainty view and verify token coloring based on probability - Hover over tokens to see tooltip with alternatives - Click "Regenerate from here" on a token mid-response - Verify the response continues naturally from that point - Verify aborting mid-generation and regenerating works without server crash ### Automated Testing - Added tests for Claude Messages API adapter - Added tests for OpenAI Responses API adapter 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.5 Co-authored-by: Evan --- .../src/lib/components/ChatMessages.svelte | 72 +++- .../src/lib/components/TokenHeatmap.svelte | 236 +++++++++++++ dashboard/src/lib/stores/app.svelte.ts | 319 +++++++++++++++++- src/exo/master/adapters/chat_completions.py | 30 ++ src/exo/master/api.py | 15 + src/exo/shared/types/chunks.py | 9 +- src/exo/shared/types/text_generation.py | 2 + .../shared/types/worker/runner_response.py | 4 +- src/exo/worker/engines/mlx/constants.py | 2 + .../worker/engines/mlx/generator/generate.py | 75 +++- src/exo/worker/engines/mlx/utils_mlx.py | 9 + src/exo/worker/runner/runner.py | 2 + 12 files changed, 760 insertions(+), 15 deletions(-) create mode 100644 dashboard/src/lib/components/TokenHeatmap.svelte diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte index 15ea088d..44b9ec0d 100644 --- a/dashboard/src/lib/components/ChatMessages.svelte +++ b/dashboard/src/lib/components/ChatMessages.svelte @@ -6,11 +6,13 @@ deleteMessage, editAndRegenerate, regenerateLastResponse, + regenerateFromToken, setEditingImage, } from "$lib/stores/app.svelte"; import type { Message } from "$lib/stores/app.svelte"; import type { MessageAttachment } from "$lib/stores/app.svelte"; import MarkdownContent from "./MarkdownContent.svelte"; + import TokenHeatmap from "./TokenHeatmap.svelte"; interface Props { class?: string; @@ -99,6 +101,23 @@ let copiedMessageId = $state(null); let expandedThinkingMessageIds = $state>(new Set()); + // Uncertainty heatmap toggle + let heatmapMessageIds = $state>(new Set()); + + function toggleHeatmap(messageId: string) { + const next = new Set(heatmapMessageIds); + if (next.has(messageId)) { + next.delete(messageId); + } else { + next.add(messageId); + } + heatmapMessageIds = next; + } + + function isHeatmapVisible(messageId: string): boolean { + return heatmapMessageIds.has(messageId); + } + function formatTimestamp(timestamp: number): string { return new Date(timestamp).toLocaleTimeString("en-US", { hour12: false, @@ -548,13 +567,23 @@ > {:else if message.content || (loading && !message.attachments?.some((a) => a.type === "generated-image"))} - - {#if loading && !message.content} - + {#if isHeatmapVisible(message.id) && message.tokens && message.tokens.length > 0} + + regenerateFromToken(message.id, tokenIndex)} + /> + {:else} + + {#if loading && !message.content} + + {/if} {/if} {/if} @@ -629,6 +658,35 @@ {/if} + + {#if message.role === "assistant" && message.tokens && message.tokens.length > 0} + + {/if} + {#if message.role === "assistant" && isLastAssistantMessage(message.id) && !loading} + {/if} + + +
+
+
+ +{/if} + + diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 51de6c66..6fdb0c7c 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -242,6 +242,19 @@ export interface MessageAttachment { mimeType?: string; } +export interface TopLogprob { + token: string; + logprob: number; + bytes: number[] | null; +} + +export interface TokenData { + token: string; + logprob: number; + probability: number; + topLogprobs: TopLogprob[]; +} + export interface Message { id: string; role: "user" | "assistant" | "system"; @@ -253,6 +266,7 @@ export interface Message { tps?: number; // Tokens per second (for assistant messages) requestType?: "chat" | "image-generation" | "image-editing"; sourceImageDataUrl?: string; // For image editing regeneration + tokens?: TokenData[]; } export interface Conversation { @@ -540,7 +554,18 @@ class AppStore { */ private saveConversationsToStorage() { try { - localStorage.setItem(STORAGE_KEY, JSON.stringify(this.conversations)); + // Strip tokens from messages before saving to avoid bloating localStorage + const stripped = this.conversations.map((conv) => ({ + ...conv, + messages: conv.messages.map((msg) => { + if (msg.tokens) { + const { tokens: _, ...rest } = msg; + return rest; + } + return msg; + }), + })); + localStorage.setItem(STORAGE_KEY, JSON.stringify(stripped)); } catch (error) { console.error("Failed to save conversations:", error); } @@ -1445,6 +1470,213 @@ class AppStore { } } + /** + * Regenerate response from a specific token index. + * Truncates the assistant message at the given token and re-generates from there. + */ + async regenerateFromToken( + messageId: string, + tokenIndex: number, + ): Promise { + if (this.isLoading) return; + + const targetConversationId = this.activeConversationId; + if (!targetConversationId) return; + + const msgIndex = this.messages.findIndex((m) => m.id === messageId); + if (msgIndex === -1) return; + + const msg = this.messages[msgIndex]; + if ( + msg.role !== "assistant" || + !msg.tokens || + tokenIndex >= msg.tokens.length + ) + return; + + // Keep tokens up to (not including) the specified index + const tokensToKeep = msg.tokens.slice(0, tokenIndex); + const prefixText = tokensToKeep.map((t) => t.token).join(""); + + // Remove all messages after this assistant message + this.messages = this.messages.slice(0, msgIndex + 1); + + // Update the message to show the prefix + this.messages[msgIndex].content = prefixText; + this.messages[msgIndex].tokens = tokensToKeep; + this.updateActiveConversation(); + + // Set up for continuation - modify the existing message in place + this.isLoading = true; + this.currentResponse = prefixText; + this.ttftMs = null; + this.tps = null; + this.totalTokens = tokensToKeep.length; + + try { + // Build messages for API - include the partial assistant message + const systemPrompt = { + role: "system" as const, + content: + "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process.", + }; + + const apiMessages = [ + systemPrompt, + ...this.messages.map((m) => { + let msgContent = m.content; + if (m.attachments) { + for (const attachment of m.attachments) { + if (attachment.type === "text" && attachment.content) { + msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``; + } + } + } + return { role: m.role, content: msgContent }; + }), + ]; + + const modelToUse = this.getModelForRequest(); + if (!modelToUse) { + throw new Error("No model available"); + } + + const requestStartTime = performance.now(); + let firstTokenTime: number | null = null; + let tokenCount = tokensToKeep.length; + + const response = await fetch("/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: modelToUse, + messages: apiMessages, + stream: true, + logprobs: true, + top_logprobs: 5, + }), + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error(`API error: ${response.status} - ${errorText}`); + } + + const reader = response.body?.getReader(); + if (!reader) throw new Error("No response body"); + + let fullContent = prefixText; + const collectedTokens: TokenData[] = [...tokensToKeep]; + + interface ChatCompletionChunk { + choices?: Array<{ + delta?: { content?: string }; + logprobs?: { + content?: Array<{ + token: string; + logprob: number; + top_logprobs?: Array<{ + token: string; + logprob: number; + bytes: number[] | null; + }>; + }>; + }; + }>; + } + + await this.parseSSEStream( + reader, + targetConversationId, + (parsed) => { + const choice = parsed.choices?.[0]; + const delta = choice?.delta?.content; + + // Collect logprobs data + const logprobsContent = choice?.logprobs?.content; + if (logprobsContent) { + for (const item of logprobsContent) { + collectedTokens.push({ + token: item.token, + logprob: item.logprob, + probability: Math.exp(item.logprob), + topLogprobs: (item.top_logprobs || []).map((t) => ({ + token: t.token, + logprob: t.logprob, + bytes: t.bytes, + })), + }); + } + } + + if (delta) { + if (firstTokenTime === null) { + firstTokenTime = performance.now(); + this.ttftMs = firstTokenTime - requestStartTime; + } + + tokenCount += 1; + this.totalTokens = tokenCount; + + if (firstTokenTime !== null && tokenCount > tokensToKeep.length) { + const elapsed = performance.now() - firstTokenTime; + this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000; + } + + fullContent += delta; + const { displayContent, thinkingContent } = + this.stripThinkingTags(fullContent); + + if (this.activeConversationId === targetConversationId) { + this.currentResponse = displayContent; + } + + // Update existing message in place + this.updateConversationMessage( + targetConversationId, + messageId, + (m) => { + m.content = displayContent; + m.thinking = thinkingContent || undefined; + m.tokens = [...collectedTokens]; + }, + ); + this.syncActiveMessagesIfNeeded(targetConversationId); + this.persistConversation(targetConversationId); + } + }, + ); + + // Final update + if (this.conversationExists(targetConversationId)) { + const { displayContent, thinkingContent } = + this.stripThinkingTags(fullContent); + this.updateConversationMessage(targetConversationId, messageId, (m) => { + m.content = displayContent; + m.thinking = thinkingContent || undefined; + m.tokens = [...collectedTokens]; + if (this.ttftMs !== null) m.ttftMs = this.ttftMs; + if (this.tps !== null) m.tps = this.tps; + }); + this.syncActiveMessagesIfNeeded(targetConversationId); + this.persistConversation(targetConversationId); + } + } catch (error) { + console.error("Error regenerating from token:", error); + if (this.conversationExists(targetConversationId)) { + this.updateConversationMessage(targetConversationId, messageId, (m) => { + m.content = `${prefixText}\n\nError: ${error instanceof Error ? error.message : "Unknown error"}`; + }); + this.syncActiveMessagesIfNeeded(targetConversationId); + this.persistConversation(targetConversationId); + } + } finally { + this.isLoading = false; + this.currentResponse = ""; + this.saveConversationsToStorage(); + } + } + /** * Helper method to regenerate a chat completion response */ @@ -1513,6 +1745,8 @@ class AppStore { model: modelToUse, messages: apiMessages, stream: true, + logprobs: true, + top_logprobs: 5, }), }); @@ -1527,16 +1761,49 @@ class AppStore { } let streamedContent = ""; + const collectedTokens: TokenData[] = []; interface ChatCompletionChunk { - choices?: Array<{ delta?: { content?: string } }>; + choices?: Array<{ + delta?: { content?: string }; + logprobs?: { + content?: Array<{ + token: string; + logprob: number; + top_logprobs?: Array<{ + token: string; + logprob: number; + bytes: number[] | null; + }>; + }>; + }; + }>; } await this.parseSSEStream( reader, targetConversationId, (parsed) => { - const delta = parsed.choices?.[0]?.delta?.content; + const choice = parsed.choices?.[0]; + const delta = choice?.delta?.content; + + // Collect logprobs data + const logprobsContent = choice?.logprobs?.content; + if (logprobsContent) { + for (const item of logprobsContent) { + collectedTokens.push({ + token: item.token, + logprob: item.logprob, + probability: Math.exp(item.logprob), + topLogprobs: (item.top_logprobs || []).map((t) => ({ + token: t.token, + logprob: t.logprob, + bytes: t.bytes, + })), + }); + } + } + if (delta) { streamedContent += delta; const { displayContent, thinkingContent } = @@ -1554,6 +1821,7 @@ class AppStore { (msg) => { msg.content = displayContent; msg.thinking = thinkingContent || undefined; + msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); @@ -1572,6 +1840,7 @@ class AppStore { (msg) => { msg.content = displayContent; msg.thinking = thinkingContent || undefined; + msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); @@ -1914,6 +2183,8 @@ class AppStore { messages: apiMessages, temperature: 0.7, stream: true, + logprobs: true, + top_logprobs: 5, }), }); @@ -1930,14 +2201,48 @@ class AppStore { let streamedContent = ""; interface ChatCompletionChunk { - choices?: Array<{ delta?: { content?: string } }>; + choices?: Array<{ + delta?: { content?: string }; + logprobs?: { + content?: Array<{ + token: string; + logprob: number; + top_logprobs?: Array<{ + token: string; + logprob: number; + bytes: number[] | null; + }>; + }>; + }; + }>; } + const collectedTokens: TokenData[] = []; + await this.parseSSEStream( reader, targetConversationId, (parsed) => { - const tokenContent = parsed.choices?.[0]?.delta?.content; + const choice = parsed.choices?.[0]; + const tokenContent = choice?.delta?.content; + + // Collect logprobs data + const logprobsContent = choice?.logprobs?.content; + if (logprobsContent) { + for (const item of logprobsContent) { + collectedTokens.push({ + token: item.token, + logprob: item.logprob, + probability: Math.exp(item.logprob), + topLogprobs: (item.top_logprobs || []).map((t) => ({ + token: t.token, + logprob: t.logprob, + bytes: t.bytes, + })), + }); + } + } + if (tokenContent) { // Track first token for TTFT if (firstTokenTime === null) { @@ -1973,6 +2278,7 @@ class AppStore { (msg) => { msg.content = displayContent; msg.thinking = thinkingContent || undefined; + msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); @@ -1997,6 +2303,7 @@ class AppStore { (msg) => { msg.content = displayContent; msg.thinking = thinkingContent || undefined; + msg.tokens = [...collectedTokens]; // Store performance metrics on the message if (this.ttftMs !== null) { msg.ttftMs = this.ttftMs; @@ -2693,6 +3000,8 @@ export const editMessage = (messageId: string, newContent: string) => export const editAndRegenerate = (messageId: string, newContent: string) => appStore.editAndRegenerate(messageId, newContent); export const regenerateLastResponse = () => appStore.regenerateLastResponse(); +export const regenerateFromToken = (messageId: string, tokenIndex: number) => + appStore.regenerateFromToken(messageId, tokenIndex); // Conversation actions export const conversations = () => appStore.conversations; diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py index e144696b..3e013079 100644 --- a/src/exo/master/adapters/chat_completions.py +++ b/src/exo/master/adapters/chat_completions.py @@ -14,6 +14,8 @@ from exo.shared.types.api import ( ErrorInfo, ErrorResponse, FinishReason, + Logprobs, + LogprobsContentItem, StreamingChoiceResponse, ToolCall, ) @@ -81,6 +83,8 @@ def chat_request_to_text_generation( chat_template_messages=chat_template_messages if chat_template_messages else None, + logprobs=request.logprobs or False, + top_logprobs=request.top_logprobs, ) @@ -88,6 +92,19 @@ def chunk_to_response( chunk: TokenChunk, command_id: CommandId ) -> ChatCompletionResponse: """Convert a TokenChunk to a streaming ChatCompletionResponse.""" + # Build logprobs if available + logprobs: Logprobs | None = None + if chunk.logprob is not None: + logprobs = Logprobs( + content=[ + LogprobsContentItem( + token=chunk.text, + logprob=chunk.logprob, + top_logprobs=chunk.top_logprobs or [], + ) + ] + ) + return ChatCompletionResponse( id=command_id, created=int(time.time()), @@ -96,6 +113,7 @@ def chunk_to_response( StreamingChoiceResponse( index=0, delta=ChatCompletionMessage(role="assistant", content=chunk.text), + logprobs=logprobs, finish_reason=chunk.finish_reason, ) ], @@ -162,6 +180,7 @@ async def collect_chat_response( """Collect all token chunks and return a single ChatCompletionResponse.""" text_parts: list[str] = [] tool_calls: list[ToolCall] = [] + logprobs_content: list[LogprobsContentItem] = [] model: str | None = None finish_reason: FinishReason | None = None error_message: str | None = None @@ -176,6 +195,14 @@ async def collect_chat_response( if isinstance(chunk, TokenChunk): text_parts.append(chunk.text) + if chunk.logprob is not None: + logprobs_content.append( + LogprobsContentItem( + token=chunk.text, + logprob=chunk.logprob, + top_logprobs=chunk.top_logprobs or [], + ) + ) if isinstance(chunk, ToolCallChunk): tool_calls.extend( @@ -208,6 +235,9 @@ async def collect_chat_response( content=combined_text, tool_calls=tool_calls if tool_calls else None, ), + logprobs=Logprobs(content=logprobs_content) + if logprobs_content + else None, finish_reason=finish_reason, ) ], diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 9bd8cbcf..0ad5454c 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -627,6 +627,11 @@ class API: self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, ) return await collect_chat_response( @@ -1183,6 +1188,11 @@ class API: self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, ) return await collect_claude_response( @@ -1210,6 +1220,11 @@ class API: self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, ) return await collect_responses_response( diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index e96dbc9d..5fe9eb1c 100644 --- a/src/exo/shared/types/chunks.py +++ b/src/exo/shared/types/chunks.py @@ -2,7 +2,12 @@ from collections.abc import Generator from typing import Any, Literal from exo.shared.models.model_cards import ModelId -from exo.shared.types.api import GenerationStats, ImageGenerationStats, Usage +from exo.shared.types.api import ( + GenerationStats, + ImageGenerationStats, + TopLogprobItem, + Usage, +) from exo.utils.pydantic_ext import TaggedModel from .api import FinishReason @@ -20,6 +25,8 @@ class TokenChunk(BaseChunk): usage: Usage | None finish_reason: Literal["stop", "length", "content_filter"] | None = None stats: GenerationStats | None = None + logprob: float | None = None + top_logprobs: list[TopLogprobItem] | None = None class ErrorChunk(BaseChunk): diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py index 31f97a70..3e7b89fd 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -40,3 +40,5 @@ class TextGenerationTaskParams(BaseModel, frozen=True): stop: str | list[str] | None = None seed: int | None = None chat_template_messages: list[dict[str, Any]] | None = None + logprobs: bool = False + top_logprobs: int | None = None diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py index 5dfbe547..d1bea77e 100644 --- a/src/exo/shared/types/worker/runner_response.py +++ b/src/exo/shared/types/worker/runner_response.py @@ -6,6 +6,7 @@ from exo.shared.types.api import ( GenerationStats, ImageGenerationStats, ToolCallItem, + TopLogprobItem, Usage, ) from exo.utils.pydantic_ext import TaggedModel @@ -22,7 +23,8 @@ class TokenizedResponse(BaseRunnerResponse): class GenerationResponse(BaseRunnerResponse): text: str token: int - # logprobs: list[float] | None = None # too big. we can change to be top-k + logprob: float | None = None + top_logprobs: list[TopLogprobItem] | None = None finish_reason: FinishReason | None = None stats: GenerationStats | None = None usage: Usage | None diff --git a/src/exo/worker/engines/mlx/constants.py b/src/exo/worker/engines/mlx/constants.py index dbffdfa0..86a663e4 100644 --- a/src/exo/worker/engines/mlx/constants.py +++ b/src/exo/worker/engines/mlx/constants.py @@ -11,5 +11,7 @@ QUANTIZE_MODEL_MODE: str | None = "affine" CACHE_GROUP_SIZE: int = 64 KV_CACHE_BITS: int | None = None +DEFAULT_TOP_LOGPROBS: int = 5 + # TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True TRUST_REMOTE_CODE: bool = True diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index a38d70c5..67a31ae0 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -12,6 +12,7 @@ from exo.shared.types.api import ( FinishReason, GenerationStats, PromptTokensDetails, + TopLogprobItem, Usage, ) from exo.shared.types.common import ModelId @@ -23,7 +24,12 @@ from exo.shared.types.worker.runner_response import ( ) from exo.worker.engines.mlx import Model from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_cache -from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS +from exo.worker.engines.mlx.constants import ( + DEFAULT_TOP_LOGPROBS, + KV_BITS, + KV_GROUP_SIZE, + MAX_TOKENS, +) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, mx_barrier, @@ -155,6 +161,60 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]: return eos +def extract_top_logprobs( + logprobs: mx.array, + tokenizer: TokenizerWrapper, + top_logprobs: int, + selected_token: int, +) -> tuple[float, list[TopLogprobItem]]: + """Extract the selected token's logprob and top alternative tokens. + + Args: + logprobs: Full vocabulary logprobs array from MLX + tokenizer: Tokenizer for decoding token IDs to strings + top_logprobs: Number of top alternatives to return + selected_token: The token ID that was actually sampled + + Returns: + Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives) + """ + # Get the logprob of the selected token + selected_logprob = float(logprobs[selected_token].item()) + + # Get top indices (most probable tokens) + # mx.argpartition gives indices that would partition the array + # We negate logprobs since argpartition finds smallest, and we want largest + top_logprobs = min(top_logprobs, logprobs.shape[0]) # Don't exceed vocab size + top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs] + + # Get the actual logprob values for these indices + top_values = logprobs[top_indices] + + # Sort by logprob (descending) for consistent ordering + sort_order = mx.argsort(-top_values) + top_indices = top_indices[sort_order] + top_values = top_values[sort_order] + + # Convert to list of TopLogprobItem + top_logprob_items: list[TopLogprobItem] = [] + for i in range(top_logprobs): + token_id = int(top_indices[i].item()) + token_logprob = float(top_values[i].item()) + # Decode token ID to string + token_str = tokenizer.decode([token_id]) + # Get byte representation + token_bytes = list(token_str.encode("utf-8")) + top_logprob_items.append( + TopLogprobItem( + token=token_str, + logprob=token_logprob, + bytes=token_bytes, + ) + ) + + return selected_logprob, top_logprob_items + + def mlx_generate( model: Model, tokenizer: TokenizerWrapper, @@ -296,9 +356,22 @@ def mlx_generate( ), ) + # Extract logprobs from the full vocabulary logprobs array + logprob: float | None = None + top_logprobs: list[TopLogprobItem] | None = None + if task.logprobs: + logprob, top_logprobs = extract_top_logprobs( + logprobs=out.logprobs, + tokenizer=tokenizer, + top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS, + selected_token=out.token, + ) + yield GenerationResponse( text=text, token=out.token, + logprob=logprob, + top_logprobs=top_logprobs, finish_reason=finish_reason, stats=stats, usage=usage, diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index e12aa185..4f1140fb 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -459,6 +459,12 @@ def apply_chat_template( continue formatted_messages.append({"role": msg.role, "content": msg.content}) + # For assistant prefilling, append content after templating to avoid a closing turn token. + partial_assistant_content: str | None = None + if formatted_messages and formatted_messages[-1].get("role") == "assistant": + partial_assistant_content = cast(str, formatted_messages[-1].get("content", "")) + formatted_messages = formatted_messages[:-1] + prompt: str = tokenizer.apply_chat_template( formatted_messages, tokenize=False, @@ -466,6 +472,9 @@ def apply_chat_template( tools=task_params.tools, ) + if partial_assistant_content: + prompt += partial_assistant_content + logger.info(prompt) return prompt diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 109ea219..b0e655cc 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -344,6 +344,8 @@ def main( usage=response.usage, finish_reason=response.finish_reason, stats=response.stats, + logprob=response.logprob, + top_logprobs=response.top_logprobs, ), ) ) From 3a9baeb9db469975f9c760960bf0d1fcb187d2ce Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Tue, 3 Feb 2026 22:49:47 +0000 Subject: [PATCH 2/2] EXO: add CLI flags for root install/uninstall The macOS app required user interaction via AppleScript prompts to install or uninstall network configuration components, making automated deployments difficult. Added --install and --uninstall command line flags that execute the network setup scripts directly when running as root, bypassing GUI prompts. Created a new main.swift entry point that parses CLI arguments and delegates to NetworkSetupHelper's new direct execution methods. This enables headless installation via `sudo EXO --install` for automated deployment scenarios while preserving the existing GUI behavior when launched normally. Test plan: - Deployed to a machine that didn't have the content installed. Got blocked on the popup and EXO never launched. - Relaunched EXO, confirmed it still never starts because of the popup. - Ran `sudo /Applications/EXO.app/Contents/MacOS/EXO --install` - Launched EXO - the API started as expected. - Ran `sudo /Applications/EXO.app/Contents/MacOS/EXO --uninstall` - Launched EXO - got the popup. --- app/EXO/EXO/EXOApp.swift | 1 - app/EXO/EXO/Services/NetworkSetupHelper.swift | 55 ++++++++++++ app/EXO/EXO/main.swift | 85 +++++++++++++++++++ 3 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 app/EXO/EXO/main.swift diff --git a/app/EXO/EXO/EXOApp.swift b/app/EXO/EXO/EXOApp.swift index 7669c408..3aff58c3 100644 --- a/app/EXO/EXO/EXOApp.swift +++ b/app/EXO/EXO/EXOApp.swift @@ -14,7 +14,6 @@ import SwiftUI import UserNotifications import os.log -@main struct EXOApp: App { @StateObject private var controller: ExoProcessController @StateObject private var stateService: ClusterStateService diff --git a/app/EXO/EXO/Services/NetworkSetupHelper.swift b/app/EXO/EXO/Services/NetworkSetupHelper.swift index 82cb82d4..5428ee8e 100644 --- a/app/EXO/EXO/Services/NetworkSetupHelper.swift +++ b/app/EXO/EXO/Services/NetworkSetupHelper.swift @@ -288,6 +288,61 @@ enum NetworkSetupHelper { """ } + /// Direct install without GUI (requires root). + /// Returns true on success, false on failure. + static func installDirectly() -> Bool { + let script = makeInstallerScript() + return runShellDirectly(script) + } + + /// Direct uninstall without GUI (requires root). + /// Returns true on success, false on failure. + static func uninstallDirectly() -> Bool { + let script = makeUninstallScript() + return runShellDirectly(script) + } + + /// Run a shell script directly via Process (no AppleScript, requires root). + /// Returns true on success, false on failure. + private static func runShellDirectly(_ script: String) -> Bool { + let process = Process() + process.executableURL = URL(fileURLWithPath: "/bin/bash") + process.arguments = ["-c", script] + + let outputPipe = Pipe() + let errorPipe = Pipe() + process.standardOutput = outputPipe + process.standardError = errorPipe + + do { + try process.run() + process.waitUntilExit() + + let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile() + let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile() + + if let output = String(data: outputData, encoding: .utf8), !output.isEmpty { + print(output) + } + if let errorOutput = String(data: errorData, encoding: .utf8), !errorOutput.isEmpty { + fputs(errorOutput, stderr) + } + + if process.terminationStatus == 0 { + logger.info("Shell script completed successfully") + return true + } else { + logger.error("Shell script failed with exit code \(process.terminationStatus)") + return false + } + } catch { + logger.error( + "Failed to run shell script: \(error.localizedDescription, privacy: .public)") + fputs("Error: \(error.localizedDescription)\n", stderr) + return false + } + } + private static func runShellAsAdmin(_ script: String) throws { let escapedScript = script diff --git a/app/EXO/EXO/main.swift b/app/EXO/EXO/main.swift new file mode 100644 index 00000000..9383981f --- /dev/null +++ b/app/EXO/EXO/main.swift @@ -0,0 +1,85 @@ +// +// main.swift +// EXO +// +// Created by Jake Hillion on 2026-02-03. +// + +import Foundation + +/// Command line options for the EXO app +enum CLICommand { + case install + case uninstall + case help + case none +} + +/// Parse command line arguments to determine the CLI command +func parseArguments() -> CLICommand { + let args = CommandLine.arguments + if args.contains("--help") || args.contains("-h") { + return .help + } + if args.contains("--install") { + return .install + } + if args.contains("--uninstall") { + return .uninstall + } + return .none +} + +/// Print usage information +func printUsage() { + let programName = (CommandLine.arguments.first as NSString?)?.lastPathComponent ?? "EXO" + print( + """ + Usage: \(programName) [OPTIONS] + + Options: + --install Install EXO network configuration (requires root) + --uninstall Uninstall EXO network configuration (requires root) + --help, -h Show this help message + + When run without options, starts the normal GUI application. + + Examples: + sudo \(programName) --install Install network components as root + sudo \(programName) --uninstall Remove network components as root + """) +} + +/// Check if running as root +func isRunningAsRoot() -> Bool { + return getuid() == 0 +} + +// Main entry point +let command = parseArguments() + +switch command { +case .help: + printUsage() + exit(0) + +case .install: + if !isRunningAsRoot() { + fputs("Error: --install requires root privileges. Run with sudo.\n", stderr) + exit(1) + } + let success = NetworkSetupHelper.installDirectly() + exit(success ? 0 : 1) + +case .uninstall: + if !isRunningAsRoot() { + fputs("Error: --uninstall requires root privileges. Run with sudo.\n", stderr) + exit(1) + } + let success = NetworkSetupHelper.uninstallDirectly() + exit(success ? 0 : 1) + +case .none: + // Start normal GUI application + EXOApp.main() +}