Feat/tracer update (#3)

NoahKex · web-flow · commit 541f53db1c90 · 2025-04-08T09:43:15.000+08:00
* chore(@cozeloop/ai): adjust code details

* feat(@cozeloop/ai): support reporting the baggage field

* chore(@cozeloop/ai): trace demo
diff --git a/examples/cozeloop-ai-node/src/api/api-client.ts b/examples/cozeloop-ai-node/src/api/api-client.ts
@@ -23,7 +23,6 @@ export async function run() {
     // axiosOptions: {},
     /** Custom headers */
     // headers: {};
-    headers: { 'x-tt-env': 'boe_commercial' }, // TODO: remove
   });
 
   const resp = await apiClient.post<{ code: number }>('/v1/loop/prompts/mget', {
diff --git a/examples/cozeloop-ai-node/src/prompt/hub.ts b/examples/cozeloop-ai-node/src/prompt/hub.ts
@@ -9,7 +9,6 @@ export async function run() {
     apiClient: {
       // baseURL: 'api_base_url',
       // token: 'your_api_token',
-      headers: { 'x-tt-env': 'boe_commercial' }, // TODO: remove
     },
   });
 
diff --git a/examples/cozeloop-ai-node/src/tracer/README.md b/examples/cozeloop-ai-node/src/tracer/README.md
@@ -17,7 +17,6 @@ loopTracer.initialize({
   apiClient: {
     // baseURL: 'https://api.coze.cn',
     // token: 'your_api_token',
-    headers: { 'x-tt-env': 'boe_commercial' }, // TODO: remove
   },
   /** Allow ultra long text report */
   // ultraLargeReport: true,
diff --git a/examples/cozeloop-ai-node/src/tracer/basic.ts b/examples/cozeloop-ai-node/src/tracer/basic.ts
@@ -0,0 +1,40 @@
+import { cozeLoopTracer } from '@cozeloop/ai';
+
+import { doSomething } from './utils';
+
+export async function runCustom() {
+  // Wrap any function to make it traceable
+  await cozeLoopTracer.traceable(
+    async parentSpan => {
+      // Manually set input
+      cozeLoopTracer.setInput(parentSpan, 'xxx');
+
+      // Invoke any function, if it throws error, error will be caught and automatically set span as error
+      const result = await doSomething();
+
+      // Or, you can manually set error
+      cozeLoopTracer.setError(parentSpan, 'custom error message');
+
+      // You can also trace nested span, the parent-child relationship of span will be automatically concatenated
+      await cozeLoopTracer.traceable(
+        async childSpan => {
+          // Set custom tags
+          childSpan.setAttribute('custom-tag', 'xxx');
+
+          await doSomething();
+        },
+        {
+          name: 'TestCustomChildSpan',
+          type: 'MyCustomType',
+        },
+      );
+
+      // Automatically set return value as output
+      return result;
+    },
+    {
+      name: 'TestCustomParentSpan',
+      type: 'MyCustomType',
+    },
+  );
+}
diff --git a/examples/cozeloop-ai-node/src/tracer/index.ts b/examples/cozeloop-ai-node/src/tracer/index.ts
@@ -1,23 +1,24 @@
 import { cozeLoopTracer } from '@cozeloop/ai';
 
-import { runModel, runRoot, runCustom } from './simple';
+import { runRoot } from './root';
 import { runMultiModality } from './multi-modality';
+import { runModel } from './llm';
 import { runLargeText } from './large-text';
-
-// initialize tracer globally
-cozeLoopTracer.initialize({
-  /** workspace id, use process.env.COZELOOP_WORKSPACE_ID when unprovided */
-  // workspaceId: 'your_workspace_id',
-  apiClient: {
-    // baseURL: 'https://api.coze.cn',
-    // token: 'your_api_token',
-    headers: { 'x-tt-env': 'boe_commercial' }, // TODO: remove
-  },
-  /** Allow ultra long text report */
-  ultraLargeReport: true,
-});
+import { runCustom } from './basic';
 
 export async function run() {
+  // initialize tracer globally
+  cozeLoopTracer.initialize({
+    /** workspace id, use process.env.COZELOOP_WORKSPACE_ID when unprovided */
+    // workspaceId: 'your_workspace_id',
+    apiClient: {
+      // baseURL: 'https://api.coze.cn',
+      // token: 'your_api_token',
+    },
+    /** Allow ultra long text report */
+    ultraLargeReport: true,
+  });
+
   await Promise.all([
     runRoot(),
     runCustom(),
diff --git a/examples/cozeloop-ai-node/src/tracer/large-text.ts b/examples/cozeloop-ai-node/src/tracer/large-text.ts
@@ -21,7 +21,7 @@ export async function runLargeText() {
   await cozeLoopTracer.traceable(
     async span => {
       // Reporting of ultra-long texts will only take effect when the
-      // input satisfies the LoopTraceLLMCallInput structure
+      // input / output satisfies the LoopTraceLLMCallInput / LoopTraceLLMCallOutput structure
       const input: LoopTraceLLMCallInput = {
         messages: [
           {
diff --git a/examples/cozeloop-ai-node/src/tracer/llm.ts b/examples/cozeloop-ai-node/src/tracer/llm.ts
@@ -1,85 +1,11 @@
-import { setTimeout } from 'node:timers/promises';
-
+import { SpanKind } from '@cozeloop/ai';
 import {
   COZELOOP_TRACE_BUSINESS_TAGS,
   cozeLoopTracer,
   type LoopTraceLLMCallInput,
-  type LoopTraceLLMCallOutput,
-  SpanKind,
 } from '@cozeloop/ai';
 
-async function doSomething() {
-  await setTimeout(2000, 'result');
-}
-
-async function fakeLLMCall(): Promise<LoopTraceLLMCallOutput> {
-  return await setTimeout(2000, {
-    choices: [
-      {
-        index: 0,
-        finish_reason: 'stop',
-        message: {
-          role: 'assistant',
-          content: "hi, I'm xx model",
-        },
-      },
-    ],
-  });
-}
-
-export async function runRoot() {
-  // We recommend concatenating a complete user request into a trace,
-  // so the recommended approach is to report a root span at the entrance of the entire execution
-  await cozeLoopTracer.traceable(
-    async () => {
-      // execute your method
-      const result = await doSomething();
-
-      return result;
-    },
-    {
-      name: 'TestRootSpan',
-      type: 'RootSpanType',
-    },
-  );
-}
-
-export async function runCustom() {
-  // Wrap any function to make it traceable
-  await cozeLoopTracer.traceable(
-    async parentSpan => {
-      // Manually set input
-      cozeLoopTracer.setInput(parentSpan, 'xxx');
-
-      // Invoke any function, if it throws error, error will be caught and automatically set span as error
-      const result = await doSomething();
-
-      // Or, you can manually set error
-      cozeLoopTracer.setError(parentSpan, 'custom error message');
-
-      // You can also trace nested span, the parent-child relationship of span will be automatically concatenated
-      await cozeLoopTracer.traceable(
-        async childSpan => {
-          // Set custom tags
-          childSpan.setAttribute('custom-tag', 'xxx');
-
-          await doSomething();
-        },
-        {
-          name: 'TestCustomChildSpan',
-          type: 'MyCustomType',
-        },
-      );
-
-      // Automatically set return value as output
-      return result;
-    },
-    {
-      name: 'TestCustomParentSpan',
-      type: 'MyCustomType',
-    },
-  );
-}
+import { fakeLLMCall } from './utils';
 
 export async function runModel() {
   // Wrap model invoke function to make it traceable
diff --git a/examples/cozeloop-ai-node/src/tracer/multi-modality.ts b/examples/cozeloop-ai-node/src/tracer/multi-modality.ts
@@ -14,7 +14,7 @@ export async function runMultiModality() {
   await cozeLoopTracer.traceable(
     async span => {
       // Reporting of multi modality will only take effect when the
-      // input satisfies the LoopTraceLLMCallInput structure
+      // input / output satisfies the LoopTraceLLMCallInput / LoopTraceLLMCallOutput structure
       const input: LoopTraceLLMCallInput = {
         messages: [
           {
diff --git a/examples/cozeloop-ai-node/src/tracer/root.ts b/examples/cozeloop-ai-node/src/tracer/root.ts
@@ -0,0 +1,25 @@
+import { cozeLoopTracer } from '@cozeloop/ai';
+
+import { doSomething } from './utils';
+
+export async function runRoot() {
+  // We recommend concatenating a complete user request into a trace,
+  // so the recommended approach is to report a root span at the entrance of the entire execution
+  await cozeLoopTracer.traceable(
+    async () => {
+      // execute your method
+      const result = await doSomething();
+
+      return result;
+    },
+    {
+      name: 'TestRootSpan',
+      type: 'RootSpanType',
+      // you can set your own baggage fields (eg. userId),
+      // these fields will be automatically passed through and set in all sub-spans
+      userId: 'uid-123',
+      messageId: 'msg-123',
+      threadId: 'thread-123',
+    },
+  );
+}
diff --git a/examples/cozeloop-ai-node/src/tracer/utils.ts b/examples/cozeloop-ai-node/src/tracer/utils.ts
@@ -0,0 +1,22 @@
+import { setTimeout } from 'node:timers/promises';
+
+import { type LoopTraceLLMCallOutput } from '@cozeloop/ai';
+
+export async function doSomething() {
+  await setTimeout(2000, 'result');
+}
+
+export async function fakeLLMCall(): Promise<LoopTraceLLMCallOutput> {
+  return await setTimeout(2000, {
+    choices: [
+      {
+        index: 0,
+        finish_reason: 'stop',
+        message: {
+          role: 'assistant',
+          content: "hi, I'm xx model",
+        },
+      },
+    ],
+  });
+}
diff --git a/packages/cozeloop-ai/__tests__/tracer/traceable.test.ts b/packages/cozeloop-ai/__tests__/tracer/traceable.test.ts
@@ -14,15 +14,8 @@ config();
 
 cozeLoopTracer.initialize({
   apiClient: {
-    baseURL: 'https://api.coze.cn',
-    token: process.env.COZELOOP_API_TOKEN || '',
-    headers: {
-      // TODO remove
-      'x-tt-env': 'boe_commercial',
-    },
     logger: simpleConsoleLogger,
   },
-  workspaceId: '7480474966061006892',
   ultraLargeReport: true,
 });
 
@@ -136,6 +129,7 @@ describe('Test Trace Report', { timeout: 50_000 }, () => {
       {
         name: 'TestParent',
         type: SpanKind.PromptTemplate,
+        userId: 'test-user-id',
         attributes: {
           string_tag: 'string',
           boolean_tag: true,
diff --git a/packages/cozeloop-ai/src/tracer/constants/index.ts b/packages/cozeloop-ai/src/tracer/constants/index.ts
@@ -1,9 +1,22 @@
 // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+
+import { createContextKey } from '@opentelemetry/api';
+
 // SPDX-License-Identifier: MIT
 export const ROOT_SPAN_PARENT_ID = '0';
 
 export const COZELOOP_LOGGER_TRACER_TAG = 'Tracer';
 
+export const COZELOOP_TRACE_SPAN_USER_ID_KEY = createContextKey(
+  'cozeloop.trace.span.user_id',
+);
+export const COZELOOP_TRACE_SPAN_MESSAGE_ID_KEY = createContextKey(
+  'cozeloop.trace.span.message_id',
+);
+export const COZELOOP_TRACE_SPAN_THREAD_ID_KEY = createContextKey(
+  'cozeloop.trace.span.thread_id',
+);
+
 export enum COZELOOP_TRACE_IDENTIFIER {
   LOOP = 'cozeloop-js',
 }
@@ -23,16 +36,18 @@ export enum COZELOOP_TRACE_BASIC_TAGS {
   SPAN_ULTRA_LARGE_REPORT = 'cozeloop_ultra_large_report',
 
   // Reserved tags
-  SPAN_USER_ID = 'cozeloop_user_id',
-  SPAN_MESSAGE_ID = 'cozeloop_message_id',
-  SPAN_THREAD_ID = 'cozeloop_thread_id',
   SPAN_LOG_ID = 'cozeloop_span_log_id',
   SPAN_PSM = 'cozeloop_psm',
   SPAN_METHOD = 'cozeloop_method',
   SPAN_CALL_TYPE = 'cozeloop_call_type',
 }
 
 export enum COZELOOP_TRACE_BUSINESS_TAGS {
+  // Common tags
+  USER_ID = 'user_id',
+  MESSAGE_ID = 'message_id',
+  THREAD_ID = 'thread_id',
+
   // Tags for Model span
   /** The timestamp of the model's first response when using stream response. The unit is microseconds. */
   START_TIME_FIRST_RESP = 'start_time_first_resp',
diff --git a/packages/cozeloop-ai/src/tracer/exporter/loop/index.ts b/packages/cozeloop-ai/src/tracer/exporter/loop/index.ts
diff --git a/packages/cozeloop-ai/src/tracer/initialize.ts b/packages/cozeloop-ai/src/tracer/initialize.ts
diff --git a/packages/cozeloop-ai/src/tracer/types/initialize.ts b/packages/cozeloop-ai/src/tracer/types/initialize.ts
diff --git a/packages/cozeloop-ai/src/tracer/types/wrapper.ts b/packages/cozeloop-ai/src/tracer/types/wrapper.ts
diff --git a/packages/cozeloop-ai/src/tracer/wrapper.ts b/packages/cozeloop-ai/src/tracer/wrapper.ts

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ export async function runLargeText() {`
`21`	`21`	`await cozeLoopTracer.traceable(`
`22`	`22`	`async span => {`
`23`	`23`	`// Reporting of ultra-long texts will only take effect when the`
`24`		`- // input satisfies the LoopTraceLLMCallInput structure`
	`24`	`+ // input / output satisfies the LoopTraceLLMCallInput / LoopTraceLLMCallOutput structure`
`25`	`25`	`const input: LoopTraceLLMCallInput = {`
`26`	`26`	`messages: [`
`27`	`27`	`{`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ export async function runMultiModality() {`
`14`	`14`	`await cozeLoopTracer.traceable(`
`15`	`15`	`async span => {`
`16`	`16`	`// Reporting of multi modality will only take effect when the`
`17`		`- // input satisfies the LoopTraceLLMCallInput structure`
	`17`	`+ // input / output satisfies the LoopTraceLLMCallInput / LoopTraceLLMCallOutput structure`
`18`	`18`	`const input: LoopTraceLLMCallInput = {`
`19`	`19`	`messages: [`
`20`	`20`	`{`