Skip to content

Commit a0fd132

Browse files
dlantsdlants
committed
add stop message debug view, token counting impl
1 parent e844408 commit a0fd132

File tree

7 files changed

+230
-88
lines changed

7 files changed

+230
-88
lines changed

‎node/chat/chat.ts

+61-10
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,14 @@ import {
88
type Update,
99
wrapThunk,
1010
} from "../tea/tea.ts";
11-
import { d, type View } from "../tea/view.ts";
11+
import { d, withBindings, type View } from "../tea/view.ts";
1212
import * as ToolManager from "../tools/toolManager.ts";
1313
import { type Result } from "../utils/result.ts";
1414
import { Counter } from "../utils/uniqueId.ts";
1515
import type { Nvim } from "nvim-node";
1616
import type { Lsp } from "../lsp.ts";
1717
import {
18-
getClient,
18+
getClient as getProvider,
1919
type ProviderMessage,
2020
type ProviderMessageContent,
2121
type ProviderName,
@@ -24,6 +24,7 @@ import {
2424
} from "../providers/provider.ts";
2525
import { assertUnreachable } from "../utils/assertUnreachable.ts";
2626
import { DEFAULT_OPTIONS, type MagentaOptions } from "../options.ts";
27+
import { getOption } from "../nvim/nvim.ts";
2728

2829
export type Role = "user" | "assistant";
2930

@@ -103,6 +104,9 @@ export type Msg =
103104
| {
104105
type: "set-opts";
105106
options: MagentaOptions;
107+
}
108+
| {
109+
type: "show-message-debug-info";
106110
};
107111

108112
export function init({ nvim, lsp }: { nvim: Nvim; lsp: Lsp }) {
@@ -421,7 +425,7 @@ ${msg.error.stack}`,
421425
model,
422426
// eslint-disable-next-line @typescript-eslint/require-await
423427
async () => {
424-
getClient(nvim, model.activeProvider, model.options).abort();
428+
getProvider(nvim, model.activeProvider, model.options).abort();
425429
},
426430
];
427431
}
@@ -430,6 +434,10 @@ ${msg.error.stack}`,
430434
return [{ ...model, options: msg.options }];
431435
}
432436

437+
case "show-message-debug-info": {
438+
return [model, () => showDebugInfo(model)];
439+
}
440+
433441
default:
434442
assertUnreachable(msg);
435443
}
@@ -490,7 +498,7 @@ ${msg.error.stack}`,
490498
});
491499
let res;
492500
try {
493-
res = await getClient(
501+
res = await getProvider(
494502
nvim,
495503
model.activeProvider,
496504
model.options,
@@ -562,12 +570,17 @@ ${msg.error.stack}`,
562570
) % MESSAGE_ANIMATION.length
563571
]
564572
}`
565-
: d`Stopped (${model.conversation.stopReason}) [input: ${model.conversation.usage.inputTokens.toString()}, output: ${model.conversation.usage.outputTokens.toString()}${
566-
model.conversation.usage.cacheHits !== undefined &&
567-
model.conversation.usage.cacheMisses !== undefined
568-
? d`, cache hits: ${model.conversation.usage.cacheHits.toString()}, cache misses: ${model.conversation.usage.cacheMisses.toString()}`
569-
: ""
570-
}]`
573+
: withBindings(
574+
d`Stopped (${model.conversation.stopReason}) [input: ${model.conversation.usage.inputTokens.toString()}, output: ${model.conversation.usage.outputTokens.toString()}${
575+
model.conversation.usage.cacheHits !== undefined &&
576+
model.conversation.usage.cacheMisses !== undefined
577+
? d`, cache hits: ${model.conversation.usage.cacheHits.toString()}, cache misses: ${model.conversation.usage.cacheMisses.toString()}`
578+
: ""
579+
}]`,
580+
{
581+
"<CR>": () => dispatch({ type: "show-message-debug-info" }),
582+
},
583+
)
571584
: ""
572585
}${
573586
model.conversation.state == "stopped" &&
@@ -642,6 +655,44 @@ ${msg.error.stack}`,
642655
return messages.map((m) => m.message);
643656
}
644657

658+
async function showDebugInfo(model: Model) {
659+
const messages = await getMessages(model);
660+
const provider = getProvider(nvim, model.activeProvider, model.options);
661+
const params = provider.createStreamParameters(messages);
662+
const nTokens = await provider.countTokens(messages);
663+
664+
// Create a floating window
665+
const bufnr = await nvim.call("nvim_create_buf", [false, true]);
666+
await nvim.call("nvim_buf_set_option", [bufnr, "bufhidden", "wipe"]);
667+
const [editorWidth, editorHeight] = (await Promise.all([
668+
getOption("columns", nvim),
669+
await getOption("lines", nvim),
670+
])) as [number, number];
671+
const width = 80;
672+
const height = editorHeight - 20;
673+
await nvim.call("nvim_open_win", [
674+
bufnr,
675+
true,
676+
{
677+
relative: "editor",
678+
width,
679+
height,
680+
col: Math.floor((editorWidth - width) / 2),
681+
row: Math.floor((editorHeight - height) / 2),
682+
style: "minimal",
683+
border: "single",
684+
},
685+
]);
686+
687+
const lines = JSON.stringify(params, null, 2).split("\n");
688+
lines.push(`nTokens: ${nTokens}`);
689+
await nvim.call("nvim_buf_set_lines", [bufnr, 0, -1, false, lines]);
690+
691+
// Set buffer options
692+
await nvim.call("nvim_buf_set_option", [bufnr, "modifiable", false]);
693+
await nvim.call("nvim_buf_set_option", [bufnr, "filetype", "json"]);
694+
}
695+
645696
return {
646697
initModel,
647698
update,

‎node/providers/anthropic.ts

+72-51
Original file line numberDiff line numberDiff line change
@@ -47,34 +47,9 @@ export class AnthropicProvider implements Provider {
4747
}
4848
}
4949

