Просмотр исходного кода

feat(core): propagate LLM token usage from agents through pipeline

Writer sums creative + settle usage. Auditor and Reviser expose usage
from their chat calls. Runner aggregates all agent usage per chapter
and persists to ChapterMeta.tokenUsage in index.json.
Ma 7 месяцев назад
Родитель
Сommit
e13933c4ac

+ 7 - 1
packages/core/src/agents/continuity.ts

@@ -9,6 +9,11 @@ export interface AuditResult {
   readonly passed: boolean;
   readonly issues: ReadonlyArray<AuditIssue>;
   readonly summary: string;
+  readonly tokenUsage?: {
+    readonly promptTokens: number;
+    readonly completionTokens: number;
+    readonly totalTokens: number;
+  };
 }
 
 export interface AuditIssue {
@@ -296,7 +301,8 @@ ${chapterContent}`;
       ? await this.chatWithSearch(chatMessages, chatOptions)
       : await this.chat(chatMessages, chatOptions);
 
-    return this.parseAuditResult(response.content);
+    const result = this.parseAuditResult(response.content);
+    return { ...result, tokenUsage: response.usage };
   }
 
   private parseAuditResult(content: string): AuditResult {

+ 7 - 1
packages/core/src/agents/reviser.ts

@@ -15,6 +15,11 @@ export interface ReviseOutput {
   readonly updatedState: string;
   readonly updatedLedger: string;
   readonly updatedHooks: string;
+  readonly tokenUsage?: {
+    readonly promptTokens: number;
+    readonly completionTokens: number;
+    readonly totalTokens: number;
+  };
 }
 
 const MODE_DESCRIPTIONS: Record<ReviseMode, string> = {
@@ -152,7 +157,8 @@ ${chapterContent}`;
       { temperature: 0.3, maxTokens },
     );
 
