src/agents/sql.py - master - en_vui/ge

src/agents/sql.py
141 строка6 KB
Famel1x
init
11 фев 2026, 10:08
11 фев 2026, 10:081504a02
import logging
import re
from sqlalchemy import create_engine, text
from src.core.config import settings
from src.core.dto import AgentResponse, SQLResult
from src.services.rag import SchemaRAG
import time


logger = logging.getLogger("GenBI.SQLAgent")

class SQLAgent:
    def __init__(self, llm_service):
        self.llm = llm_service
        self.rag = SchemaRAG()
        # Force UTF-8 for Windows compatibility
        self.engine = create_engine(
            settings.DATABASE_URL,
            connect_args={"client_encoding": "utf8"}
        )

    def _clean_sql(self, llm_output: str) -> str:
        """
        Cleans LLM output: removes markdown code blocks and extra text.
        """
        # 1. Remove ```sql ... ``` wrappers
        pattern = r"```sql\s*(.*?)\s*```"
        match = re.search(pattern, llm_output, re.DOTALL)
        if match:
            sql = match.group(1)
        else:
            # Try generic code block
            match_generic = re.search(r"```\s*(.*?)\s*```", llm_output, re.DOTALL)
            if match_generic:
                sql = match_generic.group(1)
            else:
                sql = llm_output

        # 2. Remove comments (optional, but safer) and trim
        # sql = re.sub(r"--.*", "", sql) 
        return sql.strip()

    def _build_prompt(self, query: str, context: str, errors: list) -> str:
        # СЛОВАРЬ И ПРАВИЛА (Hardcoded Knowledge)
        schema_hints = (
            "CRITICAL DATA DICTIONARY (Translate values to English!):\n"
            "1. METRICS (column 'kpi_short_name'):\n"
            "   - 'Ипотека' -> 'Mortgage'\n"
            "   - 'Кредиты' / 'Потреб' -> 'Consumer Loans'\n"
            "   - 'Депозиты' / 'Пассивы' -> 'Corp Deposits'\n"
            "   - 'NPS' -> 'NPS'\n"
            "2. BANKS (column 'tb'):\n"
            "   - 'Уральский' -> 'Ural Bank'\n"
            "   - 'Волго-Вятский' -> 'Volgo-Vyatsky Bank'\n"
            "   - 'Московский' -> 'Moscow Bank'\n"
            "3. BRANCHES (column 'gosb'):\n"
            "   - 'Нижегородский' -> 'Nizhny Novgorod GOSB'\n"
            "   - 'Казанский' -> 'Kazan GOSB'\n"
            "4. BUSINESS LINES (column 'bp_direction'):\n"
            "   - 'Корпоративный' -> 'Corporate'\n"
            "   - 'Розничный' -> 'Retail'\n"
        )

        base = (
            f"Role: Senior BI SQL Developer.\n"
            f"Task: Generate executable SQL for PostgreSQL.\n"
            f"Target Table: 'kpi_report'.\n\n"
            f"Schema Context:\n{context}\n\n"
            f"{schema_hints}\n"
            f"CRITICAL RULES:\n"
            f"1. Database content is STRICTLY in English. ALWAYS translate Russian terms using the Dictionary above.\n"
            f"2. Use 'bp_direction' for 'block' or 'line' queries (Retail/Corporate).\n"
            f"3. Use ILIKE for text search.\n"
            f"4. Return ONLY the raw SQL code.\n\n"
            f"User Request: {query}"
        )
        
        if errors:
            base += (
                f"\n\n🛑 PREVIOUS FAILED ATTEMPTS:\n"
                f"{'; '.join(errors)}\n"
                f"FIX: Check if you used Russian values instead of English."
            )
            
        return base
    def execute(self, query: str) -> AgentResponse:
            logger.info(f"Analyzing query: '{query}'")
            
            context = self.rag.get_context(query)
            if not context:
                logger.error("RAG failed to retrieve context")
                return AgentResponse(success=False, error="Could not retrieve schema info.")

            errors = []
            for attempt in range(3):
                logger.info(f"Attempt {attempt+1}/3 generating SQL...")
                
                prompt = self._build_prompt(query, context, errors)
                
                # Call LLM
                start_gen = time.time()
                llm_response = self.llm.generate(prompt)
                gen_time = time.time() - start_gen
                
                raw_text = llm_response.content 
                sql_query = self._clean_sql(raw_text)
                
                # Логируем сгенерированный SQL (уровень INFO, чтобы ты видел это)
                logger.info(f"Generated SQL ({gen_time:.2f}s): {sql_query}")

                try:
                    # Execute
                    with self.engine.connect() as conn:
                        start_exec = time.time()
                        cursor = conn.execute(text(sql_query))
                        keys = list(cursor.keys())
                        result = cursor.fetchall()
                        exec_time = time.time() - start_exec
                        
                        data = [dict(zip(keys, row)) for row in result]
                        
                        logger.info(f"SQL Executed successfully in {exec_time:.2f}s. Rows returned: {len(data)}")
                        
                        return AgentResponse(
                            success=True,
                            payload=SQLResult(
                                query=sql_query,
                                data=data,
                                columns=keys,  
                                row_count=len(data),
                                execution_time=exec_time
                            )
                        )

                except Exception as e:
                    err_msg = str(e).split('\n')[0]
                    logger.warning(f"Attempt {attempt+1} failed. Error: {err_msg}")
                    errors.append(f"SQL: {sql_query} | Error: {err_msg}")

            logger.error("All 3 attempts failed.")
            return AgentResponse(success=False, error="Max retries exceeded. Could not generate valid SQL.")