# Script para establecer métricas baseline de tu sistema LLM
import time
import openai
from typing import Dict, List
import statistics

def measure_baseline_metrics(
    test_queries: List[str],
    expected_answers: List[str],
    model: str = "gpt-4",
    context: str = ""
) -> Dict[str, float]:
    """
    Mide métricas baseline de performance.
    
    Returns:
        Dict con latency_p50, latency_p95, accuracy, tokens_avg
    """
    latencies = []
    accuracies = []
    token_counts = []

    for query, expected in zip(test_queries, expected_answers):
        # Medir latency
        start = time.time()
        response = openai.ChatCompletion.create(
            model=model,
            messages=[
                {"role": "system", "content": context},
                {"role": "user", "content": query}
            ],
            temperature=0
        )
        latency = time.time() - start
        latencies.append(latency)

        # Contar tokens
        total_tokens = response['usage']['total_tokens']
        token_counts.append(total_tokens)

        # Medir accuracy simple (contiene respuesta esperada)
        answer = response.choices[0].message.content
        accuracy = 1.0 if expected.lower() in answer.lower() else 0.0
        accuracies.append(accuracy)

    return {
        "latency_p50": statistics.median(latencies),
        "latency_p95": statistics.quantiles(latencies, n=20)[18],  # 95th percentile
        "accuracy": sum(accuracies) / len(accuracies),
        "tokens_avg": sum(token_counts) / len(token_counts),
        "tokens_p95": statistics.quantiles(token_counts, n=20)[18]
    }

# Ejemplo de uso
test_queries = [
    "¿Cuál es el código de proyecto mencionado en el contexto?",
    "¿Qué tecnologías se mencionan?",
    # ... más queries de testing
]

expected_answers = [
    "PROJ-2847",
    "Python, LangChain",
    # ... respuestas esperadas
]

# Cargar tu contexto de producción (ej: documentos RAG)
with open("production_context.txt") as f:
    production_context = f.read()

# Medir baseline
baseline = measure_baseline_metrics(
    test_queries=test_queries,
    expected_answers=expected_answers,
    model="gpt-4",
    context=production_context
)

print("BASELINE METRICS:")
print(f"Latency P50: {baseline['latency_p50']:.2f}s")
print(f"Latency P95: {baseline['latency_p95']:.2f}s")
print(f"Accuracy: {baseline['accuracy']*100:.1f}%")
print(f"Tokens Avg: {baseline['tokens_avg']:.0f}")
print(f"Tokens P95: {baseline['tokens_p95']:.0f}")

# Ahora aplica optimizaciones y re-mide para comparar

# Complejidad computacional del self-attention
import numpy as np

def attention_complexity(n_tokens: int, d_model: int = 768, n_heads: int = 12):
    """
    Calcula complejidad computacional y memoria del self-attention.
    
    Args:
        n_tokens: Número de tokens en el contexto
        d_model: Dimensión del modelo (768 para BERT, 1024 para GPT-2, etc)
        n_heads: Número de attention heads
    
    Returns:
        Dict con métricas de complejidad
    """
    # Complejidad computacional: O(n² * d) para single-head
    # O(h * n² * d) para multi-head donde h = n_heads
    flops = n_heads * (n_tokens ** 2) * d_model

    # Memoria para almacenar matriz de atención: O(n²) por head
    attention_matrix_memory = n_heads * (n_tokens ** 2) * 4  # 4 bytes por float32

    # Memoria total incluyendo activations
    total_memory = attention_matrix_memory + (n_tokens * d_model * 4)

    return {
        "flops": flops,
        "attention_matrix_gb": attention_matrix_memory / (1024**3),
        "total_memory_gb": total_memory / (1024**3),
        "complexity_class": "O(n²)"
    }

# Ejemplo: Comparar diferentes tamaños de contexto
contexts = [4096, 32768, 131072]  # 4K, 32K, 128K tokens

print("Complejidad Self-Attention por Tamaño de Contexto:\n")
for n in contexts:
    metrics = attention_complexity(n)
    print(f"{n:,} tokens:")
    print(f"  FLOPs: {metrics['flops']:,.0f}")
    print(f"  Memoria atención: {metrics['attention_matrix_gb']:.2f} GB")
    print(f"  Memoria total: {metrics['total_memory_gb']:.2f} GB\n")

