openlayer 0.4.0 → 0.7.0

package/src/lib/integrations/langchainCallback.ts

@@ -0,0 +1,141 @@
+import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
+import { LLMResult } from '@langchain/core/dist/outputs';
+import type { Serialized } from '@langchain/core/load/serializable';
+import { AIMessage, BaseMessage, SystemMessage } from '@langchain/core/messages';
+import { addChatCompletionStepToTrace } from '../tracing/tracer';
+
+const LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP: Record<string, string> = {
+  openai: 'OpenAI',
+  'openai-chat': 'OpenAI',
+  'chat-ollama': 'Ollama',
+  vertexai: 'Google',
+};
+const PROVIDER_TO_STEP_NAME: Record<string, string> = {
+  OpenAI: 'OpenAI Chat Completion',
+  Ollama: 'Ollama Chat Completion',
+  Google: 'Google Vertex AI Chat Completion',
+};
+
+export class OpenlayerHandler extends BaseCallbackHandler {
+  name = 'OpenlayerHandler';
+  startTime: number | null = null;
+  endTime: number | null = null;
+  prompt: Array<{ role: string; content: string }> | null = null;
+  latency: number = 0;
+  provider: string | undefined;
+  model: string | null = null;
+  modelParameters: Record<string, any> | null = null;
+  promptTokens: number | null = 0;
+  completionTokens: number | null = 0;
+  totalTokens: number | null = 0;
+  output: string = '';
+  metadata: Record<string, any>;
+
+  constructor(kwargs: Record<string, any> = {}) {
+    super();
+    this.metadata = kwargs;
+  }
+  override async handleChatModelStart(
+    llm: Serialized,
+    messages: BaseMessage[][],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string,
+  ): Promise<void> {
+    this.initializeRun(extraParams || {}, metadata || {});
+    this.prompt = this.langchainMassagesToPrompt(messages);
+    this.startTime = performance.now();
+  }
+
+  private initializeRun(extraParams: Record<string, any>, metadata: Record<string, unknown>): void {
+    this.modelParameters = extraParams['invocation_params'] || {};
+
+    const provider = metadata?.['ls_provider'] as string;
+    if (provider && LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP[provider]) {
+      this.provider = LANGCHAIN_TO_OPENLAYER_PROVIDER_MAP[provider];
+    }
+    this.model = (this.modelParameters?.['model'] as string) || (metadata['ls_model_name'] as string) || null;
+    this.output = '';
+  }
+
+  private langchainMassagesToPrompt(messages: BaseMessage[][]): Array<{ role: string; content: string }> {
+    let prompt: Array<{ role: string; content: string }> = [];
+    for (const message of messages) {
+      for (const m of message) {
+        if (m instanceof AIMessage) {
+          prompt.push({ role: 'assistant', content: m.content as string });
+        } else if (m instanceof SystemMessage) {
+          prompt.push({ role: 'system', content: m.content as string });
+        } else {
+          prompt.push({ role: 'user', content: m.content as string });
+        }
+      }
+    }
+    return prompt;
+  }
+
+  override async handleLLMStart(
+    llm: Serialized,
+    prompts: string[],
+    runId: string,
+    parentRunId?: string,
+    extraParams?: Record<string, unknown>,
+    tags?: string[],
+    metadata?: Record<string, unknown>,
+    runName?: string,
+  ) {
+    this.initializeRun(extraParams || {}, metadata || {});
+    this.prompt = prompts.map((p) => ({ role: 'user', content: p }));
+    this.startTime = performance.now();
+  }
+
+  override async handleLLMEnd(output: LLMResult, runId: string, parentRunId?: string, tags?: string[]) {
+    this.endTime = performance.now();
+    this.latency = this.endTime - this.startTime!;
+    this.extractTokenInformation(output);
+    this.extractOutput(output);
+    this.addToTrace();
+  }
+
+  private extractTokenInformation(output: LLMResult) {
+    if (this.provider === 'OpenAI') {
+      this.openaiTokenInformation(output);
+    }
+  }
+
+  private openaiTokenInformation(output: LLMResult) {
+    if (output.llmOutput && 'tokenUsage' in output.llmOutput) {
+      this.promptTokens = output.llmOutput?.['tokenUsage']?.promptTokens ?? 0;
+      this.completionTokens = output.llmOutput?.['tokenUsage']?.completionTokens ?? 0;
+      this.totalTokens = output.llmOutput?.['tokenUsage']?.totalTokens ?? 0;
+    }
+  }
+
+  private extractOutput(output: LLMResult) {
+    const lastResponse = output?.generations?.at(-1)?.at(-1) ?? undefined;
+    this.output += lastResponse?.text ?? '';
+  }
+
+  private addToTrace() {
+    let name = 'Chat Completion Model';
+    if (this.provider && this.provider in PROVIDER_TO_STEP_NAME) {
+      name = PROVIDER_TO_STEP_NAME[this.provider] ?? 'Chat Completion Model';
+    }
+    addChatCompletionStepToTrace({
+      name: name,
+      inputs: { prompt: this.prompt },
+      output: this.output,
+      latency: this.latency,
+      tokens: this.totalTokens,
+      promptTokens: this.promptTokens,
+      completionTokens: this.completionTokens,
+      model: this.model,
+      modelParameters: this.modelParameters,
+      metadata: this.metadata,
+      provider: this.provider ?? '',
+    });
+  }
+}

