src/agents/agent_director.py

import re
from openai import OpenAI
from typing import List, Dict, Any, Optional, Tuple
import sys
import os
import traceback

# Add the project root to the path to ensure imports work
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# Import configuration 
from src.utils.config import CHAT_MODEL, OPENAI_API_KEY

# Import other modules needed for the agents
from src.models.retriever import Retriever

class QueryAnalyzer:
    """
    Agent responsible for analyzing and refining the user's query.
    """
    
    def __init__(self, model: str = CHAT_MODEL):
        """Initialize the query analyzer."""
        self.model = model
        self.client = OpenAI(api_key=OPENAI_API_KEY)
    
    def analyze_query(self, query: str) -> Dict[str, Any]:
        """
        Analyze the user's query to extract key information and refine it if needed.
        
        Args:
            query: The user's query
            
        Returns:
            Dictionary containing analysis results
        """
        # Create a system prompt for the query analyzer oriented toward Cathy Barriga defense
        system_prompt = (
            "Eres un analista legal especializado en defensa penal relacionado con el delito económico. Tu tarea es analizar las consultas relacionadas con "
            "el caso de Cathy Barriga para identificar:"
            "\n1. Las posibles teorías de defensa aplicables a Cathy Barriga"
            "\n2. Los elementos del tipo penal que podrían ser refutados o desacreditados"
            "\n3. Los vacíos probatorios o debilidades en la acusación"
            "\n4. Las interpretaciones alternativas de los hechos que favorezcan a la defensa"
            "\n5. Los principios jurídicos que podrían invocarse a favor de la defensa (in dubio pro reo, presunción de inocencia, etc.)"
            "\n\nProporciona tu análisis en un formato estructurado que oriente la búsqueda hacia elementos exculpatorios o mitigantes."
        )
        
        # Get analysis from the LLM
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": query}
                ],
                temperature=0.3
            )
            
            analysis = response.choices[0].message.content.strip()
            
            # Create a structured analysis
            struct_analysis = self._extract_structured_analysis(analysis, query)
            
            return {
                "original_query": query,
                "analysis": analysis,
                "structured_analysis": struct_analysis
            }
        except Exception as e:
            print(f"Error analyzing query: {e}")
            return {
                "original_query": query,
                "analysis": f"Error: {str(e)}",
                "structured_analysis": ""
            }
    
    def _extract_structured_analysis(self, analysis: str, query: str) -> str:
        """Extract a structured analysis from the raw analysis text."""
        # This would normally do more sophisticated extraction
        # For demo purposes, we'll just format it with some headers
        formatted = "## Análisis fundamental del contexto\n\n"
        formatted += "- **Dominio y expertiz**: Defensa penalista con especialidad en delito económico\n"
        formatted += "- **Consulta Original**: {query}\n"
        # formatted += "- **Enfoque de Defensa**: Estrategia de defensa para Cathy Barriga\n"
        # formatted += "- **Conceptos Clave**: Presunción de inocencia, carga de la prueba, elementos del tipo penal\n"
        return formatted


