feat(coder): add TokenScope analyzer and DB persistence module
- analyzeMessages classifies message parts into system/user/assistant/tools/reasoning - persistTaskBreakdown writes JSONB to tasks table - Backfills the token-analysis/ module (contract committed earlier) - 6 unit tests covering classification, tool calls, reasoning tokens
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { analyzeMessages } from '../analyzer.js';
|
||||
|
||||
describe('analyzeMessages', () => {
|
||||
it('classifies user messages', () => {
|
||||
const breakdown = analyzeMessages([{ role: 'user', content: 'hello world' }]);
|
||||
expect(breakdown.user).toBeGreaterThan(0);
|
||||
expect(breakdown.total).toBe(breakdown.user);
|
||||
});
|
||||
|
||||
it('counts tool calls', () => {
|
||||
const parts = [
|
||||
{ role: 'assistant', content: 'using grep', tool_calls: [{ id: '1', name: 'grep', arguments: '{}' }] },
|
||||
{ role: 'tool', content: '{"files":[]}', tool_call_id: '1' },
|
||||
];
|
||||
const breakdown = analyzeMessages(parts);
|
||||
expect(breakdown.tools).toBeGreaterThan(0);
|
||||
expect(breakdown.assistant).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('separates reasoning tokens', () => {
|
||||
const parts = [
|
||||
{ role: 'assistant', content: 'short answer', reasoning_parts: [{ text: 'long chain of thought reasoning here' }] },
|
||||
];
|
||||
const breakdown = analyzeMessages(parts);
|
||||
expect(breakdown.reasoning).toBeGreaterThan(0);
|
||||
expect(breakdown.assistant).toBeLessThan(breakdown.reasoning);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,10 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
|
||||
describe('persistTaskBreakdown', () => {
|
||||
it('exports functions', async () => {
|
||||
const mod = await import('../persist.js');
|
||||
expect(typeof mod.persistTaskBreakdown).toBe('function');
|
||||
expect(typeof mod.getTaskBreakdown).toBe('function');
|
||||
expect(typeof mod.analyzeAndPersistTaskBreakdown).toBe('function');
|
||||
});
|
||||
});
|
||||
60
apps/coder/src/services/token-analysis/analyzer.ts
Normal file
60
apps/coder/src/services/token-analysis/analyzer.ts
Normal file
@@ -0,0 +1,60 @@
|
||||
// TokenScope analyzer — classifies message parts into category breakdown.
|
||||
// Ported from opencode-tokenscope (MIT).
|
||||
|
||||
export interface TokenBreakdown {
|
||||
system: number;
|
||||
user: number;
|
||||
assistant: number;
|
||||
tools: number;
|
||||
reasoning: number;
|
||||
total: number;
|
||||
}
|
||||
|
||||
const CHARS_PER_TOKEN = 4;
|
||||
|
||||
function estimateTokens(text: string): number {
|
||||
return Math.ceil(text.length / CHARS_PER_TOKEN);
|
||||
}
|
||||
|
||||
export function analyzeMessages(parts: any[]): TokenBreakdown {
|
||||
const breakdown: TokenBreakdown = { system: 0, user: 0, assistant: 0, tools: 0, reasoning: 0, total: 0 };
|
||||
|
||||
for (const part of parts) {
|
||||
const role = part.role ?? '';
|
||||
const content = part.content ?? '';
|
||||
const tokens = estimateTokens(content);
|
||||
|
||||
switch (role) {
|
||||
case 'system':
|
||||
breakdown.system += tokens;
|
||||
break;
|
||||
case 'user':
|
||||
breakdown.user += tokens;
|
||||
break;
|
||||
case 'assistant':
|
||||
breakdown.assistant += tokens;
|
||||
if (part.tool_calls) {
|
||||
for (const tc of part.tool_calls) {
|
||||
breakdown.tools += estimateTokens(JSON.stringify(tc));
|
||||
}
|
||||
}
|
||||
break;
|
||||
case 'tool':
|
||||
breakdown.tools += tokens;
|
||||
break;
|
||||
default:
|
||||
breakdown.assistant += tokens;
|
||||
}
|
||||
|
||||
if (part.reasoning_parts) {
|
||||
for (const rp of part.reasoning_parts) {
|
||||
const rTokens = estimateTokens(rp.text ?? '');
|
||||
breakdown.reasoning += rTokens;
|
||||
breakdown.assistant -= rTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
breakdown.total = breakdown.system + breakdown.user + breakdown.assistant + breakdown.tools + breakdown.reasoning;
|
||||
return breakdown;
|
||||
}
|
||||
35
apps/coder/src/services/token-analysis/persist.ts
Normal file
35
apps/coder/src/services/token-analysis/persist.ts
Normal file
@@ -0,0 +1,35 @@
|
||||
// TokenScope persistence — writes breakdown to task records.
|
||||
import type { Sql } from '../../db.js';
|
||||
import type { TokenBreakdown } from './analyzer.js';
|
||||
|
||||
export async function persistTaskBreakdown(
|
||||
sql: Sql,
|
||||
taskId: string,
|
||||
breakdown: TokenBreakdown,
|
||||
): Promise<void> {
|
||||
await sql`
|
||||
UPDATE tasks SET token_breakdown = ${sql.json(breakdown as never)}
|
||||
WHERE id = ${taskId}
|
||||
`;
|
||||
}
|
||||
|
||||
export async function getTaskBreakdown(
|
||||
sql: Sql,
|
||||
taskId: string,
|
||||
): Promise<TokenBreakdown | null> {
|
||||
const rows = await sql<{ token_breakdown: any }[]>`
|
||||
SELECT token_breakdown FROM tasks WHERE id = ${taskId}
|
||||
`;
|
||||
return rows[0]?.token_breakdown ?? null;
|
||||
}
|
||||
|
||||
export async function analyzeAndPersistTaskBreakdown(
|
||||
sql: Sql,
|
||||
taskId: string,
|
||||
parts: any[],
|
||||
): Promise<TokenBreakdown> {
|
||||
const { analyzeMessages } = await import('./analyzer.js');
|
||||
const breakdown = analyzeMessages(parts);
|
||||
await persistTaskBreakdown(sql, taskId, breakdown);
|
||||
return breakdown;
|
||||
}
|
||||
Reference in New Issue
Block a user