package/src/lib/integrations/openAiTracer.ts

@@ -0,0 +1,137 @@
+import OpenAI from 'openai';
+import { Stream } from 'openai/src/streaming';
+import { addChatCompletionStepToTrace } from '../tracing/tracer';
+
+export function traceOpenAI(openai: OpenAI): OpenAI {
+  const createFunction = openai.chat.completions.create;
+
+  openai.chat.completions.create = async function (
+    this: typeof openai.chat.completions,
+    ...args: Parameters<typeof createFunction>
+  ): Promise<Stream<OpenAI.Chat.Completions.ChatCompletionChunk> | OpenAI.Chat.Completions.ChatCompletion> {
+    const [params, options] = args;
+    const stream = params?.stream ?? false;
+
+    try {
+      const startTime = performance.now();
+
+      // Call the original `create` function
+      let response = await createFunction.apply(this, args);
+
+      if (stream) {
+        // Handle streaming responses
+        const chunks: OpenAI.Chat.Completions.ChatCompletionChunk[] = [];
+        let collectedOutputData: any[] = [];
+        let firstTokenTime: number | undefined;
+        let completionTokens: number = 0;
+        if (isAsyncIterable(response)) {
+          async function* tracedOutputGenerator(): AsyncGenerator<
+            OpenAI.Chat.Completions.ChatCompletionChunk,
+            void,
+            unknown
+          > {
+            for await (const rawChunk of response as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>) {
+              if (chunks.length === 0) {
+                firstTokenTime = performance.now();
+              }
+              chunks.push(rawChunk);
+              const delta = rawChunk.choices[0]?.delta;
+              if (delta?.content) {
+                collectedOutputData.push(delta?.content);
+              } else if (delta?.tool_calls) {
+                const tool_call = delta.tool_calls[0];
+                if (tool_call?.function?.name) {
+                  const functionName: string =
+                    '{\n    "name": ' + '"' + tool_call.function.name + '"' + '\n    "arguments": ';
+                  collectedOutputData.push(functionName);
+                } else if (tool_call?.function?.arguments) {
+                  collectedOutputData.push(tool_call.function.arguments);
+                }
+              }
+
+              if (rawChunk.choices[0]?.finish_reason === 'tool_calls') {
+                collectedOutputData.push('\n}');
+              }
+              completionTokens += 1;
+              yield rawChunk;
+            }
+            const endTime = performance.now();
+            const traceData = {
+              name: 'OpenAI Chat Completion',
+              inputs: { prompt: params.messages },
+              output: collectedOutputData.join(''),
+              latency: endTime - startTime,
+              model: chunks[0]?.model as string,
+              modelParameters: getModelParameters(args),
+              rawOutput: chunks.map((chunk) => JSON.stringify(chunk, null, 2)).join('\n'),
+              metadata: { timeToFistToken: firstTokenTime ? firstTokenTime - startTime : null },
+              provider: 'OpenAI',
+              completionTokens: completionTokens,
+              promptTokens: 0,
+              tokens: completionTokens,
+            };
+            addChatCompletionStepToTrace(traceData);
+          }
+          return tracedOutputGenerator() as unknown as Stream<OpenAI.Chat.Completions.ChatCompletionChunk>;
+        }
+      } else {
+        // Handle non-streaming responses
+        response = response as OpenAI.Chat.Completions.ChatCompletion;
+        const completion = response.choices[0];
+        const endTime = performance.now();
+
+        let output: string = '';
+        if (completion?.message?.content) {
+          output = completion.message.content;
+        } else if (completion?.message.tool_calls) {
+          const tool_call = completion.message.tool_calls[0];
+          output = JSON.stringify(tool_call?.function, null, 2);
+        }
+
+        const traceData = {
+          name: 'OpenAI Chat Completion',
+          inputs: { prompt: params.messages },
+          output: output,
+          latency: endTime - startTime,
+          tokens: response.usage?.total_tokens ?? null,
+          promptTokens: response.usage?.prompt_tokens ?? null,
+          completionTokens: response.usage?.completion_tokens ?? null,
+          model: response.model,
+          modelParameters: getModelParameters(args),
+          rawOutput: JSON.stringify(response, null, 2),
+          metadata: {},
+          provider: 'OpenAI',
+        };
+        addChatCompletionStepToTrace(traceData);
+        return response;
+      }
+    } catch (error) {
+      console.error('Failed to trace the create chat completion request with Openlayer', error);
+      throw error;
+    }
+    // Ensure a return statement is present
+    return undefined as any;
+  } as typeof createFunction;
+
+  return openai;
+}
+
+function getModelParameters(args: any): Record<string, any> {
+  const params = args[0];
+  return {
+    frequency_penalty: params?.frequencyPenalty ?? 0,
+    logit_bias: params?.logitBias ?? null,
+    logprobs: params?.logprobs ?? false,
+    top_logprobs: params?.topLogprobs ?? null,
+    max_tokens: params?.maxTokens ?? null,
+    n: params?.n ?? 1,
+    presence_penalty: params?.presencePenalty ?? 0,
+    seed: params?.seed ?? null,
+    stop: params?.stop ?? null,
+    temperature: params?.temperature ?? 1,
+    top_p: params?.topP ?? 1,
+  };
+}
+
+const isAsyncIterable = (x: any) =>
+  x != null && typeof x === 'object' && typeof x[Symbol.asyncIterator] === 'function';

