Tristan L.

import json
import logging
import os

import boto3
from arize.otel import Endpoint
from botocore.exceptions import ClientError
from openinference.instrumentation.strands_agents import StrandsAgentsToOpenInferenceProcessor
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

logger = logging.getLogger(__name__)

ENDPOINT = Endpoint.ARIZE

def setup_tracing():
    project_name = "project_name"
    arize_secret_arn = os.getenv("ARIZE_SECRET_ARN")

    if arize_secret_arn:
        try:
            secrets_client = boto3.client('secretsmanager')
            secret_value = secrets_client.get_secret_value(SecretId=arize_secret_arn)
            secret_dict = json.loads(secret_value['SecretString'])

            logger.info(f"Arize credentials retrieved from Secrets Manager: {secret_dict}")

            # The CDK code uses 'space_id' and 'api_key' in the secret object
            if 'space_id' in secret_dict:
                os.environ['ARIZE_SPACE_ID'] = secret_dict['space_id']
            if 'api_key' in secret_dict:
                os.environ['ARIZE_API_KEY'] = secret_dict['api_key']

            logger.info("Successfully loaded Arize credentials from Secrets Manager")
        except (ClientError, json.JSONDecodeError, KeyError) as e:
            logger.error(f"Failed to load Arize credentials from Secrets Manager: {e}")

    try:
        strands_processor = StrandsAgentsToOpenInferenceProcessor()
        resource = Resource.create({
            "model_id": project_name,  ### <-- Update with your Arize AX Project Name
        })
        provider = TracerProvider(resource=resource)
        provider.add_span_processor(strands_processor)

        otel_exporter = OTLPSpanExporter(
            endpoint=ENDPOINT,
            headers={
                'space_id': os.environ.get('ARIZE_SPACE_ID'),
                'api_key': os.environ.get('ARIZE_API_KEY'),
            }
        )
        provider.add_span_processor(BatchSpanProcessor(otel_exporter))
        trace.set_tracer_provider(provider)

        logger.info(f"Tracing initialized for project: {project_name}")
    except Exception as e:
        logger.error(f"Failed to initialize tracing: {e}")
        return None