# Salida:
# 4,096 tokens:
#   FLOPs: 1,528,823,808
#   Memoria atención: 0.00 GB
#   Memoria total: 0.03 GB
#
# 32,768 tokens:
#   FLOPs: 98,304,000,000
#   Memoria atención: 0.15 GB ← 64x más que 4K
#   Memoria total: 1.15 GB
#
# 128,072 tokens:
#   FLOPs: 1,572,864,000,000
#   Memoria atención: 2.34 GB ← 1024x más que 4K
#   Memoria total: 18.34 GB

# Implementación LLMLingua para prompt compression
from llmlingua import PromptCompressor

# Inicializar compressor (usa modelo pequeño local)
compressor = PromptCompressor(
    model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
    use_llmlingua2=True,  # Usar versión 2 (mejor performance)
    device_map="cuda"  # o "cpu" si no tienes GPU
)

def compress_rag_context(
    documents: list[str],
    query: str,
    target_compression: float = 5.0  # 5x compression
) -> dict:
    """
    Comprime contexto RAG manteniendo información relevante.
    
    Args:
        documents: Lista de documentos recuperados
        query: Query del usuario
        target_compression: Factor de compresión deseado (5.0 = 5x)
    
    Returns:
        Dict con prompt original, comprimido, y métricas
    """
    # Construir contexto original
    original_context = "\n\n".join([
        f"Documento {i+1}:\n{doc}"
        for i, doc in enumerate(documents)
    ])

    original_prompt = f"""Basándote en el siguiente contexto, responde la pregunta.

Contexto:
{original_context}

Pregunta: {query}
"""

    # Comprimir usando LLMLingua
    compressed = compressor.compress_prompt(
        original_prompt,
        instruction="",  # No instruction prefix
        question=query,
        target_token=len(original_prompt.split()) // target_compression,
        condition_compare=True,  # Mejor calidad
        condition_in_question="after",
        rank_method="longllmlingua",  # Método optimizado para contextos largos
        use_sentence_level_filter=True,
        context_budget="+100",
        dynamic_context_compression_ratio=0.3,
        reorder_context="sort"  # Reordenar por relevancia
    )

    return {
        "original": original_prompt,
        "compressed": compressed["compressed_prompt"],
        "original_tokens": len(original_prompt.split()),
        "compressed_tokens": len(compressed["compressed_prompt"].split()),
        "compression_ratio": compressed["ratio"],
        "saving": compressed["saving"]
    }

# Ejemplo de uso real
documents_from_rag = [
    "El proyecto PROJ-2847 fue iniciado en Q1 2025 con un presupuesto de desarrollo de infraestructura cloud...",
    "Documentación técnica completa sobre arquitectura de microservicios y patrones de deployment...",
    "Guía de troubleshooting para errores comunes en producción incluyendo timeouts y fallos de conexión...",
    # ... 47 documentos más (context bloat típico)
]

user_query = "¿Cuál es el código de proyecto para Q1 2025?"

result = compress_rag_context(
    documents=documents_from_rag,
    query=user_query,
    target_compression=5.0  # Comprimir 5x
)

print(f"Original tokens: {result['original_tokens']}")
print(f"Compressed tokens: {result['compressed_tokens']}")
print(f"Compression ratio: {result['compression_ratio']:.1f}x")
print(f"Cost saving: {result['saving']:.1%}")
print(f"\nCompressed prompt:\n{result['compressed']}")

# Extractive compression usando sentence ranking
from sentence_transformers import SentenceTransformer, util
import numpy as np

