Skip to content

Commit 84480cf

Browse files
gustavocidornelasgustavocidornelas
committed
feat: initial working version of the tracing module
1 parent 0e5113d commit 84480cf

File tree

8 files changed

+319
-5
lines changed

8 files changed

+319
-5
lines changed

‎src/lib/.keep‎

Lines changed: 0 additions & 4 deletions
This file was deleted.

‎src/lib/core/openai-monitor.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ class OpenAIMonitor {
147147
content:
148148
role === 'user' ? `{{ message_${i} }}`
149149
: content === null || typeof content === 'undefined' ? ''
150-
: content,
150+
: String(content),
151151
role,
152152
}));
153153

‎src/lib/index.d.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
export * as core from './core/index';
2+
export * as tracing from './tracing/index';

‎src/lib/tracing/enums.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
export enum StepType {
2+
USER_CALL = 'user_call',
3+
CHAT_COMPLETION = 'chat_completion',
4+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export * as tracer from './tracer';

‎src/lib/tracing/steps.ts‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import { v4 as uuidv4 } from 'uuid';
2+
import { StepType } from './enums';
3+
4+
export interface StepData {
5+
name: string;
6+
id: string;
7+
type: StepType;
8+
inputs: any;
9+
output: any;
10+
groundTruth: any;
11+
metadata: Record<string, any>;
12+
steps: StepData[];
13+
latency: number | null;
14+
startTime: number;
15+
endTime: number | null;
16+
}
17+
18+
export interface ChatCompletionStepData extends StepData {
19+
provider: string | null;
20+
promptTokens: number | null;
21+
completionTokens: number | null;
22+
tokens: number | null;
23+
cost: number | null;
24+
model: string | null;
25+
modelParameters: Record<string, any> | null;
26+
rawOutput: string | null;
27+
}
28+
29+
export class Step {
30+
name: string;
31+
id: string;
32+
inputs: any;
33+
output: any;
34+
metadata: Record<string, any>;
35+
stepType: StepType | null = null;
36+
startTime: number;
37+
endTime: number | null = null;
38+
groundTruth: any = null;
39+
latency: number | null = null;
40+
steps: Step[] = [];
41+
42+
constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
43+
this.name = name;
44+
this.id = uuidv4();
45+
this.inputs = inputs;
46+
this.output = output;
47+
this.metadata = metadata;
48+
49+
this.startTime = Date.now();
50+
}
51+
52+
addNestedStep(nestedStep: Step): void {
53+
this.steps.push(nestedStep);
54+
}
55+
56+
log(data: Partial<Record<keyof this, any>>): void {
57+
Object.keys(data).forEach((key) => {
58+
if (key in this) {
59+
// @ts-ignore
60+
this[key] = data[key];
61+
}
62+
});
63+
}
64+
65+
toJSON(): StepData {
66+
return {
67+
name: this.name,
68+
id: this.id,
69+
type: this.stepType!,
70+
inputs: this.inputs,
71+
output: this.output,
72+
groundTruth: this.groundTruth,
73+
metadata: this.metadata,
74+
steps: this.steps.map((nestedStep) => nestedStep.toJSON()),
75+
latency: this.latency,
76+
startTime: this.startTime,
77+
endTime: this.endTime,
78+
};
79+
}
80+
}
81+
82+
export class UserCallStep extends Step {
83+
constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
84+
super(name, inputs, output, metadata);
85+
this.stepType = StepType.USER_CALL;
86+
}
87+
}
88+
89+
export class ChatCompletionStep extends Step {
90+
provider: string | null = null;
91+
promptTokens: number | null = null;
92+
completionTokens: number | null = null;
93+
tokens: number | null = null;
94+
cost: number | null = null;
95+
model: string | null = null;
96+
modelParameters: Record<string, any> | null = null;
97+
rawOutput: string | null = null;
98+
99+
constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
100+
super(name, inputs, output, metadata);
101+
this.stepType = StepType.CHAT_COMPLETION;
102+
}
103+
104+
override toJSON(): ChatCompletionStepData {
105+
const stepData = super.toJSON();
106+
return {
107+
...stepData,
108+
provider: this.provider,
109+
promptTokens: this.promptTokens,
110+
completionTokens: this.completionTokens,
111+
tokens: this.tokens,
112+
cost: this.cost,
113+
model: this.model,
114+
modelParameters: this.modelParameters,
115+
rawOutput: this.rawOutput,
116+
};
117+
}
118+
}

