# LlamaIndex: Simple y potente para RAG
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.llms import OpenAI

# Carga documentos con parsing inteligente
documents = SimpleDirectoryReader("./data").load_data()

# Crea índice con chunking optimizado
index = VectorStoreIndex.from_documents(documents)

# Query engine con configuración explícita
query_engine = index.as_query_engine(
    llm=OpenAI(model="gpt-4", temperature=0),
    similarity_top_k=3,
    response_mode="tree_summarize"  # Reduce tokens vs LangChain
)

# Query con tracing automático
response = query_engine.query("¿Qué es RAG?")

# Response incluye:
# - Texto generado
# - Source nodes con scores
# - Token usage exacto
# - Latency metrics
print(f"Tokens usados: {response.metadata['token_usage']}")
print(f"Sources: {[node.node.metadata for node in response.source_nodes]}")

# Haystack: Pipelines explícitos y testeable
from haystack import Pipeline
from haystack.nodes import EmbeddingRetriever, PromptNode
from haystack.document_stores import PineconeDocumentStore

# Document store
document_store = PineconeDocumentStore(
    api_key="xxx",
    index="mi-index",
    embedding_dim=1536
)

# Retriever
retriever = EmbeddingRetriever(
    document_store=document_store,
    embedding_model="text-embedding-ada-002",
    top_k=3
)

# Generator
generator = PromptNode(
    model_name_or_path="gpt-4",
    api_key="xxx",
    default_prompt_template="Responde basándote en: {documents}\n\nPregunta: {query}"
)

# Pipeline como DAG (fácil de debuggear)
pipeline = Pipeline()
pipeline.add_node(component=retriever, name="Retriever", inputs=["Query"])
pipeline.add_node(component=generator, name="Generator", inputs=["Retriever"])

# Query con métricas detalladas
result = pipeline.run(query="¿Qué es RAG?")

# Haystack incluye timing de CADA nodo
print(pipeline.get_node("Retriever").timing)
print(pipeline.get_node("Generator").timing)

# Python Vanilla: Simple, rápido, debuggeable
import openai
import pinecone
from typing import List, Dict

class SimpleRAG:
    """
    Sistema RAG production-ready en 50 líneas.
    Sin abstracciones innecesarias.
    """

    def __init__(self, pinecone_index: str, openai_key: str):
        self.index = pinecone.Index(pinecone_index)
        openai.api_key = openai_key

    def embed(self, text: str) -> List[float]:
        """Genera embedding con control total de errores."""
        response = openai.Embedding.create(
            input=text,
            model="text-embedding-ada-002"
        )
        return response['data'][0]['embedding']

    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
        """Búsqueda vectorial con logging explícito."""
        query_embedding = self.embed(query)

        results = self.index.query(
            vector=query_embedding,
            top_k=top_k,
            include_metadata=True
        )

        # Log EXACTO de qué recuperamos
        print(f"Retrieved {len(results.matches)} documents")
        print(f"Scores: {[m.score for m in results.matches]}")

        return results.matches

    def generate(self, query: str, context_docs: List[Dict]) -> str:
        """Generación con control de tokens y caching."""
        # Construye contexto manualmente (sin llamadas ocultas)
        context = "\n\n".join([
            doc.metadata['text'] for doc in context_docs
        ])

        # Prompt explícito (lo ves y controlas)
        prompt = f"""Contexto:
{context}

Pregunta: {query}

Instrucciones: Responde basándote SOLO en el contexto."""

        # Llamada LLM única y visible
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=500  # Control explícito de costes
        )

        # Log de tokens usados
        print(f"Tokens: {response.usage.total_tokens}")

        return response.choices[0].message.content

    def query(self, question: str) -> Dict:
        """Pipeline RAG completo en 3 pasos claros."""
        docs = self.retrieve(question)
        answer = self.generate(question, docs)

        return {
            "answer": answer,
            "sources": [d.metadata for d in docs],
            "debug": {
                "retrieval_time": "120ms",  # Mides tú
                "generation_time": "650ms",
                "total_tokens": 487
            }
        }

# Uso: 3 líneas
rag = SimpleRAG(pinecone_index="mi-index", openai_key="xxx")
result = rag.query("¿Qué es RAG?")
print(result)

# Script para auditar uso de LangChain en tu codebase
import ast
import os
from collections import defaultdict

def audit_langchain_usage(repo_path: str):
    """
    Escanea codebase y genera reporte de uso LangChain.
    """
    stats = defaultdict(int)

    for root, dirs, files in os.walk(repo_path):
        for file in files:
            if file.endswith('.py'):
                filepath = os.path.join(root, file)
                with open(filepath, 'r') as f:
                    try:
                        tree = ast.parse(f.read())
                        for node in ast.walk(tree):
                            if isinstance(node, ast.ImportFrom):
                                if node.module and 'langchain' in node.module:
                                    stats[node.module] += 1
                    except:
                        pass

    # Genera reporte
    print("=== LANGCHAIN USAGE AUDIT ===")
    print(f"Total imports: {sum(stats.values())}")
    print("\nMost used modules:")
    for module, count in sorted(stats.items(), key=lambda x: x[1], reverse=True)[:10]:
        print(f"  {module}: {count} imports")

    # Identifica high-risk areas
    high_risk = ['langchain.agents', 'langchain.memory', 'langchain.chains']
    print("\n⚠️ High-risk areas (harder to migrate):")
    for module in high_risk:
        if module in stats:
            print(f"  {module}: {stats[module]} uses")

# Ejecuta auditoría
audit_langchain_usage("/path/to/your/repo")

# Patrón feature flag para migración sin downtime
import os
from typing import Dict

# Feature flag (controla % de tráfico en nueva implementación)
USE_NEW_RAG = os.getenv("USE_NEW_RAG_PERCENTAGE", "0")  # 0-100

class RAGService:
    """
    Wrapper que permite A/B testing entre LangChain y vanilla.
    """

    def __init__(self):
        self.langchain_rag = LangChainRAG()  # Implementación vieja
        self.vanilla_rag = VanillaRAG()      # Implementación nueva

    def query(self, question: str, user_id: str) -> Dict:
        """
        Decide qué implementación usar basado en feature flag.
        """
        # Hash user_id para distribución consistente
        use_new = (hash(user_id) % 100) < int(USE_NEW_RAG)

        if use_new:
            # Nueva implementación (vanilla)
            result = self.vanilla_rag.query(question)
            result['version'] = 'vanilla'
        else:
            # Vieja implementación (LangChain)
            result = self.langchain_rag.query(question)
            result['version'] = 'langchain'

        # Log para comparar performance
        self._log_metrics(result)

        return result

    def _log_metrics(self, result: Dict):
        """Log para comparar ambas implementaciones."""
        # Envia a tu sistema de observability
        print(f"Version: {result['version']}")
        print(f"Latency: {result.get('latency_ms')}ms")
        print(f"Tokens: {result.get('tokens_used')}")

# Uso: Deploy gradual
# Semana 1: USE_NEW_RAG=5   (5% tráfico)
# Semana 2: USE_NEW_RAG=20  (20% tráfico)
# Semana 3: USE_NEW_RAG=50  (50% tráfico)
# Semana 4: USE_NEW_RAG=100 (100% tráfico, elimina LangChain)