joelgilbert/NL2SQL

1"""2Agent 3: Explainer - Generate natural language explanations of query results.3"""4 5import logging6import json7from typing import List, Dict8from groq import Groq9from tenacity import retry, stop_after_attempt, wait_exponential10 11from config.settings import settings12from config.prompts import PromptTemplates13 14logger = logging.getLogger(__name__)15 16 17class ExplainerAgent:18    """Agent responsible for explaining query results in natural language."""19    20    def __init__(self):21        """Initialize explainer agent with Groq API client."""22        self.client = Groq(api_key=settings.api.groq_api_key)23        self.model = "llama-3.1-8b-instant"24    25    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))26    def explain_results(self, question: str, sql: str, results: List[Dict]) -> str:27        """28        Generate natural language explanation of query results.29        30        Args:31            question: Original user question32            sql: SQL query that was executed33            results: Query results as list of dictionaries34            35        Returns:36            Natural language explanation in markdown format37        """38        try:39            # Limit result size in prompt to avoid token limits40            sample_results = results[:10] if len(results) > 10 else results41            42            # Convert results to JSON string43            results_json = json.dumps(sample_results, indent=2, default=str)44            45            prompt = PromptTemplates.results_explanation_prompt(46                question=question,47                sql=sql,48                results=results_json,49                result_count=len(results)50            )51            52            completion = self.client.chat.completions.create(53                model=self.model,54                messages=[55                    {"role": "system", "content": PromptTemplates.RESULTS_EXPLANATION_SYSTEM},56                    {"role": "user", "content": prompt}57                ],58                temperature=0.5,59                max_tokens=80060            )61            62            explanation = completion.choices[0].message.content63            64            logger.info("Generated results explanation")65            return explanation66            67        except Exception as e:68            logger.error(f"Failed to generate explanation: {e}")69            # Fallback to basic summary70            return self._fallback_explanation(question, results)71    72    def _fallback_explanation(self, question: str, results: List[Dict]) -> str:73        """74        Generate basic explanation when API fails.75        76        Args:77            question: Original user question78            results: Query results79            80        Returns:81            Basic explanation text82        """83        row_count = len(results)84        85        explanation = f"### Query Results\n\n"86        explanation += f"Found **{row_count}** results for your question: \"{question}\"\n\n"87        88        if row_count > 0:89            # Show column names90            columns = list(results[0].keys())91            explanation += f"**Columns:** {', '.join(columns)}\n\n"92            93            # Basic insights94            explanation += "**Key Insights:**\n"95            explanation += f"- Total records returned: {row_count}\n"96            97            # If there's a numeric column, try to show sum/avg98            for col in columns:99                if row_count > 0 and isinstance(results[0].get(col), (int, float)):100                    total = sum(row.get(col, 0) for row in results if row.get(col) is not None)101                    avg = total / row_count102                    explanation += f"- Total {col}: {total:.2f}\n"103                    explanation += f"- Average {col}: {avg:.2f}\n"104                    break105        else:106            explanation += "No results found matching your criteria.\n"107        108        return explanation109    110    def generate_insights(self, results: List[Dict]) -> List[str]:111        """112        Analyze result data for patterns and insights.113        114        Args:115            results: Query results116            117        Returns:118            List of insight strings119        """120        insights = []121        122        if not results:123            insights.append("No data available for analysis")124            return insights125        126        row_count = len(results)127        insights.append(f"Total records: {row_count}")128        129        # Analyze numeric columns130        columns = list(results[0].keys())131        132        for col in columns:133            # Check if column is numeric134            sample_value = results[0].get(col)135            136            if isinstance(sample_value, (int, float)):137                values = [row.get(col) for row in results if row.get(col) is not None]138                139                if values:140                    total = sum(values)141                    avg = total / len(values)142                    max_val = max(values)143                    min_val = min(values)144                    145                    insights.append(f"{col}: min={min_val}, max={max_val}, avg={avg:.2f}")146        147        return insights148    149    def format_explanation(self, summary: str, insights: List[str]) -> str:150        """151        Combine summary and insights into formatted markdown.152        153        Args:154            summary: Main summary text155            insights: List of key insights156            157        Returns:158            Formatted explanation in markdown159        """160        explanation = f"{summary}\n\n"161        162        if insights:163            explanation += "### Key Insights\n\n"164            for insight in insights:165                explanation += f"- {insight}\n"166        167        return explanation168 169 170# Global explainer agent instance171explainer = ExplainerAgent()172