def extractive_compress(
    documents: list[str],
    query: str,
    target_sentences: int = 10
) -> str:
    """
    Extrae las N oraciones más relevantes del contexto.
    
    Args:
        documents: Documentos originales
        query: Query del usuario
        target_sentences: Número de oraciones a mantener
    
    Returns:
        Contexto comprimido con solo oraciones más relevantes
    """
    # Cargar modelo de embeddings
    model = SentenceTransformer('all-MiniLM-L6-v2')

    # Dividir documentos en oraciones
    sentences = []
    for doc in documents:
        doc_sentences = doc.split('. ')
        sentences.extend([s.strip() + '.' for s in doc_sentences if len(s.strip()) > 10])

    # Generar embeddings
    query_embedding = model.encode(query, convert_to_tensor=True)
    sentence_embeddings = model.encode(sentences, convert_to_tensor=True)

    # Calcular similaridad cosine
    similarities = util.cos_sim(query_embedding, sentence_embeddings)[0]

    # Rankear y seleccionar top N
    top_indices = np.argsort(similarities.cpu().numpy())[::-1][:target_sentences]
    top_sentences = [sentences[i] for i in sorted(top_indices)]

    # Reconstruir contexto comprimido
    compressed_context = ' '.join(top_sentences)

    return compressed_context

# Ejemplo
compressed = extractive_compress(
    documents=documents_from_rag,
    query="¿Cuál es el código de proyecto para Q1 2025?",
    target_sentences=10  # Solo 10 oraciones más relevantes
)

print(f"Compressed context ({len(compressed.split())} tokens):\n{compressed}")

# Implementación Claude Prompt Caching para RAG
import anthropic

client = anthropic.Anthropic(api_key="your-api-key")

def rag_with_caching(
    knowledge_base: str,  # Tu documentación/KB completa (cacheable)
    user_query: str  # Query específico del usuario (no cacheable)
) -> str:
    """
    RAG system usando Claude con prompt caching.
    
    La KB se cachea (10x más barato en requests subsecuentes),
    solo pagas full price por el user query nuevo cada vez.
    """
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=[
            {
                "type": "text",
                "text": "Eres un asistente experto. Responde basándote solo en el contexto proporcionado."
            },
            {
                "type": "text",
                "text": f"Contexto (Knowledge Base):\n\n{knowledge_base}",
                "cache_control": {"type": "ephemeral"}  # ← Esto cachea el KB
            }
        ],
        messages=[
            {
                "role": "user",
                "content": user_query
            }
        ]
    )

    # Inspeccionar uso de cache
    usage = response.usage
    print(f"Input tokens: {usage.input_tokens}")
    print(f"Cache creation tokens: {usage.cache_creation_input_tokens}")
    print(f"Cache read tokens: {usage.cache_read_input_tokens}")

    return response.content[0].text

# Ejemplo real: Sistema de soporte técnico
kb_documentation = """
[Tu documentación técnica completa aquí - puede ser 50K+ tokens]

# Troubleshooting Guide

## Error: Connection Timeout
Causa: Network issues o firewall blocking...
Solución: Verificar...

## Error: Authentication Failed
Causa: Credentials incorrectas o token expirado...
Solución: Regenerar token...

[... miles de líneas más de documentación ...]
"""

# Primera query - crea cache (costo full)
answer1 = rag_with_caching(
    knowledge_base=kb_documentation,
    user_query="¿Cómo soluciono connection timeout?"
)
# Output: Cache creation tokens: 50,000 (costo normal)

# Segunda query - usa cache (10x más barato)
answer2 = rag_with_caching(
    knowledge_base=kb_documentation,  # ← MISMO KB, cache hit
    user_query="¿Qué hacer si authentication falla?"
)
# Output: Cache read tokens: 50,000 (10x más barato)
#         Input tokens: 15 (solo el query nuevo, costo normal)

# RAG Híbrido: Combina retrieval + long context + caching + compression
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from sentence_transformers import CrossEncoder
import anthropic

