feat: increase complexity of agents

src/agents/agent_director.py

@@ -1,3 +1,4 @@
+import re
 from openai import OpenAI
 from typing import List, Dict, Any, Optional, Tuple
 import sys
@@ -299,6 +300,7 @@ class AgentDirector:
         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):
@@ -306,6 +308,281 @@ class AgentDirector:
         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.
@@ -350,19 +627,31 @@ class AgentDirector:
             
             results["query_analysis"] = query_analysis
             
-            # Step 2: Retrieve relevant chunks
+            # 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:
-                retrieved_chunks = self.retriever.retrieve(query, self.top_k)
+                # 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))
+                    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:
@@ -431,11 +720,36 @@ class AgentDirector:
                 
                 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, aggregated_context)
+                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

src/agents/legal_agent.py

@@ -70,16 +70,25 @@ class LegalAgent:
             # 2. Format the context for the LLM
             context = self.retriever.get_formatted_context(retrieved_chunks)
             
-            # 3. Build the system prompt
-            system_prompt = self._get_legal_system_prompt(context, query)
+            # 3. Extract key concepts for deeper analysis
+            key_concepts = self._extract_key_concepts(query, context)
             
-            # 4. Prepare messages for the chat completion
+            # 4. Enhance context with concept highlighting and reorganization
+            enhanced_context = self._enhance_context_with_concepts(context, key_concepts, query)
+            
+            # 5. Add metacognitive reflection to guide reasoning
+            final_context = self._add_strategic_guidance(query, enhanced_context, key_concepts) + enhanced_context
+            
+            # 6. Build the system prompt
+            system_prompt = self._get_legal_system_prompt(final_context, query)
+            
+            # 7. Prepare messages for the chat completion
             messages = [
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": query}
             ]
             
