fix: emit Langfuse generations for root RLM calls

dist/src/langfuse.d.ts

@@ -12,6 +12,19 @@ export interface LangfuseConfig {
     fetchImpl?: typeof fetch;
     flushTimeoutMs?: number;
 }
+export interface RootGenerationStartData {
+    name: string;
+    input: unknown;
+    model: string;
+    iteration: number;
+}
+export interface RootGenerationEndData {
+    output: unknown;
+    durationMs: number;
+    usage?: UsageStats;
+    isError?: boolean;
+    errorMessage?: string;
+}
 export declare class LangfuseTraceRecorder {
     private host;
     private publicKey;
@@ -29,6 +42,8 @@ export declare class LangfuseTraceRecorder {
         userId?: string;
         metadata?: Record<string, unknown>;
     }): void;
+    rootGenerationStart(data: RootGenerationStartData): string;
+    rootGenerationEnd(generationId: string, data: RootGenerationEndData): void;
     childStart(data: {
         parentRunId: string;
         childRunId?: string;

src/langfuse.ts

@@ -23,6 +23,21 @@ interface LangfuseEvent {
   body: Record<string, unknown>;
 }
 
+export interface RootGenerationStartData {
+  name: string;
+  input: unknown;
+  model: string;
+  iteration: number;
+}
+
+export interface RootGenerationEndData {
+  output: unknown;
+  durationMs: number;
+  usage?: UsageStats;
+  isError?: boolean;
+  errorMessage?: string;
+}
+
 export class LangfuseTraceRecorder {
   private host: string | null;
   private publicKey: string | null;
@@ -62,6 +77,64 @@ export class LangfuseTraceRecorder {
     });
   }
 
+  rootGenerationStart(data: RootGenerationStartData): string {
+    const generationId = randomUUID();
+    if (!this.enabled) return generationId;
+    this.enqueue("generation-create", {
+      id: generationId,
+      traceId: this.traceId,
+      name: data.name,
+      model: data.model,
+      input: data.input,
+      startTime: new Date().toISOString(),
+      metadata: {
+        event: "root_generation_start",
+        iteration: data.iteration,
+      },
+    });
+    return generationId;
+  }
+
+  rootGenerationEnd(generationId: string, data: RootGenerationEndData): void {
+    if (!this.enabled) return;
+    const input = data.usage?.inputTokens ?? 0;
+    const output = data.usage?.outputTokens ?? 0;
+    const cacheRead = data.usage?.cacheReadTokens ?? 0;
+    const cacheWrite = data.usage?.cacheWriteTokens ?? 0;
+    this.enqueue("generation-update", {
+      id: generationId,
+      output: data.output,
+      endTime: new Date().toISOString(),
+      level: data.isError ? "ERROR" : "DEFAULT",
+      statusMessage: data.errorMessage,
+      usage: {
+        input,
+        output,
+        total: input + output + cacheRead + cacheWrite,
+      },
+      usageDetails: {
+        input,
+        output,
+        cache_read: cacheRead,
+        cache_write: cacheWrite,
+        total: input + output + cacheRead + cacheWrite,
+      },
+      costDetails: {
+        total: data.usage?.totalCost ?? 0,
+      },
+      metadata: {
+        event: "root_generation_end",
+        duration_ms: data.durationMs,
+        input_tokens: input,
+        output_tokens: output,
+        cache_read_tokens: cacheRead,
+        cache_write_tokens: cacheWrite,
+        total_cost: data.usage?.totalCost ?? 0,
+        llm_calls: data.usage?.llmCalls ?? 0,
+      },
+    });
+  }
+
   childStart(data: {
     parentRunId: string;
     childRunId?: string;

src/rlm.ts

@@ -426,6 +426,12 @@ export async function rlmLoop(
 
       // Call LLM
       const llmStartMs = Date.now();
+      const generationId = langfuse.rootGenerationStart({
+        name: `Model call — root iteration ${iteration + 1}`,
+        input: messages,
+        model: `${config.model.provider}/${config.model.model}`,
+        iteration,
+      });
       const response = await llmComplete(messages, config.model, {
         signal: abortController.signal,
         cacheConfig,
@@ -434,6 +440,11 @@ export async function rlmLoop(
         geminiConfig: config.gemini,
       });
       const llmDurationMs = Date.now() - llmStartMs;
+      langfuse.rootGenerationEnd(generationId, {
+        output: response.text,
+        durationMs: llmDurationMs,
+        usage: response.usage,
+      });
       mergeUsage(usage, response.usage);
       budget.record(response.usage.inputTokens, response.usage.outputTokens, response.usage.totalCost);
 
@@ -664,7 +675,7 @@ export async function rlmLoop(
       logVerbose(actualIterations, `${reason}, forcing final answer`);
     }
 
-    const forcedResult = await forceFinalAnswer(messages, config, usage, abortController.signal, cacheConfig);
+    const forcedResult = await forceFinalAnswer(messages, config, usage, abortController.signal, cacheConfig, langfuse, actualIterations);
     return finalize(forcedResult, actualIterations);
   } catch (err: unknown) {
     clearTimeout(timeoutHandle);
@@ -697,7 +708,9 @@ async function forceFinalAnswer(
   config: RlmxConfig,
   usage: UsageStats,
   signal?: AbortSignal,
-  cacheConfig?: CacheLLMConfig
+  cacheConfig?: CacheLLMConfig,
+  langfuse?: LangfuseTraceRecorder,
+  iteration = 0
 ): Promise<string> {
   const forceMessages: ChatMessage[] = [
     ...messages,
@@ -708,13 +721,27 @@ async function forceFinalAnswer(
     },
   ];
 
+  const generationId = langfuse?.rootGenerationStart({
+    name: "Model call — forced final answer",
+    input: forceMessages,
+    model: `${config.model.provider}/${config.model.model}`,
+    iteration,
+  });
+  const llmStartMs = Date.now();
   const response = await llmComplete(forceMessages, config.model, {
     signal,
     cacheConfig,
     thinkingLevel: config.gemini.thinkingLevel,
     outputSchema: config.output.schema,
     geminiConfig: config.gemini,
   });
+  if (generationId) {
+    langfuse?.rootGenerationEnd(generationId, {
+      output: response.text,
+      durationMs: Date.now() - llmStartMs,
+      usage: response.usage,
+    });
+  }
   mergeUsage(usage, response.usage);
   return response.text;
 }