simstudioai--sim/apps/sim/lib/tokenization/streaming.ts

/**
 * Streaming-specific tokenization helpers
 */

import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { calculateStreamingCost } from '@/lib/tokenization/calculators'
import { TOKENIZATION_CONFIG } from '@/lib/tokenization/constants'
import {
  extractTextContent,
  hasRealCostData,
  hasRealTokenData,
  isTokenizableBlockType,
  logTokenizationDetails,
} from '@/lib/tokenization/utils'
import type { BlockLog } from '@/executor/types'

const logger = createLogger('StreamingTokenization')

/**
 * Processes a block log and adds tokenization data if needed
 */
export function processStreamingBlockLog(log: BlockLog, streamedContent: string): boolean {
  // Check if this block should be tokenized
  if (!isTokenizableBlockType(log.blockType)) {
    return false
  }

  // Check if we already have meaningful token/cost data
  if (hasRealTokenData(log.output?.tokens) && hasRealCostData(log.output?.cost)) {
    return false
  }

  // Skip recalculation if cost was explicitly set by the billing layer (e.g. BYOK zero cost)
  if (log.output?.cost?.pricing) {
    return false
  }

  // Check if we have content to tokenize
  if (!streamedContent?.trim()) {
    return false
  }

  try {
    // Determine model to use
    const model = getModelForBlock(log)

    // Prepare input text from log
    const inputText = extractTextContent(log.input)

    // Calculate streaming cost
    const systemPrompt =
      typeof log.input?.systemPrompt === 'string' ? log.input.systemPrompt : undefined
    const context = typeof log.input?.context === 'string' ? log.input.context : undefined
    const messages = Array.isArray(log.input?.messages)
      ? (log.input.messages as Array<{ role: string; content: string }>)
      : undefined
    const result = calculateStreamingCost(
      model,
      inputText,
      streamedContent,
      systemPrompt,
      context,
      messages
    )

    // Update the log output with tokenization data
    if (!log.output) {
      log.output = {}
    }

    log.output.tokens = result.tokens
    log.output.cost = result.cost
    log.output.model = result.model

    logTokenizationDetails(`Streaming tokenization completed for ${log.blockType}`, {
      blockId: log.blockId,
      blockType: log.blockType,
      model: result.model,
      provider: result.provider,
      inputLength: inputText.length,
      outputLength: streamedContent.length,
      tokens: result.tokens,
      cost: result.cost,
      method: result.method,
    })

    return true
  } catch (error) {
    logger.error(`Streaming tokenization failed for block ${log.blockId}`, {
      blockType: log.blockType,
      error: toError(error).message,
      contentLength: streamedContent?.length || 0,
    })

    // Don't throw - graceful degradation
    return false
  }
}

/**
 * Determines the appropriate model for a block
 */
function getModelForBlock(log: BlockLog): string {
  // Try to get model from output first
  if (log.output?.model?.trim()) {
    return log.output.model
  }

  // Try to get model from input
  const inputModel = log.input?.model
  if (typeof inputModel === 'string' && inputModel.trim()) {
    return inputModel
  }

  // Use block type specific defaults
  const blockType = log.blockType
  if (blockType === 'agent' || blockType === 'router' || blockType === 'evaluator') {
    return TOKENIZATION_CONFIG.defaults.model
  }

  // Final fallback
  return TOKENIZATION_CONFIG.defaults.model
}

/**
 * Processes multiple block logs for streaming tokenization
 */
export function processStreamingBlockLogs(
  logs: BlockLog[],
  streamedContentMap: Map<string, string>
): number {
  let processedCount = 0

  for (const log of logs) {
    const content = streamedContentMap.get(log.blockId)
    if (content && processStreamingBlockLog(log, content)) {
      processedCount++
    }
  }

  logger.info(`Streaming tokenization summary`, {
    totalLogs: logs.length,
    processedBlocks: processedCount,
    streamedBlocks: streamedContentMap.size,
  })

  return processedCount
}