allow setting of telemetry options

packages/core/src/agent/agent.test.ts

@@ -1,6 +1,8 @@
 import { createOpenAI } from '@ai-sdk/openai';
+import { InMemorySpanExporter, SimpleSpanProcessor } from '@opentelemetry/sdk-trace-base';
+import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
 import { config } from 'dotenv';
-import { describe, it, expect, vi } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { z } from 'zod';
 
 import { TestIntegration } from '../integration/openapi-toolset.mock';
@@ -275,4 +277,42 @@ describe('agent', () => {
     expect(sanitizedMessages).not.toContainEqual(toolCallThree);
     expect(sanitizedMessages).toHaveLength(2);
   });
+
+  it('should use telemetry options when generating a response', async () => {
+    const electionAgent = new Agent({
+      name: 'US Election agent',
+      instructions: 'You know about the past US elections',
+      model: openai('gpt-4o'),
+    });
+
+    const memoryExporter = new InMemorySpanExporter();
+    const tracerProvider = new NodeTracerProvider({
+      spanProcessors: [new SimpleSpanProcessor(memoryExporter)],
+    });
+    tracerProvider.register();
+
+    const mastra = new Mastra({
+      agents: { electionAgent },
+      telemetry: {
+        enabled: true,
+        serviceName: 'test-service',
+        export: {
+          type: 'custom',
+          exporter: memoryExporter,
+        },
+      },
+    });
+    const agentOne = mastra.getAgent('electionAgent');
+
+    const response = await agentOne.generate('Who won the 2016 US presidential election?', {
+      telemetryOptions: { functionId: 'test-function-id', metadata: { test: 'test' } },
+    });
+
+    const spans = memoryExporter.getFinishedSpans();
+    const aiSpan = spans.find(span => span.name === 'ai.generateText');
+    expect(aiSpan).toBeDefined();
+    expect(aiSpan?.attributes['ai.telemetry.metadata.test']).toBe('test');
+    expect(aiSpan?.attributes['resource.name']).toBe('test-function-id');
+    await tracerProvider.shutdown();
+  });
 });

packages/core/src/agent/index.ts

@@ -4,23 +4,23 @@ import type {
   CoreMessage,
   CoreToolMessage,
   CoreUserMessage,
+  LanguageModelV1,
   TextPart,
   ToolCallPart,
   UserContent,
-  LanguageModelV1,
 } from 'ai';
 import { randomUUID } from 'crypto';
 import type { JSONSchema7 } from 'json-schema';
-import { z } from 'zod';
 import type { ZodSchema } from 'zod';
+import { z } from 'zod';
 
 import type { MastraPrimitives } from '../action';
 import { MastraBase } from '../base';
 import type { Metric } from '../eval';
 import { AvailableHooks, executeHook } from '../hooks';
 import type { GenerateReturn, StreamReturn } from '../llm';
-import { MastraLLM } from '../llm/model';
 import type { MastraLLMBase } from '../llm/model';
+import { MastraLLM } from '../llm/model';
 import { LogLevel, RegisteredLogger } from '../logger';
 import type { MastraMemory } from '../memory/memory';
 import type { MemoryConfig, StorageThreadType } from '../memory/types';
@@ -777,6 +777,7 @@ export class Agent<
       output = 'text',
       temperature,
       toolChoice = 'auto',
+      telemetryOptions,
     }: AgentGenerateOptions<Z> = {},
   ): Promise<GenerateReturn<Z>> {
     let messagesToUse: CoreMessage[] = [];
@@ -824,6 +825,7 @@ export class Agent<
         runId: runIdToUse,
         temperature,
         toolChoice,
+        telemetryOptions,
       });
 
       const outputText = result.text;
@@ -843,6 +845,7 @@ export class Agent<
       runId: runIdToUse,
       temperature,
       toolChoice,
+      telemetryOptions,
     });
 
     const outputText = JSON.stringify(result.object);
@@ -867,6 +870,7 @@ export class Agent<
       output = 'text',
       temperature,
       toolChoice = 'auto',