-            # 5. Generate the answer
+            # 8. Generate the answer
             response = self.client.chat.completions.create(
                 model=self.model,
                 messages=messages,
@@ -89,10 +98,11 @@ class LegalAgent:
             
             answer = response.choices[0].message.content
             
-            # 6. Build result dictionary
+            # 9. Build result dictionary
             result = {
                 "answer": answer,
-                "sources": [chunk.get('source', 'desconocido') for chunk in retrieved_chunks[:5]]
+                "sources": [chunk.get('source', 'desconocido') for chunk in retrieved_chunks[:5]],
+                "key_concepts": key_concepts
             }
             
             return result
@@ -105,4 +115,174 @@ class LegalAgent:
             return {
                 "answer": f"I apologize, but I encountered an error while processing your legal query. This is likely due to the test mode of the application. In a full deployment, I would provide a detailed legal analysis based on relevant precedents and statutes.\n\nYour query was: {query}\n\nError details: {str(e)}",
                 "sources": ["Error - no sources available"]
-            } 
+            }
+            
+    def _extract_key_concepts(self, query: str, context: str) -> List[str]:
+        """
+        Extract key legal concepts from query and context to guide analysis.
+        
+        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 especializado en derecho penal y defensa legal. "
+                "Tu tarea es identificar y extraer los conceptos legales clave, fundamentos jurídicos "
+                "y argumentos principales de la información proporcionada. Enfócate específicamente en "
+                "conceptos relevantes para la defensa de Cathy Barriga. "
+                "Extrae entre 5-8 conceptos clave organizados por orden de importancia para la defensa."
+            )
+            
+            # 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 legal: {query}\n\nContexto legal:\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 reasonable number of concepts
+            if len(concepts) > 10:
+                concepts = concepts[:10]
+            elif len(concepts) < 3 and len(concepts_text) > 50:
+                # If 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]
+            
+            return concepts
+        except Exception as e:
+            print(f"Error extracting key concepts: {e}")
+            return []  # Return empty list if extraction fails
+            
+    def _enhance_context_with_concepts(self, context: str, concepts: List[str], query: str) -> str:
+        """
+        Enhance context by emphasizing key concepts and improving structure.
+        
+        Args:
+            context: The original context
+            concepts: Key legal concepts extracted
+            query: The user's query
+            
+        Returns:
+            Enhanced context
+        """
+        try:
+            if not concepts:
+                return context
+                
+            system_prompt = (
+                "Eres un experto legal especializado en estructurar información para maximizar su utilidad "
+                "en estrategias de defensa penal. Tu tarea es reorganizar el contexto proporcionado para:"
+                "\n1. Enfatizar los conceptos clave identificados relevantes para la defensa de Cathy Barriga"
+                "\n2. Organizar la información en una estructura lógica y progresiva"
+                "\n3. Conectar elementos relacionados que puedan fortalecer la defensa"
+                "\n4. Incluir un resumen inicial que integre los conceptos fundamentales"
+                "\n5. Destacar elementos exculpatorios o mitigantes"
+                "\nMantén TODA la información relevante del contexto original."
+            )
+            
+            # Create a concepts section
+            concepts_text = "\n".join([f"- {concept}" for concept in concepts])
+            
+            # For large contexts, process a reasonable portion
+            if len(context) > 6000:
+                context_preview = context[:6000]
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": f"Consulta legal: {query}\n\nConceptos clave a enfatizar:\n{concepts_text}\n\n"
+                                              f"Contexto a mejorar (extracto):\n\n{context_preview}"}
+                    ],
+                    temperature=0.3
+                )
+                
+                enhanced_preview = response.choices[0].message.content.strip()
+                
+                # Add notice about partial enhancement
+                return enhanced_preview + "\n\n--- CONTEXTO ADICIONAL ---\n\n" + context[6000:]
+            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 legal: {query}\n\nConceptos clave a enfatizar:\n{concepts_text}\n\n"
+                                              f"Contexto a mejorar:\n\n{context}"}
+                    ],
+                    temperature=0.3
+                )
+                
+                return response.choices[0].message.content.strip()
+        except Exception as e:
+            print(f"Error enhancing context: {e}")
+            return context  # Return original context if enhancement fails
+            
+    def _add_strategic_guidance(self, query: str, context: str, concepts: List[str]) -> str:
+        """
+        Add strategic guidance to direct reasoning and defense strategy.
+        
+        Args:
+            query: The user's query
+            context: The enhanced context
+            concepts: Key legal concepts
+            
+        Returns:
+            Strategic guidance section to prepend to context
+        """
+        try:
+            system_prompt = (
+                "Eres un estratega legal de elite especializado en defensa penal. Tu función es proporcionar "
+                "una guía estratégica que:"
+                "\n1. Identifique las principales líneas argumentativas disponibles para la defensa de Cathy Barriga"
+                "\n2. Señale conexiones no evidentes entre elementos del contexto que podrían fortalecer la defensa"
+                "\n3. Identifique debilidades en la argumentación acusatoria"
+                "\n4. Sugiera enfoques interpretativos alternativos favorables"
+                "\n5. Destaque principios legales fundamentales aplicables a la defensa"
+                "\nEsta guía orientará la estrategia de defensa."
+            )
+            
+            # Use 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 para la defensa de Cathy Barriga: {query}\n\n"
+                                          f"Conceptos legales clave:\n{concepts_text}\n\n"
+                                          f"Contexto (extracto):\n{context_preview}"}
+                ],
+                temperature=0.4
+            )
+            
+            guidance = response.choices[0].message.content.strip()
+            
+            # Format the guidance as a distinct section
+            formatted_guidance = f"## Guía Estratégica para la Defensa\n\n{guidance}\n\n"
+            formatted_guidance += "## Análisis Detallado\n\n"
+            
+            return formatted_guidance
+        except Exception as e:
+            print(f"Error generating strategic guidance: {e}")
+            return ""  # Return empty string if guidance generation fails