@@ -10,10 +10,13 @@ import type { ResponseCache } from '../caching/types'
1010import { type EntityType , VALID_LINK_TYPES } from '../search/types'
1111import {
1212 type ActivityCategory ,
13+ addLlmUsage ,
1314 type CandidateMessage ,
1415 type ClassifiedActivity ,
1516 type ClassifierConfig ,
1617 calculateCombinedScore ,
18+ EMPTY_LLM_USAGE ,
19+ type LlmUsage ,
1720 type Result
1821} from '../types'
1922import { DEFAULT_MODELS } from './models'
@@ -27,6 +30,18 @@ import { callProviderWithFallbacks } from './providers'
2730import type { ParsedClassification } from './response-parser'
2831import { countTokens , MAX_BATCH_TOKENS } from './tokenizer'
2932
33+ /** Result from classifyBatch including activities and usage data */
34+ export interface ClassifyBatchResult {
35+ readonly activities : ClassifiedActivity [ ]
36+ readonly usage : LlmUsage
37+ }
38+
39+ /** Result from classifyMessages including activities and total usage */
40+ export interface ClassifyMessagesResult {
41+ readonly activities : ClassifiedActivity [ ]
42+ readonly usage : LlmUsage
43+ }
44+
3045export { createSmartBatches , groupCandidatesByProximity } from './batching'
3146export type { ResolvedModel } from './models'
3247export {
@@ -149,13 +164,14 @@ function toClassifiedActivity(
149164 * @param config Classifier configuration
150165 * @param cache Optional response cache
151166 * @param promptType Optional prompt type override (auto-detected if not specified)
167+ * @returns Result with activities AND usage data (usage is 0 for cache hits)
152168 */
153169export async function classifyBatch (
154170 candidates : readonly CandidateMessage [ ] ,
155171 config : ClassifierConfig ,
156172 cache ?: ResponseCache ,
157173 promptType ?: PromptType
158- ) : Promise < Result < ClassifiedActivity [ ] > > {
174+ ) : Promise < Result < ClassifyBatchResult > > {
159175 const prompt = buildClassificationPrompt (
160176 candidates ,
161177 {
@@ -191,7 +207,7 @@ export async function classifyBatch(
191207
192208 try {
193209 const expectedIds = candidates . map ( ( c ) => c . messageId )
194- const parsed = parseClassificationResponse ( responseResult . value , expectedIds )
210+ const parsed = parseClassificationResponse ( responseResult . value . text , expectedIds )
195211
196212 // Map responses to candidates, filtering out invalid classifications
197213 const suggestions : ClassifiedActivity [ ] = [ ]
@@ -206,7 +222,10 @@ export async function classifyBatch(
206222 }
207223 }
208224
209- return { ok : true , value : suggestions }
225+ return {
226+ ok : true ,
227+ value : { activities : suggestions , usage : responseResult . value . usage }
228+ }
210229 } catch ( error ) {
211230 const message = error instanceof Error ? error . message : String ( error )
212231 return {
@@ -233,6 +252,12 @@ function createBatches(
233252 return batches
234253}
235254
255+ /** Internal result from processBatches */
256+ interface ProcessBatchesResult {
257+ activities : ClassifiedActivity [ ]
258+ usage : LlmUsage
259+ }
260+
236261/**
237262 * Process batches of a specific type.
238263 */
@@ -244,8 +269,9 @@ async function processBatches(
244269 model : string ,
245270 batchIndexOffset : number ,
246271 totalBatches : number
247- ) : Promise < Result < ClassifiedActivity [ ] > > {
272+ ) : Promise < Result < ProcessBatchesResult > > {
248273 const results : ClassifiedActivity [ ] = [ ]
274+ let totalUsage : LlmUsage = EMPTY_LLM_USAGE
249275
250276 for ( let i = 0 ; i < batches . length ; i ++ ) {
251277 const batch = batches [ i ]
@@ -281,14 +307,15 @@ async function processBatches(
281307 config . onBatchComplete ?.( {
282308 batchIndex : globalBatchIndex ,
283309 totalBatches,
284- activityCount : result . value . length ,
310+ activityCount : result . value . activities . length ,
285311 durationMs
286312 } )
287313
288- results . push ( ...result . value )
314+ results . push ( ...result . value . activities )
315+ totalUsage = addLlmUsage ( totalUsage , result . value . usage )
289316 }
290317
291- return { ok : true , value : results }
318+ return { ok : true , value : { activities : results , usage : totalUsage } }
292319}
293320
294321/**
@@ -300,13 +327,13 @@ async function processBatches(
300327 * @param candidates Candidate messages to classify
301328 * @param config Classifier configuration
302329 * @param cache Optional response cache to prevent duplicate API calls
303- * @returns Classified suggestions or error
330+ * @returns Classified activities with aggregated usage data, or error
304331 */
305332export async function classifyMessages (
306333 candidates : readonly CandidateMessage [ ] ,
307334 config : ClassifierConfig ,
308335 cache ?: ResponseCache
309- ) : Promise < Result < ClassifiedActivity [ ] > > {
336+ ) : Promise < Result < ClassifyMessagesResult > > {
310337 const batchSize = config . batchSize ?? DEFAULT_BATCH_SIZE
311338 const model = config . model ?? DEFAULT_MODELS [ config . provider ]
312339
@@ -319,6 +346,7 @@ export async function classifyMessages(
319346 const totalBatches = suggestionBatches . length + agreementBatches . length
320347
321348 const results : ClassifiedActivity [ ] = [ ]
349+ let totalUsage : LlmUsage = EMPTY_LLM_USAGE
322350
323351 // Process suggestion batches first
324352 if ( suggestionBatches . length > 0 ) {
@@ -334,7 +362,8 @@ export async function classifyMessages(
334362 if ( ! suggestionResult . ok ) {
335363 return suggestionResult
336364 }
337- results . push ( ...suggestionResult . value )
365+ results . push ( ...suggestionResult . value . activities )
366+ totalUsage = addLlmUsage ( totalUsage , suggestionResult . value . usage )
338367 }
339368
340369 // Process agreement batches
@@ -351,10 +380,11 @@ export async function classifyMessages(
351380 if ( ! agreementResult . ok ) {
352381 return agreementResult
353382 }
354- results . push ( ...agreementResult . value )
383+ results . push ( ...agreementResult . value . activities )
384+ totalUsage = addLlmUsage ( totalUsage , agreementResult . value . usage )
355385 }
356386
357- return { ok : true , value : results }
387+ return { ok : true , value : { activities : results , usage : totalUsage } }
358388}
359389
360390/**
0 commit comments