Add incremental chat completion traces

Author: jahooma

cli/src/hooks/helpers/__tests__/send-message.test.ts

@@ -1,7 +1,12 @@
 import * as fs from 'fs'
 import path from 'path'
+import { randomUUID } from 'node:crypto'
 
-import { getCurrentChatDir, getMostRecentChatDir, getProjectDataDir } from '../project-files'
+import {
+  getCurrentChatDir,
+  getMostRecentChatDir,
+  getProjectDataDir,
+} from '../project-files'
 import { logger } from './logger'
 
 import type { ChatMessage, ContentBlock } from '../types/chat'
@@ -21,9 +26,9 @@ type SavedChatState = {
  */
 function extractToggleIds(blocks: ContentBlock[] | undefined): string[] {
   if (!blocks) return []
-  
+
   const ids: string[] = []
-  
+
   for (const block of blocks) {
     if (block.type === 'agent') {
       ids.push(block.agentId)
@@ -33,7 +38,7 @@ function extractToggleIds(blocks: ContentBlock[] | undefined): string[] {
       ids.push(block.toolCallId)
     }
   }
-  
+
   return ids
 }
 
@@ -42,11 +47,11 @@ function extractToggleIds(blocks: ContentBlock[] | undefined): string[] {
  */
 export function getAllToggleIdsFromMessages(messages: ChatMessage[]): string[] {
   const ids: string[] = []
-  
+
   for (const message of messages) {
     ids.push(...extractToggleIds(message.blocks))
   }
-  
+
   return ids
 }
 
@@ -69,11 +74,14 @@ export function getChatMessagesPath(): string {
 /**
  * Save both the RunState and ChatMessage[] to disk
  */
-export function saveChatState(runState: RunState, messages: ChatMessage[]): void {
+export function saveChatState(
+  runState: RunState,
+  messages: ChatMessage[],
+): void {
   try {
     const runStatePath = getRunStatePath()
     const messagesPath = getChatMessagesPath()
-    
+
     fs.writeFileSync(runStatePath, JSON.stringify(runState, null, 2))
     fs.writeFileSync(messagesPath, JSON.stringify(messages, null, 2))
   } catch (error) {
@@ -92,14 +100,19 @@ export function saveChatState(runState: RunState, messages: ChatMessage[]): void
  * recently modified chat directory is used.
  * Returns null if no previous chat exists or files can't be parsed.
  */
-export function loadMostRecentChatState(chatId?: string): SavedChatState | null {
+export function loadMostRecentChatState(
+  chatId?: string,
+): SavedChatState | null {
   try {
     let chatDir: string | null = null
 
     if (chatId && chatId.trim().length > 0) {
       const baseDir = path.join(getProjectDataDir(), 'chats')
       const candidateDir = path.join(baseDir, chatId.trim())
-      if (fs.existsSync(candidateDir) && fs.statSync(candidateDir).isDirectory()) {
+      if (
+        fs.existsSync(candidateDir) &&
+        fs.statSync(candidateDir).isDirectory()
+      ) {
         chatDir = candidateDir
       } else {
         logger.debug(
@@ -133,12 +146,18 @@ export function loadMostRecentChatState(chatId?: string): SavedChatState | null
     const messagesContent = fs.readFileSync(messagesPath, 'utf8')
 
     const runState = JSON.parse(runStateContent) as RunState
+    runState.traceSessionId ??= randomUUID()
     const messages = JSON.parse(messagesContent) as ChatMessage[]
 
     const resolvedChatId = path.basename(chatDir)
 
     logger.info(
-      { runStatePath, messagesPath, messageCount: messages.length, chatId: resolvedChatId },
+      {
+        runStatePath,
+        messagesPath,
+        messageCount: messages.length,
+        chatId: resolvedChatId,
+      },
       'Loaded chat state from chat directory',
     )
 
@@ -161,18 +180,15 @@ export function clearChatState(): void {
   try {
     const runStatePath = getRunStatePath()
     const messagesPath = getChatMessagesPath()
-    
+
     if (fs.existsSync(runStatePath)) {
       fs.unlinkSync(runStatePath)
     }
     if (fs.existsSync(messagesPath)) {
       fs.unlinkSync(messagesPath)
     }
-    
-    logger.debug(
-      { runStatePath, messagesPath },
-      'Cleared chat state files'
-    )
+
+    logger.debug({ runStatePath, messagesPath }, 'Cleared chat state files')
   } catch (error) {
     logger.error(
       {

cli/src/utils/run-state-storage.ts

@@ -21,3 +21,35 @@ export type InsertMessageBigqueryFn = (params: {
   dataset?: string
   logger: Logger
 }) => Promise<boolean>
+
+export type ChatCompletionTraceRow = {
+  id: string
+  user_id: string
+  client_id?: string | null
+  trace_session_id: string
+  trace_lineage_id: string
+  run_id: string
+  agent_id: string
+  created_at: Date
+  model: string
+  cost_mode?: string | null
+  request: unknown
+  message_count: number
+  message_start_index: number
+  message_delta_count: number
+  previous_message_count?: number | null
+  common_prefix_length: number
+  cache_hit: boolean
+  full_snapshot: boolean
+  messages: unknown[]
+  delta_message_hashes: string[]
+  tool_count: number
+  tools?: unknown[] | null
+  tools_omitted: boolean
+}
+
+export type InsertChatCompletionTraceBigqueryFn = (params: {
+  row: ChatCompletionTraceRow
+  dataset?: string
+  logger: Logger
+}) => Promise<boolean>

common/src/types/contracts/bigquery.ts

@@ -2,17 +2,26 @@ import { IS_PROD } from '@codebuff/common/env'
 import { getErrorObject } from '@codebuff/common/util/error'
 import { BigQuery } from '@google-cloud/bigquery'
 
-import { MESSAGE_SCHEMA, RELABELS_SCHEMA, TRACES_SCHEMA } from './schema'
+import {
+  CHAT_COMPLETION_TRACES_SCHEMA,
+  MESSAGE_SCHEMA,
+  RELABELS_SCHEMA,
+  TRACES_SCHEMA,
+} from './schema'
 
 import type { BaseTrace, GetRelevantFilesTrace, Relabel, Trace } from './schema'
-import type { MessageRow } from '@codebuff/common/types/contracts/bigquery'
+import type {
+  ChatCompletionTraceRow,
+  MessageRow,
+} from '@codebuff/common/types/contracts/bigquery'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 
 const DATASET = IS_PROD ? 'codebuff_data' : 'codebuff_data_dev'
 
 const TRACES_TABLE = 'traces'
 const RELABELS_TABLE = 'relabels'
 const MESSAGE_TABLE = 'message'
+const CHAT_COMPLETION_TRACES_TABLE = 'chat_completion_traces'
 
 // Create a single BigQuery client instance to be used by all functions
 let client: BigQuery | null = null
@@ -77,6 +86,17 @@ export async function setupBigQuery({
         fields: ['user_id'],
       },
     })
+    await ds.table(CHAT_COMPLETION_TRACES_TABLE).get({
+      autoCreate: true,
+      schema: CHAT_COMPLETION_TRACES_SCHEMA,
+      timePartitioning: {
+        type: 'MONTH',
+        field: 'created_at',
+      },
+      clustering: {
+        fields: ['user_id', 'trace_session_id', 'trace_lineage_id'],
+      },
+    })
   } catch (error) {
     const err = error as Error & { code?: string; details?: unknown }
     logger.error(
@@ -94,6 +114,53 @@ export async function setupBigQuery({
   }
 }
 
+export async function insertChatCompletionTraceBigquery({
+  row,
+  dataset,
+  logger,
+}: {
+  row: ChatCompletionTraceRow
+  dataset?: string
+  logger: Logger
+}) {
+  const resolvedDataset = dataset ?? DATASET
+  try {
+    await getClient()
+      .dataset(resolvedDataset)
+      .table(CHAT_COMPLETION_TRACES_TABLE)
+      .insert({
+        ...row,
+        request: JSON.stringify(row.request),
+        messages: JSON.stringify(row.messages),
+        delta_message_hashes: JSON.stringify(row.delta_message_hashes),
+        tools: row.tools ? JSON.stringify(row.tools) : null,
+      })
+
+    logger.debug(
+      {
+        traceId: row.id,
+        userId: row.user_id,
+        clientId: row.client_id,
+        traceSessionId: row.trace_session_id,
+        traceLineageId: row.trace_lineage_id,
+        runId: row.run_id,
+        messageStartIndex: row.message_start_index,
+        messageDeltaCount: row.message_delta_count,
+        fullSnapshot: row.full_snapshot,
+      },
+      'Inserted chat completion trace into BigQuery',
+    )
+    return true
+  } catch (error) {
+    logger.error(
+      { error: getErrorObject(error), traceId: row.id },
+      'Failed to insert chat completion trace into BigQuery',
+    )
+
+    return false
+  }
+}
+
 export async function insertMessageBigquery({
   row,
   dataset,

packages/bigquery/src/client.ts

@@ -143,3 +143,31 @@ export const MESSAGE_SCHEMA: TableSchema = {
     { name: 'cache_read_input_tokens', type: 'INTEGER', mode: 'NULLABLE' },
   ],
 }
+
+export const CHAT_COMPLETION_TRACES_SCHEMA: TableSchema = {
+  fields: [
+    { name: 'id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'user_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'client_id', type: 'STRING', mode: 'NULLABLE' },
+    { name: 'trace_session_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'trace_lineage_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'run_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'agent_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'created_at', type: 'TIMESTAMP', mode: 'REQUIRED' },
+    { name: 'model', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'cost_mode', type: 'STRING', mode: 'NULLABLE' },
+    { name: 'request', type: 'JSON', mode: 'REQUIRED' },
+    { name: 'message_count', type: 'INTEGER', mode: 'REQUIRED' },
+    { name: 'message_start_index', type: 'INTEGER', mode: 'REQUIRED' },
+    { name: 'message_delta_count', type: 'INTEGER', mode: 'REQUIRED' },
+    { name: 'previous_message_count', type: 'INTEGER', mode: 'NULLABLE' },
+    { name: 'common_prefix_length', type: 'INTEGER', mode: 'REQUIRED' },
+    { name: 'cache_hit', type: 'BOOLEAN', mode: 'REQUIRED' },
+    { name: 'full_snapshot', type: 'BOOLEAN', mode: 'REQUIRED' },
+    { name: 'messages', type: 'JSON', mode: 'REQUIRED' },
+    { name: 'delta_message_hashes', type: 'JSON', mode: 'REQUIRED' },
+    { name: 'tool_count', type: 'INTEGER', mode: 'REQUIRED' },
+    { name: 'tools', type: 'JSON', mode: 'NULLABLE' },
+    { name: 'tools_omitted', type: 'BOOLEAN', mode: 'REQUIRED' },
+  ],
+}

packages/bigquery/src/schema.ts

@@ -62,6 +62,7 @@ export function selectHighestPriorityKnowledgeFile(
 export type RunState = {
   sessionState?: SessionState
   output: AgentOutput
+  traceSessionId: string
 }
 
 export type InitialSessionStateOptions = {
@@ -630,6 +631,7 @@ export async function generateInitialRunState({
   fs: CodebuffFileSystem
 }): Promise<RunState> {
   return {
+    traceSessionId: crypto.randomUUID(),
     sessionState: await initialSessionState({
       cwd,
       skillsDir,