modules/scenario_contextualization/integration/enhanced_rag_engine.py

"""
Enhanced RAG Engine
Integrates scenario contextualization into RAG queries
"""
from typing import Tuple, Optional
from openai import OpenAI

from src.rag_query import RAGQueryEngine
from .feature_extractor import ADASFeatureExtractor
from ..retrieval.scenario_retriever import ScenarioRetriever
from ..formatting.constructive_formatter import ConstructiveFormatter


class EnhancedAnswer:
    """Enhanced answer with scenario context"""
    def __init__(
        self,
        answer: str,
        sources: str,
        scenarios_html: Optional[str] = None,
        scenario_count: int = 0
    ):
        self.answer = answer
        self.sources = sources
        self.scenarios_html = scenarios_html
        self.scenario_count = scenario_count


class EnhancedRAGEngine:
    """
    Enhanced RAG engine with scenario contextualization integration
    """
    def __init__(
        self,
        base_rag_engine: RAGQueryEngine,
        scenario_retriever: ScenarioRetriever,
        feature_extractor: ADASFeatureExtractor,
        formatter: ConstructiveFormatter
    ):
        self.base_rag = base_rag_engine
        self.scenario_retriever = scenario_retriever
        self.feature_extractor = feature_extractor
        self.formatter = formatter
    
    def query(
        self,
        query: str,
        user_id: Optional[str] = None,
        user_context: Optional[dict] = None
    ) -> EnhancedAnswer:
        """
        Execute enhanced query
        
        Args:
            query: User query
            user_id: User ID (optional, for personalization)
            user_context: User context (optional)
        
        Returns:
            EnhancedAnswer: Contains standard answer and scenario context
        """
        # 1. Standard RAG query
        base_answer, sources = self.base_rag.query(query)
        
        # 2. Extract ADAS features
        adas_features = self.feature_extractor.extract(query)
        
        # 3. Retrieve relevant scenarios (if related features found)
        scenarios_html = None
        scenario_count = 0
        if adas_features:
            try:
                ranked_scenarios = self.scenario_retriever.retrieve(
                    query=query,
                    adas_features=adas_features,
                    max_results=3,
                    user_context=user_context
                )
                
                if ranked_scenarios:
                    # 4. Format scenarios
                    scenarios_html = self.formatter.format_scenarios_for_ui(ranked_scenarios)
                    scenario_count = len(ranked_scenarios)
            except Exception as e:
                print(f"⚠️ Error retrieving scenarios: {e}")
                # Continue execution, scenario retrieval failure doesn't affect standard answer
        
        return EnhancedAnswer(
            answer=base_answer,
            sources=sources,
            scenarios_html=scenarios_html,
            scenario_count=scenario_count
        )