class HybridRAGSystem:
    """
    Sistema RAG híbrido production-ready que combina:
    1. Semantic search (recupera top 50 chunks)
    2. Re-ranking (refina a top 5 más relevantes)
    3. Prompt compression (reduce 5x)
    4. Context caching (Claude, 10x cheaper)
    5. Long context selectivo (para multi-hop si necesario)
    """

    def __init__(
        self,
        pinecone_index: str,
        reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    ):
        # Vector store para semantic search
        self.vectorstore = Pinecone.from_existing_index(
            index_name=pinecone_index,
            embedding=OpenAIEmbeddings()
        )

        # Re-ranker para afinar resultados
        self.reranker = CrossEncoder(reranker_model)

        # Cliente Claude con caching
        self.claude = anthropic.Anthropic()

    def retrieve_and_rerank(
        self,
        query: str,
        initial_k: int = 50,  # Retrieve mucho inicialmente
        final_k: int = 5  # Refinar a top 5 después de re-ranking
    ) -> list[str]:
        """
        Paso 1: Semantic search + Re-ranking.
        """
        # Semantic search: top 50 chunks por cosine similarity
        initial_docs = self.vectorstore.similarity_search(query, k=initial_k)

        # Re-ranking: usar cross-encoder para scoring más preciso
        doc_texts = [doc.page_content for doc in initial_docs]
        rerank_scores = self.reranker.predict([
            [query, doc] for doc in doc_texts
        ])

        # Ordenar por re-rank score y tomar top 5
        ranked_indices = sorted(
            range(len(rerank_scores)),
            key=lambda i: rerank_scores[i],
            reverse=True
        )[:final_k]

        top_docs = [doc_texts[i] for i in ranked_indices]
        return top_docs

    def compress_context(
        self,
        documents: list[str],
        target_compression: float = 5.0
    ) -> str:
        """
        Paso 2: Prompt compression (5x).
        (Aquí usarías LLMLingua, simplificado por brevedad)
        """
        # En producción: usar LLMLingua aquí
        # Por ahora, concatenación simple
        context = "\n\n".join([
            f"Documento {i+1}:\n{doc}"
            for i, doc in enumerate(documents)
        ])

        # Simular compression (en producción usar LLMLingua)
        # compressed = llmlingua.compress(context, target_compression)

        return context  # En producción retornar compressed

    def query_with_caching(
        self,
        context: str,
        query: str,
        use_long_context: bool = False
    ) -> str:
        """
        Paso 3: Query LLM con context caching.
        
        Args:
            context: Contexto comprimido y re-ranked
            query: User query
            use_long_context: Si True, usa long context approach (>32K tokens)
        """
        # Determinar estrategia basado en context size
        context_tokens = len(context.split())

        if use_long_context and context_tokens > 32000:
            # Estrategia long context: dar todo pero con caching
            system_prompt = f"""Eres un asistente experto. Analiza el siguiente contexto extenso y responde la pregunta del usuario.

Contexto completo (cached):
{context}"""
        else:
            # Estrategia RAG estándar: contexto optimizado
            system_prompt = f"""Eres un asistente experto. Responde basándote en los documentos más relevantes proporcionados.

Documentos relevantes (cached):
{context}"""

        response = self.claude.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=2048,
            system=[
                {
                    "type": "text",
                    "text": system_prompt,
                    "cache_control": {"type": "ephemeral"}  # Cache el contexto
                }
            ],
            messages=[
                {"role": "user", "content": query}
            ]
        )

        return response.content[0].text

    def answer_query(
        self,
        query: str,
        use_long_context: bool = False
    ) -> dict:
        """
        Pipeline completo: Retrieve → Re-rank → Compress → Cache → Answer
        """
        # 1. Retrieve + Re-rank
        top_docs = self.retrieve_and_rerank(query, initial_k=50, final_k=5)

        # 2. Compress context
        context = self.compress_context(top_docs, target_compression=5.0)

        # 3. Query con caching
        answer = self.query_with_caching(context, query, use_long_context)

        return {
            "answer": answer,
            "sources": top_docs,
            "context_tokens": len(context.split())
        }

# Ejemplo de uso
rag_system = HybridRAGSystem(pinecone_index="knowledge-base-prod")

result = rag_system.answer_query(
    query="¿Cuáles son las mejores prácticas para reducir latency en LLMs?",
    use_long_context=False  # RAG optimizado
)

print(result["answer"])
print(f"\nTokens procesados: {result['context_tokens']}")
print(f"Fuentes: {len(result['sources'])} documentos")

# Middleware de monitoring para tracking completo de métricas LLM
import time
import tiktoken
from prometheus_client import Histogram, Counter, Gauge
import structlog

logger = structlog.get_logger()

