import numpy as np
from scipy import stats
from scipy.spatial.distance import jensenshannon


# ALGORITMO #1: Population Stability Index (PSI)
# Mejor para: Detecting embedding distribution shifts
# Threshold: PSI >0.1 investigate, >0.25 critical

def calculate_psi(baseline, current, bins=10):
    """
    Calcula PSI entre dos distribuciones.
    
    Args:
        baseline: Array embeddings baseline period (e.g., primera semana producción)
        current: Array embeddings current period (última semana)
        bins: Número de bins para discretizar distribución
        
    Returns:
        psi_score: Float, donde >0.1 indica drift significativo
    """
    # Discretizar en bins
    breakpoints = np.quantile(baseline, np.linspace(0, 1, bins + 1))
    baseline_counts = np.histogram(baseline, bins=breakpoints)[0]
    current_counts = np.histogram(current, bins=breakpoints)[0]

    # Calcular proporciones (evitar división por cero)
    baseline_props = baseline_counts / len(baseline)
    current_props = current_counts / len(current)

    baseline_props = np.where(baseline_props == 0, 0.0001, baseline_props)
    current_props = np.where(current_props == 0, 0.0001, current_props)

    # PSI formula
    psi = np.sum(
        (current_props - baseline_props) * np.log(current_props / baseline_props)
    )

    return psi


# Ejemplo uso:
baseline_embeddings = np.random.normal(0, 1, 1000)      # Semana 1 producción
current_embeddings = np.random.normal(0.3, 1.2, 1000)  # Semana 4 (drift!)

psi_score = calculate_psi(baseline_embeddings, current_embeddings)
print(f"PSI Score: {psi_score:.4f}")

if psi_score > 0.25:
    print("🚨 CRITICAL DRIFT DETECTED - Re-index vector DB urgente")
elif psi_score > 0.1:
    print("⚠️ Moderate drift - Investigar root cause")


# ALGORITMO #2: Kolmogorov-Smirnov Test (KS Test)
# Mejor para: Detecting changes en retrieval score distributions
# Threshold: p-value 

{
"trace_id": "abc123",
"timestamp": "2025-12-17T10:23:41Z",
"input": {
    "query": "¿Dónde está mi pedido #12345?",
    "user_id": "user_789",
    "session_context": "3rd message in conversation"
},
"reasoning_steps": [
    {
    "step": 1,
    "thought": "Usuario pregunta por pedido específico (#12345). Necesito verificar status actual.",
    "available_tools": ["Search", "QueryDB", "CheckOrderStatus", "Escalate", "SendEmail"],
    "tool_selected": "CheckOrderStatus",
    "selection_reasoning": "Query contiene order number explícito. CheckOrderStatus es tool específico para tracking.",
    "confidence_score": 0.95
    },
    {
    "step": 2,
    "action": "CheckOrderStatus(order_id='12345')",
    "observation": "Order #12345 - Status: In Transit, ETA: 2025-12-19, Carrier: UPS",
    "latency_ms": 850
    },
    {
    "step": 3,
    "thought": "Tengo información completa. Puedo responder directamente sin escalation.",
    "final_response": "Tu pedido #12345 está en tránsito con UPS. Llegará el 19 de diciembre.",
    "hallucination_check": "PASS - Información verificada de DB",
    "pii_detected": "order_number",
    "safety_score": 1.0
    }
],
"evaluation": {
    "goal_achieved": true,
    "reasoning_quality": 0.92,
    "tool_selection_accuracy": 1.0,
    "output_relevance": 0.98,
    "response_time_ms": 1850,
    "tokens_consumed": 234,
    "cost_usd": 0.00047
},
"metadata": {
    "model": "gpt-4-turbo",
    "framework": "langchain",
    "agent_version": "v2.3.1"
}
}

# OpenTelemetry core
opentelemetry-api==1.22.0
opentelemetry-sdk==1.22.0
opentelemetry-instrumentation==0.43b0

