Repository: mastra-ai/mastra
File: observability/arize/src/openInferenceOTLPExporter.ts (TypeScript)

Code before:
const GEN_AI_USAGE_AUDIO_INPUT_TOKENS = 'gen_ai.usage.audio_input_tokens';
const GEN_AI_USAGE_AUDIO_OUTPUT_TOKENS = 'gen_ai.usage.audio_output_tokens';

const MASTRA_GENERAL_PREFIX = 'mastra.';
const MASTRA_METADATA_PREFIX = 'mastra.metadata.';
const MASTRA_MODEL_STEP_INPUT = 'mastra.model_step.input';
const MASTRA_MODEL_STEP_OUTPUT = 'mastra.model_step.output';
const MASTRA_MODEL_CHUNK_OUTPUT = 'mastra.model_chunk.output';
const MASTRA_SPAN_TYPE = 'mastra.span.type';

/**
 * Maps Mastra span types to OpenInference span kinds.
 *
 * The @arizeai/openinference-genai library defaults all spans to LLM kind,
 * which is incorrect for workflow, agent, tool, and other non-LLM spans.
 * This mapping overrides the span kind based on the mastra.span.type attribute.
 */
const SPAN_TYPE_TO_KIND: Record<string, OpenInferenceSpanKind> = {
  // Workflow spans -> CHAIN
  workflow_run: OpenInferenceSpanKind.CHAIN,
  workflow_step: OpenInferenceSpanKind.CHAIN,
  workflow_conditional: OpenInferenceSpanKind.CHAIN,
  workflow_conditional_eval: OpenInferenceSpanKind.CHAIN,
  workflow_parallel: OpenInferenceSpanKind.CHAIN,
  workflow_loop: OpenInferenceSpanKind.CHAIN,
  workflow_sleep: OpenInferenceSpanKind.CHAIN,
  workflow_wait_event: OpenInferenceSpanKind.CHAIN,
  // Model spans -> LLM
  model_generation: OpenInferenceSpanKind.LLM,
  model_step: OpenInferenceSpanKind.LLM,
  model_chunk: OpenInferenceSpanKind.LLM,
  // Tool spans -> TOOL
  tool_call: OpenInferenceSpanKind.TOOL,
  mcp_tool_call: OpenInferenceSpanKind.TOOL,
  // Agent spans -> AGENT
  agent_run: OpenInferenceSpanKind.AGENT,
  // Other spans -> CHAIN
  processor_run: OpenInferenceSpanKind.CHAIN,
  generic: OpenInferenceSpanKind.CHAIN,
};

/**
 * Converts GenAI usage metrics to OpenInference LLM token count attributes.
 * Maps from OTEL GenAI semantic conventions to OpenInference semantic conventions.
 *
 * @param attributes - The span attributes containing GenAI usage metrics
 * @returns OpenInference token count attributes
 */
function convertUsageMetricsToOpenInference(attributes: Record<string, any>): Record<string, any> {
  const result: Record<string, any> = {};


Review comment (suggestion):
can this map just have the LLM, TOOL, and AGENT types, and then we default to CHAIN for the others?

If we add new span types, I don't want to have to change every exporter.

Code after:
const GEN_AI_USAGE_AUDIO_INPUT_TOKENS = 'gen_ai.usage.audio_input_tokens';
const GEN_AI_USAGE_AUDIO_OUTPUT_TOKENS = 'gen_ai.usage.audio_output_tokens';

const MASTRA_GENERAL_PREFIX = 'mastra.';
const MASTRA_METADATA_PREFIX = 'mastra.metadata.';
const MASTRA_MODEL_STEP_INPUT = 'mastra.model_step.input';
const MASTRA_MODEL_STEP_OUTPUT = 'mastra.model_step.output';
const MASTRA_MODEL_CHUNK_OUTPUT = 'mastra.model_chunk.output';
const MASTRA_SPAN_TYPE = 'mastra.span.type';

/**
 * Maps Mastra span types to OpenInference span kinds for proper trace categorization.
 *
 * Only non-CHAIN types are mapped here - all other span types default to CHAIN.
 */
const SPAN_TYPE_TO_KIND: Record<string, OpenInferenceSpanKind> = {
  // Model spans -> LLM
  model_generation: OpenInferenceSpanKind.LLM,
  model_step: OpenInferenceSpanKind.LLM,
  model_chunk: OpenInferenceSpanKind.LLM,
  // Tool spans -> TOOL
  tool_call: OpenInferenceSpanKind.TOOL,
  mcp_tool_call: OpenInferenceSpanKind.TOOL,
  // Agent spans -> AGENT
  agent_run: OpenInferenceSpanKind.AGENT,
};

/**
 * Converts GenAI usage metrics to OpenInference LLM token count attributes.
 * Maps from OTEL GenAI semantic conventions to OpenInference semantic conventions.
 *
 * @param attributes - The span attributes containing GenAI usage metrics
 * @returns OpenInference token count attributes
 */
function convertUsageMetricsToOpenInference(attributes: Record<string, any>): Record<string, any> {
  const result: Record<string, any> = {};

  const inputTokens = attributes[ATTR_GEN_AI_USAGE_INPUT_TOKENS];
  const outputTokens = attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS];

  // Core token counts
  if (inputTokens !== undefined) {
    result[LLM_TOKEN_COUNT_PROMPT] = inputTokens;
  }
  if (outputTokens !== undefined) {
    result[LLM_TOKEN_COUNT_COMPLETION] = outputTokens;
  }

  // Total tokens (compute if we have both input and output)
  if (inputTokens !== undefined && outputTokens !== undefined) {
    result[LLM_TOKEN_COUNT_TOTAL] = inputTokens + outputTokens;