from langchain.memory import ConversationBufferMemory, VectorStoreRetrieverMemory
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
import pinecone

# Inicializar Pinecone para memoria vectorial
pinecone.init(api_key="tu-api-key", environment="us-west1-gcp")
index = pinecone.Index("conversation-memory")

# Configurar embeddings
embeddings = OpenAIEmbeddings()

# Memory system híbrido: buffer reciente + vector store para histórico
class HybridMemorySystem:
    """
    Sistema de memoria híbrido que combina:
    - Buffer memory para contexto inmediato (últimos 5-10 turnos)
    - Vector store memory para retrieval de contexto histórico relevante
    """
    
    def __init__(self, buffer_size=10):
        # Memoria de buffer para contexto reciente
        self.buffer_memory = ConversationBufferMemory(
            memory_key="recent_history",
            return_messages=True,
            output_key="output",
            input_key="input"
        )
        
        # Memoria vectorial para histórico
        vectorstore = Pinecone(index, embeddings.embed_query, "text")
        self.vector_memory = VectorStoreRetrieverMemory(
            retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
            memory_key="relevant_history"
        )
    
    def save_context(self, inputs, outputs):
        """Guarda contexto en ambos sistemas de memoria"""
        self.buffer_memory.save_context(inputs, outputs)
        self.vector_memory.save_context(inputs, outputs)
    
    def load_context(self, query):
        """
        Recupera contexto híbrido:
        - Últimos N turnos del buffer
        - Top K turnos históricos relevantes del vector store
        """
        recent = self.buffer_memory.load_memory_variables({})
        relevant = self.vector_memory.load_memory_variables({"prompt": query})
        return {
            "recent_context": recent["recent_history"],
            "relevant_context": relevant["relevant_history"]
        }

# Uso en production
memory = HybridMemorySystem(buffer_size=10)

# Guardar interacción
memory.save_context(
    {"input": "¿Cuáles son los términos de mi póliza de seguro?"},
    {"output": "Tu póliza cubre daños por inundación hasta el 80% del valor..."}
)

# Recuperar para nueva query (obtiene últimos 10 turnos + 3 más relevantes del histórico)
context = memory.load_context("¿Y qué pasa con daños por terremoto?")

from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
from sentence_transformers import CrossEncoder

class ContextualRetrievalRAG:
    """
    Implementación de Contextual Retrieval (Anthropic approach)
    Mejora el RAG tradicional añadiendo contexto explícito a cada chunk 
    antes de embedding, reduciendo failures de retrieval en 49-67%.
    """
    
    def __init__(self, documents, index_name="contextual-rag"):
        self.embeddings = OpenAIEmbeddings()
        self.llm = OpenAI(temperature=0)
        
        # Paso 1: Generar chunks contextualizados
        contextualized_chunks = self._contextualize_chunks(documents)
        
        # Paso 2: Crear embeddings y almacenar en vector DB
        self.vectorstore = Pinecone.from_texts(
            texts=contextualized_chunks,
            embedding=self.embeddings,
            index_name=index_name
        )
    
    def _contextualize_chunks(self, documents):
        """
        Añade contexto situacional a cada chunk usando el LLM
        En lugar de chunk aislado: "El deducible es de 500 euros."
        Genera chunk contextualizado: "Este documento describe póliza de 
        seguro de hogar. Sección de costos: El deducible es de 500 euros."
        """
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=800,
            chunk_overlap=100
        )
        
        contextualized = []
        for doc in documents:
            chunks = splitter.split_text(doc["content"])
            
            # Para cada chunk, generar contexto situacional
            for chunk in chunks:
                # Prompt para contextualizar
                context_prompt = f"""Dado el documento completo y este chunk específico, 
                genera 1-2 frases de contexto que sitúen este chunk.
                
                Documento: {doc["title"]}
                Tipo: {doc["type"]}
                Chunk: {chunk}
                
                Contexto situacional (1-2 frases):"""
                
                situational_context = self.llm(context_prompt).strip()
                
                # Combinar contexto + chunk original
                contextualized_chunk = f"{situational_context}\n\n{chunk}"
                contextualized.append(contextualized_chunk)
        
        return contextualized
    
    def retrieve_with_reranking(self, query, k=20, top_n=5):
        """
        Retrieval en dos fases:
        1. Semantic search inicial (top 20)
        2. Reranking con cross-encoder (top 5 final)
        Esto logra 67% reduction en retrieval failures según Anthropic
        """
        # Fase 1: Retrieval inicial
        initial_results = self.vectorstore.similarity_search(query, k=k)
        
        # Fase 2: Reranking con modelo especializado
        reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
        # Calcular scores de reranking
        pairs = [[query, doc.page_content] for doc in initial_results]
        scores = reranker.predict(pairs)
        
        # Ordenar por score y tomar top N
        ranked_results = sorted(
            zip(initial_results, scores),
            key=lambda x: x[1],
            reverse=True
        )[:top_n]
        
        return [doc for doc, score in ranked_results]

