Merge pull request #23 from CoolSpring8/token-probability

Implement token-level probability tracking and reroll UI

Author: CoolSpring8

bun.lock

@@ -13,6 +13,7 @@
 		"@ai-sdk/react": "^2.0.98",
 		"@base-ui-components/react": "^1.0.0-beta.4",
 		"@built-in-ai/core": "^2.0.1",
+		"@floating-ui/react": "^0.27.16",
 		"@mantine/core": "^7.8.1",
 		"@mantine/hooks": "^7.8.1",
 		"@xyflow/react": "^12.9.1",

package.json

@@ -87,6 +87,8 @@ const App = () => {
 		activeModel,
 		openAIProvider,
 		getBuiltInChatModel,
+		baseURL: activeProvider?.config.baseURL ?? "",
+		apiKey: activeProvider?.config.apiKey ?? "",
 	});
 
 	const {
@@ -113,17 +115,21 @@ const App = () => {
 		isTreeEmpty,
 		createSystemMessage,
 		setActiveTarget,
+		rerollFromToken,
 	} = useConversationController({
 		defaultSystemPrompt,
 		ensureChatReady,
 	});
 
 	const {
 		textContent,
-		setTextContent,
+		overwriteTextContent,
 		isGenerating: isTextGenerating,
 		predict,
 		cancel,
+		tokenLogprobs: textTokenLogprobs,
+		rerollFromToken: rerollTextFromToken,
+		seedText: textSeed,
 	} = useTextCompletion({
 		ensureCompletionReady,
 	});
@@ -167,8 +173,8 @@ const App = () => {
 	const handleClearConversation = useCallback(() => {
 		clearConversation();
 		cancel();
-		setTextContent("");
-	}, [cancel, clearConversation, setTextContent]);
+		overwriteTextContent("");
+	}, [cancel, clearConversation, overwriteTextContent]);
 
 	const handleImportPreparation = useCallback(() => {
 		abortActiveStreams();
@@ -217,6 +223,7 @@ const App = () => {
 								onEditCancel={cancelEdit}
 								onPromptDirtyChange={setIsPromptDirty}
 								resetSignal={resetSignal}
+								onTokenReroll={rerollFromToken}
 							/>
 						</div>
 					) : view === "diagram" ? (
@@ -238,10 +245,14 @@ const App = () => {
 									: undefined
 							}
 							onChange={(value) => {
-								setTextContent(value);
+								overwriteTextContent(value);
 							}}
 							onPredict={predict}
 							onCancel={cancel}
+							tokenLogprobs={textTokenLogprobs}
+							onTokenReroll={rerollTextFromToken}
+							showTokenOverlay
+							generatedPrefix={textSeed}
 						/>
 					)}
 					<SettingsModal open={isSettingsOpen} onClose={onSettingsClose} />

src/App.tsx

@@ -1,6 +1,8 @@
 import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
 import type { ModelInfo } from "../types";
 
+export const OPENAI_COMPATIBLE_PROVIDER_NAME = "user-openai-compatible";
+
 const normalizeBaseURL = (baseURL: string) =>
 	baseURL.trim().replace(/\/+$/, "");
 
@@ -17,7 +19,7 @@ export const buildOpenAICompatibleProvider = ({
 	}
 	return createOpenAICompatible({
 		baseURL: normalizedBaseURL,
-		name: "user-openai-compatible",
+		name: OPENAI_COMPATIBLE_PROVIDER_NAME,
 		apiKey: apiKey || "_PLACEHOLDER_",
 	});
 };