package/src/lib/tracing/enums.ts

@@ -0,0 +1,4 @@
+export enum StepType {
+  USER_CALL = 'user_call',
+  CHAT_COMPLETION = 'chat_completion',
+}

package/src/lib/tracing/index.d.ts

@@ -0,0 +1 @@
+export * as tracer from './tracer';

package/src/lib/tracing/steps.ts

@@ -0,0 +1,118 @@
+import { v4 as uuidv4 } from 'uuid';
+import { StepType } from './enums';
+
+export interface StepData {
+  name: string;
+  id: string;
+  type: StepType;
+  inputs: any;
+  output: any;
+  groundTruth: any;
+  metadata: Record<string, any>;
+  steps: StepData[];
+  latency: number | null;
+  startTime: number;
+  endTime: number | null;
+}
+
+export interface ChatCompletionStepData extends StepData {
+  provider: string | null;
+  promptTokens: number | null;
+  completionTokens: number | null;
+  tokens: number | null;
+  cost: number | null;
+  model: string | null;
+  modelParameters: Record<string, any> | null;
+  rawOutput: string | null;
+}
+
+export class Step {
+  name: string;
+  id: string;
+  inputs: any;
+  output: any;
+  metadata: Record<string, any>;
+  stepType: StepType | null = null;
+  startTime: number;
+  endTime: number | null = null;
+  groundTruth: any = null;
+  latency: number | null = null;
+  steps: Step[] = [];
+
+  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
+    this.name = name;
+    this.id = uuidv4();
+    this.inputs = inputs;
+    this.output = output;
+    this.metadata = metadata;
+
+    this.startTime = Date.now();
+  }
+
+  addNestedStep(nestedStep: Step): void {
+    this.steps.push(nestedStep);
+  }
+
+  log(data: Partial<Record<keyof this, any>>): void {
+    Object.keys(data).forEach((key) => {
+      if (key in this) {
+        // @ts-ignore
+        this[key] = data[key];
+      }
+    });
+  }
+
+  toJSON(): StepData {
+    return {
+      name: this.name,
+      id: this.id,
+      type: this.stepType!,
+      inputs: this.inputs,
+      output: this.output,
+      groundTruth: this.groundTruth,
+      metadata: this.metadata,
+      steps: this.steps.map((nestedStep) => nestedStep.toJSON()),
+      latency: this.latency,
+      startTime: this.startTime,
+      endTime: this.endTime,
+    };
+  }
+}
+
+export class UserCallStep extends Step {
+  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
+    super(name, inputs, output, metadata);
+    this.stepType = StepType.USER_CALL;
+  }
+}
+
+export class ChatCompletionStep extends Step {
+  provider: string | null = null;
+  promptTokens: number | null = null;
+  completionTokens: number | null = null;
+  tokens: number | null = null;
+  cost: number | null = null;
+  model: string | null = null;
+  modelParameters: Record<string, any> | null = null;
+  rawOutput: string | null = null;
+
+  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
+    super(name, inputs, output, metadata);
+    this.stepType = StepType.CHAT_COMPLETION;
+  }
+
+  override toJSON(): ChatCompletionStepData {
+    const stepData = super.toJSON();
+    return {
+      ...stepData,
+      provider: this.provider,
+      promptTokens: this.promptTokens,
+      completionTokens: this.completionTokens,
+      tokens: this.tokens,
+      cost: this.cost,
+      model: this.model,
+      modelParameters: this.modelParameters,
+      rawOutput: this.rawOutput,
+    };
+  }
+}

