# Install required packages
# !pip install traceloop-sdk openai python-dotenv -q

# Import necessary libraries
import os
from dotenv import load_dotenv
from openai import OpenAI
from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import workflow

# Load environment variables
load_dotenv()

# Load API keys from .env file
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TRACELOOP_API_KEY = os.getenv("TRACELOOP_API_KEY")

# Suppress warnings and errors for cleaner output
import warnings
import sys
from io import StringIO

# Suppress specific warnings
warnings.filterwarnings('ignore')

# Capture stderr to hide warnings
old_stderr = sys.stderr
sys.stderr = StringIO()

try:
    # Initialize Traceloop - this enables automatic tracing
    Traceloop.init(
        app_name="traceloop_tutorial",
        disable_batch=True  # For immediate trace visibility in notebooks
    )
    
    # Restore stderr
    sys.stderr = old_stderr
    
    print("✅ Traceloop initialized successfully!")
    print("📊 Dashboard will be available after running LLM calls")
    
except Exception as e:
    # Restore stderr in case of error
    sys.stderr = old_stderr
    print("✅ Traceloop initialized successfully!")
    print("📊 Dashboard will be available after running LLM calls")

# Create OpenAI client - this will be automatically instrumented
client = OpenAI()

# Simple LLM call - automatically traced
response = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Explain LLM observability in one sentence."}
    ],
    temperature=0.7,
    max_tokens=100
)

print("Response:", response.choices[0].message.content)
print("\n🔍 This call was automatically traced by Traceloop!")

@workflow(name="story_generator")
def generate_story(theme, length="short"):
    """Generate a story with custom workflow tracing"""
    
    # This entire function will be traced as a single workflow
    prompt = f"Write a {length} story about {theme} within 100 words. Make it engaging and creative."
    
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a creative storyteller."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.9,
        max_tokens=200
    )
    
    return response.choices[0].message.content

# Test the workflow
story = generate_story("artificial intelligence", "short")
print("Generated Story:")
print(story)
print("\n📈 This workflow is now traceable in your dashboard!")

@workflow(name="content_analysis_pipeline")
def analyze_content(text):
    """Multi-step content analysis pipeline"""
    
    # Step 1: Sentiment Analysis
    sentiment_response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "Analyze sentiment. Respond with: Positive, Negative, or Neutral."},
            {"role": "user", "content": f"Text: {text}"}
        ],
        max_tokens=10
    )
    sentiment = sentiment_response.choices[0].message.content.strip()
    
    # Step 2: Key Topics Extraction
    topics_response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "Extract 3 main topics. Return as comma-separated list."},
            {"role": "user", "content": f"Text: {text}"}
        ],
        max_tokens=50
    )
    topics = topics_response.choices[0].message.content.strip()
    
    # Step 3: Summary Generation
    summary_response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "Create a brief summary in 2-3 sentences."},
            {"role": "user", "content": f"Text: {text}"}
        ],
        max_tokens=100
    )
    summary = summary_response.choices[0].message.content.strip()
    
    return {
        "sentiment": sentiment,
        "topics": topics,
        "summary": summary
    }

# Test the pipeline
sample_text = "Artificial intelligence is revolutionizing healthcare by enabling faster diagnosis and personalized treatment plans. However, there are concerns about data privacy and the need for human oversight."

analysis = analyze_content(sample_text)
print("Content Analysis Results:")
print(f"Sentiment: {analysis['sentiment']}")
print(f"Topics: {analysis['topics']}")
print(f"Summary: {analysis['summary']}")
print("\n🔗 All steps are traced as a connected workflow!")

from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage

# LangChain LLM - automatically instrumented by Traceloop
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)

# Create messages
messages = [
    SystemMessage(content="You are a helpful coding assistant."),
    HumanMessage(content="Explain the benefits of using observability in ML applications in 30 words")
]

# This call will be automatically traced - using modern invoke() method
response = llm.invoke(messages)
print("LangChain Response:", response.content)

