scripts/test_tracing_live.py

import sys
import os
import logging
from typing import List

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("test_trace_submission")

# Add project root to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from aibrahim.config import CONFIG
from aibrahim.infra.tracing import TRACER
from langchain_core.messages import HumanMessage
from langchain_community.chat_models import ChatOpenAI

def run_test_trace():
    logger.info("Initializing test trace...")
    
    # Ensure config has tracing enabled
    if not CONFIG.enable_tracing:
        logger.error("Tracing is disabled in config.")
        return

    # Get Langfuse callback handler from TRACER which handles env setup
    try:
        # Debug: Check explicit env vars before getting callbacks
        # We simulate what TRACER does to see if it works here
        import os
        pk = CONFIG.langfuse_public_key
        sk = CONFIG.langfuse_secret_key
        host = CONFIG.langfuse_host
        
        logger.info(f"Debug - Config PK: {pk[:10]}...")
        logger.info(f"Debug - Config SK length: {len(sk)}")
        
        # We must rely on TRACER setting the env, but we can verify if it did?
        # TRACER is global and initialized at import time.
        # So it should have already set os.environ if enabled.
        
        env_sk = os.environ.get("LANGFUSE_SECRET_KEY", "NOT_SET")
        logger.info(f"Debug - Env SK after TRACER init: {env_sk[:5]}... ({len(env_sk)} chars)")
        logger.info(f"Debug - Env SK Repr: {repr(env_sk)}")
        
        callbacks = TRACER.get_callbacks(["test-trace-live"])
        logger.info("Retrieved callbacks from TRACER.")
        
    except Exception as e:
        logger.error(f"Failed to get callbacks: {e}")
        return

    if not callbacks:
        logger.error("No callbacks returned.")
        return

    logger.info("Langfuse callback handler ready.")

    # Initialize a simple model (using OpenAI as per project dependencies)
    try:
        model = ChatOpenAI(
            model="gpt-3.5-turbo",
            temperature=0,
            openai_api_key=os.getenv("OPENAI_API_KEY")
        )

        logger.info("Invoking model with tracing...")
        response = model.invoke(
            [HumanMessage(content="Return the word 'Pong' to verify tracing connectivity.")],
            config={"callbacks": callbacks}
        )
        
        logger.info(f"Model response: {response.content}")
        logger.info("---------------------------------------------------")
        logger.info("✅ Trace generated successfully!")
        
        # Checking if explicit handler needs flushing
        if callbacks and hasattr(callbacks[0], "langfuse"):
            logger.info("Flushing langfuse client...")
            callbacks[0].langfuse.flush()
            logger.info("Flush complete.")
            
    except Exception as e:
        logger.error(f"Failed to run model invocation: {e}")

if __name__ == "__main__":
    run_test_trace()