50-
async sendMessage(
51-
messages: Array<ProviderMessage>,
52-
onText: (text: string) => void,
53-
onError: (error: Error) => void,
54-
): Promise<{
55-
toolRequests: Result<ToolManager.ToolRequest, { rawRequest: unknown }>[];
56-
stopReason: StopReason;
57-
usage: Usage;
58-
}> {
59-
const buf: string[] = [];
60-
let flushInProgress: boolean = false;
61-
62-
const flushBuffer = () => {
63-
if (buf.length && !flushInProgress) {
64-
const text = buf.join("");
65-
buf.splice(0);
66-
67-
flushInProgress = true;
68-
69-
try {
70-
onText(text);
71-
} finally {
72-
flushInProgress = false;
73-
setInterval(flushBuffer, 1);
74-
}
75-
}
76-
};
77-
50+
createStreamParameters(
51+
messages: ProviderMessage[],
52+
): Anthropic.Messages.MessageStreamParams {
7853
const anthropicMessages = messages.map((m): MessageParam => {
7954
let content: Anthropic.Messages.ContentBlockParam[];
8055
if (typeof m.content == "string") {
@@ -127,31 +102,77 @@ export class AnthropicProvider implements Provider {
127102
},
128103
);
129104

105+
return {
106+
messages: anthropicMessages,
107+
model: this.options.model,
108+
max_tokens: 4096,
109+
system: [
110+
{
111+
type: "text",
112+
text: DEFAULT_SYSTEM_PROMPT,
113+
// the prompt appears in the following order:
114+
// tools
115+
// system
116+
// messages
117+
// This ensures the tools + system prompt (which is approx 1400 tokens) is cached.
118+
cache_control:
119+
cacheControlItemsPlaced < 4 ? { type: "ephemeral" } : null,
120+
},
121+
],
122+
tool_choice: {
123+
type: "auto",
124+
disable_parallel_tool_use: false,
125+
},
126+
tools,
127+
};
128+
}
129+
130+
async countTokens(messages: Array<ProviderMessage>): Promise<number> {
131+
const params = this.createStreamParameters(messages);
132+
const lastMessage = params.messages[params.messages.length - 1];
133+
if (!lastMessage || lastMessage.role != "user") {
134+
params.messages.push({ role: "user", content: "test" });
135+
}
136+
const res = await this.client.messages.countTokens({
137+
messages: params.messages,
138+
model: params.model,
139+
system: params.system as Anthropic.TextBlockParam[],
140+
tools: params.tools as Anthropic.Tool[],
141+
});
142+
return res.input_tokens;
143+
}
144+
145+
async sendMessage(
146+
messages: Array<ProviderMessage>,
147+
onText: (text: string) => void,
148+
onError: (error: Error) => void,
149+
): Promise<{
150+
toolRequests: Result<ToolManager.ToolRequest, { rawRequest: unknown }>[];
151+
stopReason: StopReason;
152+
usage: Usage;
153+
}> {
154+
const buf: string[] = [];
155+
let flushInProgress: boolean = false;
156+
157+
const flushBuffer = () => {
158+
if (buf.length && !flushInProgress) {
159+
const text = buf.join("");
160+
buf.splice(0);
161+
162+
flushInProgress = true;
163+
164+
try {
165+
onText(text);
166+
} finally {
167+
flushInProgress = false;
168+
setInterval(flushBuffer, 1);
169+
}
170+
}
171+
};
172+
130173
try {
131174
this.request = this.client.messages
132-
.stream({
133-
messages: anthropicMessages,
134-
model: this.options.model,
135-
max_tokens: 4096,
136-
system: [
137-
{
138-
type: "text",
139-
text: DEFAULT_SYSTEM_PROMPT,
140-
// the prompt appears in the following order:
141-
// tools
142-
// system
143-
// messages
144-
// This ensures the tools + system prompt (which is approx 1400 tokens) is cached.
145-
cache_control:
146-
cacheControlItemsPlaced < 4 ? { type: "ephemeral" } : null,
147-
},
148-
],
149-
tool_choice: {
150-
type: "auto",
151-
disable_parallel_tool_use: false,
152-
},
153-
tools,
154-
})
175+
.stream(this.createStreamParameters(messages))
155176
.on("text", (text: string) => {
156177
buf.push(text);
157178
flushBuffer();

‎node/providers/mock.ts

+9
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,15 @@ export class MockProvider implements Provider {
3939
}
4040
}
4141

42+
createStreamParameters(messages: Array<ProviderMessage>): unknown {
43+
return messages;
44+
}
45+
46+
// eslint-disable-next-line @typescript-eslint/require-await
47+
async countTokens(messages: Array<ProviderMessage>): Promise<number> {
48+
return messages.length;
49+
}
50+
4251
async sendMessage(
4352
messages: Array<ProviderMessage>,
4453
onText: (text: string) => void,

0 commit comments

Comments
 (0)