package/src/lib/tracing/tracer.ts

@@ -0,0 +1,204 @@
+// tracing/tracer.ts
+
+import { Trace } from './traces';
+import { Step, UserCallStep, ChatCompletionStep } from './steps';
+import { StepType } from './enums';
+import Openlayer from '../../index';
+
+let currentTrace: Trace | null = null;
+
+const publish = process.env['OPENLAYER_DISABLE_PUBLISH'] != 'true';
+let client: Openlayer | null = null;
+if (publish) {
+  console.debug('Publishing is enabled');
+  client = new Openlayer();
+}
+
+function getCurrentTrace(): Trace | null {
+  return currentTrace;
+}
+
+function setCurrentTrace(trace: Trace | null) {
+  currentTrace = trace;
+}
+
+// Function to create a new step
+const stepStack: Step[] = [];
+
+function createStep(
+  name: string,
+  stepType: StepType = StepType.USER_CALL,
+  inputs?: any,
+  output?: any,
+  metadata: Record<string, any> | null = null,
+): [Step, () => void] {
+  metadata = metadata || {};
+  let newStep: Step;
+  if (stepType === StepType.CHAT_COMPLETION) {
+    newStep = new ChatCompletionStep(name, inputs, output, metadata);
+  } else {
+    newStep = new UserCallStep(name, inputs, output, metadata);
+  }
+  newStep.startTime = performance.now();
+
+  const parentStep = getCurrentStep();
+  const isRootStep = parentStep === null;
+
+  if (isRootStep) {
+    console.debug('Starting a new trace...');
+    console.debug(`Adding step ${name} as the root step`);
+    const currentTrace = new Trace();
+    setCurrentTrace(currentTrace);
+    currentTrace.addStep(newStep);
+  } else {
+    console.debug(`Adding step ${name} as a nested step to ${parentStep!.name}`);
+    currentTrace = getCurrentTrace()!;
+    parentStep!.addNestedStep(newStep);
+  }
+
+  stepStack.push(newStep);
+
+  const endStep = () => {
+    newStep.endTime = performance.now();
+    newStep.latency = newStep.endTime - newStep.startTime;
+
+    stepStack.pop(); // Remove the current step from the stack
+
+    if (isRootStep) {
+      console.debug('Ending the trace...');
+      const traceData = getCurrentTrace();
+      // Post process trace and get the input variable names
+      const { traceData: processedTraceData, inputVariableNames } = postProcessTrace(traceData!);
+
+      if (publish && process.env['OPENLAYER_INFERENCE_PIPELINE_ID']) {
+        client!.inferencePipelines.data.stream(process.env['OPENLAYER_INFERENCE_PIPELINE_ID'], {
+          config: {
+            outputColumnName: 'output',
+            inputVariableNames: inputVariableNames,
+            groundTruthColumnName: 'groundTruth',
+            latencyColumnName: 'latency',
+            costColumnName: 'cost',
+            timestampColumnName: 'inferenceTimestamp',
+            inferenceIdColumnName: 'inferenceId',
+            numOfTokenColumnName: 'tokens',
+          },
+          rows: [processedTraceData],
+        });
+      }
+
+      // Reset the entire trace state
+      setCurrentTrace(null);
+      stepStack.length = 0; // Clear the step stack
+    } else {
+      console.debug(`Ending step ${name}`);
+    }
+  };
+
+  return [newStep, endStep];
+}
+
+function getCurrentStep(): Step | null | undefined {
+  const currentStep = stepStack.length > 0 ? stepStack[stepStack.length - 1] : null;
+  return currentStep;
+}
+
+function getParamNames(func: Function): string[] {
+  const STRIP_COMMENTS = /((\/\/.*$)|(\/\*[\s\S]*?\*\/))/gm;
+  const ARGUMENT_NAMES = /([^\s,]+)/g;
+  const fnStr = func.toString().replace(STRIP_COMMENTS, '');
+  const result = fnStr.slice(fnStr.indexOf('(') + 1, fnStr.indexOf(')')).match(ARGUMENT_NAMES);
+  return result || [];
+}
+
+// Higher-order function to trace synchronous or asynchronous functions
+function trace(fn: Function, stepType: StepType = StepType.USER_CALL, stepName?: string): Function {
+  return async function (...args: any[]) {
+    const name = stepName || fn.name;
+    const paramNames = getParamNames(fn);
+    const inputs = Object.fromEntries(paramNames.map((name, index) => [name, args[index]]));
+    const [step, endStep] = createStep(name, stepType, args);
+
+    try {
+      const result = await fn(...args);
+      step.log({ inputs, output: result });
+      return result;
+    } catch (error: any) {
+      step.log({ inputs, metadata: { error: error.message } });
+      throw error;
+    } finally {
+      endStep();
+    }
+  };
+}
+
+export function addChatCompletionStepToTrace({
+  name,
+  inputs,
+  output,
+  latency,
+  tokens = null,
+  promptTokens = null,
+  completionTokens = null,
+  model = null,
+  modelParameters = null,
+  rawOutput = null,
+  metadata = {},
+  provider = 'OpenAI',
+}: {
+  name: string;
+  inputs: any;
+  output: any;
+  latency: number;
+  tokens?: number | null;
+  promptTokens?: number | null;
+  completionTokens?: number | null;
+  model?: string | null;
+  modelParameters?: Record<string, any> | null;
+  rawOutput?: string | null;
+  metadata?: Record<string, any>;
+  provider?: string;
+}) {
+  const [step, endStep] = createStep(name, StepType.CHAT_COMPLETION, inputs, output, metadata);
+
+  if (step instanceof ChatCompletionStep) {
+    step.provider = provider;
+    step.promptTokens = promptTokens;
+    step.completionTokens = completionTokens;
+    step.tokens = tokens;
+    step.model = model;
+    step.modelParameters = modelParameters;
+    step.rawOutput = rawOutput;
+    step.latency = latency;
+  }
+
+  step.log({ inputs, output, metadata });
+  endStep();
+}
+
+function postProcessTrace(traceObj: Trace): { traceData: any; inputVariableNames: string[] } {
+  const rootStep = traceObj.steps[0];
+
+  const input_variables = rootStep!.inputs;
+  const inputVariableNames = input_variables ? Object.keys(input_variables) : [];
+
+  const processed_steps = traceObj.toJSON();
+
+  const traceData = {
+    inferenceTimestamp: Date.now(),
+    inferenceId: rootStep!.id.toString(),
+    output: rootStep!.output,
+    groundTruth: rootStep!.groundTruth,
+    latency: rootStep!.latency,
+    cost: 0, // fix
+    tokens: 0, // fix
+    steps: processed_steps,
+  };
+
+  if (input_variables) {
+    Object.assign(traceData, input_variables);
+  }
+
+  return { traceData, inputVariableNames };
+}
+
+export default trace;

