Skip to content

Commit c803ebc

Browse files
ahmadbilaldevsaqibameen
ahmadbilaldev
and
saqibameen
authored
📦 NEW: Tracing support (#108)
* 📦 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 --------- Co-authored-by: Saqib Ameen <[email protected]>
1 parent 691d31c commit c803ebc

File tree

5 files changed

+553
-73
lines changed

5 files changed

+553
-73
lines changed

‎examples/nodejs/workflows/workflows.ts‎

Lines changed: 74 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,85 @@
1-
// Experimental upcoming beta AI primitve.
2-
// Please refer to the documentation for more information: https://langbase.com/docs for more information.
3-
1+
// Test script for the simplified proxy approach
42
import 'dotenv/config';
5-
import {Langbase, Workflow} from 'langbase';
3+
import {Langbase} from 'langbase';
64

5+
// Create Langbase instance
76
const langbase = new Langbase({
87
apiKey: process.env.LANGBASE_API_KEY!,
98
});
109

1110
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-
},
11+
// Create a workflow with debug mode enabled
12+
const workflow = langbase.workflow({
13+
name: 'simplified-proxy-test',
14+
debug: true, // Enable debug logging
3515
});
3616

37-
console.log(result['completion']);
17+
try {
18+
// STEP 1: Call langbase.agent.run but don't return its result directly
19+
const step1Result = await workflow.step({
20+
id: 'call-but-return-custom',
21+
run: async () => {
22+
// Return custom result instead
23+
return {
24+
customField: 'Custom result from simplified proxy',
25+
timestamp: new Date().toISOString(),
26+
};
27+
},
28+
});
29+
30+
// STEP 2: Return agent.run result directly
31+
const step2Result = await workflow.step({
32+
id: 'return-agent-run-directly',
33+
run: async () => {
34+
// Call Langbase API and return the result directly
35+
return langbase.agent.run({
36+
model: 'openai:gpt-4o-mini',
37+
apiKey: process.env.OPENAI_API_KEY!,
38+
instructions: 'Be brief and concise.',
39+
input: 'What is 2+2?',
40+
stream: false,
41+
});
42+
},
43+
});
44+
45+
// STEP 3: Make multiple Langbase calls in one step
46+
const step3Result = await workflow.step({
47+
id: 'multiple-calls',
48+
run: async () => {
49+
// First call
50+
const call1 = await langbase.agent.run({
51+
model: 'openai:gpt-4o-mini',
52+
apiKey: process.env.OPENAI_API_KEY!,
53+
instructions: 'Be brief.',
54+
input: 'First proxy test',
55+
stream: false,
56+
});
57+
58+
// Second call with different method
59+
const call2 = await langbase.agent.run({
60+
model: 'openai:gpt-4o-mini',
61+
apiKey: process.env.OPENAI_API_KEY!,
62+
instructions: 'Be brief.',
63+
input: 'Second proxy test',
64+
stream: false,
65+
});
66+
67+
// Return combined result
68+
return {
69+
summary: 'Multiple calls completed with simplified proxy',
70+
calls: 2,
71+
firstOutput: call1.output,
72+
secondOutput: call2.output,
73+
};
74+
},
75+
});
76+
} catch (error) {
77+
console.error('❌ Workflow error:', error);
78+
} finally {
79+
// End the workflow to show trace report
80+
workflow.end();
81+
}
3882
}
3983

40-
main();
84+
// Run the test
85+
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: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
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+
if (type === 'workflow') {
56+
trace = {
57+
workflow: {
58+
createdAt,
59+
id: traceId,
60+
agentWorkflowId: process.env.LANGBASE_AGENT_ID || '',
61+
name: traceData.name || '',
62+
startTime: Date.now(),
63+
steps: [],
64+
},
65+
entityAuthId: '',
66+
};
67+
} else if (type === 'agent') {
68+
trace = {agent: {...traceData, createdAt, id: traceId}};
69+
} else if (type === 'chunk') {
70+
trace = {chunk: {...traceData, createdAt, id: traceId}};
71+
} else if (type === 'memory') {
72+
trace = {memory: {...traceData, createdAt, id: traceId}};
73+
} else if (type === 'parse') {
74+
trace = {parse: {...traceData, createdAt, id: traceId}};
75+
} else if (type === 'embed') {
76+
trace = {embed: {...traceData, createdAt, id: traceId}};
77+
} else {
78+
throw new Error('Unknown trace type');
79+
}
80+
this.traces.set(traceId, trace);
81+
return traceId;
82+
}
83+
84+
addStep(traceId: string, step: StepTrace) {
85+
const trace = this.traces.get(traceId);
86+
if (trace && 'workflow' in trace) {
87+
trace.workflow.steps.push(step);
88+
}
89+
}
90+
91+
endTrace(traceId: string) {
92+
const trace = this.traces.get(traceId);
93+
if (trace && 'workflow' in trace) {
94+
trace.workflow.endTime = Date.now();
95+
trace.workflow.duration =
96+
trace.workflow.endTime - trace.workflow.startTime;
97+
}
98+
}
99+
100+
getTrace(traceId: string): PrimitiveTrace | undefined {
101+
return this.traces.get(traceId);
102+
}
103+
104+
printTrace(traceId: string) {
105+
const trace = this.traces.get(traceId);
106+
if (!trace) return;
107+
if ('workflow' in trace) {
108+
const wf = trace.workflow;
109+
const duration = wf.endTime
110+
? wf.endTime - wf.startTime
111+
: Date.now() - wf.startTime;
112+
console.log('\n📊 Workflow Trace:');
113+
console.log(`Name: ${wf.name}`);
114+
console.log(`Duration: ${duration}ms`);
115+
console.log(`Start Time: ${new Date(wf.startTime).toISOString()}`);
116+
if (wf.endTime) {
117+
console.log(`End Time: ${new Date(wf.endTime).toISOString()}`);
118+
}
119+
console.log('\nSteps:');
120+
wf.steps.forEach(step => {
121+
console.log(`\n Step: ${step.name}`);
122+
console.log(` Duration: ${step.duration}ms`);
123+
if (step.traces && step.traces.length > 0) {
124+
console.log(` Traces:`, step.traces);
125+
}
126+
console.log(` Output:`, step.output);
127+
});
128+
} else {
129+
console.log('\n📊 Primitive Trace:');
130+
console.dir(trace, {depth: 4});
131+
}
132+
}
133+
}

0 commit comments

Comments
 (0)