diff --git a/CHANGELOG.md b/CHANGELOG.md index 18b71a0..c69ad29 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## 0.1.3 + +### Changed + +- Set `stream_options.include_usage` on streaming requests so oMLX returns + token counts in the final SSE chunk. +- Emit token usage as a `LanguageModelDataPart` (mime type `usage`) so the + Copilot Chat context window indicator shows real usage instead of 0%. + Requires VS Code 1.120 or newer (microsoft/vscode#315394). + ## 0.1.2 ### Changed diff --git a/package-lock.json b/package-lock.json index 7d15084..5dac08e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,21 +1,21 @@ { "name": "omlx-copilot-chat", - "version": "0.1.2", + "version": "0.1.3", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "omlx-copilot-chat", - "version": "0.1.2", + "version": "0.1.3", "license": "MIT", "devDependencies": { "@types/node": "^22.0.0", - "@types/vscode": "^1.104.0", + "@types/vscode": "^1.120.0", "@vscode/vsce": "^3.6.0", "typescript": "^5.9.0" }, "engines": { - "vscode": "^1.104.0" + "vscode": "^1.120.0" } }, "node_modules/@azu/format-text": { @@ -572,9 +572,9 @@ "license": "MIT" }, "node_modules/@types/vscode": { - "version": "1.118.0", - "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.118.0.tgz", - "integrity": "sha512-Ah6eTlqDcwIMELEVwQMO++rJAFBRz/oLluLD/vWdYrH1KuI9kfpaM+7pg0OvvascgcJy+ghLCERAYouM4QbzGw==", + "version": "1.125.0", + "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.125.0.tgz", + "integrity": "sha512-0icm/ZQAaism87P0ekHqi4/Ju9du+Tm0RUW+y7vqRsxY2cY0FNRX1nAnaW7nT6npPt2tfHiheZ55Zm9UhqonFA==", "dev": true, "license": "MIT" }, diff --git a/package.json b/package.json index e9ffc32..9f31ccd 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "omlx-copilot-chat", "displayName": "oMLX Copilot Chat", "description": "Use oMLX OpenAI-compatible models from the VS Code Copilot Chat model picker.", - "version": "0.1.2", + "version": "0.1.3", "publisher": "techopolis", "preview": true, "license": "MIT", @@ -15,7 +15,7 @@ }, "homepage": "https://github.com/mikedoise/oMLX-Copilot-Chat#readme", "engines": { - "vscode": "^1.104.0" + "vscode": "^1.120.0" }, "categories": [ "AI", @@ -118,7 +118,7 @@ }, "devDependencies": { "@types/node": "^22.0.0", - "@types/vscode": "^1.104.0", + "@types/vscode": "^1.120.0", "@vscode/vsce": "^3.6.0", "typescript": "^5.9.0" } diff --git a/src/provider/OmlxClient.ts b/src/provider/OmlxClient.ts index e3f2a3e..f4881a4 100644 --- a/src/provider/OmlxClient.ts +++ b/src/provider/OmlxClient.ts @@ -8,12 +8,14 @@ import { OpenAIModelListResponse, OpenAITool, OpenAIToolCall, - OpenAIToolCallDelta + OpenAIToolCallDelta, + OpenAIUsage } from './openAICompatTypes'; export type OmlxStreamPart = | { readonly type: 'text'; readonly value: string } - | { readonly type: 'toolCall'; readonly callId: string; readonly name: string; readonly input: object }; + | { readonly type: 'toolCall'; readonly callId: string; readonly name: string; readonly input: object } + | { readonly type: 'usage'; readonly usage: OpenAIUsage }; export class OmlxPromptTooLongError extends Error { constructor( @@ -71,6 +73,7 @@ export class OmlxClient { model, messages, stream: true, + stream_options: { include_usage: true }, temperature: options.temperature, max_tokens: options.maxTokens, tools: options.tools, @@ -114,6 +117,12 @@ export class OmlxClient { const line = buffer.slice(0, newlineIndex); buffer = buffer.slice(newlineIndex + 1); const fragment = parseSseLine(line); + // oMLX emits the usage chunk (choices: [], usage: {...}) immediately + // before the [DONE] sentinel when stream_options.include_usage is set, + // so it is read here as a normal chunk ahead of the done check below. + if (fragment.usage) { + yield { type: 'usage', usage: fragment.usage }; + } if (fragment.done) { yield* completeToolCalls(toolCallAccumulators); return; @@ -199,6 +208,7 @@ interface ParsedSseLine { readonly done: boolean; readonly text?: string; readonly toolCalls: readonly OpenAIToolCallDelta[]; + readonly usage?: OpenAIUsage; } function parseSseLine(line: string): ParsedSseLine { @@ -217,7 +227,8 @@ function parseSseLine(line: string): ParsedSseLine { return { done: false, text: delta?.content ?? undefined, - toolCalls: delta?.tool_calls ?? [] + toolCalls: delta?.tool_calls ?? [], + usage: chunk.usage ?? undefined }; } diff --git a/src/provider/OmlxLanguageModelProvider.ts b/src/provider/OmlxLanguageModelProvider.ts index c246668..3fb7a56 100644 --- a/src/provider/OmlxLanguageModelProvider.ts +++ b/src/provider/OmlxLanguageModelProvider.ts @@ -81,8 +81,13 @@ export class OmlxLanguageModelProvider implements vscode.LanguageModelChatProvid for await (const part of client.streamChatCompletion(model.id, convertedMessages, requestOptions, token)) { if (part.type === 'text') { progress.report(new vscode.LanguageModelTextPart(part.value)); - } else { + } else if (part.type === 'toolCall') { progress.report(new vscode.LanguageModelToolCallPart(part.callId, part.name, part.input)); + } else if (part.type === 'usage') { + this.output.appendLine( + `Usage ${model.id}: prompt_tokens=${part.usage.prompt_tokens}, completion_tokens=${part.usage.completion_tokens}, total_tokens=${part.usage.total_tokens}, cached_tokens=${part.usage.prompt_tokens_details?.cached_tokens ?? 0}` + ); + progress.report(vscode.LanguageModelDataPart.json(part.usage, 'usage')); } } } catch (error) { diff --git a/src/provider/openAICompatTypes.ts b/src/provider/openAICompatTypes.ts index 29fcee6..db9c3c3 100644 --- a/src/provider/openAICompatTypes.ts +++ b/src/provider/openAICompatTypes.ts @@ -55,6 +55,7 @@ export interface OpenAIChatCompletionRequest { readonly model: string; readonly messages: readonly OpenAIChatMessage[]; readonly stream: boolean; + readonly stream_options?: { readonly include_usage: boolean }; readonly temperature?: number; readonly max_tokens?: number; readonly tools?: readonly OpenAITool[]; @@ -88,6 +89,15 @@ export interface OpenAIChatCompletionResponse { }[]; } +export interface OpenAIUsage { + readonly prompt_tokens: number; + readonly completion_tokens: number; + readonly total_tokens: number; + readonly prompt_tokens_details?: { + readonly cached_tokens?: number; + }; +} + export interface OpenAIChatCompletionChunk { readonly choices?: readonly { readonly delta?: { @@ -95,6 +105,7 @@ export interface OpenAIChatCompletionChunk { readonly tool_calls?: readonly OpenAIToolCallDelta[]; }; }[]; + readonly usage?: OpenAIUsage | null; } export interface OpenAIToolCallDelta {