package/src/lib/tracing/traces.ts

@@ -0,0 +1,19 @@
+import { Step } from './steps';
+
+export class Trace {
+  public steps: Step[];
+  private currentStep: Step | null;
+
+  constructor() {
+    this.steps = [];
+    this.currentStep = null;
+  }
+
+  public addStep(step: Step): void {
+    this.steps.push(step);
+  }
+
+  public toJSON(): Array<Record<string, any>> {
+    return this.steps.map((step) => step.toJSON());
+  }
+}

package/src/resources/index.ts

@@ -1,7 +1,12 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 export { Commits } from './commits/commits';
-export { InferencePipelines } from './inference-pipelines/inference-pipelines';
+export {
+  InferencePipelineRetrieveResponse,
+  InferencePipelineUpdateResponse,
+  InferencePipelineUpdateParams,
+  InferencePipelines,
+} from './inference-pipelines/inference-pipelines';
 export {
   ProjectCreateResponse,
   ProjectListResponse,
@@ -9,3 +14,4 @@ export {
   ProjectListParams,
   Projects,
 } from './projects/projects';
+export { Storage } from './storage/storage';

package/src/resources/inference-pipelines/data.ts

@@ -6,7 +6,7 @@ import * as DataAPI from './data';
 
 export class Data extends APIResource {
   /**
-   * Create an inference data point in an inference pipeline.
+   * Publish an inference data point to an inference pipeline.
    */
   stream(
     inferencePipelineId: string,

package/src/resources/inference-pipelines/index.ts

@@ -1,6 +1,11 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 export { DataStreamResponse, DataStreamParams, Data } from './data';
-export { InferencePipelines } from './inference-pipelines';
+export {
+  InferencePipelineRetrieveResponse,
+  InferencePipelineUpdateResponse,
+  InferencePipelineUpdateParams,
+  InferencePipelines,
+} from './inference-pipelines';
 export { RowUpdateResponse, RowUpdateParams, Rows } from './rows';
 export { TestResultListResponse, TestResultListParams, TestResults } from './test-results';