class ContextAggregator:
    """
    Agent responsible for aggregating and organizing retrieved document chunks.
    """
    
    def __init__(self, model: str = CHAT_MODEL):
        """Initialize the context aggregator."""
        self.model = model
        self.client = OpenAI(api_key=OPENAI_API_KEY)
    
    def aggregate_context(self, query: str, retrieved_chunks: List[Dict[str, Any]]) -> str:
        """
        Aggregate retrieved chunks into a coherent context.
        
        Args:
            query: The user's query
            retrieved_chunks: List of retrieved document chunks
            
        Returns:
            String containing the organized context
        """
        # If small number of chunks, use a simpler approach
        if len(retrieved_chunks) <= 10:
            # For small number of chunks, just organize them
            chunk_contents = [
                {
                    'source': chunk.get('source', 'unknown'),
                    'content': chunk.get('text', chunk.get('chunk', "No content available")),
                    'is_summary': False
                }
                for chunk in retrieved_chunks
            ]
            return self._organize_content(query, chunk_contents)
        else:
            # Group chunks by source
            sources = {}
            for chunk in retrieved_chunks:
                source = chunk.get('source', 'unknown')
                if source not in sources:
                    sources[source] = []
                sources[source].append(chunk)
            
            # Create summaries for each source
            summaries = []
            for source, chunks in sources.items():
                summary = self._summarize_chunks(source, chunks, query)
                summaries.append(summary)
            
            # Aggregate the summaries and individual chunks
            aggregated_context = self._organize_content(query, summaries)
            return aggregated_context
    
    def _summarize_chunks(self, source: str, chunks: List[Dict[str, Any]], query: str) -> Dict[str, Any]:
        """Summarize a group of chunks from the same source."""
        # Combine chunks into a single text, handling different chunk formats
        try:
            chunks_text = "\n\n".join([chunk.get('text', chunk.get('chunk', "No content available")) for chunk in chunks])
        except Exception as e:
            print(f"Error combining chunks: {e}")
            # Fallback to a safer method
            chunks_text = ""
            for chunk in chunks:
                try:
                    if isinstance(chunk, dict):
                        chunk_content = chunk.get('text', chunk.get('chunk', "No content available"))
                        chunks_text += chunk_content + "\n\n"
                    else:
                        chunks_text += str(chunk) + "\n\n"
                except Exception as chunk_e:
                    print(f"Error processing individual chunk: {chunk_e}")
                    continue
        
        # Create a prompt for summarization oriented toward defense
        system_prompt = (
            "Eres un experto en derecho penal especializado en organizar información para la defensa legal. Tu tarea es resumir, con un approach divide & conquer, "
            "los documentos proporcionados priorizando información que pueda ser útil para la defensa de Cathy Barriga. Enfócate en:"
            "\n1. Elementos exculpatorios o que cuestionen la culpabilidad"
            "\n2. Inconsistencias o debilidades en la evidencia de la acusación"
            "\n3. Interpretaciones alternativas de los hechos que favorezcan a la defendida"
            "\n4. Precedentes legales que podrían apoyar la defensa"
            "\n5. Circunstancias atenuantes o justificativas"
            "\n6. Vicios procedimentales que puedan ser alegados"
            "\nMantén la precisión fáctica pero organiza la información para construir la mejor argumentación posible."
        )
        
        # Get summary from the LLM
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Consulta relacionada: {query}\n\nDocumentos de {source}:\n\n{chunks_text}"}
                ],
                temperature=0.3
            )
            
            summary = response.choices[0].message.content.strip()
            
            return {
                'source': source,
                'content': summary,
                'is_summary': True,
                'num_chunks': len(chunks)
            }
        except Exception as e:
            print(f"Error summarizing chunks from {source}: {e}")
            return {
                'source': source,
                'content': f"Error summarizing content: {str(e)}",
                'is_summary': True,
                'num_chunks': len(chunks)
            }
    
    def _organize_content(self, query: str, contents: List[Dict[str, Any]]) -> str:
        """Organize content items into a coherent structure."""
        # Simple organization - separate summaries and regular chunks
        organized_text = f"# Relevant Legal Context for: {query}\n\n"
        
        # Add summaries first
        summaries = [item for item in contents if item.get('is_summary', False)]
        if summaries:
            organized_text += "## Summaries of Key Sources\n\n"
            for summary in summaries:
                organized_text += f"### {summary['source']}\n"
                organized_text += f"{summary['content']}\n\n"
        
        # Add individual chunks
        individual_chunks = [item for item in contents if not item.get('is_summary', False)]
        if individual_chunks:
            organized_text += "## Detalles Relevantes adicionales\n\n"
            for chunk in individual_chunks:
                organized_text += f"### De {chunk['source']}\n"
                organized_text += f"{chunk['content']}\n\n"
        
        return organized_text


