diff --git a/packages/sample-app/package.json b/packages/sample-app/package.json index 00eb5743..526df3ad 100644 --- a/packages/sample-app/package.json +++ b/packages/sample-app/package.json @@ -15,6 +15,7 @@ "run:gemini": "npm run build && node dist/src/vertexai/gemini.js", "run:palm2": "npm run build && node dist/src/vertexai/palm2.js", "run:decorators": "npm run build && node dist/src/sample_decorators.js", + "run:associations": "npm run build && node dist/src/sample_associations.js", "run:with": "npm run build && node dist/src/sample_with.js", "run:prompt_mgmt": "npm run build && node dist/src/sample_prompt_mgmt.js", "run:vercel": "npm run build && node dist/src/sample_vercel_ai.js", @@ -43,6 +44,7 @@ "run:mcp": "npm run build && node dist/src/sample_mcp.js", "run:mcp:real": "npm run build && node dist/src/sample_mcp_real.js", "run:mcp:working": "npm run build && node dist/src/sample_mcp_working.js", + "run:chatbot": "npm run build && node dist/src/sample_chatbot_interactive.js", "dev:image_generation": "pnpm --filter @traceloop/instrumentation-openai build && pnpm --filter @traceloop/node-server-sdk build && npm run build && node dist/src/sample_openai_image_generation.js", "lint": "eslint .", "lint:fix": "eslint . --fix" diff --git a/packages/sample-app/src/sample_associations.ts b/packages/sample-app/src/sample_associations.ts new file mode 100644 index 00000000..0779b331 --- /dev/null +++ b/packages/sample-app/src/sample_associations.ts @@ -0,0 +1,178 @@ +import * as traceloop from "@traceloop/node-server-sdk"; +import OpenAI from "openai"; + +// Initialize Traceloop +traceloop.initialize({ + appName: "associations_demo", + apiKey: process.env.TRACELOOP_API_KEY, + disableBatch: true, +}); + +const openai = new OpenAI(); + +/** + * Sample chatbot that demonstrates the Associations API. + * This example shows how to track conversations, users, and sessions + * across multiple LLM interactions. + */ +class ChatbotWithAssociations { + constructor( + private conversationId: string, + private userId: string, + private sessionId: string, + ) {} + + /** + * Process a multi-turn conversation with associations + */ + @traceloop.workflow({ name: "chatbot_conversation" }) + async handleConversation() { + console.log("\n=== Starting Chatbot Conversation ==="); + console.log(`Conversation ID: ${this.conversationId}`); + console.log(`User ID: ${this.userId}`); + console.log(`Session ID: ${this.sessionId}\n`); + + // Set standard associations at the beginning of the conversation + // These will be automatically attached to all spans within this context + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, this.conversationId], + [traceloop.AssociationProperty.USER_ID, this.userId], + [traceloop.AssociationProperty.SESSION_ID, this.sessionId], + ]); + + // Use withAssociationProperties to add custom properties + // Custom properties (like chat_subject) will be prefixed with traceloop.association.properties + return traceloop.withAssociationProperties( + { chat_subject: "general" }, + async () => { + // First message + const greeting = await this.sendMessage( + "Hello! What's the weather like today?", + ); + console.log(`Bot: ${greeting}\n`); + + // Second message in the same conversation + const followup = await this.sendMessage( + "What should I wear for that weather?", + ); + console.log(`Bot: ${followup}\n`); + + // Third message + const final = await this.sendMessage("Thanks for the advice!"); + console.log(`Bot: ${final}\n`); + + return { + greeting, + followup, + final, + }; + }, + ); + } + + /** + * Send a single message - this is a task within the workflow + */ + @traceloop.task({ name: "send_message" }) + private async sendMessage(userMessage: string): Promise { + console.log(`User: ${userMessage}`); + + const completion = await openai.chat.completions.create({ + messages: [{ role: "user", content: userMessage }], + model: "gpt-3.5-turbo", + }); + + return completion.choices[0].message.content || "No response"; + } +} + +/** + * Simulate a customer service scenario with multiple customers + */ +async function customerServiceDemo() { + return traceloop.withWorkflow( + { name: "customer_service_scenario" }, + async () => { + console.log("\n=== Customer Service Scenario ===\n"); + + // Customer 1 + traceloop.Associations.set([ + [traceloop.AssociationProperty.CUSTOMER_ID, "cust-001"], + [traceloop.AssociationProperty.USER_ID, "agent-alice"], + ]); + + const customer1Response = await openai.chat.completions.create({ + messages: [ + { + role: "user", + content: "I need help with my order #12345", + }, + ], + model: "gpt-3.5-turbo", + }); + + console.log("Customer 1 (cust-001):"); + console.log( + `Response: ${customer1Response.choices[0].message.content}\n`, + ); + + // Customer 2 - Update associations for new customer + traceloop.Associations.set([ + [traceloop.AssociationProperty.CUSTOMER_ID, "cust-002"], + [traceloop.AssociationProperty.USER_ID, "agent-bob"], + ]); + + const customer2Response = await openai.chat.completions.create({ + messages: [ + { + role: "user", + content: "How do I return an item?", + }, + ], + model: "gpt-3.5-turbo", + }); + + console.log("Customer 2 (cust-002):"); + console.log( + `Response: ${customer2Response.choices[0].message.content}\n`, + ); + }, + ); +} + +/** + * Main execution + */ +async function main() { + console.log("============================================"); + console.log("Traceloop Associations API Demo"); + console.log("============================================"); + + try { + // Example 1: Multi-turn chatbot conversation with custom properties + const chatbot = new ChatbotWithAssociations( + "conv-abc-123", // conversation_id + "user-alice-456", // user_id + "session-xyz-789", // session_id + ); + + await chatbot.handleConversation(); + + // Example 2: Customer service with multiple customers + await customerServiceDemo(); + + console.log("\n=== Demo Complete ==="); + console.log( + "Check your Traceloop dashboard to see the associations attached to traces!", + ); + console.log( + "You can filter and search by conversation_id, user_id, session_id, customer_id, or custom properties like chat_subject.", + ); + } catch (error) { + console.error("Error running demo:", error); + process.exit(1); + } +} + +// Run the demo +main(); diff --git a/packages/sample-app/src/sample_chatbot_interactive.ts b/packages/sample-app/src/sample_chatbot_interactive.ts new file mode 100644 index 00000000..ccd86e7f --- /dev/null +++ b/packages/sample-app/src/sample_chatbot_interactive.ts @@ -0,0 +1,284 @@ +import * as traceloop from "@traceloop/node-server-sdk"; +import { openai } from "@ai-sdk/openai"; +import { streamText, CoreMessage, tool } from "ai"; +import * as readline from "readline"; +import { z } from "zod"; + +import "dotenv/config"; + +traceloop.initialize({ + appName: "sample_chatbot_interactive", + disableBatch: true, +}); + +// ANSI color codes for terminal output +const colors = { + reset: "\x1b[0m", + bright: "\x1b[1m", + dim: "\x1b[2m", + cyan: "\x1b[36m", + green: "\x1b[32m", + yellow: "\x1b[33m", + blue: "\x1b[34m", + magenta: "\x1b[35m", +}; + +class InteractiveChatbot { + private conversationHistory: CoreMessage[] = []; + private rl: readline.Interface; + private conversationId: string; + private userId: string; + + constructor() { + this.rl = readline.createInterface({ + input: process.stdin, + output: process.stdout, + prompt: `${colors.cyan}${colors.bright}You: ${colors.reset}`, + }); + // Generate unique IDs for this session + this.conversationId = `conv-${Date.now()}`; + this.userId = `user-${Math.random().toString(36).substring(7)}`; + } + + @traceloop.task({ name: "summarize_interaction" }) + async generateSummary( + userMessage: string, + assistantResponse: string, + ): Promise { + console.log( + `\n${colors.yellow}▼ SUMMARY${colors.reset} ${colors.dim}TASK${colors.reset}`, + ); + + const summaryResult = await streamText({ + model: openai("gpt-4o-mini"), + messages: [ + { + role: "system", + content: + "Create a very brief title (3-6 words) that summarizes this conversation exchange. Only return the title, nothing else.", + }, + { + role: "user", + content: `User: ${userMessage}\n\nAssistant: ${assistantResponse}`, + }, + ], + experimental_telemetry: { isEnabled: true }, + }); + + let summary = ""; + for await (const chunk of summaryResult.textStream) { + summary += chunk; + } + + const cleanSummary = summary.trim().replace(/^["']|["']$/g, ""); + console.log(`${colors.dim}${cleanSummary}${colors.reset}`); + + return cleanSummary; + } + + @traceloop.workflow({ name: "chat_interaction" }) + async processMessage(userMessage: string): Promise { + // Set associations for tracing + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, this.conversationId], + [traceloop.AssociationProperty.USER_ID, this.userId], + ]); + + // Add user message to history + this.conversationHistory.push({ + role: "user", + content: userMessage, + }); + + console.log(`\n${colors.green}${colors.bright}Assistant: ${colors.reset}`); + + // Stream the response + const result = await streamText({ + model: openai("gpt-4o"), + messages: [ + { + role: "system", + content: + "You are a helpful AI assistant with access to tools. Use the available tools when appropriate to provide accurate information. Provide clear, concise, and friendly responses.", + }, + ...this.conversationHistory, + ], + tools: { + calculator: tool({ + description: + "Perform mathematical calculations. Supports basic arithmetic operations.", + parameters: z.object({ + expression: z + .string() + .describe( + "The mathematical expression to evaluate (e.g., '2 + 2' or '10 * 5')", + ), + }), + execute: async ({ expression }) => { + try { + // Simple safe eval for basic math (only allow numbers and operators) + const sanitized = expression.replace(/[^0-9+\-*/().\s]/g, ""); + const result = eval(sanitized); + console.log( + `\n${colors.yellow}🔧 Calculator: ${expression} = ${result}${colors.reset}`, + ); + return { result, expression }; + } catch (error) { + return { error: "Invalid mathematical expression" }; + } + }, + }), + getCurrentWeather: tool({ + description: + "Get the current weather for a location. Use this when users ask about weather conditions.", + parameters: z.object({ + location: z + .string() + .describe("The city and country, e.g., 'London, UK'"), + }), + execute: async ({ location }) => { + console.log( + `\n${colors.yellow}🔧 Weather: Checking weather for ${location}${colors.reset}`, + ); + // Simulated weather data + const weatherConditions = [ + "sunny", + "cloudy", + "rainy", + "partly cloudy", + ]; + const condition = + weatherConditions[ + Math.floor(Math.random() * weatherConditions.length) + ]; + const temperature = Math.floor(Math.random() * 30) + 10; // 10-40°C + return { + location, + temperature: `${temperature}°C`, + condition, + humidity: `${Math.floor(Math.random() * 40) + 40}%`, + }; + }, + }), + getTime: tool({ + description: + "Get the current date and time. Use this when users ask about the current time or date.", + parameters: z.object({ + timezone: z + .string() + .optional() + .describe("Optional timezone (e.g., 'America/New_York')"), + }), + execute: async ({ timezone }) => { + const now = new Date(); + const options: Intl.DateTimeFormatOptions = { + timeZone: timezone, + dateStyle: "full", + timeStyle: "long", + }; + const formatted = now.toLocaleString("en-US", options); + console.log( + `\n${colors.yellow}🔧 Time: ${formatted}${colors.reset}`, + ); + return { + datetime: formatted, + timestamp: now.toISOString(), + timezone: timezone || "local", + }; + }, + }), + }, + maxSteps: 5, + experimental_telemetry: { isEnabled: true }, + }); + + let fullResponse = ""; + for await (const chunk of result.textStream) { + process.stdout.write(chunk); + fullResponse += chunk; + } + + console.log("\n"); + + // Wait for the full response to complete to get all messages including tool calls + const finalResult = await result.response; + + // Add all response messages (including tool calls and results) to history + // This ensures the conversation history includes the complete interaction + for (const message of finalResult.messages) { + this.conversationHistory.push(message); + } + + // Generate summary for this interaction + await this.generateSummary(userMessage, fullResponse); + + return fullResponse; + } + + clearHistory(): void { + this.conversationHistory = []; + console.log( + `\n${colors.magenta}✓ Conversation history cleared${colors.reset}\n`, + ); + } + + async start(): Promise { + console.log( + `${colors.bright}${colors.blue}╔════════════════════════════════════════════════════════════╗`, + ); + console.log( + `║ Interactive AI Chatbot with Traceloop ║`, + ); + console.log( + `╚════════════════════════════════════════════════════════════╝${colors.reset}\n`, + ); + console.log( + `${colors.dim}Commands: /exit (quit) | /clear (clear history)${colors.reset}\n`, + ); + + this.rl.prompt(); + + this.rl.on("line", async (input: string) => { + const trimmedInput = input.trim(); + + if (!trimmedInput) { + this.rl.prompt(); + return; + } + + if (trimmedInput === "/exit") { + console.log(`\n${colors.magenta}Goodbye! 👋${colors.reset}\n`); + this.rl.close(); + process.exit(0); + } + + if (trimmedInput === "/clear") { + this.clearHistory(); + this.rl.prompt(); + return; + } + + try { + await this.processMessage(trimmedInput); + } catch (error) { + console.error( + `\n${colors.bright}Error:${colors.reset} ${error instanceof Error ? error.message : String(error)}\n`, + ); + } + + this.rl.prompt(); + }); + + this.rl.on("close", () => { + console.log(`\n${colors.magenta}Goodbye! 👋${colors.reset}\n`); + process.exit(0); + }); + } +} + +async function main() { + const chatbot = new InteractiveChatbot(); + await chatbot.start(); +} + +main().catch(console.error); diff --git a/packages/traceloop-sdk/package.json b/packages/traceloop-sdk/package.json index 3cdba376..daaa8872 100644 --- a/packages/traceloop-sdk/package.json +++ b/packages/traceloop-sdk/package.json @@ -57,6 +57,7 @@ "dependencies": { "@google-cloud/opentelemetry-cloud-trace-exporter": "^3.0.0", "@opentelemetry/api": "^1.9.0", + "@opentelemetry/context-async-hooks": "^2.0.0", "@opentelemetry/core": "^2.0.1", "@opentelemetry/exporter-trace-otlp-proto": "^0.203.0", "@opentelemetry/instrumentation": "^0.203.0", diff --git a/packages/traceloop-sdk/src/lib/node-server-sdk.ts b/packages/traceloop-sdk/src/lib/node-server-sdk.ts index 4236a04c..997f760d 100644 --- a/packages/traceloop-sdk/src/lib/node-server-sdk.ts +++ b/packages/traceloop-sdk/src/lib/node-server-sdk.ts @@ -64,6 +64,7 @@ export { getTraceloopTracer } from "./tracing/tracing"; export * from "./tracing/decorators"; export * from "./tracing/manual"; export * from "./tracing/association"; +export * from "./tracing/associations"; export * from "./tracing/custom-metric"; export * from "./tracing/span-processor"; export * from "./prompts"; diff --git a/packages/traceloop-sdk/src/lib/tracing/association.ts b/packages/traceloop-sdk/src/lib/tracing/association.ts index 07f6c850..44d71b5e 100644 --- a/packages/traceloop-sdk/src/lib/tracing/association.ts +++ b/packages/traceloop-sdk/src/lib/tracing/association.ts @@ -14,8 +14,16 @@ export function withAssociationProperties< return fn.apply(thisArg, args); } + // Get existing associations from context and merge with new properties + const existingAssociations = context + .active() + .getValue(ASSOCATION_PROPERTIES_KEY) as Record | undefined; + const mergedAssociations = existingAssociations + ? { ...existingAssociations, ...properties } + : properties; + const newContext = context .active() - .setValue(ASSOCATION_PROPERTIES_KEY, properties); + .setValue(ASSOCATION_PROPERTIES_KEY, mergedAssociations); return context.with(newContext, fn, thisArg, ...args); } diff --git a/packages/traceloop-sdk/src/lib/tracing/associations.ts b/packages/traceloop-sdk/src/lib/tracing/associations.ts new file mode 100644 index 00000000..422e9a9e --- /dev/null +++ b/packages/traceloop-sdk/src/lib/tracing/associations.ts @@ -0,0 +1,90 @@ +import { trace, context as otelContext } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { ASSOCATION_PROPERTIES_KEY } from "./tracing"; + +/** + * Standard association properties for tracing. + */ +export enum AssociationProperty { + CONVERSATION_ID = "conversation_id", + CUSTOMER_ID = "customer_id", + USER_ID = "user_id", + SESSION_ID = "session_id", +} + +/** + * Set of standard association property keys (without prefix). + * Use this to check if a property should be set directly or with the TRACELOOP_ASSOCIATION_PROPERTIES prefix. + */ +export const STANDARD_ASSOCIATION_PROPERTIES = new Set( + Object.values(AssociationProperty), +); + +/** + * Type alias for a single association + */ +export type Association = [AssociationProperty, string]; + +/** + * Class for managing trace associations. + */ +export class Associations { + /** + * Set associations that will be added directly to all spans in the current context. + * + * @param associations - An array of [property, value] tuples + * + * @example + * // Single association + * Associations.set([[AssociationProperty.CONVERSATION_ID, "conv-123"]]); + * + * // Multiple associations + * Associations.set([ + * [AssociationProperty.USER_ID, "user-456"], + * [AssociationProperty.SESSION_ID, "session-789"] + * ]); + */ + static set(associations: Association[]): void { + // Get current associations from context or create empty object + const existingAssociations = otelContext + .active() + .getValue(ASSOCATION_PROPERTIES_KEY) as + | Record + | undefined; + const currentAssociations: Record = existingAssociations + ? { ...existingAssociations } + : {}; + + // Update associations with new values + for (const [prop, value] of associations) { + currentAssociations[prop] = value; + } + + // Store associations in context + const newContext = otelContext + .active() + .setValue(ASSOCATION_PROPERTIES_KEY, currentAssociations); + + // Set the new context as active using the context manager + // This is the equivalent of Python's attach(set_value(...)) + const contextManager = (otelContext as any)["_getContextManager"](); + if ( + contextManager && + contextManager instanceof AsyncLocalStorageContextManager + ) { + // For AsyncLocalStorageContextManager, we need to use the internal _asyncLocalStorage + const storage = (contextManager as any)._asyncLocalStorage; + if (storage) { + storage.enterWith(newContext); + } + } + + // Also set directly on the current span (use newContext after enterWith) + const span = trace.getSpan(newContext); + if (span && span.isRecording()) { + for (const [prop, value] of associations) { + span.setAttribute(prop, value); + } + } + } +} diff --git a/packages/traceloop-sdk/test/associations.test.ts b/packages/traceloop-sdk/test/associations.test.ts new file mode 100644 index 00000000..4959b953 --- /dev/null +++ b/packages/traceloop-sdk/test/associations.test.ts @@ -0,0 +1,389 @@ +/* + * Copyright Traceloop + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from "assert"; + +import type * as OpenAIModule from "openai"; + +import * as traceloop from "../src"; + +import { Polly, setupMocha as setupPolly } from "@pollyjs/core"; +import NodeHttpAdapter from "@pollyjs/adapter-node-http"; +import FetchAdapter from "@pollyjs/adapter-fetch"; +import FSPersister from "@pollyjs/persister-fs"; +import { SpanAttributes } from "@traceloop/ai-semantic-conventions"; +import { initializeSharedTraceloop, getSharedExporter } from "./test-setup"; + +const memoryExporter = getSharedExporter(); + +Polly.register(NodeHttpAdapter); +Polly.register(FetchAdapter); +Polly.register(FSPersister); + +describe("Test Associations API", () => { + let openai: OpenAIModule.OpenAI; + + setupPolly({ + adapters: ["node-http", "fetch"], + persister: "fs", + recordIfMissing: process.env.RECORD_MODE === "NEW", + recordFailedRequests: true, + mode: process.env.RECORD_MODE === "NEW" ? "record" : "replay", + matchRequestsBy: { + headers: false, + url: { + protocol: true, + hostname: true, + pathname: true, + query: false, + }, + }, + logging: true, + }); + + before(async function () { + if (process.env.RECORD_MODE !== "NEW") { + process.env.OPENAI_API_KEY = "test"; + } + + // Use shared initialization to avoid conflicts with other test suites + initializeSharedTraceloop(); + + // Initialize OpenAI after Polly is set up + const openAIModule: typeof OpenAIModule = await import("openai"); + openai = new openAIModule.OpenAI(); + }); + + beforeEach(function () { + const { server } = this.polly as Polly; + server.any().on("beforePersist", (_req, recording) => { + recording.request.headers = recording.request.headers.filter( + ({ name }: { name: string }) => name !== "authorization", + ); + }); + }); + + afterEach(async () => { + memoryExporter.reset(); + }); + + it("should set single association on spans", async () => { + const result = await traceloop.withWorkflow( + { name: "test_single_association" }, + async () => { + // Set a single association + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, "conv-123"], + ]); + + const chatCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Tell me a joke" }], + model: "gpt-3.5-turbo", + }); + + return chatCompletion.choices[0].message.content; + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const workflowSpan = spans.find( + (span) => span.name === "test_single_association.workflow", + ); + const chatSpan = spans.find((span) => span.name === "openai.chat"); + + assert.ok(result); + assert.ok(workflowSpan); + assert.ok(chatSpan); + + // Check that the association is set on both workflow and LLM spans + assert.strictEqual(workflowSpan.attributes["conversation_id"], "conv-123"); + assert.strictEqual( + chatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.conversation_id` + ], + "conv-123", + ); + }); + + it("should set multiple associations on spans", async () => { + const result = await traceloop.withWorkflow( + { name: "test_multiple_associations" }, + async () => { + // Set multiple associations + traceloop.Associations.set([ + [traceloop.AssociationProperty.USER_ID, "user-456"], + [traceloop.AssociationProperty.SESSION_ID, "session-789"], + ]); + + const chatCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Tell me a fact" }], + model: "gpt-3.5-turbo", + }); + + return chatCompletion.choices[0].message.content; + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const workflowSpan = spans.find( + (span) => span.name === "test_multiple_associations.workflow", + ); + const chatSpan = spans.find((span) => span.name === "openai.chat"); + + assert.ok(result); + assert.ok(workflowSpan); + assert.ok(chatSpan); + + // Check that both associations are set + assert.strictEqual(workflowSpan.attributes["user_id"], "user-456"); + assert.strictEqual(workflowSpan.attributes["session_id"], "session-789"); + assert.strictEqual( + chatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id` + ], + "user-456", + ); + assert.strictEqual( + chatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id` + ], + "session-789", + ); + }); + + it("should allow updating associations mid-workflow", async () => { + const result = await traceloop.withWorkflow( + { name: "test_update_associations" }, + async () => { + // Set initial association + traceloop.Associations.set([ + [traceloop.AssociationProperty.SESSION_ID, "session-initial"], + ]); + + const firstCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "First message" }], + model: "gpt-3.5-turbo", + }); + + // Update association + traceloop.Associations.set([ + [traceloop.AssociationProperty.SESSION_ID, "session-updated"], + ]); + + const secondCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Second message" }], + model: "gpt-3.5-turbo", + }); + + return { + first: firstCompletion.choices[0].message.content, + second: secondCompletion.choices[0].message.content, + }; + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const firstChatSpan = spans.find( + (span) => + span.name === "openai.chat" && + span.attributes[`${SpanAttributes.ATTR_GEN_AI_PROMPT}.0.content`] === + "First message", + ); + const secondChatSpan = spans.find( + (span) => + span.name === "openai.chat" && + span.attributes[`${SpanAttributes.ATTR_GEN_AI_PROMPT}.0.content`] === + "Second message", + ); + + assert.ok(result); + assert.ok(firstChatSpan); + assert.ok(secondChatSpan); + + // First span should have initial value + assert.strictEqual( + firstChatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id` + ], + "session-initial", + ); + + // Second span should have updated value + assert.strictEqual( + secondChatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id` + ], + "session-updated", + ); + }); + + it("should work with all AssociationProperty types", async () => { + const result = await traceloop.withWorkflow( + { name: "test_all_property_types" }, + async () => { + // Set all association property types + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, "conv-abc"], + [traceloop.AssociationProperty.CUSTOMER_ID, "customer-def"], + [traceloop.AssociationProperty.USER_ID, "user-ghi"], + [traceloop.AssociationProperty.SESSION_ID, "session-jkl"], + ]); + + const chatCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Test all properties" }], + model: "gpt-3.5-turbo", + }); + + return chatCompletion.choices[0].message.content; + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const chatSpan = spans.find((span) => span.name === "openai.chat"); + + assert.ok(result); + assert.ok(chatSpan); + + // Check all property types are set (standard properties without prefix) + assert.strictEqual(chatSpan.attributes["conversation_id"], "conv-abc"); + assert.strictEqual(chatSpan.attributes["customer_id"], "customer-def"); + assert.strictEqual(chatSpan.attributes["user_id"], "user-ghi"); + assert.strictEqual(chatSpan.attributes["session_id"], "session-jkl"); + }); + + it("should propagate associations to all child spans", async () => { + const result = await traceloop.withWorkflow( + { name: "test_child_propagation" }, + async () => { + // Set associations at the workflow level + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, "conv-propagate"], + [traceloop.AssociationProperty.USER_ID, "user-propagate"], + ]); + + // Call a child task + const taskResult = await traceloop.withTask( + { name: "subtask" }, + async () => { + const chatCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Child task message" }], + model: "gpt-3.5-turbo", + }); + return chatCompletion.choices[0].message.content; + }, + ); + + return taskResult; + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const workflowSpan = spans.find( + (span) => span.name === "test_child_propagation.workflow", + ); + const taskSpan = spans.find((span) => span.name === "subtask.task"); + const chatSpan = spans.find( + (span) => + span.name === "openai.chat" && + span.attributes[`${SpanAttributes.ATTR_GEN_AI_PROMPT}.0.content`] === + "Child task message", + ); + + assert.ok(result); + assert.ok(workflowSpan); + assert.ok(taskSpan); + assert.ok(chatSpan); + + // All spans should have the associations (standard properties without prefix) + assert.strictEqual( + workflowSpan.attributes["conversation_id"], + "conv-propagate", + ); + assert.strictEqual(workflowSpan.attributes["user_id"], "user-propagate"); + + assert.strictEqual( + taskSpan.attributes["conversation_id"], + "conv-propagate", + ); + assert.strictEqual(taskSpan.attributes["user_id"], "user-propagate"); + + assert.strictEqual( + chatSpan.attributes["conversation_id"], + "conv-propagate", + ); + assert.strictEqual(chatSpan.attributes["user_id"], "user-propagate"); + }); + + it("should merge associations from Associations.set() and withAssociationProperties()", async () => { + const result = await traceloop.withWorkflow( + { name: "test_merge_associations" }, + async () => { + // Set standard associations + traceloop.Associations.set([ + [traceloop.AssociationProperty.CONVERSATION_ID, "conv-merge"], + [traceloop.AssociationProperty.USER_ID, "user-merge"], + ]); + + // Add custom properties via withAssociationProperties + return await traceloop.withAssociationProperties( + { custom_field: "custom-value" }, + async () => { + const chatCompletion = await openai.chat.completions.create({ + messages: [{ role: "user", content: "Test merge" }], + model: "gpt-3.5-turbo", + }); + return chatCompletion.choices[0].message.content; + }, + ); + }, + ); + + await traceloop.forceFlush(); + const spans = memoryExporter.getFinishedSpans(); + + const chatSpan = spans.find( + (span) => + span.name === "openai.chat" && + span.attributes[`${SpanAttributes.ATTR_GEN_AI_PROMPT}.0.content`] === + "Test merge", + ); + + assert.ok(result); + assert.ok(chatSpan); + + // Standard properties should be without prefix + assert.strictEqual(chatSpan.attributes["conversation_id"], "conv-merge"); + assert.strictEqual(chatSpan.attributes["user_id"], "user-merge"); + + // Custom property should have prefix + assert.strictEqual( + chatSpan.attributes[ + `${SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.custom_field` + ], + "custom-value", + ); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2c521e3a..a4e41dbe 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -822,6 +822,9 @@ importers: '@opentelemetry/api': specifier: ^1.9.0 version: 1.9.0 + '@opentelemetry/context-async-hooks': + specifier: ^2.0.0 + version: 2.0.1(@opentelemetry/api@1.9.0) '@opentelemetry/core': specifier: ^2.0.1 version: 2.0.1(@opentelemetry/api@1.9.0) @@ -14375,7 +14378,7 @@ snapshots: isstream: 0.1.2 jsonwebtoken: 9.0.3 mime-types: 2.1.35 - retry-axios: 2.6.0(axios@1.13.2(debug@4.4.3)) + retry-axios: 2.6.0(axios@1.13.2) tough-cookie: 4.1.4 transitivePeerDependencies: - supports-color @@ -16317,7 +16320,7 @@ snapshots: onetime: 5.1.2 signal-exit: 3.0.7 - retry-axios@2.6.0(axios@1.13.2(debug@4.4.3)): + retry-axios@2.6.0(axios@1.13.2): dependencies: axios: 1.13.2(debug@4.4.3)