# Exporters para diferentes backends
opentelemetry-exporter-otlp-proto-grpc==1.22.0  # Tempo/Jaeger
opentelemetry-exporter-prometheus==0.43b0       # Prometheus

# Instrumentación automática para frameworks comunes
opentelemetry-instrumentation-requests==0.43b0
opentelemetry-instrumentation-urllib3==0.43b0

# LangChain integration (oficial desde v0.1.0)
langchain>=0.1.0
langchain-openai>=0.0.5

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 opentelemetry.sdk.resources import Resource


def setup_observability(service_name: str = "agentic-ai-app"):
    """
    Configura OpenTelemetry para agentic AI observability.
    
    Backend: Grafana Tempo (OTLP receiver en localhost:4317)
    Metrics: Prometheus (scraping /metrics endpoint)
    """
    # Resource attributes (identifican tu servicio)
    resource = Resource.create({
        "service.name": service_name,
        "service.version": "1.0.0",
        "deployment.environment": "production",
        "ai.framework": "langchain",
        "ai.model.provider": "openai"
    })

    # Configurar tracer provider
    tracer_provider = TracerProvider(resource=resource)

    # OTLP exporter para Tempo/Jaeger (gRPC endpoint)
    otlp_exporter = OTLPSpanExporter(
        endpoint="http://localhost:4317",  # Tempo gRPC receiver
        insecure=True  # En prod usar TLS
    )

    # Batch span processor (agrupa spans antes de exportar para efficiency)
    span_processor = BatchSpanProcessor(
        otlp_exporter,
        max_queue_size=2048,
        max_export_batch_size=512,
        export_timeout_millis=30000
    )

    tracer_provider.add_span_processor(span_processor)

    # Setear como tracer global
    trace.set_tracer_provider(tracer_provider)

    print(f"✅ OpenTelemetry configurado para {service_name}")
    print(f"📊 Traces exportando a: http://localhost:4317")
    print(f"🔍 Visualiza en Grafana: http://localhost:3000")


# Inicializar al startup de aplicación
if __name__ == "__main__":
    setup_observability("customer-support-agent")

from opentelemetry import trace
from langchain.agents import AgentExecutor, create_react_agent
from langchain_openai import ChatOpenAI
from langchain.tools import Tool
import time


# Obtener tracer (ya configurado en otel_config.py)
tracer = trace.get_tracer(__name__)