‎src/lib/tracing/tracer.ts‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
// tracing/tracer.ts
2+
3+
import { Trace } from './traces';
4+
import { Step, UserCallStep, ChatCompletionStep } from './steps';
5+
import { StepType } from './enums';
6+
import Openlayer from '../../index';
7+
8+
let currentTrace: Trace | null = null;
9+
10+
const publish = process.env['OPENLAYER_DISABLE_PUBLISH'] != 'true';
11+
let client: Openlayer | null = null;
12+
if (publish) {
13+
console.log('Publishing is enabled');
14+
client = new Openlayer();
15+
}
16+
17+
function getCurrentTrace(): Trace | null {
18+
return currentTrace;
19+
}
20+
21+
function setCurrentTrace(trace: Trace | null) {
22+
currentTrace = trace;
23+
}
24+
25+
// Function to create a new step
26+
const stepStack: Step[] = [];
27+
28+
function createStep(
29+
name: string,
30+
stepType: StepType = StepType.USER_CALL,
31+
inputs?: any,
32+
output?: any,
33+
metadata: Record<string, any> | null = null,
34+
): [Step, () => void] {
35+
metadata = metadata || {};
36+
let newStep: Step;
37+
if (stepType === StepType.CHAT_COMPLETION) {
38+
newStep = new ChatCompletionStep(name, inputs, output, metadata);
39+
} else {
40+
newStep = new UserCallStep(name, inputs, output, metadata);
41+
}
42+
newStep.startTime = performance.now();
43+
44+
const parentStep = getCurrentStep();
45+
const isRootStep = parentStep === null;
46+
47+
if (isRootStep) {
48+
console.log('Starting a new trace...');
49+
console.log(`Adding step ${name} as the root step`);
50+
const currentTrace = new Trace();
51+
setCurrentTrace(currentTrace);
52+
currentTrace.addStep(newStep);
53+
} else {
54+
console.log(`Adding step ${name} as a nested step to ${parentStep!.name}`);
55+
currentTrace = getCurrentTrace()!;
56+
parentStep!.addNestedStep(newStep);
57+
}
58+
59+
stepStack.push(newStep);
60+
61+
const endStep = () => {
62+
newStep.endTime = performance.now();
63+
newStep.latency = newStep.endTime - newStep.startTime;
64+
65+
stepStack.pop(); // Remove the current step from the stack
66+
67+
if (isRootStep) {
68+
console.log('Ending the trace...');
69+
const traceData = getCurrentTrace();
70+
// Post process trace and get the input variable names
71+
const { traceData: processedTraceData, inputVariableNames } = postProcessTrace(traceData!);
72+
console.log('Processed trace data:', JSON.stringify(processedTraceData, null, 2));
73+
console.log('Input variable names:', inputVariableNames);
74+
75+
if (publish && process.env['OPENLAYER_INFERENCE_PIPELINE_ID']) {
76+
client!.inferencePipelines.data.stream(process.env['OPENLAYER_INFERENCE_PIPELINE_ID'], {
77+
config: {
78+
outputColumnName: 'output',
79+
inputVariableNames: inputVariableNames,
80+
groundTruthColumnName: 'groundTruth',
81+
latencyColumnName: 'latency',
82+
costColumnName: 'cost',
83+
timestampColumnName: 'inferenceTimestamp',
84+
inferenceIdColumnName: 'inferenceId',
85+
numOfTokenColumnName: 'tokens',
86+
},
87+
rows: [processedTraceData],
88+
});
89+
}
90+
console.log('Trace data ready for upload:', JSON.stringify(traceData, null, 2));
91+
92+
// Reset the entire trace state
93+
setCurrentTrace(null);
94+
stepStack.length = 0; // Clear the step stack
95+
} else {
96+
console.log(`Ending step ${name}`);
97+
}
98+
};
99+
100+
return [newStep, endStep];
101+
}
102+
103+
function getCurrentStep(): Step | null | undefined {
104+
const currentStep = stepStack.length > 0 ? stepStack[stepStack.length - 1] : null;
105+
return currentStep;
106+
}
107+
108+
function getParamNames(func: Function): string[] {
109+
const STRIP_COMMENTS = /((\/\/.*$)|(\/\*[\s\S]*?\*\/))/gm;
110+
const ARGUMENT_NAMES = /([^\s,]+)/g;
111+
const fnStr = func.toString().replace(STRIP_COMMENTS, '');
112+
const result = fnStr.slice(fnStr.indexOf('(') + 1, fnStr.indexOf(')')).match(ARGUMENT_NAMES);
113+
return result || [];
114+
}
115+
116+
// Higher-order function to trace synchronous or asynchronous functions
117+
function trace(fn: Function, stepType: StepType = StepType.USER_CALL, stepName?: string): Function {
118+
return async function (...args: any[]) {
119+
const name = stepName || fn.name;
120+
const paramNames = getParamNames(fn);
121+
const inputs = Object.fromEntries(paramNames.map((name, index) => [name, args[index]]));
122+
const [step, endStep] = createStep(name, stepType, args);
123+
124+
try {
125+
const result = await fn(...args);
126+
step.log({ inputs, output: result });
127+
return result;
128+
} catch (error: any) {
129+
step.log({ inputs, metadata: { error: error.message } });
130+
throw error;
131+
} finally {
132+
endStep();
133+
}
134+
};
135+
}
136+
137+
// Example usage of specialized function to add a chat completion step
138+
export function addChatCompletionStepToTrace(
139+
name: string,
140+
inputs: any,
141+
output: any,
142+
metadata?: Record<string, any>,
143+
) {
144+
const [step, endStep] = createStep(name, StepType.CHAT_COMPLETION, inputs, output, metadata);
145+
step.log({ inputs, output });
146+
endStep();
147+
}
148+
149+
function postProcessTrace(traceObj: Trace): { traceData: any; inputVariableNames: string[] } {
150+
const rootStep = traceObj.steps[0];
151+
152+
const input_variables = rootStep!.inputs;
153+
const inputVariableNames = input_variables ? Object.keys(input_variables) : [];
154+
155+
const processed_steps = traceObj.toJSON();
156+
157+
const traceData = {
158+
inferenceTimestamp: Date.now(),
159+
inferenceId: rootStep!.id.toString(),
160+
output: rootStep!.output,
161+
groundTruth: rootStep!.groundTruth,
162+
latency: rootStep!.latency,
163+
cost: 0, // fix
164+
tokens: 0, // fix
165+
steps: processed_steps,
166+
};
167+
168+
if (input_variables) {
169+
Object.assign(traceData, input_variables);
170+
}
171+
172+
return { traceData, inputVariableNames };
173+
}
174+
175+
export default trace;

‎src/lib/tracing/traces.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import { Step } from './steps';
2+
3+
export class Trace {
4+
public steps: Step[];
5+
private currentStep: Step | null;
6+
7+
constructor() {
8+
this.steps = [];
9+
this.currentStep = null;
10+
}
11+
12+
public addStep(step: Step): void {
13+
this.steps.push(step);
14+
}
15+
16+
public toJSON(): Array<Record<string, any>> {
17+
return this.steps.map((step) => step.toJSON());
18+
}
19+
}

0 commit comments

Comments
 (0)