+      telemetryOptions,
     }: AgentStreamOptions<Z> = {},
   ): Promise<StreamReturn<Z>> {
     const runIdToUse = runId || randomUUID();
@@ -930,6 +934,7 @@ export class Agent<
         maxSteps,
         runId: runIdToUse,
         toolChoice,
+        telemetryOptions,
       }) as unknown as StreamReturn<Z>;
     }
 
@@ -959,6 +964,7 @@ export class Agent<
       maxSteps,
       runId: runIdToUse,
       toolChoice,
+      telemetryOptions,
     }) as unknown as StreamReturn<Z>;
   }
 }

packages/core/src/agent/types.ts

@@ -1,4 +1,4 @@
-import type { LanguageModelV1 } from 'ai';
+import type { LanguageModelV1, TelemetrySettings } from 'ai';
 import type { JSONSchema7 } from 'json-schema';
 import type { ZodSchema } from 'zod';
 
@@ -42,6 +42,7 @@ export interface AgentGenerateOptions<Z extends ZodSchema | JSONSchema7 | undefi
   output?: OutputType | Z;
   temperature?: number;
   toolChoice?: 'auto' | 'required';
+  telemetryOptions?: TelemetrySettings;
 }
 
 export interface AgentStreamOptions<Z extends ZodSchema | JSONSchema7 | undefined = undefined> {
@@ -57,4 +58,5 @@ export interface AgentStreamOptions<Z extends ZodSchema | JSONSchema7 | undefine
   output?: OutputType | Z;
   temperature?: number;
   toolChoice?: 'auto' | 'required';
+  telemetryOptions?: TelemetrySettings;
 }

packages/core/src/base.ts

@@ -1,5 +1,5 @@
+import type { BaseLogMessage, Logger, LogLevel } from './logger';
 import { createLogger, RegisteredLogger } from './logger';
