Streaming Tracking

import time
import openai
from ldai.tracker import TokenUsage

def call_streaming_with_tracking(ai_config, user_prompt: str) -> str | None:
    if not ai_config.enabled:
        return None

    tracker = ai_config.create_tracker()
    start_time = time.time()
    first_token_time = None

    try:
        stream = openai.chat.completions.create(
            model=ai_config.model.name,
            messages=[
                {"role": "system", "content": ai_config.messages[0].content},
                {"role": "user", "content": user_prompt},
            ],
            stream=True,
            stream_options={"include_usage": True},  # Required to get usage in final chunk
        )

        response_text = ""
        final_usage = None
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                if first_token_time is None:
                    first_token_time = time.time()
                    tracker.track_time_to_first_token(
                        int((first_token_time - start_time) * 1000)
                    )
                response_text += chunk.choices[0].delta.content
            if getattr(chunk, "usage", None):
                final_usage = chunk.usage

        tracker.track_duration(int((time.time() - start_time) * 1000))
        tracker.track_success()

        if final_usage:
            tracker.track_tokens(TokenUsage(
                total=final_usage.total_tokens,
                input=final_usage.prompt_tokens,
                output=final_usage.completion_tokens,
            ))

        return response_text

    except Exception:
        tracker.track_error()
        raise

import tiktoken
from ldai.tracker import TokenUsage

def estimate_tokens(model_name: str, prompt: str, response: str) -> TokenUsage:
    try:
        enc = tiktoken.encoding_for_model(model_name)
    except KeyError:
        enc = tiktoken.get_encoding("cl100k_base")
    input_tokens = len(enc.encode(prompt))
    output_tokens = len(enc.encode(response))
    return TokenUsage(
        total=input_tokens + output_tokens,
        input=input_tokens,
        output=output_tokens,
    )

import { OpenAI } from 'openai';

const client = new OpenAI();

async function callStreamingWithTracking(
  aiConfig: LDAICompletionConfig,
  userPrompt: string,
): Promise<string | null> {
  if (!aiConfig.enabled) return null;

  const tracker = aiConfig.createTracker();
  const startTime = Date.now();
  let firstTokenTime: number | null = null;

  try {
    const stream = await client.chat.completions.create({
      model: aiConfig.model!.name,
      messages: [
        ...aiConfig.messages,
        { role: 'user', content: userPrompt },
      ],
      stream: true,
      stream_options: { include_usage: true },
    });

    let responseText = '';
    let finalUsage: OpenAI.CompletionUsage | undefined;

    for await (const chunk of stream) {
      const delta = chunk.choices[0]?.delta?.content;
      if (delta) {
        if (firstTokenTime === null) {
          firstTokenTime = Date.now();
          tracker.trackTimeToFirstToken(firstTokenTime - startTime);
        }
        responseText += delta;
      }
      if (chunk.usage) {
        finalUsage = chunk.usage;
      }
    }

    tracker.trackDuration(Date.now() - startTime);
    tracker.trackSuccess();

    if (finalUsage) {
      tracker.trackTokens({
        total: finalUsage.total_tokens,
        input: finalUsage.prompt_tokens,
        output: finalUsage.completion_tokens,
      });
    }

    return responseText;
  } catch (err) {
    tracker.trackError();
    throw err;
  }
}

const tracker = aiConfig.createTracker();
const response = await tracker.trackStreamMetricsOf(
  (chunks) => {
    // Extract usage from the final chunk
    const final = chunks[chunks.length - 1];
    return {
      success: true,
      tokens: {
        total: final.usage?.total_tokens ?? 0,
        input: final.usage?.prompt_tokens ?? 0,
        output: final.usage?.completion_tokens ?? 0,
      },
    };
  },
  () => client.chat.completions.create({ /* ... */, stream: true, stream_options: { include_usage: true } }),
);