# Prometheus metrics
llm_request_duration = Histogram(
    'llm_request_duration_seconds',
    'LLM request duration',
    buckets=[0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0]
)

llm_tokens_total = Histogram(
    'llm_tokens_total',
    'Total tokens per request',
    labelnames=['type'],  # input/output
    buckets=[100, 500, 1000, 5000, 10000, 32000, 64000, 128000]
)

llm_cost_total = Counter(
    'llm_cost_dollars_total',
    'Cumulative API cost'
)

llm_errors_total = Counter(
    'llm_errors_total',
    'Total LLM errors',
    labelnames=['error_type']
)

llm_cache_hits = Counter(
    'llm_cache_hits_total',
    'Total cache hits'
)

class LLMMonitoringMiddleware:
    """
    Middleware para tracking exhaustivo de métricas LLM.
    Wrap todas las llamadas a LLM APIs.
    """

    def __init__(self, pricing_model: dict):
        """
        Args:
            pricing_model: Dict mapeando modelo → (input_cost, output_cost) por 1M tokens
        """
        self.pricing = pricing_model
        self.encoder = tiktoken.encoding_for_model("gpt-4")  # Para contar tokens

    def track_llm_call(
        self,
        llm_function,
        model: str,
        messages: list,
        **kwargs
    ):
        """
        Wrapper que trackea métricas de cualquier llamada LLM.
        
        Usage:
            response = middleware.track_llm_call(
                openai.ChatCompletion.create,
                model="gpt-4",
                messages=[...]
            )
        """
        # Contar tokens input
        input_text = "\n".join([m['content'] for m in messages])
        input_tokens = len(self.encoder.encode(input_text))

        # Medir latency
        start_time = time.time()

        try:
            # Ejecutar llamada real
            response = llm_function(model=model, messages=messages, **kwargs)
            latency = time.time() - start_time

            # Extraer métricas de response
            output_tokens = response['usage']['completion_tokens']
            total_tokens = response['usage']['total_tokens']

            # Detectar cache hit (si disponible en response)
            cache_hit = response.get('usage', {}).get('cache_read_input_tokens', 0) > 0

            # Calcular costo
            input_cost_per_1m, output_cost_per_1m = self.pricing.get(model, (0, 0))
            cost = (
                (input_tokens / 1_000_000) * input_cost_per_1m +
                (output_tokens / 1_000_000) * output_cost_per_1m
            )

            # Registrar métricas
            llm_request_duration.observe(latency)
            llm_tokens_total.labels(type='input').observe(input_tokens)
            llm_tokens_total.labels(type='output').observe(output_tokens)
            llm_cost_total.inc(cost)

            if cache_hit:
                llm_cache_hits.inc()

            # Structured logging
            logger.info(
                "llm_request_complete",
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                latency=latency,
                cost=cost,
                cache_hit=cache_hit
            )

            # Alertar si thresholds excedidos
            if total_tokens > 32000:
                logger.warning(
                    "high_token_usage",
                    tokens=total_tokens,
                    threshold=32000,
                    message="Entering context rot danger zone"
                )

            if latency > 5.0:
                logger.warning(
                    "high_latency",
                    latency=latency,
                    threshold=5.0,
                    message="Latency exceeds UX acceptable threshold"
                )

            return response

        except Exception as e:
            latency = time.time() - start_time

            # Track error
            error_type = type(e).__name__
            llm_errors_total.labels(error_type=error_type).inc()

            logger.error(
                "llm_request_failed",
                model=model,
                input_tokens=input_tokens,
                latency=latency,
                error=str(e),
                error_type=error_type
            )

            raise

# Ejemplo de uso
pricing = {
    "gpt-4": (30, 60),  # Estimado: input/output cost por 1M tokens (actualizar con pricing real)
    "claude-3-5-sonnet": (15, 75),
    "gemini-1.5-pro": (10, 30)
}

middleware = LLMMonitoringMiddleware(pricing_model=pricing)

# Wrap tus llamadas LLM
response = middleware.track_llm_call(
    openai.ChatCompletion.create,
    model="gpt-4",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Explain context rot"}
    ]
)

# Ahora todas las métricas están siendo tracked automáticamente
# Exponer endpoint /metrics para Prometheus scraping