# Uso en production
documents = [
    {
        "title": "Póliza Seguro Hogar",
        "type": "insurance",
        "content": "..."
    },
    {
        "title": "Términos y Condiciones", 
        "type": "legal",
        "content": "..."
    }
]

rag = ContextualRetrievalRAG(documents)

# Recuperar contexto relevante con reranking
relevant_chunks = rag.retrieve_with_reranking(
    query="¿Cuál es el deducible para daños por agua?",
    k=20,
    top_n=5
)

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

class ContextCompressor:
    """
    Comprime contexto largo manteniendo información relevante
    Técnicas:
    1. Summarization: Resumir contexto histórico
    2. Extraction: Extraer solo facts relevantes
    3. Structured notes: Convertir a formato estructurado compacto
    """
    
    def __init__(self):
        self.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
    
    def compress_conversation_history(self, history, max_tokens=500):
        """
        Comprime historial de conversación a resumen estructurado
        Entrada: 15 turnos de conversación (3000 tokens)
        Salida: Resumen estructurado (500 tokens)
        Ratio compresión: 6x
        """
        prompt = PromptTemplate(
            input_variables=["history", "max_tokens"],
            template="""Resume el siguiente historial de conversación en máximo {max_tokens} tokens.
            
Mantén SOLO información factual crítica. Usa formato estructurado:
- Tema principal:
- Hechos clave:
- Decisiones tomadas:
- Preguntas sin resolver:

Historial: {history}

Resumen estructurado:"""
        )
        
        chain = LLMChain(llm=self.llm, prompt=prompt)
        compressed = chain.run(history=history, max_tokens=max_tokens)
        return compressed
    
    def extract_relevant_facts(self, context, query):
        """
        Extrae solo facts del contexto relevantes para la query actual
        En lugar de pasar 5 chunks completos (2000 tokens), 
        extrae 10-15 facts específicos (300 tokens)
        """
        prompt = PromptTemplate(
            input_variables=["context", "query"],
            template="""Del siguiente contexto, extrae SOLO los facts específicamente 
relevantes para responder la query. Formato bullet points.

Query: {query}

Contexto: {context}

Facts relevantes (bullet points):"""
        )
        
        chain = LLMChain(llm=self.llm, prompt=prompt)
        facts = chain.run(context=context, query=query)
        return facts
    
    def sliding_window_compression(self, messages, window_size=10, summary_size=3):
        """
        Técnica de ventana deslizante con compresión:
        - Mantiene últimos N mensajes completos (window_size)
        - Comprime mensajes antiguos a resumen
        - Total tokens: constante independiente de largo conversación
        """
        if len(messages) 

from langgraph.graph import Graph, END
from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain.llms import OpenAI
from langchain.prompts import ChatPromptTemplate