class InstrumentedAgentExecutor:
    """
    Wrapper para LangChain AgentExecutor con full OpenTelemetry instrumentation.
    Implementa semantic conventions GenAI.
    """

    def __init__(self, agent_executor: AgentExecutor):
        self.agent_executor = agent_executor

    def invoke(self, input_data: dict):
        """
        Ejecuta agent con full tracing de reasoning pathway.
        """
        # ROOT SPAN: Representa interacción completa usuario → agent → respuesta
        with tracer.start_as_current_span(
            "agent.interaction",
            kind=trace.SpanKind.SERVER
        ) as root_span:

            # Semantic conventions GenAI attributes
            root_span.set_attribute("ai.operation.name", "agent_reasoning")
            root_span.set_attribute("ai.request.model", "gpt-4-turbo")
            root_span.set_attribute("ai.request.temperature", 0.0)
            root_span.set_attribute("ai.request.max_tokens", 2000)

            # Input attributes
            user_query = input_data.get("input", "")
            root_span.set_attribute("ai.prompt", user_query[:200])  # Truncar largo
            root_span.set_attribute("ai.prompt.length", len(user_query))

            start_time = time.time()

            try:
                # SPAN HIJO: Agent reasoning loop
                with tracer.start_as_current_span("agent.reasoning_loop") as reasoning_span:
                    reasoning_span.set_attribute("ai.agent.framework", "langchain")
                    reasoning_span.set_attribute("ai.agent.pattern", "react")

                    # Ejecutar agent (LangChain maneja reasoning internamente)
                    result = self.agent_executor.invoke(input_data)

                    # Capturar intermediate_steps (thought → action → observation)
                    if "intermediate_steps" in result:
                        self._trace_reasoning_steps(result["intermediate_steps"])

                # Output attributes
                output = result.get("output", "")
                root_span.set_attribute("ai.response", output[:200])
                root_span.set_attribute("ai.response.length", len(output))

                # Métricas de performance
                latency_ms = (time.time() - start_time) * 1000
                root_span.set_attribute("ai.response.latency_ms", latency_ms)

                # Token tracking (si disponible del LLM)
                if hasattr(result, "token_usage"):
                    root_span.set_attribute(
                        "ai.usage.prompt_tokens",
                        result.token_usage.prompt_tokens
                    )
                    root_span.set_attribute(
                        "ai.usage.completion_tokens",
                        result.token_usage.completion_tokens
                    )
                    root_span.set_attribute(
                        "ai.usage.total_tokens",
                        result.token_usage.total_tokens
                    )

                # Status success
                root_span.set_status(trace.Status(trace.StatusCode.OK))

                return result

            except Exception as e:
                # Capturar errores como span events
                root_span.record_exception(e)
                root_span.set_status(
                    trace.Status(trace.StatusCode.ERROR, str(e))
                )
                raise

    def _trace_reasoning_steps(self, intermediate_steps):
        """
        Crea spans para cada step del reasoning pathway
        (thought → action → observation).
        """
        for idx, (agent_action, observation) in enumerate(intermediate_steps):
            with tracer.start_as_current_span(f"agent.step.{idx+1}") as step_span:
                # Tool selection
                step_span.set_attribute("ai.tool.name", agent_action.tool)
                step_span.set_attribute(
                    "ai.tool.input",
                    str(agent_action.tool_input)[:100]
                )

                # Reasoning (log del agent sobre por qué eligió esta tool)
                if hasattr(agent_action, "log"):
                    step_span.set_attribute("ai.reasoning", agent_action.log[:200])

                # Observation (resultado de tool execution)
                step_span.set_attribute("ai.observation", str(observation)[:200])

                # Event para marcar tool invocation
                step_span.add_event(
                    "tool.invocation",
                    attributes={
                        "tool.name": agent_action.tool,
                        "tool.success": observation is not None
                    }
                )


# EJEMPLO USO
if __name__ == "__main__":
    # Setup OpenTelemetry (del paso anterior)
    from otel_config import setup_observability
    setup_observability("customer-support-agent")

    # Crear LangChain agent normal
    llm = ChatOpenAI(model="gpt-4-turbo", temperature=0)

    tools = [
        Tool(
            name="Search",
            func=lambda x: f"Documentación encontrada: {x}",
            description="Busca en knowledge base"
        ),
        Tool(
            name="QueryDB",
            func=lambda x: f"Resultado DB: {x}",
            description="Query customer database"
        )
    ]

    agent = create_react_agent(llm, tools, prompt_template)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    # Wrap con instrumentación
    instrumented_agent = InstrumentedAgentExecutor(agent_executor)

    # Ejecutar con full tracing
    result = instrumented_agent.invoke({
        "input": "¿Cuál es el status del pedido #12345?"
    })

    print(f"✅ Respuesta: {result['output']}")
    print(f"📊 Traces disponibles en Grafana Tempo")

version: '3.8'

