Implement token-level probability tracking and reroll UI#23
Implement token-level probability tracking and reroll UI#23CoolSpring8 merged 22 commits intomainfrom
Conversation
…hance reroll functionality
…xpose seed text for completion visualization
…es and improved generation control
…ontrol of hovered tokens in the inline
…rating` state in `useTextCompletion` hook
…in streams and UI
…and disable token reroll for reasoning content
…its direct text styling for consistency
…gprobStream.ts` with a new `openaiLogprobs.ts` module and updating its usage
…stAnimationFrame`
… improving hover stability and interaction
…over state management with delayed opening and separate highlight/tooltip indices
…ges and fix `isGenerating` state reset race condition
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds token-level logprob capture, parsing, streaming, UI heatmaps, and interactive token rerolling; extends provider readiness with OpenAI-compatible baseURL/apiKey; propagates tokenLogprobs through AI streaming, tree/state, hooks, and components; introduces a token inline renderer and heatmap settings. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant UI as User/UI
participant Msg as MessageItem
participant Chat as ChatView
participant Controller as useConversationController
participant Provider as OpenAI-compatible Provider
participant Tree as Conversation Tree
UI->>Msg: select token alternative
Msg->>Chat: onRerollToken(tokenIdx, alternative)
Chat->>Controller: rerollFromToken(messageId, tokenIdx, alternative)
activate Controller
Controller->>Provider: stream with seed (replacement) + logprob options
Provider-->>Controller: streamed chunks (text + logprobs)
Controller->>Controller: parseChatLogprobsChunk -> StreamChunk
Controller->>Tree: appendToNode(delta {content/reasoning, tokenLogprobs})
Tree-->>Controller: updated node
Controller-->>Chat: updated message + tokenLogprobs
deactivate Controller
Chat->>Msg: re-render with new tokenLogprobs
sequenceDiagram
autonumber
participant User as User/UI
participant App as App.tsx
participant TextHook as useTextCompletion
participant Provider as OpenAI-compatible Provider
participant Parser as openaiLogprobs
User->>App: request text completion
App->>TextHook: predict/runGeneration()
activate TextHook
TextHook->>Provider: streamText(messages, {logprobs: true})
Provider-->>TextHook: chunks (raw + logprobs)
TextHook->>Parser: parseCompletionLogprobsChunk(rawChunk)
Parser-->>TextHook: StreamChunk {content, tokenLogprobs}
TextHook->>TextHook: update tokenLogprobs state and seedText
deactivate TextHook
TextHook-->>App: {textContent, tokenLogprobs, seedText}
App->>TextCompletionView: render overlay with TokenInlineRenderer
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes
Possibly related PRs
Suggested labels
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @CoolSpring8, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a powerful new feature set centered around token probabilities, significantly enhancing user interaction and control over AI model outputs. Users can now visually inspect the probability of each generated token and, for OpenAI-compatible models, actively reroll responses from any point, offering a deeper understanding and more granular influence on the generation process. This is complemented by new display settings for customization. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and well-implemented feature for token probability visualization and interaction. The code is generally of high quality, with good attention to detail, especially in the complex UI components for displaying token information. I've provided a few suggestions to improve type safety, address potential edge cases with token parsing, and refactor some duplicated logic for better maintainability. Overall, this is an excellent addition to the project.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/hooks/useConversationController.ts (1)
163-190: Missing dependency inareMessagesEqual.
areMessagesEqualnow callsareTokenLogprobsEqual(line 179), butareTokenLogprobsEqualis not included in the dependency array (line 189). WhileareTokenLogprobsEqualhas an empty dependency array and is stable, it should still be listed for correctness and to avoid future issues if its dependencies change.}, - [areContentsEqual], + [areContentsEqual, areTokenLogprobsEqual], );
🧹 Nitpick comments (4)
src/state/useSettingsStore.ts (1)
11-20: Token heatmap settings wiring looks consistentState, persistence, and hydration for
enableTokenHeatmap/heatmapThemeare coherent and backward compatible; defaults are applied correctly when older stored settings are missing these fields. As an optional cleanup, you could centralize the heatmap theme union into a sharedHeatmapThemetype to avoid repeating the string literals across files.Also applies to: 22-35, 52-60, 71-77, 111-122, 135-145
src/components/SettingsModal.tsx (1)
38-38: Display tab and heatmap controls are wired correctlyThe new “Display” tab correctly binds to
enableTokenHeatmap/heatmapThemeand uses reasonable copy and defaults. As an optional improvement, consider centralizing the heatmap theme values (label/value pairs) in a shared config or type to keep Settings, the store, andTokenInlineRendererin sync if you add themes later.Also applies to: 40-58, 97-101, 370-395, 397-449, 760-777, 790-795
src/hooks/useConversationController.ts (1)
363-500: Token reroll implementation is well-structured.The
handleRerollFromTokenfunction:
- Properly validates readiness and provider compatibility
- Correctly aborts active streams before starting a new one
- Builds seed tokens with the replacement injected at the correct position
- Separates reasoning vs content tokens for proper prefix construction
- Handles errors gracefully with appropriate status transitions
One observation: the
temperature: 0.3at line 435 is hardcoded. Consider exposing this as a configurable option if users need different creativity levels for rerolls.src/ai/openaiLogprobs.ts (1)
92-127: Consider type refinement for ModelMessage casting.The
as unknown as ModelMessagecasts work but are somewhat brittle. If theailibrary'sModelMessagetype evolves, these casts could silently break.Consider creating a local type or using a more specific type assertion if the
ailibrary provides better typing options:// Alternative: explicit type with the properties you need type ChatMessage = { role: string; content: MessageContent };However, this is low priority if the current approach is working reliably with the
ailibrary version in use.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
bun.lockis excluded by!**/*.lock
📒 Files selected for processing (17)
package.json(1 hunks)src/App.tsx(5 hunks)src/ai/openaiCompatible.ts(2 hunks)src/ai/openaiLogprobs.ts(1 hunks)src/ai/sendMessage.ts(4 hunks)src/components/ChatView.tsx(6 hunks)src/components/MessageItem.tsx(3 hunks)src/components/SettingsModal.tsx(6 hunks)src/components/TextCompletionView.tsx(3 hunks)src/components/TokenInlineRenderer.tsx(1 hunks)src/hooks/useConversationController.ts(6 hunks)src/hooks/useProviderReadiness.ts(5 hunks)src/hooks/useTextCompletion.ts(3 hunks)src/state/useSettingsStore.ts(6 hunks)src/tree/types.ts(2 hunks)src/tree/useConversationTree.ts(10 hunks)src/types.ts(3 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
!(**/*.md)
📄 CodeRabbit inference engine (AGENTS.md)
Indentation: use one tab (Biome default) except for Markdown files where space indentation is used
Files:
package.jsonsrc/tree/types.tssrc/state/useSettingsStore.tssrc/types.tssrc/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/tree/useConversationTree.tssrc/components/TokenInlineRenderer.tsxsrc/ai/openaiCompatible.tssrc/components/TextCompletionView.tsxsrc/ai/sendMessage.tssrc/hooks/useProviderReadiness.tssrc/components/MessageItem.tsxsrc/ai/openaiLogprobs.tssrc/hooks/useTextCompletion.tssrc/App.tsxsrc/hooks/useConversationController.ts
**/*.{ts,tsx,js,jsx,json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Biome for formatting. Run
bunx biome check --write .to format and organize imports
Files:
package.jsonsrc/tree/types.tssrc/state/useSettingsStore.tssrc/types.tssrc/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/tree/useConversationTree.tssrc/components/TokenInlineRenderer.tsxsrc/ai/openaiCompatible.tssrc/components/TextCompletionView.tsxsrc/ai/sendMessage.tssrc/hooks/useProviderReadiness.tssrc/components/MessageItem.tsxsrc/ai/openaiLogprobs.tssrc/hooks/useTextCompletion.tssrc/App.tsxsrc/hooks/useConversationController.ts
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: Use TypeScript and React 18 for development
Use tsc for type-checking. Runbunx tsc --noEmitto check for type errors
Files:
src/tree/types.tssrc/state/useSettingsStore.tssrc/types.tssrc/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/tree/useConversationTree.tssrc/components/TokenInlineRenderer.tsxsrc/ai/openaiCompatible.tssrc/components/TextCompletionView.tsxsrc/ai/sendMessage.tssrc/hooks/useProviderReadiness.tssrc/components/MessageItem.tsxsrc/ai/openaiLogprobs.tssrc/hooks/useTextCompletion.tssrc/App.tsxsrc/hooks/useConversationController.ts
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx,js,jsx}: Never commit API keys. Keys are user-provided and stored client-side (IndexedDB). Avoid server code or bundling secrets intodist/
Keep code concise and self-documenting
Files:
src/tree/types.tssrc/state/useSettingsStore.tssrc/types.tssrc/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/tree/useConversationTree.tssrc/components/TokenInlineRenderer.tsxsrc/ai/openaiCompatible.tssrc/components/TextCompletionView.tsxsrc/ai/sendMessage.tssrc/hooks/useProviderReadiness.tssrc/components/MessageItem.tsxsrc/ai/openaiLogprobs.tssrc/hooks/useTextCompletion.tssrc/App.tsxsrc/hooks/useConversationController.ts
src/components/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Components under
src/components/should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Files:
src/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/components/TokenInlineRenderer.tsxsrc/components/TextCompletionView.tsxsrc/components/MessageItem.tsx
src/**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
Prefer Tailwind utilities for CSS; add global styles in
src/index.cssonly when necessary
Files:
src/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/components/TokenInlineRenderer.tsxsrc/components/TextCompletionView.tsxsrc/components/MessageItem.tsxsrc/App.tsx
**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
When using
borderin Tailwind, also includeborder-solid(and set a color) to avoid browser defaults
Files:
src/components/ChatView.tsxsrc/components/SettingsModal.tsxsrc/components/TokenInlineRenderer.tsxsrc/components/TextCompletionView.tsxsrc/components/MessageItem.tsxsrc/App.tsx
🧠 Learnings (3)
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Providers only in `index.tsx` (e.g., `MantineProvider`). Settings state lives in `src/state/useSettingsStore.ts`, conversation state in `src/tree/useConversationTree.ts` and `src/tree/types.ts`, and AI provider helpers in `src/ai/openaiCompatible.ts`
Applied to files:
src/state/useSettingsStore.tssrc/components/SettingsModal.tsxsrc/tree/useConversationTree.tssrc/ai/sendMessage.tssrc/hooks/useProviderReadiness.tssrc/ai/openaiLogprobs.tssrc/hooks/useTextCompletion.tssrc/App.tsxsrc/hooks/useConversationController.ts
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Applies to src/components/**/*.{ts,tsx} : Components under `src/components/` should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Applied to files:
src/components/SettingsModal.tsxsrc/components/TokenInlineRenderer.tsxsrc/components/TextCompletionView.tsx
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Shared conversation/tree logic belongs in the store (`src/tree/` and `src/state/`), not in individual components
Applied to files:
src/tree/useConversationTree.ts
🧬 Code graph analysis (11)
src/tree/types.ts (1)
src/types.ts (1)
TokenLogprob(23-28)
src/state/useSettingsStore.ts (2)
src/types.ts (1)
BuiltInAvailability(56-61)src/constants/storageKeys.ts (1)
settingsKey(1-1)
src/components/ChatView.tsx (1)
src/types.ts (2)
TokenAlternative(30-33)MessageContentPart(10-12)
src/tree/useConversationTree.ts (2)
src/types.ts (2)
TokenLogprob(23-28)MessageContent(14-14)src/tree/types.ts (1)
NodeID(3-3)
src/components/TokenInlineRenderer.tsx (2)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/state/useSettingsStore.ts (1)
useSettingsStore(52-275)
src/components/TextCompletionView.tsx (1)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)
src/ai/sendMessage.ts (3)
src/types.ts (1)
TokenLogprob(23-28)src/ai/openaiLogprobs.ts (3)
toModelMessages(92-127)buildChatLogprobOptions(297-302)parseChatLogprobsChunk(129-226)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/ai/openaiLogprobs.ts (1)
src/types.ts (2)
TokenLogprob(23-28)Message(16-21)
src/hooks/useTextCompletion.ts (3)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/ai/openaiLogprobs.ts (2)
buildCompletionLogprobOptions(304-309)parseCompletionLogprobsChunk(228-295)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/App.tsx (1)
src/hooks/useConversationController.ts (1)
useConversationController(28-529)
src/hooks/useConversationController.ts (4)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/tree/useConversationTree.ts (1)
useConversationTree(224-592)src/ai/openaiLogprobs.ts (3)
toModelMessages(92-127)buildChatLogprobOptions(297-302)parseChatLogprobsChunk(129-226)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
🔇 Additional comments (34)
src/ai/sendMessage.ts (6)
1-14: LGTM!The import reorganization properly brings in the necessary types and utilities for the dual-path streaming logic with token logprobs support.
28-35: LGTM!The
appendToNodedelta signature cleanly extends to supporttokenLogprobsalongside existingcontentandreasoningfields, maintaining consistent optional semantics.
74-100: LGTM!The branching logic for determining whether to append to an existing assistant message or create a new one is well-structured. The optional chaining on
lastMessage?.roleproperly guards against undefined, ensuring theshouldAppendToAssistantcondition safely short-circuits before accessinglastMessage._metadata.uuidinside the conditional block.
115-149: LGTM!The built-in provider path correctly handles text and reasoning deltas without token logprobs (as expected for browser-based AI). The conditional prefix injection for
browser-aiis properly scoped.
150-192: LGTM!The OpenAI-compatible streaming path properly:
- Enables raw chunk access with
includeRawChunks: true- Uses
buildChatLogprobOptionsfor provider-specific logprob configuration- Safely guards against undefined/empty logprobs with optional chaining
- Handles stream errors comprehensively (string, Error instance, and fallback)
200-205: LGTM!The conditional cleanup pattern correctly prevents race conditions by only clearing the generating state if this stream is still the latest one. This ensures that a newer stream won't have its state inadvertently cleared by an older finishing stream.
src/tree/types.ts (2)
1-1: LGTM!Import correctly brings in
TokenLogprobalongside the existingMessageContentimport.
15-15: LGTM!The optional
tokenLogprobsfield cleanly extendsTreeNodeto carry token-level probability data, aligning with the broader token tracking feature without breaking existing node handling.src/ai/openaiCompatible.ts (2)
4-5: LGTM!Extracting the provider name into a constant establishes a single source of truth, which is now correctly consumed by other modules (
sendMessage.ts, hooks) per the AI summary. This aligns with the codebase convention of placing AI provider helpers in this file. Based on learnings.
22-22: LGTM!Using the constant here ensures consistency with all other usages of the provider name across the codebase.
src/types.ts (3)
7-7: LGTM!Adding
tokenLogprobstoMessageMetadataenables token-level data to flow through the message layer, complementing theTreeNodeextension insrc/tree/types.ts.
23-33: LGTM!The
TokenLogprobandTokenAlternativeinterfaces are well-designed:
TokenLogprob.probabilityis correctly optional (some providers may not include probability for the selected token)TokenAlternative.probabilityis required (alternatives always have probabilities for ranking)- The
segmentfield enables distinguishing content vs reasoning tokens for extended thinking models
68-79: LGTM!Extending
ChatProviderReadyforopenai-compatiblewithbaseURLandapiKeyenables downstream features (like token reroll) to reconstruct API calls without additional lookups.package.json (1)
16-16: @floating-ui/react version is current. The specified version^0.27.16matches the latest stable release on npm, making this dependency addition appropriate.src/hooks/useProviderReadiness.ts (1)
12-20: Base URL and API key propagation is correct
baseURL/apiKeyare correctly added to the readiness options, included in memo deps, and exposed only on the"openai-compatible"branches for chat and completion, which matches the intended usage.Also applies to: 22-30, 41-47, 57-65, 67-89
src/components/ChatView.tsx (1)
7-12: Per-message token visibility and reroll wiring look solidThe
tokenViewStatesmap, cleanup effect, andhandleTokenViewChange/handleTokenRerollcallbacks are all correctly keyed onmessage._metadata.uuidand integrate cleanly withMessageItemprops. The newonTokenRerollsurface is well-shaped.Also applies to: 33-37, 53-59, 104-113, 129-152, 269-295
src/App.tsx (1)
56-66: End-to-end token probability / reroll wiring in App looks correctThe additional
baseURL/apiKeythreading intouseProviderReadiness, the newrerollFromTokenand text-completion fields from the hooks, and their propagation intoChatViewandTextCompletionVieware all coherently wired and respect existing provider-kind constraints.Also applies to: 84-92, 94-120, 124-135, 173-178, 211-227, 238-256
src/hooks/useTextCompletion.ts (1)
5-15: Token logprob streaming and reroll logic inuseTextCompletionlooks solidThe new
tokenLogprobs/seedRefflow, includingrunGeneration,predict’s alignment+pruning logic,rerollFromToken, andoverwriteTextContent, is internally consistent and should keep the overlay in sync with the underlying text while handling edits, aborts, and rerolls correctly.Also applies to: 23-28, 29-94, 96-175, 188-244
src/components/MessageItem.tsx (4)
46-55: Well-implemented controlled/uncontrolled pattern.The token view state correctly supports both controlled (via
showTokens+onShowTokensChange) and uncontrolled (via internal state) modes, following React best practices.
63-84: Token mapping logic is sound.The
buildSelectHandlercorrectly maps the local token index (within the filteredcontentTokenEntriesorreasoningTokenEntries) back to the global index intokenLogprobs. The guard at line 79 handles edge cases appropriately.
157-163: Token view toggle button correctly gated.The toggle only renders when token data exists, and the icon/title appropriately reflects the current state.
184-191: Token rendering integration looks correct.
TokenInlineRendereris correctly used for both reasoning (withdisabledto prevent rerolls) and content (withonSelectAlternativefor rerolls). The fallback to Markdown rendering when not in token view maintains backward compatibility.Also applies to: 196-203
src/hooks/useConversationController.ts (2)
117-161: Token comparison logic is thorough.The
areTokenLogprobsEqualhelper correctly handles identity checks, null/undefined cases, length mismatches, and deep comparison of token properties including alternatives.
329-334: Correctly clears tokenLogprobs on edit.Setting
tokenLogprobs: undefinedwhen replacing an edited node ensures stale token data doesn't persist after content changes.src/tree/useConversationTree.ts (5)
35-47: Type definitions correctly extended for tokenLogprobs.The delta and function signatures are properly updated to support optional
tokenLogprobsthroughout the tree operations.
166-178: Conditional tokenLogprobs serialization is correct.The
toMessagefunction properly includestokenLogprobsin_metadataonly when present, avoiding unnecessary property pollution.
361-389: Token append logic is correct.
appendToNodeproperly spreads existing tokens with new ones, preserving order and handling the empty/undefined cases appropriately.
449-472: Clone preserves token data correctly.
cloneNodenow copiestokenLogprobsfrom the source node, maintaining data integrity during duplication.
560-590: Snapshot import handles tokenLogprobs.The
importSnapshotfunction correctly preservestokenLogprobswhen restoring conversation state from a snapshot.src/ai/openaiLogprobs.ts (5)
10-33: Logprob to probability conversion is mathematically correct.
Math.exp(logprob)correctly converts log-probabilities to probabilities. TheNumber.isFinitefilter appropriately handles edge cases like-Infinitylogprobs.
35-44: Fallback token handling is sound.Ensuring the chosen token always appears in alternatives (even with probability 0 if not returned by the API) prevents UI issues when displaying token choices.
129-226: Chat logprobs parsing handles multiple API formats.The function correctly handles:
- Both array and record formats for
top_logprobs- String and array formats for
reasoning_content- Fallback paths when primary logprobs data is unavailable
The complexity is warranted given the API format variations.
228-295: Completion logprobs parsing is consistent with chat parsing.The structure mirrors
parseChatLogprobsChunkappropriately, with the key difference being thetextfield instead ofdelta.contentand hardcodedsegment: "content"which is correct for text completions.
297-309: Provider options builders are clean and reusable.The functions correctly construct provider-specific options for enabling logprobs in API requests.
CoolSpring8
changed the title
CoolSpring8
changed the title
…lity functions and improved type safety
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
src/ai/openaiLogprobs.ts(1 hunks)src/components/SettingsModal.tsx(7 hunks)src/components/TextCompletionView.tsx(3 hunks)src/components/TokenInlineRenderer.tsx(1 hunks)src/hooks/useConversationController.ts(6 hunks)src/hooks/useTextCompletion.ts(3 hunks)src/state/useSettingsStore.ts(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- src/components/TokenInlineRenderer.tsx
- src/components/SettingsModal.tsx
🧰 Additional context used
📓 Path-based instructions (7)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: Use TypeScript and React 18 for development
Use tsc for type-checking. Runbunx tsc --noEmitto check for type errors
Files:
src/state/useSettingsStore.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/openaiLogprobs.tssrc/components/TextCompletionView.tsx
!(**/*.md)
📄 CodeRabbit inference engine (AGENTS.md)
Indentation: use one tab (Biome default) except for Markdown files where space indentation is used
Files:
src/state/useSettingsStore.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/openaiLogprobs.tssrc/components/TextCompletionView.tsx
**/*.{ts,tsx,js,jsx,json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Biome for formatting. Run
bunx biome check --write .to format and organize imports
Files:
src/state/useSettingsStore.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/openaiLogprobs.tssrc/components/TextCompletionView.tsx
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx,js,jsx}: Never commit API keys. Keys are user-provided and stored client-side (IndexedDB). Avoid server code or bundling secrets intodist/
Keep code concise and self-documenting
Files:
src/state/useSettingsStore.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/openaiLogprobs.tssrc/components/TextCompletionView.tsx
src/components/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Components under
src/components/should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Files:
src/components/TextCompletionView.tsx
src/**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
Prefer Tailwind utilities for CSS; add global styles in
src/index.cssonly when necessary
Files:
src/components/TextCompletionView.tsx
**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
When using
borderin Tailwind, also includeborder-solid(and set a color) to avoid browser defaults
Files:
src/components/TextCompletionView.tsx
🧠 Learnings (3)
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Providers only in `index.tsx` (e.g., `MantineProvider`). Settings state lives in `src/state/useSettingsStore.ts`, conversation state in `src/tree/useConversationTree.ts` and `src/tree/types.ts`, and AI provider helpers in `src/ai/openaiCompatible.ts`
Applied to files:
src/state/useSettingsStore.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/openaiLogprobs.ts
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Applies to src/components/**/*.{ts,tsx} : Components under `src/components/` should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Applied to files:
src/components/TextCompletionView.tsx
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Applies to **/*.{ts,tsx} : Use TypeScript and React 18 for development
Applied to files:
src/components/TextCompletionView.tsx
🧬 Code graph analysis (5)
src/state/useSettingsStore.ts (2)
src/types.ts (2)
ProviderEntry(44-54)BuiltInAvailability(56-61)src/constants/storageKeys.ts (1)
settingsKey(1-1)
src/hooks/useTextCompletion.ts (3)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/ai/openaiLogprobs.ts (2)
buildCompletionLogprobOptions(325-330)parseCompletionLogprobsChunk(254-316)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/hooks/useConversationController.ts (3)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/ai/openaiLogprobs.ts (3)
toModelMessages(123-158)buildChatLogprobOptions(318-323)parseChatLogprobsChunk(160-252)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/ai/openaiLogprobs.ts (1)
src/types.ts (2)
TokenLogprob(23-28)Message(16-21)
src/components/TextCompletionView.tsx (1)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)
🔇 Additional comments (24)
src/state/useSettingsStore.ts (7)
11-17: LGTM! Clean type definition.The heatmap theme constant and type are well-structured. Using
as constensures literal types for the theme options, and deriving the type from the constant maintains a single source of truth.
23-24: LGTM! Proper storage schema extension.The new heatmap settings are correctly added to the persisted settings interface with appropriate types.
34-35: LGTM! Consistent state API extension.The new heatmap fields and setter methods follow the established pattern for settings state. The async signatures align with other persistence operations.
Also applies to: 40-41
60-75: LGTM! Proper persistence implementation.The new heatmap settings are correctly extracted from state and included in the persisted object, following the same pattern as existing settings.
81-82: LGTM! Sensible default values.The defaults (enableTokenHeatmap: false, heatmapTheme: "traffic-light") are appropriate. Making the heatmap opt-in is a good UX choice.
121-128: LGTM! Correct setter implementation.Both setters follow the established pattern: update in-memory state, then persist. The implementation is consistent with other settings like
setEnableBeforeUnloadWarning.
149-150: LGTM! Backward-compatible hydration.The hydration logic correctly handles existing users who don't have these settings stored yet, falling back to sensible defaults that match the initial state.
src/ai/openaiLogprobs.ts (6)
16-39: LGTM!The normalization logic correctly handles both array and object formats, properly converts log probabilities to probabilities using
Math.exp, and safely filters out non-finite values.
41-50: LGTM!The fallback logic correctly ensures that the selected token always appears in the alternatives list, using probability 0 when it's not already present.
123-158: LGTM!The function now properly handles content type conversions without unsafe type assertions. System and assistant messages are correctly converted to text-only content via helper functions, while user messages preserve multipart structure including images. The past review concern has been addressed.
160-252: LGTM!The parsing logic correctly handles both content and reasoning segments, properly allows empty string tokens (as some models may produce them), and includes appropriate fallback logic to reconstruct text from tokens when delta content is missing.
254-316: LGTM!The completion logprob parsing correctly handles token extraction and includes proper fallback logic to reconstruct text from tokens when
choice.textis unavailable.
318-330: LGTM!The option builder functions correctly construct provider-specific logprob configurations with sensible defaults.
src/hooks/useTextCompletion.ts (4)
29-94: LGTM!The generation function properly manages streaming, abort controllers, and token logprob updates. Error handling correctly distinguishes between aborts and actual errors.
96-142: LGTM!The predict function cleanly validates token reusability in a single pass and branches appropriately. The duplicate logic mentioned in past reviews has been consolidated.
144-153: LGTM!The cancel function properly aborts the active controller and resets generation state.
192-199: LGTM!The overwrite function correctly resets the seed and clears token state when manually replacing content.
src/components/TextCompletionView.tsx (3)
1-32: LGTM!Imports and props are correctly defined. The past review concerns about React type imports and unused props have been resolved.
33-92: LGTM!Mouse handling and RAF cleanup are properly implemented. The requestAnimationFrame is correctly cancelled on mousemove, mouse leave, and component unmount. Past review concerns have been addressed.
94-178: LGTM!The two-layer rendering approach correctly overlays tokens while keeping the textarea interactive. The
showTokenOverlayprop is properly wired into the rendering logic.src/hooks/useConversationController.ts (4)
117-161: LGTM!The token logprob equality checker correctly performs deep comparison of tokens, probabilities, segments, and alternatives arrays.
329-349: LGTM!Correctly clears
tokenLogprobswhen finishing an edit, as the original token metadata is no longer valid after content modification.
363-500: LGTM!The token reroll implementation is comprehensive and correct. It properly validates provider compatibility, constructs seed tokens with correct segment metadata, separates reasoning and content segments, and streams the regenerated response with logprob tracking. Error handling and cleanup are properly implemented.
527-527: LGTM!The
rerollFromTokenfunction is correctly exposed in the public API.
… context handling
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
src/hooks/useTextCompletion.ts (1)
100-142: Token alignment / reuse logic inpredictlooks correct; consider extracting if reused elsewhereThe
canReuseTokens+currentCheckwalk correctly ensures:
- Only prefixes where
seedRef.current + concatenated tokensremain a prefix oftextContentare reused.- Pure tail deletions truncate
tokenLogprobs(viavalidTokenCount) and continue generation.- Any mid-span edit or seed misalignment resets tokens and treats
textContentas a fresh seed.If you have similar alignment logic in the chat completion hook or other flows, consider extracting this into a shared helper to avoid future drift and keep the behavior consistent across modes.
🧹 Nitpick comments (1)
src/hooks/useTextCompletion.ts (1)
155-191: Reroll behavior and segment preservation look good; verify intended handling of suffix/user editsThe
rerollFromTokenimplementation is solid:
- Safely early-returns on out-of-range
tokenIndex.- Cancels in‑flight generation when
isGeneratingis true.- Rebuilds alternatives with the new choice first and removes duplicate tokens.
- Preserves
segmentviasegment: target.segment, resolving the previous missing-segment issue.- Recomputes
newSeedfromseedRef.current + prefixTokens + replacementand updatestokenLogprobsaccordingly before regenerating.One behavior to double‑check: by rebuilding
newSeedfromseedRef.currentand the prefix tokens, then callingsetTextContent(newSeed), any manual user edits after the rerolled token (or outside the tracked token region) are discarded. If the UI only exposes reroll for purely model-generated spans this is fine; otherwise, you may want to confirm this is the desired UX, or consider rebasing rerolls relative to the currenttextContentinstead ofseedRef.current.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/hooks/useTextCompletion.ts(3 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: Use TypeScript and React 18 for development
Use tsc for type-checking. Runbunx tsc --noEmitto check for type errors
Files:
src/hooks/useTextCompletion.ts
!(**/*.md)
📄 CodeRabbit inference engine (AGENTS.md)
Indentation: use one tab (Biome default) except for Markdown files where space indentation is used
Files:
src/hooks/useTextCompletion.ts
**/*.{ts,tsx,js,jsx,json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Biome for formatting. Run
bunx biome check --write .to format and organize imports
Files:
src/hooks/useTextCompletion.ts
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx,js,jsx}: Never commit API keys. Keys are user-provided and stored client-side (IndexedDB). Avoid server code or bundling secrets intodist/
Keep code concise and self-documenting
Files:
src/hooks/useTextCompletion.ts
🧠 Learnings (1)
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Providers only in `index.tsx` (e.g., `MantineProvider`). Settings state lives in `src/state/useSettingsStore.ts`, conversation state in `src/tree/useConversationTree.ts` and `src/tree/types.ts`, and AI provider helpers in `src/ai/openaiCompatible.ts`
Applied to files:
src/hooks/useTextCompletion.ts
🔇 Additional comments (2)
src/hooks/useTextCompletion.ts (2)
5-28: Imports and new token/seed state wiring look consistentThe added imports for OpenAI-compatible logprob handling and the new
tokenLogprobs/seedRefstate are coherently used below; types line up with later usage and there are no apparent typing or lifecycle issues here.
29-94: Streaming + logprob aggregation and abort handling are robustThe
runGenerationimplementation correctly:
- Guards against concurrent runs (
hasActiveGeneration+abortControllerRef).- Updates
seedRefonly when requested viaupdateSeed.- Streams text deltas into
textContentand pushes parsedtokenLogprobsincrementally.- Distinguishes between user-initiated aborts (no toast) and real errors (logs + toast) while ensuring
isGeneratingandabortControllerRefare reset infinally.No changes requested here.
…nt race conditions
|
@codex review |
There was a problem hiding this comment.
Code Review
This is an excellent pull request that introduces a sophisticated set of features for token-level interaction, including probability visualization and rerolling. The implementation is well-thought-out, especially the complex UI for text completion and the intelligent logic in the useTextCompletion hook. The code is generally of high quality.
My review includes a few suggestions to improve robustness and maintainability:
- Increasing type safety when parsing raw API stream chunks to prevent potential runtime errors.
- Refactoring duplicated stream-processing logic to improve maintainability.
- Replacing a hardcoded color value with a theme-based Tailwind class for better consistency.
- Adding comments to a particularly complex function to improve its readability.
Overall, this is a fantastic contribution.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…ntroducing utility functions for improved readability and maintainability
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
src/ai/streamUtils.ts (1)
34-50: Theif-ifpattern allows processing multiple part types in a single iteration.The use of sequential
ifstatements rather thanif-else-ifmeans all conditions are evaluated for each part. This appears intentional to allow arawpart to potentially also be anerrorpart, but in practice theailibrary'sStreamParttypes are mutually exclusive. Consider usingelse ifor aswitchstatement for clarity and marginal performance benefit.for await (const part of fullStream) { - if (part.type === "text-delta" && part.text) { + switch (part.type) { + case "text-delta": + if (part.text) append({ content: part.text }); + break; + case "reasoning-delta": + if (part.text) append({ reasoning: part.text }); + break; + case "raw": + if (parseRawChunk) { + const chunk = parseRawChunk(part.rawValue); + if (chunk?.tokenLogprobs?.length) { + append({ tokenLogprobs: chunk.tokenLogprobs }); + } + } + break; + case "error": + throw new Error(toErrorMessage(part.error)); + } - append({ content: part.text }); - } - if (part.type === "reasoning-delta" && part.text) { - append({ reasoning: part.text }); - } - if (part.type === "raw" && parseRawChunk) { - const chunk = parseRawChunk(part.rawValue); - if (chunk?.tokenLogprobs?.length) { - append({ tokenLogprobs: chunk.tokenLogprobs }); - } - } - if (part.type === "error") { - throw new Error(toErrorMessage(part.error)); - } }src/hooks/useConversationController.ts (2)
427-462: Consider avoiding the detached async IIFE pattern.The
void (async () => { ... })();pattern fires off an async operation without awaiting it, which meanshandleRerollFromTokenreturns theassistantIdbefore the streaming completes. While this may be intentional for UI responsiveness, it makes error handling less predictable for callers. The error is caught internally and shown via toast, but the caller has no way to know if the operation succeeded.If this is intentional, consider adding a comment explaining why the fire-and-forget pattern is used here.
436-436: Hardcoded temperature value.The temperature
0.3is hardcoded here and insendMessage.ts. Consider extracting this to a shared constant or making it configurable to maintain consistency and ease future adjustments.src/hooks/useTextCompletion.ts (1)
58-79: Stream processing is inlined here but usesprocessFullStreamelsewhere.Unlike
sendMessage.tsanduseConversationController.tswhich use the sharedprocessFullStreamutility, this file manually iterates over the stream. Consider usingprocessFullStreamfor consistency, though the inline approach here allows direct access tosetTextContentandsetTokenLogprobsvia Immer drafts.src/components/TextCompletionView.tsx (1)
147-147: Minor: Consider adding focus ring styles for accessibility.The textarea has
ring-0which removes the focus ring. While the caret provides some focus indication, a subtle focus ring improves accessibility for keyboard users. As per coding guidelines, prefer Tailwind utilities for CSS.- className="absolute inset-0 h-full w-full resize-none overflow-hidden bg-transparent p-3 font-sans text-base leading-8 text-transparent caret-slate-900 focus:outline-none dark:text-transparent dark:caret-slate-100 border-none outline-none ring-0" + className="absolute inset-0 h-full w-full resize-none overflow-hidden bg-transparent p-3 font-sans text-base leading-8 text-transparent caret-slate-900 focus:outline-none focus:ring-2 focus:ring-blue-500/20 dark:text-transparent dark:caret-slate-100 border-none outline-none"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/ai/openaiLogprobs.ts(1 hunks)src/ai/sendMessage.ts(4 hunks)src/ai/streamUtils.ts(1 hunks)src/components/TextCompletionView.tsx(3 hunks)src/hooks/useConversationController.ts(6 hunks)src/hooks/useTextCompletion.ts(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- src/ai/openaiLogprobs.ts
🧰 Additional context used
📓 Path-based instructions (7)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: Use TypeScript and React 18 for development
Use tsc for type-checking. Runbunx tsc --noEmitto check for type errors
Files:
src/ai/streamUtils.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/sendMessage.tssrc/components/TextCompletionView.tsx
!(**/*.md)
📄 CodeRabbit inference engine (AGENTS.md)
Indentation: use one tab (Biome default) except for Markdown files where space indentation is used
Files:
src/ai/streamUtils.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/sendMessage.tssrc/components/TextCompletionView.tsx
**/*.{ts,tsx,js,jsx,json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Biome for formatting. Run
bunx biome check --write .to format and organize imports
Files:
src/ai/streamUtils.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/sendMessage.tssrc/components/TextCompletionView.tsx
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx,js,jsx}: Never commit API keys. Keys are user-provided and stored client-side (IndexedDB). Avoid server code or bundling secrets intodist/
Keep code concise and self-documenting
Files:
src/ai/streamUtils.tssrc/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/sendMessage.tssrc/components/TextCompletionView.tsx
src/components/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Components under
src/components/should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Files:
src/components/TextCompletionView.tsx
src/**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
Prefer Tailwind utilities for CSS; add global styles in
src/index.cssonly when necessary
Files:
src/components/TextCompletionView.tsx
**/*.{tsx,css}
📄 CodeRabbit inference engine (AGENTS.md)
When using
borderin Tailwind, also includeborder-solid(and set a color) to avoid browser defaults
Files:
src/components/TextCompletionView.tsx
🧠 Learnings (2)
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Providers only in `index.tsx` (e.g., `MantineProvider`). Settings state lives in `src/state/useSettingsStore.ts`, conversation state in `src/tree/useConversationTree.ts` and `src/tree/types.ts`, and AI provider helpers in `src/ai/openaiCompatible.ts`
Applied to files:
src/hooks/useTextCompletion.tssrc/hooks/useConversationController.tssrc/ai/sendMessage.ts
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Applies to src/components/**/*.{ts,tsx} : Components under `src/components/` should include their own UI behavior (hover/edit states, popovers, menus) rather than being purely presentational
Applied to files:
src/components/TextCompletionView.tsx
🧬 Code graph analysis (4)
src/ai/streamUtils.ts (2)
src/types.ts (1)
TokenLogprob(23-28)src/ai/openaiLogprobs.ts (1)
StreamChunk(4-8)
src/hooks/useTextCompletion.ts (3)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)src/ai/openaiLogprobs.ts (2)
buildCompletionLogprobOptions(442-447)parseCompletionLogprobsChunk(371-433)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/ai/sendMessage.ts (4)
src/types.ts (1)
TokenLogprob(23-28)src/ai/streamUtils.ts (1)
processFullStream(24-51)src/ai/openaiLogprobs.ts (3)
toModelMessages(123-158)buildChatLogprobOptions(435-440)parseChatLogprobsChunk(295-369)src/ai/openaiCompatible.ts (1)
OPENAI_COMPATIBLE_PROVIDER_NAME(4-4)
src/components/TextCompletionView.tsx (1)
src/types.ts (2)
TokenLogprob(23-28)TokenAlternative(30-33)
🔇 Additional comments (21)
src/ai/streamUtils.ts (2)
1-15: LGTM!The type definitions for
StreamPartandStreamAppenderare well-structured and align with theStreamChunkandTokenLogprobtypes defined in the related modules.
17-22: LGTM!The
toErrorMessagehelper provides sensible fallback error messaging with proper type narrowing.src/hooks/useConversationController.ts (4)
118-162: LGTM!The
areTokenLogprobsEqualcomparison function is thorough, correctly handling identity checks, nullability, length mismatches, and deep comparison of token properties and alternatives.
178-191: LGTM!The
areMessagesEqualfunction is correctly extended to includetokenLogprobscomparison, and the dependency array is updated accordingly.
330-350: LGTM!Clearing
tokenLogprobswhen editing a message is correct behavior since the edited content invalidates the original token probability data.
364-478: LGTM overall for the reroll implementation.The
handleRerollFromTokenfunction correctly:
- Validates provider compatibility (OpenAI-compatible only)
- Aborts active streams before starting new generation
- Reconstructs seed tokens with the replacement
- Separates reasoning vs content segments
- Handles abort and error states appropriately
The logic for building
seedTokens,seedText,seedReasoning, andseedContentis sound.src/ai/sendMessage.ts (6)
29-36: LGTM!The
appendToNodedelta type is correctly extended to includetokenLogprobsfor token-level probability tracking.
79-95: Verify the "continue generation" logic is intentional.When
hasContentis false and the last message is from the assistant, the code reuses the existing assistant node (line 85). This enables a "continue generation" flow where the user can extend the assistant's response without adding a new user message.Ensure this behavior is documented and that callers understand that sending an empty prompt after an assistant message will continue generating into that same assistant node rather than creating a new conversation turn.
100-102: LGTM!Filtering out
toolmessages before streaming is correct since the model providers don't need tool messages in this context.
105-130: LGTM!The built-in provider streaming path correctly uses
processFullStreamwithout the raw chunk parser since the built-in model doesn't provide logprobs.
131-148: LGTM!The OpenAI-compatible streaming path correctly enables raw chunks and uses
parseChatLogprobsChunkto extract token logprobs. The use ofprocessFullStreamconsolidates the stream processing logic as intended.
156-161: LGTM!The finalization logic correctly only clears
isGeneratingwhen this stream is the latest, preventing race conditions with concurrent streams.src/hooks/useTextCompletion.ts (5)
29-94: LGTM!The
runGenerationfunction correctly manages the abort controller lifecycle, handles streaming with token logprobs, and ensures proper cleanup in the finally block.
96-144: LGTM!The
predictfunction logic for determining whether to reuse existing token probabilities is well-documented with inline comments explaining the two paths: continuing generation with existing tokens vs. treating current text as a new seed.
157-193: LGTM!The
rerollFromTokenfunction correctly preserves thesegmentproperty (line 175) as flagged in a past review. The alternatives array is properly reordered to place the replacement first.
195-205: LGTM!The
overwriteTextContentfunction correctly cancels any in-flight generation before resetting state, addressing the past review concern.
216-216: ReturningseedRef.currentyields a snapshot, not reactive state.
seedText: seedRef.currentreturns the ref's value at the time of render, but since it's not React state, changes toseedRef.currentwon't trigger re-renders. Consumers may see stale values. If reactivity is needed, consider using state instead of a ref, or document thatseedTextis a snapshot.src/components/TextCompletionView.tsx (4)
1-32: LGTM!The imports are correct with
MouseEventproperly imported from React. The new props (tokenLogprobs,onTokenReroll,showTokenOverlay,generatedPrefix) are well-typed and destructured with sensible defaults.
44-76: LGTM!The RAF-based hit-testing for hover detection is well-implemented:
- Properly cancels previous animation frames before scheduling new ones
- Temporarily disables pointer events on textarea to hit-test the overlay
- Correctly handles the token menu case to avoid changing selection
- Restores pointer events after hit-testing
78-92: LGTM!The RAF cleanup is properly handled in both
handleMouseLeave(for immediate cleanup) and theuseEffectcleanup function (for unmount), addressing the past review concern about stale callbacks firing after unmount.
106-151: LGTM!The layered UI approach with a relative overlay div and absolute textarea is a clever solution for combining token visualization with text editing. The overlay dictates height while the transparent textarea captures input.
…f streaming text and token log probabilities
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/hooks/useTextCompletion.ts (1)
88-136: Consider adding explanatory comments for the token reuse strategy.The predict function implements sophisticated token reuse logic that correctly handles various editing scenarios. However, as noted in previous reviews, the strategy isn't immediately obvious. Adding brief comments explaining the high-level approach (what canReuseTokens means, why there are two generation paths) would improve maintainability.
Example:
const predict = useCallback(async () => { if (isGenerating) { return; } + // Strategy: Reuse existing token probabilities if the user only appended text. + // This allows continuing generation without restarting. If the user edited + // the generated portion, treat the entire current text as a new seed. const canReuseTokens =
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/hooks/useTextCompletion.ts(3 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: Use TypeScript and React 18 for development
Use tsc for type-checking. Runbunx tsc --noEmitto check for type errors
Files:
src/hooks/useTextCompletion.ts
!(**/*.md)
📄 CodeRabbit inference engine (AGENTS.md)
Indentation: use one tab (Biome default) except for Markdown files where space indentation is used
Files:
src/hooks/useTextCompletion.ts
**/*.{ts,tsx,js,jsx,json}
📄 CodeRabbit inference engine (AGENTS.md)
Use Biome for formatting. Run
bunx biome check --write .to format and organize imports
Files:
src/hooks/useTextCompletion.ts
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx,js,jsx}: Never commit API keys. Keys are user-provided and stored client-side (IndexedDB). Avoid server code or bundling secrets intodist/
Keep code concise and self-documenting
Files:
src/hooks/useTextCompletion.ts
🧠 Learnings (1)
📚 Learning: 2025-11-26T05:58:11.231Z
Learnt from: CR
Repo: CoolSpring8/iaslate PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-11-26T05:58:11.231Z
Learning: Providers only in `index.tsx` (e.g., `MantineProvider`). Settings state lives in `src/state/useSettingsStore.ts`, conversation state in `src/tree/useConversationTree.ts` and `src/tree/types.ts`, and AI provider helpers in `src/ai/openaiCompatible.ts`
Applied to files:
src/hooks/useTextCompletion.ts
🔇 Additional comments (6)
src/hooks/useTextCompletion.ts (6)
5-15: LGTM: Imports align with token tracking requirements.All new imports are necessary for the token-level probability tracking and logprob parsing functionality.
26-28: LGTM: State initialization supports token tracking.The tokenLogprobs state and seedRef are appropriately initialized for tracking per-token probabilities and managing the regeneration seed.
30-86: Well-structured generation orchestration.The runGeneration function properly manages abort controllers, error handling, and cleanup. The early return on line 36 when a generation is already active prevents race conditions but silently ignores subsequent calls - this appears intentional and reasonable for preventing concurrent generations.
149-185: LGTM: Token reroll logic is correct.The rerollFromToken function properly handles token replacement, including the segment preservation that was flagged in previous reviews (line 167). The function correctly cancels active generations, updates alternatives, and restarts generation from the modified point.
187-197: LGTM: Overwrite handles cancellation correctly.The overwriteTextContent function properly cancels any in-flight generation before resetting state (lines 189-191), addressing the race condition concern from previous reviews. This ensures old streamed tokens don't mix with the new content.
199-209: LGTM: Expanded API supports token tracking and reroll UI.The return value appropriately exposes all functionality required for token-level probability tracking and interactive rerolling, aligning with the PR objectives.
1 participant
Summary by CodeRabbit
New Features
Other
✏️ Tip: You can customize this high-level summary in your review settings.