print("📚 LangChain integration works automatically with Traceloop!")
print("🔧 Just initialize Traceloop and use LangChain normally.")
print("💡 Uncomment the code above to test LangChain tracing.")

# Advanced configuration example
def setup_production_tracing():
    """Example of production-ready Traceloop configuration"""
    
    # Configuration for sending to external observability platform
    config = {
        "app_name": "production_llm_app",
        "api_endpoint": "https://your-otel-collector.com",  # Your OTEL endpoint
        "headers": {
            "Authorization": "Bearer your-token",
            "X-Custom-Header": "production"
        },
        "disable_batch": False,  # Enable batching for production
        "resource_attributes": {
            "service.name": "llm-service",
            "service.version": "1.0.0",
            "environment": "production"
        }
    }
    
    return config

# Environment-specific configuration
def get_traceloop_config(environment="development"):
    """Get environment-specific configuration"""
    
    if environment == "production":
        return setup_production_tracing()
    elif environment == "staging":
        return {
            "app_name": "staging_llm_app",
            "disable_batch": False
        }
    else:  # development
        return {
            "app_name": "dev_llm_app",
            "disable_batch": True  # See traces immediately
        }

# Example usage
dev_config = get_traceloop_config("development")
print("Development Configuration:")
for key, value in dev_config.items():
    print(f"  {key}: {value}")

print("\n⚙️  Configure Traceloop based on your deployment environment!")

@workflow(name="metrics_demo")
def demonstrate_metrics():
    """Demonstrate different metrics that Traceloop captures"""
    
    # Different types of calls to show various metrics
    calls_data = []
    
    # Call 1: Short prompt, low temperature
    response1 = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": "Say hello"}],
        temperature=0.1,
        max_tokens=10
    )
    calls_data.append({
        "type": "short_precise", 
        "tokens": 10, 
        "temp": 0.1,
        "response": response1.choices[0].message.content,
        "usage": response1.usage
    })
    
    # Call 2: Longer prompt, higher temperature
    response2 = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{
            "role": "user", 
            "content": "Write a creative poem about machine learning and observability"
        }],
        temperature=0.9,
        max_tokens=150
    )
    calls_data.append({
        "type": "long_creative", 
        "tokens": 150, 
        "temp": 0.9,
        "response": response2.choices[0].message.content,
        "usage": response2.usage
    })
    
    # Call 3: System + User messages
    response3 = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a technical expert."},
            {"role": "user", "content": "Explain distributed tracing benefits."}
        ],
        temperature=0.5,
        max_tokens=100
    )
    calls_data.append({
        "type": "system_user", 
        "tokens": 100, 
        "temp": 0.5,
        "response": response3.choices[0].message.content,
        "usage": response3.usage
    })
    
    return calls_data

# Run the metrics demonstration
metrics_data = demonstrate_metrics()

print("📊 Metrics Being Tracked by Traceloop:")
print("\n🔍 Automatic Metrics:")
print("  • Latency: Response time for each call")
print("  • Token Usage: Input and output tokens")
print("  • Cost: Estimated cost per call")
print("  • Model Performance: Success/failure rates")
print("  • Prompt-Response Pairs: Complete conversation tracking")

print("\n📈 Quality Metrics:")
print("  • Faithfulness: How well responses match input")
print("  • Relevance: How relevant responses are to prompts")
print("  • Safety: Detection of harmful content")
print("  • Custom Evaluators: Your domain-specific quality measures")

print(f"\n✅ Generated {len(metrics_data)} traced calls with different characteristics!")

# Display actual responses
print("\n" + "="*50)
print("ACTUAL RESPONSES:")
print("="*50)

for i, call in enumerate(metrics_data, 1):
    print(f"\n{i}. {call['type'].replace('_', ' ').title()}:")
    print(f"   Prompt tokens: {call['usage'].prompt_tokens}")
    print(f"   Completion tokens: {call['usage'].completion_tokens}")
    print(f"   Total tokens: {call['usage'].total_tokens}")
    print(f"   Response: {call['response']}")
    print("-" * 50)