from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from traceloop.sdk import Traceloop

# Initialize with automatic instrumentation
Traceloop.init(
    app_name="my-ai-app",
    api_endpoint="https://otel-collector:4317",
    headers={"Authorization": "Bearer YOUR_TOKEN"},
)

# Now all LLM calls are automatically traced
from anthropic import Anthropic

client = Anthropic()

@trace.get_tracer(__name__).start_as_current_span("chat_handler")
def handle_chat(user_message: str, user_id: str) -> str:
    span = trace.get_current_span()
    span.set_attribute("user.id", user_id)
    span.set_attribute("user.message.length", len(user_message))
    
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=1024,
        messages=[{"role": "user", "content": user_message}]
    )
    
    # Automatically captured by OpenLLMetry:
    # - model name, prompt tokens, completion tokens
    # - latency, TTFT
    # - full prompt + response (if enabled)
    
    return response.content[0].text

from opentelemetry import metrics
from opentelemetry.sdk.metrics import MeterProvider
import time

meter = metrics.get_meter("llm-app")

# Core counters
request_counter = meter.create_counter(
    "llm.requests.total",
    description="Total LLM requests"
)

token_counter = meter.create_counter(
    "llm.tokens.total",
    description="Total tokens consumed"
)

# Histograms for latency
latency_histogram = meter.create_histogram(
    "llm.request.duration",
    description="LLM request duration in seconds",
    unit="s"
)

ttft_histogram = meter.create_histogram(
    "llm.time_to_first_token",
    description="Time to first token",
    unit="s"
)

# Cost tracking
cost_counter = meter.create_counter(
    "llm.cost.usd",
    description="Estimated cost in USD"
)

def tracked_llm_call(messages: list, model: str = "claude-sonnet-4-5"):
    start = time.time()
    
    request_counter.add(1, {"model": model, "env": "production"})
    
    response = client.messages.create(
        model=model,
        max_tokens=2048,
        messages=messages
    )
    
    duration = time.time() - start
    
    # Track metrics
    latency_histogram.record(duration, {"model": model})
    
    token_counter.add(
        response.usage.input_tokens + response.usage.output_tokens,
        {"model": model, "type": "total"}
    )
    
    # Estimate cost (claude-sonnet-4-5 pricing)
    cost = (response.usage.input_tokens * 0.000003 + 
            response.usage.output_tokens * 0.000015)
    cost_counter.add(cost, {"model": model})
    
    return response

import structlog
import json

logger = structlog.get_logger()

def log_llm_interaction(
    request_id: str,
    model: str,
    prompt_tokens: int,
    completion_tokens: int,
    latency_ms: float,
    finish_reason: str,
    user_id: str = None,
    session_id: str = None,
    # Enable prompt/response logging only in debug mode
    log_content: bool = False,
    prompt: str = None,
    response: str = None,
):
    log_data = {
        "event": "llm_call",
        "request_id": request_id,
        "model": model,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
        "latency_ms": round(latency_ms, 2),
        "finish_reason": finish_reason,
        "estimated_cost_usd": round(
            prompt_tokens * 0.000003 + completion_tokens * 0.000015, 6
        ),
    }
    
    if user_id:
        log_data["user_id"] = user_id
    if session_id:
        log_data["session_id"] = session_id
    
    # ⚠️ IMPORTANT: Be careful about PII in prompts
    if log_content and prompt:
        # Truncate and sanitize before logging
        log_data["prompt_preview"] = prompt[:200]
        log_data["response_preview"] = response[:200] if response else None
    
    logger.info("llm_interaction", **log_data)

from anthropic import Anthropic
from dataclasses import dataclass
from enum import Enum

class EvalResult(Enum):
    PASS = "pass"
    FAIL = "fail"
    UNCLEAR = "unclear"

@dataclass
class QualityScore:
    result: EvalResult
    score: float  # 0-1
    reasoning: str
    flags: list[str]

class LLMEvaluator:
    """Use a secondary LLM to evaluate primary LLM outputs."""
    
    def __init__(self):
        self.client = Anthropic()
    
    def evaluate_response(
        self,
        user_query: str,
        llm_response: str,
        expected_properties: list[str] = None
    ) -> QualityScore:
        
        eval_prompt = f"""You are a quality evaluator for an AI assistant. 
Evaluate the following AI response for quality issues.

USER QUERY:
{user_query}

AI RESPONSE:
{llm_response}

Check for:
1. Factual accuracy and hallucinations
2. Relevance to the query
3. Completeness
4. Safety (harmful content, PII exposure)
5. Helpfulness

Return a JSON object:
result"""
        
        eval_response = self.client.messages.create(
            model="claude-haiku-3-5",  # Use cheaper model for evals
            max_tokens=512,
            messages=[{"role": "user", "content": eval_prompt}]
        )
        
        result = json.loads(eval_response.content[0].text)
        return QualityScore(**result)
    
    def check_hallucination(
        self, 
        response: str, 
        source_documents: list[str]
    ) -> float:
        """Returns hallucination score: 0 = grounded, 1 = hallucinated."""
        
        sources_text = "\n\n---\n\n".join(source_documents[:5])
        
        prompt = f"""Given these source documents and an AI response, 
determine what fraction of the response's factual claims are NOT supported by the sources.

SOURCES:
{sources_text}

AI RESPONSE:
{response}

Return only a number between 0.0 (fully grounded) and 1.0 (fully hallucinated)."""
        
        result = self.client.messages.create(
            model="claude-haiku-3-5",
            max_tokens=10,
            messages=[{"role": "user", "content": prompt}]
        )
        
        try:
            return float(result.content[0].text.strip())
        except ValueError:
            return 0.5  # Unknown

import hashlib
from collections import defaultdict
from typing import Optional

class PromptAnalytics:
    """Track performance by prompt template."""
    
    def __init__(self, db_client):
        self.db = db_client
    
    def get_prompt_hash(self, template: str) -> str:
        """Create stable hash for a prompt template (before variable substitution)."""
        return hashlib.md5(template.encode()).hexdigest()[:8]
    
    async def record_interaction(
        self,
        prompt_template: str,
        variables: dict,
        response: str,
        latency_ms: float,
        tokens: int,
        eval_score: Optional[float] = None,
        error: Optional[str] = None
    ):
        template_hash = self.get_prompt_hash(prompt_template)
        
        await self.db.insert("llm_interactions", {
            "template_hash": template_hash,
            "variables_json": json.dumps(variables),
            "response_preview": response[:500],
            "latency_ms": latency_ms,
            "tokens": tokens,
            "eval_score": eval_score,
            "error": error,
            "timestamp": datetime.utcnow()
        })
    
    async def get_worst_performing_templates(self, hours: int = 24):
        """Find prompt templates with lowest quality scores."""
        return await self.db.query("""
            SELECT 
                template_hash,
                COUNT(*) as call_count,
                AVG(eval_score) as avg_quality,
                AVG(latency_ms) as avg_latency,
                SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END) as error_count
            FROM llm_interactions
            WHERE timestamp > NOW() - INTERVAL '%s hours'
            GROUP BY template_hash
            ORDER BY avg_quality ASC
            LIMIT 10
        """, hours)