Skip to content

Commit 161bd92

Browse files
nstrayernstrayer
committed
Optimize notebook mode for caching. Greatly improves speed.
1 parent 4381230 commit 161bd92

File tree

6 files changed

+122
-162
lines changed

6 files changed

+122
-162
lines changed

‎extensions/positron-assistant/src/md/prompts/chat/notebook-mode-agent.md‎

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -26,29 +26,10 @@ You MUST use notebook-specific tools. NEVER use file tools.
2626
✓ Use EditNotebookCells tool
2727
</anti-patterns>
2828

29-
<notebook-context>
29+
<notebook-context-instructions>
3030
You are assisting the user within a Jupyter notebook in Positron.
31-
32-
<notebook-info>
33-
<kernel language="{{positron.notebookContext.kernelLanguage}}" id="{{positron.notebookContext.kernelId}}"/>
34-
<cell-count total="{{positron.notebookContext.cellCount}}" selected="{{positron.notebookContext.selectedCells.length}}"/>
35-
{{@if(positron.notebookContext.allCells)}}
36-
<context-mode>Full notebook (< 20 cells, all cells provided below)</context-mode>
37-
{{#else}}
38-
<context-mode>Selected cells only (use GetNotebookCells for other cells)</context-mode>
39-
{{/if}}
40-
</notebook-info>
41-
42-
<selected-cells>
43-
{{positron.notebookSelectedCellsInfo}}
44-
</selected-cells>
45-
46-
{{@if(positron.notebookAllCellsInfo)}}
47-
{{positron.notebookAllCellsInfo}}
48-
{{/if}}
49-
50-
{{positron.notebookContextNote}}
51-
</notebook-context>
31+
The current notebook state (kernel info, cell contents, selection) is provided in a separate context message below.
32+
</notebook-context-instructions>
5233

5334
<workflows>
5435
**Analyze/explain:** Reference cells by **index** ("cell 0", "cell 3"). Use GetNotebookCells with `cellIndices` for additional cells. Check execution order [N], status, and success/failure.
@@ -65,15 +46,13 @@ You are assisting the user within a Jupyter notebook in Positron.
6546
</workflows>
6647

6748
<critical-rules>
68-
- ALWAYS reference cells by their **zero-based index** (first cell = index 0, second cell = index 1, last cell = {{positron.notebookContext.cellCount}} - 1)
69-
- Cell indices are shown in the context above (e.g., `<cell index="0">`, `<cell index="1">`)
49+
- ALWAYS reference cells by their **zero-based index** (first cell = index 0, second cell = index 1, etc.)
50+
- Cell indices are shown in the notebook context (e.g., `<cell index="0">`, `<cell index="1">`)
7051
- MUST check execution state: order [N], status (running/pending/idle), success/failure, duration
7152
- MUST consider cell dependencies before modifications/execution
7253
- **IMPORTANT:** When you add or delete cells, remember that indices shift:
7354
- Adding cell at index 2: cells 2+ become 3+
7455
- Deleting cell at index 2: cells 3+ become 2+
7556
- MUST maintain clear notebook structure with appropriate markdown documentation
76-
77-
**Notebook URI (for reference only):** {{positron.notebookContext.uri}}
7857
</critical-rules>
7958
{{/if}}

‎extensions/positron-assistant/src/md/prompts/chat/notebook-mode-ask.md‎

Lines changed: 4 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -28,33 +28,14 @@ If the user requests cell modifications or execution, explain that these require
2828
✓ Use EditNotebookCells tool
2929
</anti-patterns>
3030

31-
<notebook-context>
31+
<notebook-context-instructions>
3232
You are assisting the user within a Jupyter notebook in Positron.
33-
34-
<notebook-info>
35-
<kernel language="{{positron.notebookContext.kernelLanguage}}" id="{{positron.notebookContext.kernelId}}"/>
36-
<cell-count total="{{positron.notebookContext.cellCount}}" selected="{{positron.notebookContext.selectedCells.length}}"/>
37-
{{@if(positron.notebookContext.allCells)}}
38-
<context-mode>Full notebook (< 20 cells, all cells provided below)</context-mode>
39-
{{#else}}
40-
<context-mode>Selected cells only (use GetNotebookCells for other cells)</context-mode>
41-
{{/if}}
42-
</notebook-info>
43-
44-
<selected-cells>
45-
{{positron.notebookSelectedCellsInfo}}
46-
</selected-cells>
47-
48-
{{@if(positron.notebookAllCellsInfo)}}
49-
{{positron.notebookAllCellsInfo}}
50-
{{/if}}
51-
52-
{{positron.notebookContextNote}}
53-
</notebook-context>
33+
The current notebook state (kernel info, cell contents, selection) is provided in a separate context message below.
34+
</notebook-context-instructions>
5435

5536
<critical-rules>
5637
- ALWAYS reference cells by their **zero-based index** (first cell = index 0, second cell = index 1, etc.)
57-
- Cell indices are shown in the context above (e.g., `<cell index="0">`, `<cell index="1">`)
38+
- Cell indices are shown in the notebook context (e.g., `<cell index="0">`, `<cell index="1">`)
5839
- MUST consider notebook's execution state, cell dependencies, and execution history
5940
- MUST pay attention to cell status (selection, execution status, execution order, success/failure, duration)
6041
- Execution order numbers [N] indicate sequence in which cells were executed
@@ -68,6 +49,4 @@ You are assisting the user within a Jupyter notebook in Positron.
6849

6950
**Debug issues:** Check cell execution status, order, success/failure. Use GetCellOutputs with `cellIndex` to inspect errors/outputs. Consider cell dependencies and sequence.
7051
</workflows>
71-
72-
**Notebook URI (for reference only):** {{positron.notebookContext.uri}}
7352
{{/if}}

‎extensions/positron-assistant/src/md/prompts/chat/notebook-mode-edit.md‎

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -29,29 +29,10 @@ If the user requests cell execution, suggest switching to Agent mode for executi
2929
✓ Use EditNotebookCells tool
3030
</anti-patterns>
3131

32-
<notebook-context>
32+
<notebook-context-instructions>
3333
You are assisting the user within a Jupyter notebook in Positron with modification access.
34-
35-
<notebook-info>
36-
<kernel language="{{positron.notebookContext.kernelLanguage}}" id="{{positron.notebookContext.kernelId}}"/>
37-
<cell-count total="{{positron.notebookContext.cellCount}}" selected="{{positron.notebookContext.selectedCells.length}}"/>
38-
{{@if(positron.notebookContext.allCells)}}
39-
<context-mode>Full notebook (< 20 cells, all cells provided below)</context-mode>
40-
{{#else}}
41-
<context-mode>Selected cells only (use GetNotebookCells for other cells)</context-mode>
42-
{{/if}}
43-
</notebook-info>
44-
45-
<selected-cells>
46-
{{positron.notebookSelectedCellsInfo}}
47-
</selected-cells>
48-
49-
{{@if(positron.notebookAllCellsInfo)}}
50-
{{positron.notebookAllCellsInfo}}
51-
{{/if}}
52-
53-
{{positron.notebookContextNote}}
54-
</notebook-context>
34+
The current notebook state (kernel info, cell contents, selection) is provided in a separate context message below.
35+
</notebook-context-instructions>
5536

5637
<workflows>
5738
**Mode capabilities:** View, modify, add, delete cells. Cannot execute (Agent mode only). If execution requested: "Cannot execute in Edit mode. Switch to Agent mode to run cells."
@@ -70,8 +51,8 @@ You are assisting the user within a Jupyter notebook in Positron with modificati
7051
</workflows>
7152

7253
<critical-rules>
73-
- ALWAYS reference cells by their **zero-based index** (first cell = index 0, second cell = index 1, last cell = {{positron.notebookContext.cellCount}} - 1)
74-
- Cell indices are shown in the context above (e.g., `<cell index="0">`, `<cell index="1">`)
54+
- ALWAYS reference cells by their **zero-based index** (first cell = index 0, second cell = index 1, etc.)
55+
- Cell indices are shown in the notebook context (e.g., `<cell index="0">`, `<cell index="1">`)
7556
- MUST check execution state: order [N], status (running/pending/idle), success/failure, duration
7657
- MUST consider cell dependencies before modifications
7758
- **IMPORTANT:** When you add or delete cells, remember that indices shift:
@@ -80,7 +61,5 @@ You are assisting the user within a Jupyter notebook in Positron with modificati
8061
- When modifying cells, preserve notebook structure and maintain cell dependencies
8162
- When adding cells, choose positions that respect logical flow
8263
- When execution requested → "Cannot execute in Edit mode. Switch to Agent mode to run cells."
83-
84-
**Notebook URI (for reference only):** {{positron.notebookContext.uri}}
8564
</critical-rules>
8665
{{/if}}

‎extensions/positron-assistant/src/participants.ts‎

Lines changed: 34 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import { getCommitChanges } from './git.js';
2121
import { getEnabledTools, getPositronContextPrompts } from './api.js';
2222
import { TokenUsage } from './tokens.js';
2323
import { PromptRenderer } from './promptRender.js';
24-
import { SerializedNotebookContext, serializeNotebookContext, getAttachedNotebookContext } from './tools/notebookUtils.js';
24+
import { getAttachedNotebookContext, serializeNotebookContextAsUserMessage } from './tools/notebookUtils.js';
2525

2626
export enum ParticipantID {
2727
/** The participant used in the chat pane in Ask mode. */
@@ -328,9 +328,25 @@ abstract class PositronAssistantParticipant implements IPositronAssistantPartici
328328
addCacheControlBreakpointPartsToLastUserMessages(messages, 2);
329329

330330
// Add a user message containing context about the request, workspace, running sessions, etc.
331-
// The context message is the second-last message in the chat messages.
331+
// The context message is fairly stable during iterative workflows, so it comes before
332+
// the more volatile notebook context.
332333
const contextInfo = await attachContextInfo(messages);
333334

335+
// Add notebook context as a separate user message with cache breakpoint.
336+
// This is placed AFTER the general context message because notebook state (cells, selection)
337+
// changes more frequently than session variables. By placing the most volatile content last,
338+
// we maximize cache hits on the stable prefix (system prompt + context message).
339+
const notebookContext = await getAttachedNotebookContext(request);
340+
if (notebookContext) {
341+
const notebookContextContent = serializeNotebookContextAsUserMessage(notebookContext);
342+
const notebookMessage = vscode.LanguageModelChatMessage.User([
343+
new vscode.LanguageModelTextPart(notebookContextContent),
344+
languageModelCacheBreakpointPart(), // Cache breakpoint after notebook context
345+
]);
346+
messages.push(notebookMessage);
347+
log.debug(`[participant] Added notebook context as user message (${notebookContextContent.length} chars)`);
348+
}
349+
334350
// Add the user's prompt.
335351
// The user's prompt is the last message in the chat messages.
336352
// If this ordering changes, also update getUserPrompt for the EchoLanguageModel in extensions/positron-assistant/src/models.ts.
@@ -908,70 +924,25 @@ export class PositronAssistantEditorParticipant extends PositronAssistantPartici
908924

909925
}
910926

911-
/** The participant used in notebook inline chats. */
927+
/**
928+
* The participant used in notebook inline chats.
929+
*
930+
* Notebook context is injected as a separate user message in defaultRequestHandler()
931+
* rather than via getCustomPrompt(). This allows the system prompt to contain only
932+
* static instructions (cacheable), while dynamic notebook state is added as a user
933+
* message with its own cache breakpoint.
934+
*/
912935
export class PositronAssistantNotebookParticipant extends PositronAssistantEditorParticipant implements IPositronAssistantParticipant {
913936
id = ParticipantID.Notebook;
914937

915-
override async getCustomPrompt(request: vscode.ChatRequest): Promise<string> {
916-
// Check if notebook mode feature is enabled
917-
const notebookModeEnabled = vscode.workspace
918-
.getConfiguration('positron.assistant.notebookMode')
919-
.get('enable', false);
920-
921-
if (!notebookModeEnabled) {
922-
log.debug('[notebook participant] Notebook mode disabled via feature flag');
923-
return super.getCustomPrompt(request);
924-
}
925-
926-
// Get the active notebook context
927-
const notebookContext = await positron.notebooks.getContext();
928-
if (!notebookContext) {
929-
log.debug('[notebook participant] No notebook context available for inline chat');
930-
return super.getCustomPrompt(request);
931-
}
932-
933-
// Positron notebooks send requests with the location2 type of
934-
// ChatRequestEditorData which is because it's much easier/cleaner to
935-
// just route the requests from positron notebooks here without totally
936-
// changing the request body. Standard VS Code notebooks send We rely on
937-
// the editor being a positron notebook for notebook-wide awareness so
938-
// we can return early if it's not.
939-
if (!(request.location2 instanceof vscode.ChatRequestEditorData)) {
940-
// Fall back to non-positron-notebook aware behavior.
941-
return super.getCustomPrompt(request);
942-
}
943-
944-
const cellDoc = request.location2.document;
945-
946-
const cellUri = cellDoc.uri.toString();
947-
948-
// Get all cells from the notebook
949-
let allCells: positron.notebooks.NotebookCell[];
950-
try {
951-
allCells = await positron.notebooks.getCells(notebookContext.uri);
952-
} catch (err) {
953-
log.error('[notebook participant] Failed to get notebook cells:', err);
954-
return super.getCustomPrompt(request);
955-
}
956-
957-
// Find the current cell by matching URI
958-
const currentCell = allCells.find(c => c.id === cellUri);
959-
if (!currentCell) {
960-
log.debug('[notebook participant] Could not find current cell in notebook');
961-
return super.getCustomPrompt(request);
962-
}
963-
964-
const currentIndex = currentCell.index;
965-
966-
// Use unified serialization helper with current cell as anchor and full wrapping
967-
// This applies filtering logic (sliding window around current cell) and formats consistently with chat pane
968-
const serialized = serializeNotebookContext(notebookContext, {
969-
anchorIndex: currentIndex,
970-
wrapInNotebookContext: true
971-
});
972-
973-
const serializedContext = serialized.fullContext || '';
974-
return serializedContext;
938+
/**
939+
* Returns empty string to prevent the parent EditorParticipant from adding
940+
* cell content as if it were a standalone file. Notebook cells should be
941+
* presented as part of the notebook context (handled in defaultRequestHandler),
942+
* not as individual documents with selection/diagnostics.
943+
*/
944+
override async getCustomPrompt(_request: vscode.ChatRequest): Promise<string> {
945+
return '';
975946
}
976947
}
977948

‎extensions/positron-assistant/src/promptRender.ts‎

Lines changed: 4 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import * as positron from 'positron';
1010
import * as yaml from 'yaml';
1111
import { MARKDOWN_DIR } from './constants';
1212
import { log } from './extension.js';
13-
import { SerializedNotebookContext, serializeNotebookContext } from './tools/notebookUtils.js';
13+
import { SerializedNotebookContext } from './tools/notebookUtils.js';
1414
import * as xml from './xml.js';
1515

1616
const PROMPT_MODE_SELECTIONS_KEY = 'positron.assistant.promptModeSelections';
@@ -26,46 +26,28 @@ interface AugmentedRenderData {
2626
hasRSession: boolean;
2727
hasPythonSession: boolean;
2828
hasNotebookContext: boolean;
29-
notebookKernelInfo?: string;
30-
notebookSelectedCellsInfo?: string;
31-
notebookAllCellsInfo?: string;
32-
notebookContextNote?: string;
3329
}
3430

3531
/**
3632
* Simple template engine supporting conditionals and interpolation
3733
*/
3834
class PromptTemplateEngine {
3935
/**
40-
* Augment render data with helper properties for template conditions
36+
* Augment render data with helper properties for template conditions.
4137
*/
4238
private static augmentRenderData(data: PromptRenderData): PromptRenderData & AugmentedRenderData {
4339
const hasRSession = data.sessions?.some(s => s.languageId === 'r') ?? false;
4440
const hasPythonSession = data.sessions?.some(s => s.languageId === 'python') ?? false;
4541

46-
// Notebook context augmentation
42+
// Only track whether notebook context exists, not its content.
43+
// Dynamic notebook content is injected as a user message for caching.
4744
const hasNotebookContext = !!data.notebookContext;
48-
let notebookKernelInfo: string | undefined;
49-
let notebookSelectedCellsInfo: string | undefined;
50-
let notebookAllCellsInfo: string | undefined;
51-
let notebookContextNote: string | undefined;
52-
53-
if (data.notebookContext) {
54-
notebookKernelInfo = data.notebookContext.kernelInfo;
55-
notebookSelectedCellsInfo = data.notebookContext.selectedCellsInfo;
56-
notebookAllCellsInfo = data.notebookContext.allCellsInfo;
57-
notebookContextNote = data.notebookContext.contextNote;
58-
}
5945

6046
return {
6147
...data,
6248
hasRSession,
6349
hasPythonSession,
6450
hasNotebookContext,
65-
notebookKernelInfo,
66-
notebookSelectedCellsInfo,
67-
notebookAllCellsInfo,
68-
notebookContextNote,
6951
};
7052
}
7153

0 commit comments

Comments
 (0)