-    return this.parseOutput(response.content, gp);
+    const output = this.parseOutput(response.content, gp);
+    return { ...output, tokenUsage: response.usage };
   }
 
   private parseOutput(content: string, gp: GenreProfile): ReviseOutput {

+ 24 - 4
packages/core/src/agents/writer.ts

@@ -21,6 +21,12 @@ export interface WriteChapterInput {
   readonly temperatureOverride?: number;
 }
 
+export interface TokenUsage {
+  readonly promptTokens: number;
+  readonly completionTokens: number;
+  readonly totalTokens: number;
+}
+
 export interface WriteChapterOutput {
   readonly chapterNumber: number;
   readonly title: string;
@@ -37,6 +43,7 @@ export interface WriteChapterOutput {
   readonly updatedCharacterMatrix: string;
   readonly postWriteErrors: ReadonlyArray<PostWriteViolation>;
   readonly postWriteWarnings: ReadonlyArray<PostWriteViolation>;
+  readonly tokenUsage?: TokenUsage;
 }
 
 export class WriterAgent extends BaseAgent {
@@ -122,13 +129,14 @@ export class WriterAgent extends BaseAgent {
       ],
       { maxTokens: creativeMaxTokens, temperature: creativeTemperature },
     );
+    const creativeUsage = creativeResponse.usage;
 
     const creative = parseCreativeOutput(chapterNumber, creativeResponse.content);
 
     // ── Phase 2: State settlement (temperature 0.3) ──
     process.stderr.write(`[writer] Phase 2: state settlement for chapter ${chapterNumber} (${creative.wordCount} chars)\n`);
 
-    const settlement = await this.settle({
+    const settleResult = await this.settle({
       book,
       genreProfile,
       bookRules,
@@ -144,6 +152,8 @@ export class WriterAgent extends BaseAgent {
       characterMatrix,
       volumeOutline,
     });
+    const settlement = settleResult.settlement;
+    const settleUsage = settleResult.usage;
 
     // ── Post-write validation (regex + rule-based, zero LLM cost) ──
     const ruleViolations = validatePostWrite(creative.content, genreProfile, bookRules);
@@ -169,7 +179,13 @@ export class WriterAgent extends BaseAgent {
       }
     }
 
-    // ── Merge into WriteChapterOutput (interface unchanged) ──
+    // ── Merge into WriteChapterOutput ──
+    const tokenUsage: TokenUsage = {
+      promptTokens: creativeUsage.promptTokens + settleUsage.promptTokens,
+      completionTokens: creativeUsage.completionTokens + settleUsage.completionTokens,
+      totalTokens: creativeUsage.totalTokens + settleUsage.totalTokens,
+    };
+
     return {
       chapterNumber,
       title: creative.title,
@@ -186,6 +202,7 @@ export class WriterAgent extends BaseAgent {
       updatedCharacterMatrix: settlement.updatedCharacterMatrix,
       postWriteErrors,
       postWriteWarnings,
+      tokenUsage,
     };
   }
 
@@ -204,7 +221,7 @@ export class WriterAgent extends BaseAgent {
     readonly emotionalArcs: string;
     readonly characterMatrix: string;
     readonly volumeOutline: string;
-  }) {
+  }): Promise<{ settlement: ReturnType<typeof parseSettlementOutput>; usage: TokenUsage }> {
     const settlerSystem = buildSettlerSystemPrompt(
       params.book, params.genreProfile, params.bookRules,
     );
@@ -234,7 +251,10 @@ export class WriterAgent extends BaseAgent {
       { maxTokens: settlerMaxTokens, temperature: 0.3 },
     );
 
-    return parseSettlementOutput(response.content, params.genreProfile);
+    return {
+      settlement: parseSettlementOutput(response.content, params.genreProfile),
+      usage: response.usage,
+    };
   }
 
   async saveChapter(

+ 2 - 2
packages/core/src/index.ts

@@ -14,7 +14,7 @@ export { createLLMClient, chatCompletion, chatWithTools, type LLMClient, type LL
 // Agents
 export { BaseAgent, type AgentContext } from "./agents/base.js";
 export { ArchitectAgent, type ArchitectOutput } from "./agents/architect.js";
-export { WriterAgent, type WriteChapterInput, type WriteChapterOutput } from "./agents/writer.js";
+export { WriterAgent, type WriteChapterInput, type WriteChapterOutput, type TokenUsage } from "./agents/writer.js";
 export { ContinuityAuditor, type AuditResult, type AuditIssue } from "./agents/continuity.js";
 export { ReviserAgent, type ReviseOutput, type ReviseMode } from "./agents/reviser.js";
 export { RadarAgent, type RadarResult, type RadarRecommendation } from "./agents/radar.js";
@@ -37,7 +37,7 @@ export { fetchUrl } from "./utils/web-search.js";
 export { splitChapters, type SplitChapter } from "./utils/chapter-splitter.js";
 
 // Pipeline
-export { PipelineRunner, type PipelineConfig, type ChapterPipelineResult, type DraftResult, type ReviseResult, type TruthFiles, type BookStatusInfo, type ImportChaptersInput, type ImportChaptersResult } from "./pipeline/runner.js";
+export { PipelineRunner, type PipelineConfig, type ChapterPipelineResult, type DraftResult, type ReviseResult, type TruthFiles, type BookStatusInfo, type ImportChaptersInput, type ImportChaptersResult, type TokenUsageSummary } from "./pipeline/runner.js";
 export { Scheduler, type SchedulerConfig } from "./pipeline/scheduler.js";
 export { runAgentLoop, AGENT_TOOLS as AGENT_TOOLS, type AgentLoopOptions } from "./pipeline/agent.js";
 export { detectChapter, detectAndRewrite, loadDetectionHistory, type DetectChapterResult, type DetectAndRewriteResult } from "./pipeline/detection-runner.js";

+ 31 - 1
packages/core/src/pipeline/runner.ts

@@ -34,6 +34,12 @@ export interface PipelineConfig {
   readonly modelOverrides?: Record<string, string | AgentLLMOverride>;
 }
 
+export interface TokenUsageSummary {
+  readonly promptTokens: number;
+  readonly completionTokens: number;
+  readonly totalTokens: number;
+}
+
 export interface ChapterPipelineResult {
   readonly chapterNumber: number;
   readonly title: string;
@@ -41,6 +47,7 @@ export interface ChapterPipelineResult {
   readonly auditResult: AuditResult;
   readonly revised: boolean;
   readonly status: "ready-for-review" | "audit-failed";
+  readonly tokenUsage?: TokenUsageSummary;
 }
 
 // Atomic operation results
@@ -49,6 +56,7 @@ export interface DraftResult {
   readonly title: string;
   readonly wordCount: number;
   readonly filePath: string;
+  readonly tokenUsage?: TokenUsageSummary;
 }
 
 export interface ReviseResult {
@@ -224,6 +232,7 @@ export class PipelineRunner {
         createdAt: now,
         updatedAt: now,
         auditIssues: [],
+        ...(output.tokenUsage ? { tokenUsage: output.tokenUsage } : {}),
       };
       await this.state.saveChapterIndex(bookId, [...existingIndex, newEntry]);
 
@@ -235,7 +244,7 @@ export class PipelineRunner {
         wordCount: output.wordCount,
       });
 
-      return { chapterNumber, title: output.title, wordCount: output.wordCount, filePath };
+      return { chapterNumber, title: output.title, wordCount: output.wordCount, filePath, tokenUsage: output.tokenUsage };
     } finally {
       await releaseLock();
     }
@@ -464,6 +473,9 @@ export class PipelineRunner {
       ...(temperatureOverride ? { temperatureOverride } : {}),
     });
 
+    // Token usage accumulator
+    let totalUsage: TokenUsageSummary = output.tokenUsage ?? { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
+
     // 2a. Post-write error gate: if deterministic rules found errors, auto-fix before LLM audit
     let finalContent = output.content;
     let finalWordCount = output.wordCount;
@@ -488,6 +500,7 @@ export class PipelineRunner {
         "spot-fix",
         book.genre,
       );
+      totalUsage = PipelineRunner.addUsage(totalUsage, fixResult.tokenUsage);
       if (fixResult.revisedContent.length > 0) {
         finalContent = fixResult.revisedContent;
         finalWordCount = fixResult.wordCount;
@@ -503,6 +516,7 @@ export class PipelineRunner {
       chapterNumber,
       book.genre,
     );
+    totalUsage = PipelineRunner.addUsage(totalUsage, llmAudit.tokenUsage);
     const aiTellsResult = analyzeAITells(finalContent);
     const sensitiveWriteResult = analyzeSensitiveWords(finalContent);
     const hasBlockedWriteWords = sensitiveWriteResult.found.some((f) => f.severity === "block");
@@ -527,6 +541,7 @@ export class PipelineRunner {
           "spot-fix",
           book.genre,
         );
+        totalUsage = PipelineRunner.addUsage(totalUsage, reviseOutput.tokenUsage);
 
         if (reviseOutput.revisedContent.length > 0) {
           // Guard: reject revision if AI markers increased
@@ -551,6 +566,7 @@ export class PipelineRunner {
             book.genre,
             { temperature: 0 },
           );
+          totalUsage = PipelineRunner.addUsage(totalUsage, reAudit.tokenUsage);
           const reAITells = analyzeAITells(finalContent);
           const reSensitive = analyzeSensitiveWords(finalContent);
           const reHasBlocked = reSensitive.found.some((f) => f.severity === "block");
@@ -608,6 +624,7 @@ export class PipelineRunner {
       auditIssues: auditResult.issues.map(
         (i) => `[${i.severity}] ${i.description}`,
       ),
+      tokenUsage: totalUsage,
     };
     await this.state.saveChapterIndex(bookId, [...existingIndex, newEntry]);
 
@@ -646,6 +663,7 @@ export class PipelineRunner {
       auditResult,
       revised,
       status: auditResult.passed ? "ready-for-review" : "audit-failed",
+      tokenUsage: totalUsage,
     };
   }
 
@@ -955,6 +973,18 @@ ${matrix}`,
     }
   }
 
+  private static addUsage(
+    a: TokenUsageSummary,
+    b?: { readonly promptTokens: number; readonly completionTokens: number; readonly totalTokens: number },
+  ): TokenUsageSummary {
+    if (!b) return a;
+    return {
+      promptTokens: a.promptTokens + b.promptTokens,
+      completionTokens: a.completionTokens + b.completionTokens,
+      totalTokens: a.totalTokens + b.totalTokens,
+    };
+  }
+
   // ---------------------------------------------------------------------------
   // Helpers
   // ---------------------------------------------------------------------------