services:
  # Grafana Tempo: Distributed tracing backend
  tempo:
    image: grafana/tempo:latest
    command: ["-config.file=/etc/tempo.yaml"]
    volumes:
      - ./tempo-config.yaml:/etc/tempo.yaml
      - tempo-data:/tmp/tempo
    ports:
      - "4317:4317"  # OTLP gRPC receiver
      - "3200:3200"  # Tempo API

  # Prometheus: Metrics collection
  prometheus:
    image: prom/prometheus:latest
    volumes:
      - ./prometheus.yml:/etc/prometheus/prometheus.yml
      - prometheus-data:/prometheus
    ports:
      - "9090:9090"
    command:
      - '--config.file=/etc/prometheus/prometheus.yml'
      - '--storage.tsdb.retention.time=30d'

  # Grafana: Unified visualization
  grafana:
    image: grafana/grafana:latest
    volumes:
      - grafana-data:/var/lib/grafana
      - ./grafana-datasources.yml:/etc/grafana/provisioning/datasources/datasources.yml
    environment:
      - GF_AUTH_ANONYMOUS_ENABLED=true
      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
    ports:
      - "3000:3000"
    depends_on:
      - tempo
      - prometheus

volumes:
  tempo-data:
  prometheus-data:
  grafana-data:

from langchain_openai import ChatOpenAI
from typing import Literal


class ModelRouter:
    """
    Ruta queries a modelo apropiado basado en complexity analysis.
    Savings típicos: 40-50% en LLM costs.
    """

    # Pricing por 1M tokens (OpenAI Dec 2025)
    PRICING = {
        "gpt-4-turbo": {"input": 10.0, "output": 30.0},   # Premium
        "gpt-4o-mini": {"input": 0.15, "output": 0.60},   # Budget
        "gpt-3.5-turbo": {"input": 0.50, "output": 1.50}  # Legacy
    }

    def __init__(self):
        self.models = {
            "premium": ChatOpenAI(model="gpt-4-turbo", temperature=0),
            "standard": ChatOpenAI(model="gpt-4o-mini", temperature=0),
            "budget": ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
        }

    def route(self, query: str) -> Literal["premium", "standard", "budget"]:
        """
        Analiza complexity de query y retorna tier apropiado.
        """
        query_lower = query.lower()

        # PREMIUM (GPT-4 Turbo): Multi-step reasoning, complex analysis
        premium_indicators = [
            "analiza", "compara", "evalúa", "recomienda", "estrategia",
            "por qué", "explica detalladamente", "pros y contras"
        ]
        if any(ind in query_lower for ind in premium_indicators):
            return "premium"

        # BUDGET (GPT-3.5): Simple Q&A, fact retrieval
        budget_indicators = [
            "qué es", "cuándo", "dónde", "quién",
            "status", "check", "busca", "encuentra"
        ]
        if any(ind in query_lower for ind in budget_indicators):
            return "budget"

        # STANDARD (GPT-4o Mini): Default para queries moderadas
        return "standard"

    def estimate_cost(self, query: str, response: str, tier: str) -> float:
        """
        Estima cost de interacción en USD.
        """
        # Estimación tokens (rough: 1 token ≈ 4 chars)
        input_tokens = len(query) / 4
        output_tokens = len(response) / 4

        pricing = self.PRICING[
            "gpt-4-turbo" if tier == "premium"
            else "gpt-4o-mini" if tier == "standard"
            else "gpt-3.5-turbo"
        ]

        cost = (
            (input_tokens / 1_000_000) * pricing["input"] +
            (output_tokens / 1_000_000) * pricing["output"]
        )

        return cost


# Ejemplo uso con tracking
router = ModelRouter()

queries = [
    "¿Cuál es el status del pedido #12345?",  # → budget (simple lookup)
    "Analiza por qué nuestro churn rate aumentó 15% y recomienda 3 estrategias"  # → premium
]

for query in queries:
    tier = router.route(query)
    model = router.models[tier]
    response = model.invoke(query).content
    cost = router.estimate_cost(query, response, tier)

    print(f"Query: {query[:50]}...")
    print(f"Tier: {tier} | Cost: ${cost:.6f}")
    print(f"Savings vs always-premium: {((0.00034 - cost) / 0.00034 * 100):.1f}%\n")