class AnswerGenerator:
    """
    Agent responsible for generating comprehensive answers based on the context.
    """
    
    def __init__(self, model: str = CHAT_MODEL):
        """Initialize the answer generator."""
        self.model = model
        self.client = OpenAI(api_key=OPENAI_API_KEY)
    
    def generate_answer(self, query: str, context: str) -> str:
        """
        Generate a comprehensive answer to the user's query using the provided context.
        
        Args:
            query: The user's query
            context: The organized context
            
        Returns:
            The generated answer
        """
        # Create a system prompt for the answer generator focused on defense
        system_prompt = (
            "Eres un abogado defensor experto especializado en derecho penal, particularmente en la defensa de casos complejos. "
            "Tu objetivo es estructurar con máxima precisión y claridad lo expuesto en el contexto. Al responder preguntas:\n"
            "\n1. Basa tus respuestas exclusivamente en la información proporcionada en el contexto y en principios jurídicos de defensa penal\n"
            "\n2. Cuestiona sistemáticamente las pruebas de cargo, identificando sus debilidades metodológicas, procedimentales o interpretativas\n"
            "\n3. Explora interpretaciones alternativas de los hechos que favorezcan a la defendida\n"
            "\n4. Identifica vicios procesales que puedan ser alegados\n"
            "\n5. Construye narrativas coherentes que expliquen los hechos de manera objetiva\n"
            "\n6. Cita fuentes específicas cuando te refieres a información o argumentos clave\n"
            "\n7. Estructura tu respuesta claramente con secciones apropiadas\n"
            "\n8. Maximiza la claridad y precisión de tus respuestas en todo momento, con un enfoque lógico, penalista, algorítmico y objetivo"
            # "\n8. Mantén un tono objetivo y profesional pero claramente orientado a la defensa, evitando cualquier admisión de culpabilidad"
        )
        
        # Get answer from the LLM
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Consulta sobre la defensa de Cathy Barriga: {query}\n\nContexto:\n\n{context}"}
                ],
                temperature=0.3
            )
            
            answer = response.choices[0].message.content.strip()
            return answer
        except Exception as e:
            print(f"Error generating answer: {e}")
            return f"Error generating answer: {str(e)}"