class MultiAgentContextIsolation:
    """
    Arquitectura multi-agente con contextos aislados
    Beneficios:
    - Cada agente maneja 5-10 variables máximo (dentro del límite working memory)
    - Contexts específicos, no bloat
    - Paralelización posible
    - Failure isolation (si un agente falla, otros continúan)
    """
    
    def __init__(self):
        self.llm = OpenAI(temperature=0)
        
        # Definir agentes especializados
        self.research_agent = self._create_research_agent()
        self.analysis_agent = self._create_analysis_agent()
        self.writing_agent = self._create_writing_agent()
        
        # Crear workflow graph
        self.workflow = self._build_workflow()
    
    def _create_research_agent(self):
        """Agente especializado en retrieval de información
        Contexto aislado: solo query + retrieved docs"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Eres un agente de research especializado. 
Tu ÚNICO trabajo es recuperar información relevante de la base de conocimiento. 
NO respondas preguntas, solo retorna los 5 documentos más relevantes."""), 
            ("user", "{query}")
        ])
        
        agent = create_openai_functions_agent(
            llm=self.llm, 
            prompt=prompt, 
            tools=[self._get_retrieval_tool()]
        )
        return AgentExecutor(agent=agent, tools=[self._get_retrieval_tool()])
    
    def _create_analysis_agent(self):
        """Agente especializado en análisis de información
        Contexto aislado: query + documentos recuperados (NO historial completo)"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Eres un agente de análisis especializado. 
Analiza los documentos proporcionados y extrae insights clave. 
Retorna SOLO facts verificables, no opiniones."""), 
            ("user", "Query: {query}\nDocumentos: {documents}")
        ])
        
        agent = create_openai_functions_agent(llm=self.llm, prompt=prompt, tools=[])
        return AgentExecutor(agent=agent, tools=[])
    
    def _create_writing_agent(self):
        """Agente especializado en redacción de respuesta final
        Contexto aislado: query + analyzed facts (NO documentos crudos)"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """Eres un agente de writing especializado. 
Genera una respuesta clara y concisa basada ÚNICAMENTE en los facts proporcionados. 
Si los facts no son suficientes, di "No tengo suficiente información"."""), 
            ("user", "Query: {query}\nFacts verificados: {facts}")
        ])
        
        agent = create_openai_functions_agent(llm=self.llm, prompt=prompt, tools=[])
        return AgentExecutor(agent=agent, tools=[])
    
    def _get_retrieval_tool(self):
        """Placeholder para tool de retrieval"""
        # Implementar según tu vectorstore
        pass
    
    def _build_workflow(self):
        """Define el workflow de agentes con LangGraph
        Flow: User Query → Research Agent → Analysis Agent → Writing Agent → Response"""
        workflow = Graph()
        
        # Nodos
        workflow.add_node("research", self._research_step)
        workflow.add_node("analysis", self._analysis_step)
        workflow.add_node("writing", self._writing_step)
        
        # Edges (flujo secuencial)
        workflow.add_edge("research", "analysis")
        workflow.add_edge("analysis", "writing")
        workflow.add_edge("writing", END)
        
        # Set entry point
        workflow.set_entry_point("research")
        return workflow.compile()
    
    def _research_step(self, state):
        """Paso 1: Retrieval de información"""
        query = state["query"]
        documents = self.research_agent.invoke({"query": query})
        # Pasar SOLO documentos al siguiente step, NO todo el state
        return {"documents": documents["output"]}
    
    def _analysis_step(self, state):
        """Paso 2: Análisis de documentos"""
        query = state["query"]
        documents = state["documents"]
        facts = self.analysis_agent.invoke({
            "query": query, 
            "documents": documents
        })
        # Pasar SOLO facts extraídos, NO documentos completos
        return {"facts": facts["output"]}
    
    def _writing_step(self, state):
        """Paso 3: Generación de respuesta"""
        query = state["query"]
        facts = state["facts"]
        response = self.writing_agent.invoke({
            "query": query, 
            "facts": facts
        })
        return {"response": response["output"]}
    
    def process_query(self, query):
        """Procesa query a través del workflow multi-agente
        Cada agente maneja contexto aislado y limitado:
        - Research: solo query
        - Analysis: solo query + docs (no research prompts)
        - Writing: solo query + facts (no docs crudos ni análisis prompts)
        Total context per agent: 

from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Pinecone
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import Anthropic
from sentence_transformers import CrossEncoder
from rank_bm25 import BM25Okapi
import pinecone
import nltk

