Skip to content

Commit fb4c0a7

Browse files
ahmadbilaldevsaqibameen
ahmadbilaldev
and
saqibameen
authored
👌 IMPROVE: Tracing (#110)
* 📦 NEW: add tracing support * 📦 NEW: Extract traceid from response headers * 📦 NEW: Error and duration * 📦 NEW: Return headers from all endpoints of all primitives * 📦 NEW: traces.create fn * 📦 NEW: Use traces.creaate fn in workflow end * 🐛 FIX: Trace collector according to types * 🐛 FIX: Trace id creation * 📦 NEW: Connect trace with agent id from deployment * 🐛 FIX: Global trace collector for serverless * 📦 NEW: Make workflow name optional * 📦 NEW: Backward compat for tracing * 👌 IMPROVE: Make workflow name optional * 👌 IMPROVE: example * 👌 IMPROVE: Error handle envs * 🐛 FIX: process --------- Co-authored-by: Saqib Ameen <saqib1@ualberta.ca>
1 parent fd04351 commit fb4c0a7

5 files changed

Lines changed: 572 additions & 72 deletions

File tree

‎examples/nodejs/workflows/workflows.ts‎

Lines changed: 78 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,90 @@
1-
// Experimental upcoming beta AI primitve.
2-
// Please refer to the documentation for more information: https://langbase.com/docs for more information.
3-
41
import 'dotenv/config';
52
import {Langbase, Workflow} from 'langbase';
63

4+
// Create Langbase instance
75
const langbase = new Langbase({
86
apiKey: process.env.LANGBASE_API_KEY!,
97
});
108

119
async function main() {
12-
const {step} = new Workflow({debug: true});
13-
14-
const result = await step({
15-
id: 'sumamrize',
16-
run: async () => {
17-
return langbase.llm.run({
18-
model: 'openai:gpt-4o-mini',
19-
apiKey: process.env.OPENAI_API_KEY!,
20-
messages: [
21-
{
22-
role: 'system',
23-
content:
24-
'You are an expert summarizer. Summarize the user input.',
25-
},
26-
{
27-
role: 'user',
28-
content:
29-
'I am testing workflows. I just created an example of summarize workflow. Can you summarize this?',
30-
},
31-
],
32-
stream: false,
33-
});
34-
},
10+
// Create a workflow with debug mode enabled
11+
const workflow = langbase.workflow({
12+
name: 'Test Agent Workflow',
13+
debug: true,
3514
});
15+
// OR
16+
// const workflow = new Workflow({
17+
// name: 'Test Agent Workflow',
18+
// debug: true,
19+
// langbase,
20+
// });
21+
22+
try {
23+
// STEP 1: Call langbase.agent.run but don't return its result directly
24+
const step1Result = await workflow.step({
25+
id: 'call-but-return-custom',
26+
run: async () => {
27+
// Return custom result instead
28+
return {
29+
customField: 'Custom result from simplified proxy',
30+
timestamp: new Date().toISOString(),
31+
};
32+
},
33+
});
34+
35+
// STEP 2: Return agent.run result directly
36+
const step2Result = await workflow.step({
37+
id: 'return-agent-run-directly',
38+
run: async () => {
39+
// Call Langbase API and return the result directly
40+
return langbase.agent.run({
41+
model: 'openai:gpt-4o-mini',
42+
apiKey: process.env.OPENAI_API_KEY!,
43+
instructions: 'Be brief and concise.',
44+
input: 'What is 2+2?',
45+
stream: false,
46+
});
47+
},
48+
});
49+
50+
// STEP 3: Make multiple Langbase calls in one step
51+
const step3Result = await workflow.step({
52+
id: 'multiple-calls',
53+
run: async () => {
54+
// First call
55+
const call1 = await langbase.agent.run({
56+
model: 'openai:gpt-4o-mini',
57+
apiKey: process.env.OPENAI_API_KEY!,
58+
instructions: 'Be brief.',
59+
input: 'First proxy test',
60+
stream: false,
61+
});
62+
63+
// Second call with different method
64+
const call2 = await langbase.agent.run({
65+
model: 'openai:gpt-4o-mini',
66+
apiKey: process.env.OPENAI_API_KEY!,
67+
instructions: 'Be brief.',
68+
input: 'Second proxy test',
69+
stream: false,
70+
});
3671

37-
console.log(result['completion']);
72+
// Return combined result
73+
return {
74+
summary: 'Multiple calls completed with simplified proxy',
75+
calls: 2,
76+
firstOutput: call1.output,
77+
secondOutput: call2.output,
78+
};
79+
},
80+
});
81+
} catch (error) {
82+
console.error('❌ Workflow error:', error);
83+
} finally {
84+
// End the workflow to send trace
85+
workflow.end();
86+
}
3887
}
3988

40-
main();
89+
// Run the test
90+
main().catch(console.error);

‎packages/langbase/src/common/request.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,17 @@ export class Request {
6262
const isLllmGenerationEndpoint =
6363
GENERATION_ENDPOINTS.includes(endpoint);
6464

65+
// All endpoints should return headers if rawResponse is true
66+
if (!isLllmGenerationEndpoint && options.body?.rawResponse) {
67+
const responseData = await response.json();
68+
return {
69+
...responseData,
70+
rawResponse: {
71+
headers: Object.fromEntries(response.headers.entries()),
72+
},
73+
} as T;
74+
}
75+
6576
if (isLllmGenerationEndpoint) {
6677
const threadId = response.headers.get('lb-thread-id');
6778

‎packages/langbase/src/langbase/langbase.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import {convertDocToFormData} from '@/lib/utils/doc-to-formdata';
22
import {Request} from '../common/request';
3+
import {Workflow} from './workflows';
34

45
export type Role = 'user' | 'assistant' | 'system' | 'tool';
56

@@ -639,6 +640,12 @@ export class Langbase {
639640
};
640641
};
641642

643+
public workflow: (config: {debug?: boolean; name?: string}) => Workflow;
644+
645+
public traces: {
646+
create: (trace: any) => Promise<any>;
647+
};
648+
642649
constructor(options?: LangbaseOptions) {
643650
this.baseUrl = options?.baseUrl ?? 'https://api.langbase.com';
644651
this.apiKey = options?.apiKey ?? '';
@@ -723,6 +730,12 @@ export class Langbase {
723730
this.agent = {
724731
run: this.runAgent.bind(this),
725732
};
733+
734+
this.workflow = config => new Workflow({...config, langbase: this});
735+
736+
this.traces = {
737+
create: this.createTrace.bind(this),
738+
};
726739
}
727740

728741
private async runPipe(
@@ -1146,4 +1159,17 @@ export class Langbase {
11461159
},
11471160
});
11481161
}
1162+
1163+
/**
1164+
* Creates a new trace on Langbase.
1165+
*
1166+
* @param {any} trace - The trace data to send.
1167+
* @returns {Promise<any>} A promise that resolves to the response of the trace creation.
1168+
*/
1169+
private async createTrace(trace: any): Promise<any> {
1170+
return this.request.post({
1171+
endpoint: '/v1/traces',
1172+
body: trace,
1173+
});
1174+
}
11491175
}

‎packages/langbase/src/langbase/trace.ts‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
export interface Trace {
2+
name: string;
3+
startTime: number;
4+
endTime?: number;
5+
duration?: number;
6+
steps: StepTrace[];
7+
error?: string;
8+
}
9+
10+
export interface StepTrace {
11+
name: string;
12+
output: any;
13+
error?: string;
14+
traces: string[] | null;
15+
duration: number;
16+
startTime: number;
17+
endTime: number;
18+
}
19+
20+
export type TraceType =
21+
| 'workflow'
22+
| 'agent'
23+
| 'chunk'
24+
| 'memory'
25+
| 'parse'
26+
| 'embed';
27+
28+
export type PrimitiveTrace =
29+
| {chunk: any}
30+
| {agent: any}
31+
| {memory: any}
32+
| {parse: any}
33+
| {embed: any}
34+
| {workflow: WorkflowTrace; entityAuthId: string};
35+
36+
type WorkflowTrace = {
37+
createdAt: string;
38+
id: string;
39+
agentWorkflowId: string;
40+
name: string;
41+
startTime: number;
42+
endTime?: number;
43+
duration?: number;
44+
steps: StepTrace[];
45+
error?: string;
46+
};
47+
48+
export class TraceManager {
49+
private traces: Map<string, PrimitiveTrace> = new Map();
50+
51+
createTrace(type: TraceType, traceData: any = {}): string {
52+
const traceId = crypto.randomUUID();
53+
let trace: PrimitiveTrace;
54+
const createdAt = new Date().toISOString();
55+
const agentWorkflowId =
56+
typeof process !== 'undefined' && process.env?.LANGBASE_AGENT_ID
57+
? process.env.LANGBASE_AGENT_ID
58+
: '';
59+
60+
if (type === 'workflow') {
61+
trace = {
62+
workflow: {
63+
createdAt,
64+
id: traceId,
65+
agentWorkflowId,
66+
name: traceData.name || '',
67+
startTime: Date.now(),
68+
steps: [],
69+
},
70+
entityAuthId: '',
71+
};
72+
} else if (type === 'agent') {
73+
trace = {agent: {...traceData, createdAt, id: traceId}};
74+
} else if (type === 'chunk') {
75+
trace = {chunk: {...traceData, createdAt, id: traceId}};
76+
} else if (type === 'memory') {
77+
trace = {memory: {...traceData, createdAt, id: traceId}};
78+
} else if (type === 'parse') {
79+
trace = {parse: {...traceData, createdAt, id: traceId}};
80+
} else if (type === 'embed') {
81+
trace = {embed: {...traceData, createdAt, id: traceId}};
82+
} else {
83+
throw new Error('Unknown trace type');
84+
}
85+
this.traces.set(traceId, trace);
86+
return traceId;
87+
}
88+
89+
addStep(traceId: string, step: StepTrace) {
90+
const trace = this.traces.get(traceId);
91+
if (trace && 'workflow' in trace) {
92+
trace.workflow.steps.push(step);
93+
}
94+
}
95+
96+
endTrace(traceId: string) {
97+
const trace = this.traces.get(traceId);
98+
if (trace && 'workflow' in trace) {
99+
trace.workflow.endTime = Date.now();
100+
trace.workflow.duration =
101+
trace.workflow.endTime - trace.workflow.startTime;
102+
}
103+
}
104+
105+
getTrace(traceId: string): PrimitiveTrace | undefined {
106+
return this.traces.get(traceId);
107+
}
108+
109+
printTrace(traceId: string) {
110+
const trace = this.traces.get(traceId);
111+
if (!trace) return;
112+
if ('workflow' in trace) {
113+
const wf = trace.workflow;
114+
const duration = wf.endTime
115+
? wf.endTime - wf.startTime
116+
: Date.now() - wf.startTime;
117+
console.log('\n📊 Workflow Trace:');
118+
console.log(`Name: ${wf.name}`);
119+
console.log(`Duration: ${duration}ms`);
120+
console.log(`Start Time: ${new Date(wf.startTime).toISOString()}`);
121+
if (wf.endTime) {
122+
console.log(`End Time: ${new Date(wf.endTime).toISOString()}`);
123+
}
124+
console.log('\nSteps:');
125+
wf.steps.forEach(step => {
126+
console.log(`\n Step: ${step.name}`);
127+
console.log(` Duration: ${step.duration}ms`);
128+
if (step.traces && step.traces.length > 0) {
129+
console.log(` Traces:`, step.traces);
130+
}
131+
console.log(` Output:`, step.output);
132+
});
133+
} else {
134+
console.log('\n📊 Primitive Trace:');
135+
console.dir(trace, {depth: 4});
136+
}
137+
}
138+
}

0 commit comments

Comments
 (0)