add data to bigquery table

charleslien · charleslien · commit 230e39b82f71 · 2025-10-06T10:16:15.000-07:00
diff --git a/npm-app/src/asdf.ts b/npm-app/src/asdf.ts
@@ -10,13 +10,13 @@ const codebuffBackendProvider = createOpenAICompatible({
 })
 
 const response = streamText({
-  model: codebuffBackendProvider('anthropic/claude-sonnet-4.5'),
+  model: codebuffBackendProvider('openai/gpt-5'),
   messages: [
     {
       role: 'user',
       content:
         'This is a bunch of text just to fill out some space. Ignore this.'.repeat(
-          1000,
+          100,
         ),
     },
     {
@@ -29,6 +29,15 @@ const response = streamText({
       },
     },
   ],
+  providerOptions: {
+    codebuff: {
+      // all these get directly added to the body at the top level
+      reasoningEffort: 'low',
+      codebuff_metadata: {
+        agent_run_id: 'testing',
+      },
+    },
+  },
 })
 for await (const chunk of response.fullStream) {
   console.log('asdf', { chunk })
diff --git a/packages/bigquery/src/client.ts b/packages/bigquery/src/client.ts
@@ -1,9 +1,16 @@
 import { logger } from '@codebuff/common/util/logger'
+import { errorToObject } from '@codebuff/common/util/object'
 import { BigQuery } from '@google-cloud/bigquery'
 
-import { RELABELS_SCHEMA, TRACES_SCHEMA } from './schema'
+import { MESSAGE_SCHEMA, RELABELS_SCHEMA, TRACES_SCHEMA } from './schema'
 
-import type { BaseTrace, GetRelevantFilesTrace, Relabel, Trace } from './schema'
+import type {
+  BaseTrace,
+  GetRelevantFilesTrace,
+  MessageRow,
+  Relabel,
+  Trace,
+} from './schema'
 
 const DATASET =
   process.env.NEXT_PUBLIC_CB_ENVIRONMENT === 'prod'
@@ -12,6 +19,7 @@ const DATASET =
 
 const TRACES_TABLE = 'traces'
 const RELABELS_TABLE = 'relabels'
+const MESSAGE_TABLE = 'message'
 
 // Create a single BigQuery client instance to be used by all functions
 let client: BigQuery | null = null
@@ -26,6 +34,9 @@ function getClient(): BigQuery {
 }
 
 export async function setupBigQuery(dataset: string = DATASET) {
+  if (client) {
+    return
+  }
   try {
     client = new BigQuery()
 
@@ -55,6 +66,17 @@ export async function setupBigQuery(dataset: string = DATASET) {
         fields: ['user_id', 'agent_step_id'],
       },
     })
+    await ds.table(MESSAGE_TABLE).get({
+      autoCreate: true,
+      schema: MESSAGE_SCHEMA,
+      timePartitioning: {
+        type: 'MONTH',
+        field: 'finished_at',
+      },
+      clustering: {
+        fields: ['user_id'],
+      },
+    })
   } catch (error) {
     logger.error(
       {
@@ -71,6 +93,34 @@ export async function setupBigQuery(dataset: string = DATASET) {
   }
 }
 
+export async function insertMessage(
+  row: MessageRow,
+  dataset: string = DATASET,
+) {
+  try {
+    await getClient()
+      .dataset(dataset)
+      .table(MESSAGE_TABLE)
+      .insert({ ...row, request: JSON.stringify(row.request) })
+
+    logger.debug(
+      {
+        ...row,
+        request: undefined,
+      },
+      'Inserted message into BigQuery',
+    )
+    return true
+  } catch (error) {
+    logger.error(
+      { error: errorToObject(error), messageId: row.id },
+      'Failed to insert message into BigQuery',
+    )
+
+    return false
+  }
+}
+
 export async function insertTrace(trace: Trace, dataset: string = DATASET) {
   try {
     // Create a copy of the trace and stringify payload if needed
diff --git a/packages/bigquery/src/schema.ts b/packages/bigquery/src/schema.ts
@@ -125,3 +125,35 @@ export const RELABELS_SCHEMA: TableSchema = {
     { name: 'payload', type: 'JSON', mode: 'REQUIRED' },
   ],
 }
+
+export type MessageRow = {
+  id: string
+  user_id: string
+  finished_at: Date
+  created_at: Date
+  request: unknown
+  response: string
+  output_tokens?: number | null
+  reasoning_tokens?: number | null
+  cost?: number | null
+  upstream_inference_cost?: number | null
+  input_tokens?: number | null
+  cache_read_input_tokens?: number | null
+}
+
+export const MESSAGE_SCHEMA: TableSchema = {
+  fields: [
+    { name: 'id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'user_id', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'finished_at', type: 'TIMESTAMP', mode: 'REQUIRED' },
+    { name: 'created_at', type: 'TIMESTAMP', mode: 'REQUIRED' },
+    { name: 'request', type: 'JSON', mode: 'REQUIRED' },
+    { name: 'response', type: 'STRING', mode: 'REQUIRED' },
+    { name: 'output_tokens', type: 'INTEGER', mode: 'NULLABLE' },
+    { name: 'reasoning_tokens', type: 'INTEGER', mode: 'NULLABLE' },
+    { name: 'cost', type: 'FLOAT', mode: 'NULLABLE' },
+    { name: 'upstream_inference_cost', type: 'FLOAT', mode: 'NULLABLE' },
+    { name: 'input_tokens', type: 'INTEGER', mode: 'NULLABLE' },
+    { name: 'cache_read_input_tokens', type: 'INTEGER', mode: 'NULLABLE' },
+  ],
+}
diff --git a/web/src/app/api/v1/chat/completions/route.ts b/web/src/app/api/v1/chat/completions/route.ts
@@ -1,9 +1,10 @@
+import { getUserUsageData } from '@codebuff/billing/usage-service'
 import { NextResponse } from 'next/server'
 
 import type { NextRequest } from 'next/server'
 
 import { getUserInfoFromApiKey } from '@/db/user'
-import { handleOpenrouterStream } from '@/llm-api/openrouter'
+import { handleOpenRouterStream } from '@/llm-api/openrouter'
 import { extractApiKeyFromHeader } from '@/util/auth'
 import { errorToObject } from '@/util/error'
 import { logger } from '@/util/logger'
@@ -30,10 +31,25 @@ export async function POST(req: NextRequest) {
       )
     }
 
+    const userId = userInfo.id
+    const {
+      balance: { totalRemaining },
+      nextQuotaReset,
+    } = await getUserUsageData(userId)
+    if (totalRemaining <= 0) {
+      return NextResponse.json(
+        {
+          message: `Insufficient credits. Please add credits at ${process.env.NEXT_PUBLIC_APP_URL}/usage or wait for your next cycle to begin (${nextQuotaReset}).`,
+        },
+        { status: 402 }
+      )
+    }
+
     if (body.stream) {
       try {
-        const stream = await handleOpenrouterStream({
+        const stream = await handleOpenRouterStream({
           body,
+          userId,
         })
 
         return new NextResponse(stream, {
diff --git a/web/src/llm-api/openrouter.ts b/web/src/llm-api/openrouter.ts
@@ -1,6 +1,25 @@
+import {
+  insertMessage as insertMessageIntoBigquery,
+  setupBigQuery,
+} from '@codebuff/bigquery'
 import { env } from '@codebuff/internal/env'
 
-export async function handleOpenrouterStream({ body }: { body: any }) {
+import { OpenRouterStreamChatCompletionChunkSchema } from './type/openrouter'
+
+import type { OpenRouterStreamChatCompletionChunk } from './type/openrouter'
+
+import { errorToObject } from '@/util/error'
+import { logger } from '@/util/logger'
+
+type StreamState = { responseText: string }
+
+export async function handleOpenRouterStream({
+  body,
+  userId,
+}: {
+  body: any
+  userId: string
+}) {
   // Ensure usage tracking is enabled
   if (body.usage === undefined) {
     body.usage = {}
@@ -30,7 +49,8 @@ export async function handleOpenrouterStream({ body }: { body: any }) {
     throw new Error('Failed to get response reader')
   }
 
-  let heartbeatInterval: ReturnType<typeof setInterval>
+  let heartbeatInterval: NodeJS.Timeout
+  let state: StreamState = { responseText: '' }
 
   // Create a ReadableStream that Next.js can handle
   const stream = new ReadableStream({
@@ -67,6 +87,8 @@ export async function handleOpenrouterStream({ body }: { body: any }) {
             const line = buffer.slice(0, lineEnd + 1)
             buffer = buffer.slice(lineEnd + 1)
 
+            state = await handleLine({ userId, request: body, line, state })
+
             // Forward the line to the client
             controller.enqueue(new TextEncoder().encode(line))
 
@@ -90,3 +112,109 @@ export async function handleOpenrouterStream({ body }: { body: any }) {
 
   return stream
 }
+
+async function handleLine({
+  userId,
+  request,
+  line,
+  state,
+}: {
+  userId: string
+  request: unknown
+  line: string
+  state: StreamState
+}): Promise<StreamState> {
+  if (!line.startsWith('data: ')) {
+    return state
+  }
+
+  const raw = line.slice('data: '.length)
+  if (raw === '[DONE]\n') {
+    return state
+  }
+
+  // Parse the string into an object
+  let obj
+  try {
+    obj = JSON.parse(raw)
+  } catch (error) {
+    logger.warn(
+      `Received non-JSON OpenRouter response: ${JSON.stringify(errorToObject(error), null, 2)}`
+    )
+    return state
+  }
+
+  // Extract usage
+  const parsed = OpenRouterStreamChatCompletionChunkSchema.safeParse(obj)
+  if (!parsed.success) {
+    logger.warn(
+      `Unable to parse OpenRotuer response: ${JSON.stringify(errorToObject(parsed.error), null, 2)}`
+    )
+    return state
+  }
+
+  return await handleResponse({ userId, request, data: parsed.data, state })
+}
+
+async function handleResponse({
+  userId,
+  request,
+  data,
+  state,
+}: {
+  userId: string
+  request: unknown
+  data: OpenRouterStreamChatCompletionChunk
+  state: StreamState
+}): Promise<StreamState> {
+  state = await handleStreamChunk({ data, state })
+
+  if ('error' in data || !data.usage) {
+    // Stream not finished
+    return state
+  }
+  const usage = data.usage
+
+  // do not await this
+  setupBigQuery().then(() =>
+    insertMessageIntoBigquery({
+      id: data.id,
+      user_id: userId,
+      finished_at: new Date(),
+      created_at: new Date(data.created * 1000),
+      request,
+      response: state.responseText,
+      output_tokens: usage.completion_tokens,
+      reasoning_tokens: usage.completion_tokens_details?.reasoning_tokens,
+      cost: usage.cost,
+      upstream_inference_cost: usage.cost_details?.upstream_inference_cost,
+      input_tokens: usage.prompt_tokens,
+      cache_read_input_tokens: usage.prompt_tokens_details?.cached_tokens,
+    })
+  )
+  const openRouterCost = usage.cost ?? 0
+  const upstreamCost = usage.cost_details?.upstream_inference_cost ?? 0
+  const cost = openRouterCost + upstreamCost
+  // asdf todo: charge user
+  return state
+}
+
+async function handleStreamChunk({
+  data,
+  state,
+}: {
+  data: OpenRouterStreamChatCompletionChunk
+  state: StreamState
+}): Promise<StreamState> {
+  if ('error' in data) {
+    logger.warn({ streamChunk: data }, 'Received error from OpenRouter')
+    return state
+  }
+
+  if (!data.choices.length) {
+    logger.warn({ streamChunk: data }, 'Received empty choices from OpenRouter')
+  }
+  const choice = data.choices[0]
+  state.responseText += choice.delta?.content ?? ''
+  return state
+}
diff --git a/web/src/llm-api/type/openrouter.ts b/web/src/llm-api/type/openrouter.ts
diff --git a/web/src/util/error.ts b/web/src/util/error.ts
