Skip to content

Commit 3f71de3

Browse files
committed
feat(ollama): add generateWithUsage() and export Usage/GenerateResult types
Adds generateWithUsage() method to OllamaClient that returns token usage metadata (prompt_tokens, completion_tokens, total_tokens) alongside content. The existing generate() method delegates to generateWithUsage() internally, maintaining full backward compatibility. Also exports the Usage and GenerateResult interfaces so consumers can type their metering/billing integrations.
1 parent d4f3f66 commit 3f71de3

1 file changed

Lines changed: 25 additions & 3 deletions

File tree

packages/ollama/src/index.ts

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ interface ToolCallContent {
4444
arguments: Record<string, JsonValue | undefined>;
4545
}
4646

47-
interface Usage {
47+
export interface Usage {
4848
input: number;
4949
output: number;
5050
cacheRead: number;
@@ -226,6 +226,13 @@ export interface GenerateInput {
226226
maxTokens?: number;
227227
}
228228

229+
export interface GenerateResult {
230+
content: string;
231+
usage: Usage;
232+
model: string;
233+
stopReason: 'stop' | 'length' | 'toolUse' | 'error' | 'aborted';
234+
}
235+
229236
interface OllamaTagsResponse {
230237
models?: Array<{ name: string }>;
231238
}
@@ -310,6 +317,15 @@ export class OllamaClient {
310317
input: GenerateInput,
311318
onChunk?: (chunk: string) => void
312319
): Promise<string | void> {
320+
const result = await this.generateWithUsage(input, onChunk);
321+
if (onChunk || input.stream) return;
322+
return result.content;
323+
}
324+
325+
async generateWithUsage(
326+
input: GenerateInput,
327+
onChunk?: (chunk: string) => void
328+
): Promise<GenerateResult> {
313329
const context = legacyInputToContext(input);
314330
const model: ModelDescriptor = {
315331
id: input.model,
@@ -336,14 +352,20 @@ export class OllamaClient {
336352
onChunk?.(event.delta);
337353
}
338354
}
339-
return;
340355
}
341356

342357
const message = await response.result();
343-
return message.content
358+
const content = message.content
344359
.filter((block): block is TextContent => block.type === 'text')
345360
.map((block) => block.text)
346361
.join('');
362+
363+
return {
364+
content,
365+
usage: message.usage,
366+
model: message.model,
367+
stopReason: message.stopReason,
368+
};
347369
}
348370
}
349371

0 commit comments

Comments
 (0)