-import type { Logger, LogLevel, BaseLogMessage } from './logger';
 import type { Telemetry } from './telemetry';
 
 export class MastraBase {

packages/core/src/llm/index.ts

@@ -11,6 +11,7 @@ import type {
   LanguageModelV1,
   StreamObjectResult,
   StreamTextResult,
+  TelemetrySettings,
 } from 'ai';
 import type { JSONSchema7 } from 'json-schema';
 import type { z, ZodSchema } from 'zod';
@@ -82,6 +83,7 @@ export type LLMStreamOptions<Z extends ZodSchema | JSONSchema7 | undefined = und
   convertedTools?: Record<string, CoreTool>;
   output?: OutputType | Z;
   temperature?: number;
+  telemetryOptions?: TelemetrySettings;
 };
 
 export type LLMTextOptions = {
@@ -92,6 +94,7 @@ export type LLMTextOptions = {
   toolChoice?: 'auto' | 'required';
   maxSteps?: number;
   temperature?: number;
+  telemetryOptions?: TelemetrySettings;
 } & Run;
 
 export type LLMTextObjectOptions<T> = LLMTextOptions & {
@@ -107,6 +110,7 @@ export type LLMInnerStreamOptions = {
   maxSteps?: number;
   temperature?: number;
   toolChoice?: 'auto' | 'required';
+  telemetryOptions?: TelemetrySettings;
 } & Run;
 
 export type LLMStreamObjectOptions<T> = LLMInnerStreamOptions & {

packages/core/src/llm/model/model.ts

@@ -1,16 +1,16 @@
-import { generateText, generateObject, jsonSchema, streamText, streamObject } from 'ai';
-import type { LanguageModel, Schema, CoreMessage } from 'ai';
+import type { CoreMessage, LanguageModel, Schema } from 'ai';
+import { generateObject, generateText, jsonSchema, streamObject, streamText } from 'ai';
 import type { JSONSchema7 } from 'json-schema';
-import { z } from 'zod';
 import type { ZodSchema } from 'zod';
+import { z } from 'zod';
 
 import type {
   GenerateReturn,
-  LLMTextOptions,
   LLMInnerStreamOptions,
   LLMStreamObjectOptions,
   LLMStreamOptions,
   LLMTextObjectOptions,
+  LLMTextOptions,
   StreamReturn,
 } from '../';
 import type { MastraPrimitives } from '../../action';
@@ -114,6 +114,7 @@ export class MastraLLM extends MastraLLMBase {
     temperature,
     toolChoice = 'auto',
     onStepFinish,
+    telemetryOptions,
   }: LLMTextOptions) {
     const model = this.#model;
 
@@ -159,7 +160,10 @@ export class MastraLLM extends MastraLLMBase {
     return await generateText({
       messages,
       ...argsForExecute,
-      experimental_telemetry: this.experimental_telemetry,
+      experimental_telemetry: {
+        ...this.experimental_telemetry,
+        ...telemetryOptions,
+      },
     });
   }
 
@@ -173,6 +177,7 @@ export class MastraLLM extends MastraLLMBase {
     runId,
     temperature,
     toolChoice = 'auto',
+    telemetryOptions,
   }: LLMTextObjectOptions<T>) {
     const model = this.#model;
 
@@ -228,7 +233,10 @@ export class MastraLLM extends MastraLLMBase {
       ...argsForExecute,
       output: output as any,
       schema,
-      experimental_telemetry: this.experimental_telemetry,
+      experimental_telemetry: {
+        ...this.experimental_telemetry,
+        ...telemetryOptions,
+      },
     });
   }
 
@@ -242,6 +250,7 @@ export class MastraLLM extends MastraLLMBase {
     runId,
     temperature,
     toolChoice = 'auto',
+    telemetryOptions,
   }: LLMInnerStreamOptions) {
     const model = this.#model;
     this.logger.debug(`[LLM] - Streaming text`, {
@@ -298,7 +307,10 @@ export class MastraLLM extends MastraLLMBase {
     return await streamText({
       messages,
       ...argsForExecute,
-      experimental_telemetry: this.experimental_telemetry,
+      experimental_telemetry: {
+        ...this.experimental_telemetry,
+        ...telemetryOptions,
+      },
     });
   }
 
@@ -313,6 +325,7 @@ export class MastraLLM extends MastraLLMBase {
     runId,
     temperature,
     toolChoice = 'auto',
+    telemetryOptions,
   }: LLMStreamObjectOptions<T>) {
     const model = this.#model;
     this.logger.debug(`[LLM] - Streaming structured output`, {
@@ -384,7 +397,10 @@ export class MastraLLM extends MastraLLMBase {
       ...argsForExecute,
       output: output as any,
       schema,
-      experimental_telemetry: this.experimental_telemetry,
+      experimental_telemetry: {
+        ...this.experimental_telemetry,
+        ...telemetryOptions,
+      },
     });
   }
 
@@ -398,6 +414,7 @@ export class MastraLLM extends MastraLLMBase {
       runId,
       output = 'text',
       temperature,
+      telemetryOptions,
     }: LLMStreamOptions<Z> = {},
   ): Promise<GenerateReturn<Z>> {
     const msgs = this.convertToMessages(messages);
@@ -422,6 +439,7 @@ export class MastraLLM extends MastraLLMBase {
       tools,
       convertedTools,
       runId,
+      telemetryOptions,
     })) as unknown as GenerateReturn<Z>;
   }
 
@@ -436,6 +454,7 @@ export class MastraLLM extends MastraLLMBase {
       runId,
       output = 'text',
       temperature,
+      telemetryOptions,
     }: LLMStreamOptions<Z> = {},
   ) {
     const msgs = this.convertToMessages(messages);
@@ -450,6 +469,7 @@ export class MastraLLM extends MastraLLMBase {
         convertedTools,
         runId,
         temperature,
+        telemetryOptions,
       })) as unknown as StreamReturn<Z>;
     }
 
@@ -463,6 +483,7 @@ export class MastraLLM extends MastraLLMBase {
       convertedTools,
       runId,
       temperature,
+      telemetryOptions,
     })) as unknown as StreamReturn<Z>;
   }
 }