src/agent/sql_generator.py - master - yakobsonsa/bi_agent

src/agent/sql_generator.py
155 строк6 KB
Yakobsonsa
vers 2.0
02 фев 2026, 00:06
02 фев 2026, 00:0672f33ab
"""SQL generation service with validation."""

import logging
from typing import Optional
from src.interfaces.schema_provider import ISchemaProvider
from src.interfaces.sql_validator import ISQLValidator
from src.interfaces.llm_client import ILLMClient
from src.interfaces.conversation_memory import IConversationManager
from src.interfaces.cache import ICache
from src.db.database import DatabaseError

logger = logging.getLogger(__name__)


class SQLGenerator:
    """
    Responsible for SQL generation and validation.
    
    Single Responsibility: Generate valid SQL from natural language.
    Decoupled from BIAgent orchestration logic.
    
    Features:
    - Context-aware SQL generation
    - Multi-turn dialogue support
    - Session-based question refinement
    """
    
    def __init__(
        self,
        llm_client: ILLMClient,
        schema_provider: ISchemaProvider,
        sql_validator: ISQLValidator,
        conversation_manager: Optional[IConversationManager] = None,
        cache: Optional[ICache] = None,
        cache_ttl_seconds: int = 600
    ):
        """
        Initialize SQL generator.
        
        Args:
            llm_client: LLM client (QwenClient or compatible)
            schema_provider: Schema metadata provider
            sql_validator: SQL validation service
            conversation_manager: Optional conversation context manager
        """
        self.llm = llm_client
        self.schema_provider = schema_provider
        self.validator = sql_validator
        self.conversation_manager = conversation_manager
        self.cache = cache
        self.cache_ttl_seconds = cache_ttl_seconds
    
    async def generate(self, question: str, user_id: Optional[int] = None) -> str:
        """
        Generate validated SQL from natural language question.
        
        Supports context-aware generation for multi-turn dialogue.
        
        Process:
        1. Get conversation context (if user_id provided)
        2. Get schema context
        3. Generate SQL via LLM
        4. Normalize SQL
        5. Fix hallucinations
        6. Validate against schema
        
        Args:
            question: User's question in natural language
            user_id: Optional user ID for conversation context
        
        Returns:
            Valid SQL query string
        
        Raises:
            DatabaseError: If SQL generation or validation fails
        """
        # Get conversation context for enriched prompts
        context_prompt = ""
        if self.conversation_manager and user_id:
            context_str = self.conversation_manager.get_context(user_id, num_turns=3)
            if context_str:
                context_prompt = f"Контекст предыдущих вопросов:\n{context_str}\n\n"
        # Step 2: Get schema context
        schema_summary = self.schema_provider.get_schema_summary()
        schema_description = self.schema_provider.get_schema_description()
        schema_with_metadata = f"{schema_summary}\n\n{schema_description}"
        
        logger.info(f"Generating SQL for question: {question[:100]}...")
        
        # Step 3: Generate SQL via LLM with context
        enhanced_question = context_prompt + question
        raw_sql = await self.llm.generate_sql(enhanced_question, schema_with_metadata)
        
        if not raw_sql or not raw_sql.strip():
            raise DatabaseError("❌ Ваш вопрос не связан с анализом данных.\n\nЯ могу помочь с запросами о продажах, выручке, продуктах, городах и сотрудниках.")
        
        # Step 3: Normalize SQL
        sql_query = self.validator.normalize_sql(raw_sql)
        
        # Check if it looks like SQL
        normalized_lower = sql_query.lower()
        if "select" not in normalized_lower or "from" not in normalized_lower:
            # Check if this is a schema description request
            if self._is_schema_question(question):
                logger.info("ℹ️ Detected schema description request")
                tables = self.schema_provider.get_valid_tables()
                if tables:
                    sql_query = f"SELECT * FROM {tables[0]} LIMIT 1;"
                else:
                    raise DatabaseError("No tables available in database.")
            else:
                raise DatabaseError("❌ Ваш вопрос не связан с анализом данных.\n\nЯ могу помочь с запросами о продажах, выручке, продуктах, городах и сотрудниках.")
        
        # Step 4: Fix hallucinations
        logger.info("🧹 Checking for hallucinations in SQL...")
        fixed_sql, was_fixed = self.validator.fix_hallucinations(sql_query)
        
        if was_fixed:
            logger.info(f"🔧 Fixed hallucinated columns: {sql_query[:80]}... → {fixed_sql[:80]}...")
            sql_query = fixed_sql
        else:
            logger.info("✅ SQL is clean, no hallucinations detected")
        
        # Step 5: Validate against schema
        validation = self.validator.validate(sql_query)
        if not validation['valid']:
            error_msg = "Invalid SQL: " + "; ".join(validation['errors'])
            logger.error(f"⚠️  {error_msg}")
            raise DatabaseError(error_msg)
        
        logger.info(f"📝 Generated valid SQL: {sql_query}")
        
        return sql_query
    
    def _is_schema_question(self, question: str) -> bool:
        """
        Check if question is asking about database schema/structure.
        
        Args:
            question: User question
            
        Returns:
            True if question is about schema
        """
        question_lower = question.lower()
        schema_keywords = [
            'какие данные',
            'опиши',
            'описание',
            'структура',
            'схема',
            'что есть',
            'какие таблицы'
        ]
        return any(keyword in question_lower for keyword in schema_keywords)