class AgentDirector:
    """
    Director that coordinates the various specialized agents to process a query.
    """
    
    def __init__(self, model: str = None, top_k: int = 200, debug: bool = False):
        """
        Initialize the agent director.
        
        Args:
            model: The OpenAI chat model to use
            top_k: Number of chunks to retrieve
            debug: Whether to show detailed reasoning steps
        """
        # Ensure model is not None, default to CHAT_MODEL if not provided
        self.model = model if model is not None else CHAT_MODEL
        self.top_k = top_k
        self.debug = debug
        self.retriever = Retriever(top_k=top_k)
        self.query_analyzer = QueryAnalyzer(model=self.model)
        self.context_aggregator = ContextAggregator(model=self.model)
        self.answer_generator = AnswerGenerator(model=self.model)
        self.client = OpenAI(api_key=OPENAI_API_KEY)
        # LegalAgent will be imported on demand
    
    def _debug_print(self, message):
        """Print debug message if debug mode is enabled."""
        if self.debug:
            print(f"\n🧠 AGENT THINKING: {message}")
    
    def _enhance_query_with_insights(self, query: str, analysis: Dict[str, Any]) -> str:
        """
        Enhance the original query with insights from analysis to improve retrieval.
        
        Args:
            query: Original user query
            analysis: Query analysis results
            
        Returns:
            Enhanced query for better retrieval
        """
        try:
            system_prompt = (
                "Eres un experto en búsqueda semántica para investigaciones legales. Tu objetivo es reformular y mejorar "
                "la consulta original para maximizar la recuperación de información relevante para la defensa legal. "
                "Debes expandir la consulta con términos técnicos legales, conceptos relacionados, y posibles "
                "contraargumentos que debería contemplar la defensa."
            )
            
            analysis_text = analysis.get("analysis", "")
            
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Consulta original: {query}\n\nAnálisis de la consulta: {analysis_text}\n\n"
                                               f"Formula una consulta mejorada que maximice la recuperación de información relevante para "
                                               f"la defensa. La consulta debe ser expansiva pero precisa."}
                ],
                temperature=0.3
            )
            
            enhanced_query = response.choices[0].message.content.strip()
            
            # If response is too verbose, extract just the query part
            if len(enhanced_query.split()) > 30:
                # Try to find a clear query section
                import re
                query_matches = re.search(r"(?:consulta mejorada|consulta refinada|consulta expandida):\s*(.*?)(?:\n|$)", 
                                         enhanced_query, re.IGNORECASE)
                if query_matches:
                    enhanced_query = query_matches.group(1).strip()
            
            self._debug_print(f"Consulta original: {query}\nConsulta mejorada: {enhanced_query}")
            return enhanced_query
        except Exception as e:
            print(f"Error enhancing query: {e}")
            return query  # Fallback to original query
    
    def _extract_key_concepts(self, query: str, context: str) -> List[str]:
        """
        Extract key legal concepts from query and context.
        
        Args:
            query: The user's query
            context: The retrieved context
            
        Returns:
            List of key legal concepts
        """
        try:
            system_prompt = (
                "Eres un experto en análisis jurídico de primer nivel. Tu tarea es identificar y extraer los conceptos "
                "legales clave, fundamentos jurídicos y argumentos principales de la información proporcionada. "
                "Extrae exactamente los 5-7 conceptos más relevantes para la defensa legal, organizados por orden de importancia."
            )
            
            # Use a preview of the context to avoid token issues
            context_preview = context[:3000] if len(context) > 3000 else context
            
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Consulta: {query}\n\nContexto:\n{context_preview}"}
                ],
                temperature=0.2
            )
            
            concepts_text = response.choices[0].message.content.strip()
            
            # Extract concepts as a list
            import re
            # Look for numbered or bulleted items
            concepts = re.findall(r"(?:^|\n)(?:\d+\.\s*|\*\s*|\-\s*)(.*?)(?:$|\n)", concepts_text)
            
            # If regex didn't find structured concepts, split by newlines and filter
            if not concepts:
                concepts = [line.strip() for line in concepts_text.split("\n") 
                           if line.strip() and not line.strip().startswith("#")]
            
            # Ensure we have a reasonable number of concepts
            if len(concepts) > 10:
                concepts = concepts[:10]
            elif len(concepts) < 3 and len(concepts_text) > 50:
                # If we have too few but a lot of text, try line splitting
                concepts = [line.strip() for line in concepts_text.split("\n") if len(line.strip()) > 10][:7]
            
            self._debug_print(f"Conceptos clave extraídos: {concepts}")
            return concepts
        except Exception as e:
            print(f"Error extracting key concepts: {e}")
            return []
    
    def _refine_context_with_concepts(self, context: str, concepts: List[str], query: str) -> str:
        """
        Refine the context by emphasizing key concepts and restructuring for clarity.
        
        Args:
            context: The original context
            concepts: Key legal concepts extracted
            query: The user's query
            
        Returns:
            Refined context
        """
        try:
            if not concepts:
                return context
                
            system_prompt = (
                "Eres un experto legal especializado en reestructurar información para maximizar su utilidad en "
                "argumentación jurídica. Tu tarea es refinar y reorganizar el contexto proporcionado para:"
                "\n1. Enfatizar los conceptos clave identificados"
                "\n2. Organizar la información en una estructura coherente y progresiva"
                "\n3. Conectar elementos relacionados entre sí"
                "\n4. Incluir un resumen de alto nivel al inicio que integre los conceptos clave"
                "\n5. Añadir subtítulos informativos para mejorar la navegabilidad"
                "\nMantén TODA la información relevante del contexto original."
            )
            
            # Create a concepts section
            concepts_text = "\n".join([f"- {concept}" for concept in concepts])
            
            # Only process a portion of the context if it's very large
            if len(context) > 5000:
                # Find natural breakpoints to split the context
                chunks = []
                current_pos = 0
                while current_pos < len(context):
                    next_chunk_end = min(current_pos + 5000, len(context))
                    
                    # Try to find a natural breakpoint (like a paragraph end)
                    if next_chunk_end < len(context):
                        breakpoint_search = context[next_chunk_end-200:next_chunk_end+200]
                        paragraph_breaks = [m.start() for m in re.finditer(r'\n\n', breakpoint_search)]
                        
                        if paragraph_breaks:
                            # Find the break closest to the 5000 char mark
                            closest_break = min(paragraph_breaks, key=lambda x: abs(x - 200))
                            next_chunk_end = next_chunk_end - 200 + closest_break + 2  # +2 for the \n\n
                    
                    chunks.append(context[current_pos:next_chunk_end])
                    current_pos = next_chunk_end
                
                # Process each chunk
                refined_chunks = []
                for i, chunk in enumerate(chunks):
                    try:
                        # Use a simpler prompt for chunk refinement
                        chunk_prompt = (
                            "Reorganiza este fragmento de texto para enfatizar los conceptos clave, "
                            "manteniendo toda la información relevante. Añade conectores cuando sea necesario "
                            "para mejorar la fluidez y coherencia."
                        )
                        
                        response = self.client.chat.completions.create(
                            model=self.model,
                            messages=[
                                {"role": "system", "content": system_prompt},
                                {"role": "user", "content": f"Fragmento {i+1} de {len(chunks)}:\n\n{chunk}\n\n"
                                                         f"Conceptos clave a enfatizar:\n{concepts_text}"}
                            ],
                            temperature=0.3
                        )
                        
                        refined_chunks.append(response.choices[0].message.content.strip())
                    except Exception as e:
                        print(f"Error refining chunk {i+1}: {e}")
                        refined_chunks.append(chunk)  # Use original chunk if refinement fails
                
                # Combine refined chunks
                refined_context = "\n\n".join(refined_chunks)
                
                # Add a global summary at the beginning
                try:
                    summary_prompt = (
                        "Crea un resumen ejecutivo conciso (máximo 300 palabras) que integre todos los conceptos clave "
                        "y presente una visión de alto nivel de la información más relevante para la defensa."
                    )
                    
                    summary_response = self.client.chat.completions.create(
                        model=self.model,
                        messages=[
                            {"role": "system", "content": "Eres un experto legal que crea resúmenes ejecutivos precisos y orientados a la defensa."},
                            {"role": "user", "content": f"Consulta: {query}\n\nConceptos clave:\n{concepts_text}\n\n{summary_prompt}"}
                        ],
                        temperature=0.3
                    )
                    
                    summary = summary_response.choices[0].message.content.strip()
                    refined_context = f"# Resumen Ejecutivo\n\n{summary}\n\n# Información Detallada\n\n{refined_context}"
                except Exception as e:
                    print(f"Error creating summary: {e}")
                    # Continue without summary if it fails
                
                return refined_context
                
            else:
                # For smaller contexts, process all at once
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": f"Consulta: {query}\n\nConceptos clave a enfatizar:\n{concepts_text}\n\n"
                                                 f"Contexto a refinar:\n\n{context}"}
                    ],
                    temperature=0.3
                )
                
                return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"Error refining context: {e}")
            return context  # Return original context if refinement fails
    
    def _add_metacognitive_reflection(self, query: str, context: str, concepts: List[str]) -> str:
        """
        Add metacognitive reflections to help guide the answer generation.
        
        Args:
            query: The user's query
            context: The context
            concepts: Key legal concepts
            
        Returns:
            Metacognitive guidance for answer generation
        """
        try:
            system_prompt = (
                "Eres un abogado estratega de alto nivel capaz de analizar situaciones legales complejas desde múltiples "
                "perspectivas. Tu función es proporcionar una guía metacognitiva que:"
                "\n1. Identifique las principales líneas argumentativas disponibles para la defensa"
                "\n2. Señale conexiones no obvias entre elementos del contexto que podrían fortalecer la defensa"
                "\n3. Sugiera perspectivas alternativas que podrían no ser evidentes"
                "\n4. Identifique áreas donde un contraargumento podría ser necesario"
                "\n5. Destaque principios legales fundamentales aplicables"
                "\nEsta guía será utilizada para estructurar una respuesta legal sólida y coherente."
            )
            
            # Use just a preview of the context
            context_preview = context[:2000] if len(context) > 2000 else context
            concepts_text = "\n".join([f"- {concept}" for concept in concepts]) if concepts else "No se identificaron conceptos específicos."
            
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"Consulta de defensa: {query}\n\n"
                                              f"Conceptos clave:\n{concepts_text}\n\n"
                                              f"Contexto (extracto):\n{context_preview}"}
                ],
                temperature=0.4
            )
            
            reflection = response.choices[0].message.content.strip()
            
            # Format the reflection as a guidance section
            formatted_reflection = f"## Guía Estratégica para la Respuesta\n\n{reflection}\n\n"
            formatted_reflection += "## Contexto Completo\n\n"
            
            return formatted_reflection
        except Exception as e:
            print(f"Error generating metacognitive reflection: {e}")
            return ""  # Return empty string if reflection fails
    
    def process_query(self, query: str) -> Dict[str, Any]:
        """
        Process a user query through the agent pipeline.
        
        Args:
            query: The user's query
            
        Returns:
            Dictionary containing the results and intermediate steps
        """
        results = {
            "original_query": query,
            "model_used": self.model,
            "reasoning_steps": [] if self.debug else None
        }
        
        try:
            # Step 1: Analyze the query
            self._debug_print("Analizando la consulta...")
            print("Analizando consulta...")
            query_analysis = self.query_analyzer.analyze_query(query)
            
            if self.debug:
                # Extract key findings from analysis
                analysis_text = query_analysis.get("analysis", "")
                structured_analysis = query_analysis.get("structured_analysis", "")
                
                reasoning = f"Análisis de consulta completo:\n"
                
                # Add a simplified version of the analysis
                if structured_analysis:
                    reasoning += f"{structured_analysis}\n"
                else:
                    reasoning += f"{analysis_text[:300]}...\n"
                
                results["reasoning_steps"].append({
                    "stage": "Análisis de Consulta",
                    "reasoning": reasoning
                })
                
                self._debug_print(reasoning)
            
            results["query_analysis"] = query_analysis
            
            # Step 1.5: Enhance query based on analysis for better retrieval
            enhanced_query = self._enhance_query_with_insights(query, query_analysis)
            results["enhanced_query"] = enhanced_query
            
            if self.debug and enhanced_query != query:
                reasoning = f"Consulta mejorada para recuperación:\n{enhanced_query}\n"
                results["reasoning_steps"].append({
                    "stage": "Mejora de Consulta",
                    "reasoning": reasoning
                })
            
            # Step 2: Retrieve relevant chunks using the enhanced query
            self._debug_print("Buscando documentos relevantes...")
            print("Recuperando documentos para la defensa...")
            
            try:
                # Use the enhanced query for retrieval
                retrieved_chunks = self.retriever.retrieve(enhanced_query, self.top_k)
                results["num_chunks_retrieved"] = len(retrieved_chunks)
            except Exception as e:
                print(f"Error durante la recuperación: {e}")
                # Try a simpler approach with fewer chunks
                print("Probando con parámetros reducidos...")
                try:
                    retrieved_chunks = self.retriever.retrieve(query, min(5, self.top_k))  # Fallback to original query
                    results["num_chunks_retrieved"] = len(retrieved_chunks)
                    results["retrieval_fallback_used"] = True
                except Exception as inner_e:
                    print(f"La recuperación falló completamente: {inner_e}")
                    raise
            
            if self.debug:
                # Analyze the retrieved chunks
                num_chunks = len(retrieved_chunks)
                source_summary = {}
                
                # Count chunks per source
                for chunk in retrieved_chunks:
                    source = chunk.get('source', 'unknown')
                    if source in source_summary:
                        source_summary[source] += 1
                    else:
                        source_summary[source] = 1
                
                # Build the reasoning text
                sources_text = ", ".join([f"{src} ({count})" for src, count in source_summary.items()])
                reasoning = f"Recuperados {num_chunks} fragmentos relevantes de fuentes: {sources_text}\n"
                
                if num_chunks > 0:
                    # Add preview of top chunks
                    reasoning += f"\nVista previa de resultados prioritarios para la defensa:\n"
                    for i, chunk in enumerate(retrieved_chunks[:3]):
                        chunk_text = chunk.get('text', chunk.get('chunk', 'No content available'))
                        preview = chunk_text[:100] + "..." if len(chunk_text) > 100 else chunk_text
                        reasoning += f"{i+1}. {preview}\n"
                
                results["reasoning_steps"].append({
                    "stage": "Recuperación de Documentos",
                    "reasoning": reasoning
                })
                
                self._debug_print(reasoning)
            
            # Step 3: Aggregate and organize context
            self._debug_print("Organizando información para construir argumentos de defensa sólidos...")
            print("Agregando contexto...")
            
            try:
                aggregated_context = self.context_aggregator.aggregate_context(query, retrieved_chunks)
                results["context_length"] = len(aggregated_context)
            except Exception as e:
                print(f"Error durante la agregación de contexto: {e}")
                # Use a simple fallback context
                print("Usando agregación simple de contexto como alternativa...")
                aggregated_context = self.retriever.get_formatted_context(retrieved_chunks)
                results["context_length"] = len(aggregated_context)
                results["context_fallback_used"] = True
            
            if self.debug:
                # Analyze the context
                context_preview = aggregated_context[:200] + "..." if len(aggregated_context) > 200 else aggregated_context
                word_count = len(aggregated_context.split())
                
                reasoning = f"Organizadas {word_count} palabras de información contextual para generar precisión y claridad.\n"
                reasoning += f"Vista previa del contexto: {context_preview}\n"
                
                results["reasoning_steps"].append({
                    "stage": "Organización del Contexto",
                    "reasoning": reasoning
                })
                
                self._debug_print(reasoning)
            
            # Step 3.5: Extract key concepts and refine context
            key_concepts = self._extract_key_concepts(query, aggregated_context)
            results["key_concepts"] = key_concepts
            
            if key_concepts:
                # Refine the context based on key concepts
                refined_context = self._refine_context_with_concepts(aggregated_context, key_concepts, query)
                
                # Add metacognitive reflection for guidance
                final_context = self._add_metacognitive_reflection(query, refined_context, key_concepts) + refined_context
                
                if self.debug:
                    reasoning = "Contexto refinado basado en conceptos clave:\n"
                    reasoning += f"Conceptos identificados: {', '.join(key_concepts[:5])}"
                    if len(key_concepts) > 5:
                        reasoning += f" y {len(key_concepts) - 5} más"
                    reasoning += "\n"
                    
                    results["reasoning_steps"].append({
                        "stage": "Refinamiento de Contexto",
                        "reasoning": reasoning
                    })
            else:
                final_context = aggregated_context
            
            # Step 4: Generate the answer
            self._debug_print("Formulando respuesta basados en la evidencia organizada...")
            print("Generando respuesta...")
            try:
                answer = self.answer_generator.generate_answer(query, final_context)
            except Exception as e:
                print(f"Error durante la generación de respuesta: {e}")
                # Use legal agent as fallback
                print("Usando agente legal alternativo como respaldo...")
                # Import legal agent here to avoid circular dependencies
                try:
                    from src.agents.legal_agent import LegalAgent
                    # Create an instance of LegalAgent with defense focus
                    legal_agent = LegalAgent(model=self.model)  
                    legal_answer = legal_agent.answer_query(query, 5)  # Use just 5 chunks for fallback
                    answer = legal_answer.get("answer", "No se pudo generar un contenido preciso.")
                    results["answer_fallback_used"] = True
                except Exception as legal_error:
                    print(f"Error usando el agente legal alternativo: {legal_error}")
                    answer = "No se pudo generar un contenido preciso debido a dificultades técnicas."
                    results["answer_fallback_used"] = False
                
            results["answer"] = answer
            
            # Add sources to results
            try:
                results["sources"] = [chunk.get('source', 'unknown') for chunk in retrieved_chunks[:5]]
            except Exception as e:
                print(f"Error extrayendo fuentes: {e}")
                results["sources"] = ["Información de fuentes no disponible"]
            
            if self.debug:
                # Analyze the answer generation
                answer_preview = answer[:150] + "..." if answer else "No se generó respuesta"
                reasoning = "Respuesta generada basada en el contexto organizado.\n"
                reasoning += f"Vista previa: {answer_preview}\n"
                
                results["reasoning_steps"].append({
                    "stage": "Generación de Respuesta",
                    "reasoning": reasoning
                })
                
                self._debug_print("Generación de respuesta completa.")
            
            return results
        except Exception as e:
            error_details = traceback.format_exc()
            print(f"Error en el pipeline de agentes, usando agente legal estándar como alternativa: {e}")
            print(f"Error detallado: {error_details}")
            
            if self.debug:
                reasoning = f"Se encontró un error: {str(e)}\n"
                reasoning += "Usando agente legal estándar como alternativa."
                
                results["reasoning_steps"].append({
                    "stage": "Recuperación de Error",
                    "reasoning": reasoning
                })
                
                self._debug_print(reasoning)
            
            # Fall back to the standard legal agent
            try:
                print("Usando agente legal alternativo...")
                # Import legal agent here to avoid circular dependencies
                try:
                    from src.agents.legal_agent import LegalAgent
                    # Create an instance of LegalAgent with defense focus
                    legal_agent = LegalAgent(model=self.model)
                    legal_agent_result = legal_agent.answer_query(query, self.top_k)
                    results["error"] = str(e)
                    results["answer"] = legal_agent_result.get("answer", "No hay contenido preciso disponible del agente alternativo.")
                    
                    if "sources" in legal_agent_result:
                        results["sources"] = legal_agent_result["sources"]
                        
                    return results
                except Exception as import_error:
                    print(f"Error importando agente legal: {import_error}")
                    raise
            except Exception as fallback_error:
                # Even the fallback failed, return a simple response
                error_msg = f"Error principal: {e}\nError de alternativa: {fallback_error}"
                print(f"El agente alternativo también falló: {fallback_error}")
                results["error"] = error_msg
                results["answer"] = "Me disculpo, pero estoy teniendo dificultades técnicas procesando su consulta sobre la defensa de Cathy Barriga. Por favor, intente nuevamente más tarde."
                return results