class ProductionContextualRetrieval:
    """
    Implementación production-ready de Contextual Retrieval
    Features:
    - Contextual embeddings (usando Claude para generar contexto)
    - Contextual BM25 (keyword search contextualizado)
    - Reranking con cross-encoder
    - Hybrid retrieval (combina embeddings + BM25)
    Resultado esperado: 67% reduction en retrieval failures
    """
    
    def __init__(self, pinecone_api_key, pinecone_env, index_name):
        # Inicializar servicios
        pinecone.init(api_key=pinecone_api_key, environment=pinecone_env)
        self.index = pinecone.Index(index_name)
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
        self.claude = Anthropic(model="claude-3-haiku-20240307", temperature=0)
        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
        # Almacenar documentos originales para BM25
        self.bm25_corpus = []
        self.bm25 = None
    
    def contextualize_chunk(self, chunk, document_metadata):
        """Genera contexto situacional para un chunk usando Claude"""
        prompt = f"""Documento: {document_metadata['title']}
Tipo: {document_metadata.get('type', 'general')}
Sección: {document_metadata.get('section', 'N/A')}
Chunk específico: {chunk}

Genera 1-2 frases concisas de contexto situacional que expliquen 
de qué trata este chunk en el contexto del documento completo.

Contexto situacional:"""
        
        context = self.claude(prompt).strip()
        # Retornar chunk contextualizado
        return f"{context}\n\n{chunk}"
    
    def process_and_index_documents(self, documents):
        """Procesa documentos, genera chunks contextualizados y los indexa"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=800,
            chunk_overlap=100,
            separators=["\n\n", "\n", ". ", " ", ""]
        )
        
        all_chunks = []
        all_metadatas = []
        
        for doc in documents:
            # Split documento
            chunks = splitter.split_text(doc['content'])
            
            # Contextualizar cada chunk
            for i, chunk in enumerate(chunks):
                contextualized_chunk = self.contextualize_chunk(chunk, doc['metadata'])
                all_chunks.append(contextualized_chunk)
                all_metadatas.append({
                    **doc['metadata'],
                    'chunk_id': i,
                    'original_chunk': chunk  # Guardar chunk original también
                })
        
        # Generar embeddings y indexar en Pinecone
        vectorstore = Pinecone.from_texts(
            texts=all_chunks,
            embedding=self.embeddings,
            index_name=self.index.name,
            metadatas=all_metadatas
        )
        
        # Construir corpus BM25 con chunks contextualizados
        nltk.download('punkt', quiet=True)
        tokenized_corpus = [nltk.word_tokenize(doc.lower()) for doc in all_chunks]
        self.bm25 = BM25Okapi(tokenized_corpus)
        self.bm25_corpus = all_chunks
        
        return vectorstore
    
    def hybrid_retrieval(self, query, k=20, alpha=0.5):
        """Hybrid retrieval: combina semantic search + keyword search (BM25)"""
        import nltk
        
        # 1. Semantic search con embeddings
        embedding_results = self.vectorstore.similarity_search_with_score(query, k=k*2)
        
        # 2. BM25 keyword search
        tokenized_query = nltk.word_tokenize(query.lower())
        bm25_scores = self.bm25.get_scores(tokenized_query)
        
        # 3. Normalizar scores (0-1 range)
        def normalize_scores(scores):
            min_score = min(scores)
            max_score = max(scores)
            if max_score == min_score:
                return [0.5] * len(scores)
            return [(s - min_score) / (max_score - min_score) for s in scores]
        
        embedding_scores_normalized = normalize_scores(
            [1 - score for doc, score in embedding_results]  # Pinecone usa distance
        )
        bm25_scores_normalized = normalize_scores(bm25_scores)
        
        # 4. Calcular scores híbridos
        hybrid_results = []
        for i, (doc, _) in enumerate(embedding_results):
            hybrid_score = (
                alpha * bm25_scores_normalized[i] +
                (1 - alpha) * embedding_scores_normalized[i]
            )
            hybrid_results.append((doc, hybrid_score))
        
        # 5. Ordenar por hybrid score y retornar top K
        hybrid_results.sort(key=lambda x: x[1], reverse=True)
        return hybrid_results[:k]
    
    def retrieve_with_reranking(self, query, k_initial=20, k_final=5):
        """Retrieval en dos fases para 67% reduction en failures"""
        # Fase 1: Hybrid retrieval
        initial_results = self.hybrid_retrieval(query, k=k_initial)
        
        # Fase 2: Reranking
        pairs = [[query, doc.page_content] for doc, score in initial_results]
        rerank_scores = self.reranker.predict(pairs)
        
        # Combinar con scores híbridos (weighted average)
        final_results = []
        for i, (doc, hybrid_score) in enumerate(initial_results):
            # 70% reranking score + 30% hybrid score
            final_score = 0.7 * rerank_scores[i] + 0.3 * hybrid_score
            final_results.append((doc, final_score))
        
        # Ordenar por final score
        final_results.sort(key=lambda x: x[1], reverse=True)
        return final_results[:k_final]

# Uso en production
retrieval_system = ProductionContextualRetrieval(
    pinecone_api_key="tu-api-key",
    pinecone_env="us-west1-gcp",
    index_name="contextual-rag-prod"
)

# Indexar documentos (hacer UNA VEZ, no por query)
documents = [
    {
        'content': "...",  # Contenido completo del documento
        'metadata': {
            'title': "Póliza de Seguro de Hogar Estándar 2025",
            'type': "insurance_policy",
            'section': "coverages"
        }
    }
    # ... más documentos
]

vectorstore = retrieval_system.process_and_index_documents(documents)

# Retrieval para queries de usuarios
query = "¿Cuál es el deducible para daños por agua en mi póliza de hogar?"
results = retrieval_system.retrieve_with_reranking(
    query, k_initial=20, k_final=5
)

# Results ahora tiene los 5 chunks más relevantes
for doc, score in results:
    print(f"Score: {score:.3f}")
    print(f"Content: {doc.page_content[:200]}...")
    print("---")

from guardrails import Guard
from guardrails.hub import CompetitorCheck, ToxicLanguage, SemanticSimilarity
from langchain.llms import OpenAI
from sentence_transformers import CrossEncoder
import numpy as np

class HallucinationDetectionSystem:
    """
    Sistema multi-layer de detección de hallucinations
    Layers:
    1. Semantic similarity check (response vs retrieved context)
    2. Source verification (claims verificables en contexto)
    3. Confidence scoring (LLM-as-judge approach)
    4. Toxic/competitor content check
    Si ANY layer falla → trigger re-generation o human review
    """
    
    def __init__(self):
        self.llm = OpenAI(temperature=0)
        self.semantic_model = CrossEncoder('cross-encoder/stsb-roberta-large')
        
        # Configurar guardrails
        self.guard = Guard().use_many(
            ToxicLanguage(threshold=0.5, validation_method="sentence"),
            CompetitorCheck(competitors=["CompetitorX", "CompetitorY"])
        )
    
    def semantic_similarity_check(self, response, context_chunks, threshold=0.6):
        """Layer 1: Verifica que response esté semánticamente alineada con contexto"""
        # Calcular similarity entre response y cada chunk de contexto
        pairs = [[response, chunk] for chunk in context_chunks]
        scores = self.semantic_model.predict(pairs)
        max_similarity = np.max(scores)
        return (max_similarity >= threshold, max_similarity)
    
    def source_verification_check(self, response, context_chunks):
        """Layer 2: Extrae claims factuales del response y verifica si están en contexto"""
        # Extraer claims
        claims_prompt = f"""Extrae TODOS los claims factuales específicos de este texto. 
Lista solo facts verificables (números, nombres, fechas, etc).

Texto: {response}

Claims factuales (uno por línea):"""
        
        claims_text = self.llm(claims_prompt).strip()
        claims = [c.strip() for c in claims_text.split('\n') if c.strip()]
        
        if not claims:
            return (True, 1.0)  # No factual claims = can't hallucinate
        
        # Verificar cada claim
        verified_count = 0
        for claim in claims:
            verify_prompt = f"""¿Este claim está verificado por ALGUNO de estos contextos?
            
Claim: {claim}
Contextos: {chr(10).join([f"- {chunk}" for chunk in context_chunks])}

Responde SOLO: SÍ o NO"""
            
            verification = self.llm(verify_prompt).strip().upper()
            if "SÍ" in verification or "SI" in verification:
                verified_count += 1
        
        verification_rate = verified_count / len(claims)
        return (verification_rate >= 0.8, verification_rate)  # 80% threshold
    
    def confidence_scoring_llm_judge(self, query, response, context_chunks):
        """Layer 3: LLM-as-judge approach para confidence scoring"""
        judge_prompt = f"""Evalúa esta respuesta en una escala de 0-100:

QUERY: {query}

CONTEXTO DISPONIBLE: {chr(10).join([f"- {chunk[:200]}..." for chunk in context_chunks])}

RESPUESTA GENERADA: {response}

Criterios de evaluación:
- ¿Responde la query completamente? (0-30 puntos)
- ¿Está basada SOLO en el contexto proporcionado? (0-40 puntos)
- ¿Evita inventar información no presente en contexto? (0-30 puntos)

Retorna SOLO el score numérico (0-100) y una línea de justificación.

Score: