Session 4: Query Enhancement & Context Augmentation - Bridging the Semantic Gap¶

# Demonstrating the semantic gap

class SemanticGapAnalyzer:
    """Analyze and demonstrate semantic gaps in retrieval."""

    def __init__(self, embedding_model):
        self.embedding_model = embedding_model

    def analyze_query_document_gap(self, query: str, documents: List[str]) -> Dict:
        """Analyze semantic gaps between queries and documents."""

        # Embed query and documents
        query_embedding = self.embedding_model.encode([query])
        doc_embeddings = self.embedding_model.encode(documents)

        # Calculate similarities
        similarities = []
        for doc_emb in doc_embeddings:
            similarity = np.dot(query_embedding[0], doc_emb) / (
                np.linalg.norm(query_embedding[0]) * np.linalg.norm(doc_emb)
            )
            similarities.append(similarity)

        return {
            'query': query,
            'avg_similarity': np.mean(similarities),
            'max_similarity': np.max(similarities),
            'min_similarity': np.min(similarities),
            'gap_analysis': self._analyze_gap(query, documents, similarities)
        }

# Core HyDE setup

from typing import List, Dict, Any, Optional
import openai
from sentence_transformers import SentenceTransformer
import numpy as np

class HyDEQueryEnhancer:
    """Hypothetical Document Embeddings for query enhancement."""

    def __init__(self, llm_model, embedding_model, temperature: float = 0.7):
        self.llm_model = llm_model
        self.embedding_model = embedding_model
        self.temperature = temperature

        # Different HyDE templates for various query types
        self.hyde_templates = {
            'factual': self._factual_hyde_template,
            'procedural': self._procedural_hyde_template,
            'analytical': self._analytical_hyde_template,
            'creative': self._creative_hyde_template
        }

    def enhance_query_with_hyde(self, query: str,
                               query_type: str = 'factual',
                               num_hypotheticals: int = 3) -> Dict[str, Any]:
        """Generate hypothetical documents for query enhancement."""

        print(f"Generating HyDE for query: {query[:100]}...")

        # Step 1: Classify query type if not provided
        if query_type == 'auto':
            query_type = self._classify_query_type(query)

        # Step 2: Generate hypothetical documents
        hypothetical_docs = self._generate_hypothetical_documents(
            query, query_type, num_hypotheticals
        )

        # Step 3: Create enhanced embeddings
        enhanced_embeddings = self._create_enhanced_embeddings(
            query, hypothetical_docs
        )

        # Step 4: Return comprehensive enhancement data
        return {
            'original_query': query,
            'query_type': query_type,
            'hypothetical_documents': hypothetical_docs,
            'enhanced_embedding': enhanced_embeddings['combined'],
            'original_embedding': enhanced_embeddings['original'],
            'hyde_embeddings': enhanced_embeddings['hyde_embeddings'],
            'confidence_score': self._calculate_enhancement_confidence(
                enhanced_embeddings
            )
        }

    def _classify_query_type(self, query: str) -> str:
        """Classify query type for appropriate HyDE template selection."""

        classification_prompt = f"""
        Classify the following query into one of these categories:
        1. factual - Questions seeking specific facts or information
        2. procedural - Questions asking how to do something or step-by-step processes
        3. analytical - Questions requiring analysis, comparison, or interpretation
        4. creative - Questions requiring creative or open-ended responses

        Query: {query}

        Return only the category name:
        """

        try:
            response = self.llm_model.predict(classification_prompt).strip().lower()
            if response in self.hyde_templates:
                return response
        except Exception as e:
            print(f"Classification error: {e}")

        # Default to factual
        return 'factual'

    def _factual_hyde_template(self, query: str) -> str:
        """Template for factual query types."""
        return f"""
        Write a detailed, informative document that comprehensively answers this question: {query}

        The document should:
        - Provide specific facts and data points
        - Include relevant context and background information
        - Use authoritative tone and precise language
        - Cover multiple aspects of the topic
        - Include examples where relevant

        Document:
        """

    def _procedural_hyde_template(self, query: str) -> str:
        """Template for procedural query types."""
        return f"""
        Write a detailed how-to guide that explains: {query}

        The guide should:
        - Provide step-by-step instructions
        - Include prerequisite requirements
        - Mention common pitfalls and how to avoid them
        - Explain the reasoning behind each step
        - Include tips for success

        Guide:
        """

    def _analytical_hyde_template(self, query: str) -> str:
        """Template for analytical query types."""
        return f"""
        Write an analytical document that thoroughly examines: {query}

        The analysis should:
        - Present multiple perspectives or approaches
        - Compare and contrast different options
        - Discuss pros and cons
        - Provide evidence-based reasoning
        - Draw insightful conclusions

        Analysis:
        """

    def _generate_hypothetical_documents(self, query: str,
                                       query_type: str,
                                       num_docs: int) -> List[Dict]:
        """Generate multiple hypothetical documents with variations."""

        base_template = self.hyde_templates[query_type]
        hypothetical_docs = []

        for i in range(num_docs):
            # Add variation to temperature for diversity
            varied_temperature = self.temperature + (i * 0.1)
            varied_temperature = min(varied_temperature, 1.0)

            try:
                # Generate hypothetical document
                prompt = base_template(query)

                # Use varied temperature for diversity
                doc_text = self.llm_model.predict(
                    prompt,
                    temperature=varied_temperature
                )

                hypothetical_docs.append({
                    'document': doc_text.strip(),
                    'temperature': varied_temperature,
                    'variation': i + 1,
                    'word_count': len(doc_text.split()),
                    'quality_score': self._assess_document_quality(doc_text, query)
                })

                print(f"Generated hypothetical document {i+1}/{num_docs}")

            except Exception as e:
                print(f"Error generating document {i+1}: {e}")
                continue

        # Sort by quality score
        hypothetical_docs.sort(key=lambda x: x['quality_score'], reverse=True)

        return hypothetical_docs

    def _create_enhanced_embeddings(self, query: str,
                                  hypothetical_docs: List[Dict]) -> Dict:
        """Create enhanced embeddings combining query and hypothetical docs."""

        # Original query embedding
        original_embedding = self.embedding_model.encode([query])[0]

        # Hypothetical document embeddings
        hyde_texts = [doc['document'] for doc in hypothetical_docs]
        hyde_embeddings = self.embedding_model.encode(hyde_texts)

        # Combine embeddings using weighted average
        weights = self._calculate_document_weights(hypothetical_docs)

        # Weighted combination
        weighted_hyde = np.average(hyde_embeddings, axis=0, weights=weights)

        # Combine original query with weighted hypotheticals
        # Give more weight to original query to preserve intent
        combined_embedding = (0.3 * original_embedding +
                            0.7 * weighted_hyde)

        # Normalize the combined embedding
        combined_embedding = combined_embedding / np.linalg.norm(combined_embedding)

        return {
            'original': original_embedding,
            'hyde_embeddings': hyde_embeddings,
            'weighted_hyde': weighted_hyde,
            'combined': combined_embedding,
            'weights': weights
        }

    def _calculate_document_weights(self, hypothetical_docs: List[Dict]) -> List[float]:
        """Calculate weights for hypothetical documents based on quality."""

        quality_scores = [doc['quality_score'] for doc in hypothetical_docs]

        # Softmax normalization for weights
        exp_scores = np.exp(np.array(quality_scores))
        weights = exp_scores / np.sum(exp_scores)

        return weights.tolist()

# Advanced query expansion system

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import nltk
from nltk.corpus import wordnet
from collections import defaultdict

class IntelligentQueryExpander:
    """Advanced query expansion using multiple strategies."""

    def __init__(self, llm_model, domain_corpus: Optional[List[str]] = None):
        self.llm_model = llm_model
        self.domain_corpus = domain_corpus

        # Initialize expansion strategies
        self.expansion_strategies = {
            'synonym': self._synonym_expansion,
            'semantic': self._semantic_expansion,
            'contextual': self._contextual_expansion,
            'domain_specific': self._domain_specific_expansion
        }

        # Domain-specific TF-IDF if corpus provided
        if domain_corpus:
            self.domain_tfidf = TfidfVectorizer(
                max_features=10000,
                stop_words='english',
                ngram_range=(1, 3)
            )
            self.domain_tfidf.fit(domain_corpus)

    def expand_query(self, query: str,
                    strategies: List[str] = ['semantic', 'contextual'],
                    max_expansions: int = 5) -> Dict[str, Any]:
        """Expand query using multiple strategies."""

        expansion_results = {}
        combined_expansions = set()

        # Apply each expansion strategy
        for strategy in strategies:
            if strategy in self.expansion_strategies:
                expansions = self.expansion_strategies[strategy](
                    query, max_expansions
                )
                expansion_results[strategy] = expansions
                combined_expansions.update(expansions)

        # Create final expanded query
        expanded_query = self._create_expanded_query(
            query, list(combined_expansions)
        )

        return {
            'original_query': query,
            'expansions_by_strategy': expansion_results,
            'all_expansions': list(combined_expansions),
            'expanded_query': expanded_query,
            'expansion_count': len(combined_expansions)
        }

    def _semantic_expansion(self, query: str, max_expansions: int) -> List[str]:
        """Generate semantic expansions using LLM understanding."""

        semantic_prompt = f"""
        Given this query, generate {max_expansions} semantically related terms or phrases that would help find relevant information:

        Query: {query}

        Requirements:
        1. Include synonyms and related concepts
        2. Add domain-specific terminology if applicable
        3. Include both broader and narrower terms
        4. Focus on terms likely to appear in relevant documents

        Return only the expanded terms, one per line:
        """

        try:
            response = self.llm_model.predict(semantic_prompt)
            expansions = [
                term.strip()
                for term in response.strip().split('\n')
                if term.strip() and not term.strip().startswith(('-', '*', '•'))
            ]
            return expansions[:max_expansions]

        except Exception as e:
            print(f"Semantic expansion error: {e}")
            return []

    def _contextual_expansion(self, query: str, max_expansions: int) -> List[str]:
        """Generate contextual reformulations of the query."""

        reformulation_prompt = f"""
        Reformulate this query in {max_expansions} different ways that express the same information need:

        Original Query: {query}

        Create variations that:
        1. Use different phrasing and vocabulary
        2. Approach the question from different angles
        3. Include both specific and general formulations
        4. Maintain the original intent and meaning

        Reformulations:
        """

        try:
            response = self.llm_model.predict(reformulation_prompt)
            reformulations = [
                reform.strip().rstrip('.')
                for reform in response.strip().split('\n')
                if reform.strip() and '?' in reform or len(reform.split()) > 3
            ]
            return reformulations[:max_expansions]

        except Exception as e:
            print(f"Contextual expansion error: {e}")
            return []

# Multi-query generation system

class MultiQueryGenerator:
    """Generate multiple query perspectives for comprehensive retrieval."""

    def __init__(self, llm_model):
        self.llm_model = llm_model

        self.query_perspectives = {
            'decomposition': self._decompose_complex_query,
            'specificity_levels': self._generate_specificity_variants,
            'temporal_variants': self._generate_temporal_variants,
            'perspective_shifts': self._generate_perspective_variants,
            'domain_focused': self._generate_domain_variants
        }

    def generate_multi_query_set(self, query: str,
                               perspectives: List[str] = None,
                               total_queries: int = 8) -> Dict[str, Any]:
        """Generate comprehensive query set from multiple perspectives."""

        if perspectives is None:
            perspectives = ['decomposition', 'specificity_levels', 'perspective_shifts']

        all_queries = {'original': query}
        generation_metadata = {}

        # Distribute query generation across perspectives
        queries_per_perspective = total_queries // len(perspectives)
        remaining_queries = total_queries % len(perspectives)

        for i, perspective in enumerate(perspectives):
            num_queries = queries_per_perspective
            if i < remaining_queries:
                num_queries += 1

            generated = self.query_perspectives[perspective](query, num_queries)
            all_queries[perspective] = generated

            generation_metadata[perspective] = {
                'count': len(generated),
                'method': perspective
            }

        # Flatten and deduplicate
        flattened_queries = self._flatten_and_deduplicate(all_queries)

        return {
            'original_query': query,
            'query_variants': flattened_queries,
            'queries_by_perspective': all_queries,
            'generation_metadata': generation_metadata,
            'total_variants': len(flattened_queries)
        }

    def _decompose_complex_query(self, query: str, num_queries: int) -> List[str]:
        """Decompose complex queries into simpler sub-questions."""

        decomposition_prompt = f"""
        Break down this complex query into {num_queries} simpler, focused sub-questions that together would fully answer the original question:

        Complex Query: {query}

        Requirements:
        1. Each sub-question should be independently searchable
        2. Sub-questions should cover different aspects of the main query
        3. Avoid redundancy between sub-questions
        4. Maintain logical flow and completeness

        Sub-questions:
        """

        try:
            response = self.llm_model.predict(decomposition_prompt)
            sub_questions = [
                q.strip().rstrip('?') + '?'
                for q in response.strip().split('\n')
                if q.strip() and ('?' in q or len(q.split()) > 3)
            ]

            # Ensure we have question marks
            sub_questions = [
                q if q.endswith('?') else q + '?'
                for q in sub_questions
            ]

            return sub_questions[:num_queries]

        except Exception as e:
            print(f"Decomposition error: {e}")
            return []

    def _generate_specificity_variants(self, query: str, num_queries: int) -> List[str]:
        """Generate queries at different levels of specificity."""

        specificity_prompt = f"""
        Generate {num_queries} variants of this query at different levels of specificity:

        Original Query: {query}

        Create variants that range from:
        1. Very broad/general versions
        2. Medium specificity versions
        3. Very specific/detailed versions

        Each variant should maintain the core intent but adjust the scope:
        """

        try:
            response = self.llm_model.predict(specificity_prompt)
            variants = [
                variant.strip()
                for variant in response.strip().split('\n')
                if variant.strip() and len(variant.split()) > 2
            ]

            return variants[:num_queries]

        except Exception as e:
            print(f"Specificity variant error: {e}")
            return []

# Context window optimization system

class ContextWindowOptimizer:
    """Optimize context assembly for maximum information density."""

    def __init__(self, llm_tokenizer, max_context_tokens: int = 4000):
        self.tokenizer = llm_tokenizer
        self.max_context_tokens = max_context_tokens

        # Context optimization strategies
        self.optimization_strategies = {
            'relevance_ranking': self._relevance_based_selection,
            'diversity_clustering': self._diversity_based_selection,
            'hierarchical_summary': self._hierarchical_summarization,
            'semantic_compression': self._semantic_compression
        }

    def optimize_context_window(self, query: str,
                               retrieved_chunks: List[Dict],
                               strategy: str = 'relevance_ranking') -> Dict[str, Any]:
        """Optimize context window using specified strategy."""

        # Calculate available token budget
        query_tokens = len(self.tokenizer.encode(query))
        available_tokens = self.max_context_tokens - query_tokens - 200  # Buffer

        # Apply optimization strategy
        optimized_context = self.optimization_strategies[strategy](
            query, retrieved_chunks, available_tokens
        )

        return {
            'optimized_context': optimized_context['context'],
            'selected_chunks': optimized_context['chunks'],
            'context_tokens': optimized_context['token_count'],
            'efficiency_score': optimized_context['efficiency'],
            'strategy_used': strategy,
            'original_chunk_count': len(retrieved_chunks)
        }

    def _relevance_based_selection(self, query: str, chunks: List[Dict],
                                 token_budget: int) -> Dict[str, Any]:
        """Select chunks based on relevance scores and token efficiency."""

        # Calculate relevance scores and token costs
        chunk_analysis = []
        for i, chunk in enumerate(chunks):
            content = chunk['document'].page_content
            tokens = len(self.tokenizer.encode(content))
            relevance = 1 - chunk.get('similarity_score', 0.5)  # Convert distance to similarity

            # Calculate efficiency: relevance per token
            efficiency = relevance / tokens if tokens > 0 else 0

            chunk_analysis.append({
                'index': i,
                'content': content,
                'tokens': tokens,
                'relevance': relevance,
                'efficiency': efficiency,
                'metadata': chunk.get('metadata', {})
            })

        # Sort by efficiency (relevance per token)
        chunk_analysis.sort(key=lambda x: x['efficiency'], reverse=True)

        # Select chunks within token budget
        selected_chunks = []
        total_tokens = 0

        for chunk_data in chunk_analysis:
            if total_tokens + chunk_data['tokens'] <= token_budget:
                selected_chunks.append(chunk_data)
                total_tokens += chunk_data['tokens']
            else:
                break

        # Assemble context
        context_parts = []
        for chunk_data in selected_chunks:
            source = chunk_data['metadata'].get('source', 'Unknown')
            context_parts.append(f"[Source: {source}]\n{chunk_data['content']}")

        final_context = '\n\n'.join(context_parts)

        return {
            'context': final_context,
            'chunks': selected_chunks,
            'token_count': total_tokens,
            'efficiency': np.mean([c['efficiency'] for c in selected_chunks])
        }

    def _hierarchical_summarization(self, query: str, chunks: List[Dict],
                                  token_budget: int) -> Dict[str, Any]:
        """Create hierarchical summaries when context exceeds budget."""

        # Group chunks by source/topic
        chunk_groups = self._group_chunks_by_source(chunks)

        summarized_chunks = []
        total_tokens = 0

        for group_key, group_chunks in chunk_groups.items():
            # Calculate total tokens for this group
            group_content = '\n'.join([
                chunk['document'].page_content for chunk in group_chunks
            ])
            group_tokens = len(self.tokenizer.encode(group_content))

            if group_tokens > token_budget // 4:  # Group too large, summarize
                summary = self._summarize_chunk_group(query, group_chunks)
                summary_tokens = len(self.tokenizer.encode(summary))

                if total_tokens + summary_tokens <= token_budget:
                    summarized_chunks.append({
                        'content': summary,
                        'tokens': summary_tokens,
                        'type': 'summary',
                        'source_count': len(group_chunks),
                        'group_key': group_key
                    })
                    total_tokens += summary_tokens

            else:
                # Use original chunks if they fit
                for chunk in group_chunks:
                    content = chunk['document'].page_content
                    chunk_tokens = len(self.tokenizer.encode(content))

                    if total_tokens + chunk_tokens <= token_budget:
                        summarized_chunks.append({
                            'content': content,
                            'tokens': chunk_tokens,
                            'type': 'original',
                            'group_key': group_key,
                            'metadata': chunk.get('metadata', {})
                        })
                        total_tokens += chunk_tokens

        # Assemble final context
        context_parts = []
        for chunk in summarized_chunks:
            if chunk['type'] == 'summary':
                context_parts.append(f"[Summary from {chunk['source_count']} sources]\n{chunk['content']}")
            else:
                source = chunk['metadata'].get('source', chunk['group_key'])
                context_parts.append(f"[Source: {source}]\n{chunk['content']}")

        final_context = '\n\n'.join(context_parts)

        return {
            'context': final_context,
            'chunks': summarized_chunks,
            'token_count': total_tokens,
            'efficiency': total_tokens / token_budget
        }

# Advanced RAG prompt engineering

class RAGPromptEngineer:
    """Advanced prompt engineering specifically for RAG systems."""

    def __init__(self, llm_model):
        self.llm_model = llm_model

        # Domain-specific prompt templates
        self.prompt_templates = {
            'factual_qa': self._factual_qa_template,
            'analytical': self._analytical_template,
            'procedural': self._procedural_template,
            'comparative': self._comparative_template,
            'creative': self._creative_template,
            'technical': self._technical_template
        }

        # Prompt optimization techniques
        self.optimization_techniques = {
            'chain_of_thought': self._add_chain_of_thought,
            'source_weighting': self._add_source_weighting,
            'confidence_calibration': self._add_confidence_calibration,
            'multi_step_reasoning': self._add_multi_step_reasoning
        }

    def engineer_rag_prompt(self, query: str, context: str,
                           query_type: str = 'factual_qa',
                           optimizations: List[str] = ['confidence_calibration']) -> Dict[str, Any]:
        """Engineer optimized RAG prompt with specified techniques."""

        # Start with base template
        base_prompt = self.prompt_templates[query_type](query, context)

        # Apply optimization techniques
        optimized_prompt = base_prompt
        optimization_metadata = {}

        for technique in optimizations:
            if technique in self.optimization_techniques:
                result = self.optimization_techniques[technique](
                    optimized_prompt, query, context
                )
                optimized_prompt = result['prompt']
                optimization_metadata[technique] = result['metadata']

        return {
            'base_prompt': base_prompt,
            'optimized_prompt': optimized_prompt,
            'optimizations_applied': optimizations,
            'optimization_metadata': optimization_metadata,
            'query_type': query_type
        }

    def _factual_qa_template(self, query: str, context: str) -> str:
        """Optimized template for factual question answering."""

        return f"""You are an expert information analyst. Use the provided context to answer the user's question accurately and comprehensively.

CONTEXT INFORMATION:
{context}

QUESTION: {query}

INSTRUCTIONS:
1. Base your answer strictly on the provided context
2. If the context doesn't contain sufficient information, clearly state this limitation
3. Cite specific sources when making claims
4. Provide direct quotes when appropriate
5. Distinguish between facts and interpretations
6. If there are contradictions in the sources, acknowledge and explain them

ANSWER STRUCTURE:
- Start with a direct answer to the question
- Provide supporting details from the context
- Include relevant quotes with source attribution
- Note any limitations or uncertainties

ANSWER:"""

    def _add_chain_of_thought(self, base_prompt: str, query: str,
                            context: str) -> Dict[str, Any]:
        """Add chain-of-thought reasoning to prompt."""

        cot_enhancement = """
REASONING PROCESS:
Before providing your final answer, work through your reasoning step-by-step:

1. INFORMATION EXTRACTION: What key information from the context relates to the question?
2. RELEVANCE ASSESSMENT: How does each piece of information address the question?
3. SYNTHESIS: How do you combine the information to form a comprehensive answer?
4. CONFIDENCE EVALUATION: How confident are you in your answer based on the available evidence?

Let me work through this systematically:

STEP 1 - Information Extraction:
[Identify and list key relevant information from context]

STEP 2 - Relevance Assessment:
[Evaluate how each piece of information addresses the question]

STEP 3 - Synthesis:
[Combine information into a coherent answer]

STEP 4 - Confidence Evaluation:
[Assess confidence level and note any limitations]

FINAL ANSWER:
"""

        enhanced_prompt = base_prompt + cot_enhancement

        return {
            'prompt': enhanced_prompt,
            'metadata': {
                'technique': 'chain_of_thought',
                'added_tokens': len(cot_enhancement.split()),
                'reasoning_steps': 4
            }
        }

    def _add_confidence_calibration(self, base_prompt: str, query: str,
                                  context: str) -> Dict[str, Any]:
        """Add confidence calibration to improve answer reliability."""

        # Define confidence assessment framework
        confidence_enhancement = self._build_confidence_framework()

        # Apply confidence enhancement to base prompt
        enhanced_prompt = base_prompt + confidence_enhancement

        return {
            'prompt': enhanced_prompt,
            'metadata': {
                'technique': 'confidence_calibration',
                'added_tokens': len(confidence_enhancement.split()),
                'assessment_dimensions': 3
            }
        }

    def _build_confidence_framework(self) -> str:
        """Build structured confidence assessment framework."""

        return """
CONFIDENCE ASSESSMENT:
For your final answer, provide a confidence assessment:

1. EVIDENCE STRENGTH: Rate the strength of evidence (Strong/Medium/Weak)
   - Strong: Multiple reliable sources with consistent information
   - Medium: Some sources with mostly consistent information
   - Weak: Limited or inconsistent source information
"""

        # Continue with completeness assessment
        completeness_section = """
2. COMPLETENESS: Rate answer completeness (Complete/Partial/Incomplete)
   - Complete: Context fully addresses the question
   - Partial: Context addresses some aspects but missing key information
   - Incomplete: Context has minimal relevant information
"""

        # Add confidence scoring system
        confidence_scoring = """
3. CONFIDENCE SCORE: Provide an overall confidence score (0-100%)
   - 90-100%: Highly confident, strong evidence base
   - 70-89%: Confident, adequate evidence base
   - 50-69%: Moderately confident, limited evidence
   - Below 50%: Low confidence, insufficient evidence

Include this assessment at the end of your response:

CONFIDENCE ASSESSMENT:
- Evidence Strength: [Strong/Medium/Weak]
- Completeness: [Complete/Partial/Incomplete]
- Confidence Score: [0-100%]
- Key Limitations: [List any significant limitations or uncertainties]
"""

        return base_framework + completeness_section + confidence_scoring

# Dynamic prompt adaptation system

class DynamicPromptAdapter:
    """Adapt prompts based on context quality and query characteristics."""

    def __init__(self, llm_model):
        self.llm_model = llm_model

    def adapt_prompt_dynamically(self, query: str, context: str,
                                retrieval_metadata: Dict) -> Dict[str, Any]:
        """Dynamically adapt prompt based on context and query analysis."""

        # Step 1: Analyze input characteristics
        context_analysis = self._analyze_context_quality(context, query)
        query_analysis = self._analyze_query_characteristics(query)

        # Step 2: Select optimal strategy
        prompt_strategy = self._select_prompt_strategy(
            context_analysis, query_analysis, retrieval_metadata
        )

        # Step 3: Generate adapted prompt
        adapted_prompt = self._generate_adapted_prompt(
            query, context, prompt_strategy
        )

        return {
            'adapted_prompt': adapted_prompt,
            'context_analysis': context_analysis,
            'query_analysis': query_analysis,
            'strategy': prompt_strategy,
            'adaptation_reasoning': self._explain_adaptation_reasoning(
                context_analysis, query_analysis, prompt_strategy
            )
        }

    def _analyze_context_quality(self, context: str, query: str) -> Dict[str, Any]:
        """Analyze the quality and characteristics of retrieved context."""

        # Collect basic context metrics
        basic_metrics = self._collect_basic_metrics(context)

        # Analyze source characteristics
        source_analysis = self._analyze_context_sources(context)

        # Assess relevance using LLM
        relevance_score = self._assess_context_relevance(context, query)

        # Calculate information density
        information_density = self._calculate_information_density(context)

        # Generate overall quality classification
        quality_level = self._classify_context_quality(
            basic_metrics, source_analysis, relevance_score, information_density
        )

        return {
            **basic_metrics,
            **source_analysis,
            'relevance_score': relevance_score,
            'information_density': information_density,
            'quality_level': quality_level
        }

    def _collect_basic_metrics(self, context: str) -> Dict[str, Any]:
        """Collect basic context metrics."""

        return {
            'length': len(context.split()),
            'diversity_score': self._calculate_context_diversity(context)
        }

    def _analyze_context_sources(self, context: str) -> Dict[str, Any]:
        """Analyze source characteristics in context."""

        return {
            'source_count': context.count('[Source:'),
            'source_diversity': self._analyze_source_diversity(context)
        }

    def _assess_context_relevance(self, context: str, query: str) -> float:
        """Use LLM to assess context relevance to query."""

        relevance_prompt = f"""
        Rate the relevance of this context to the given query on a scale of 0.0 to 1.0:

        Query: {query}
        Context: {context[:1000]}...

        Consider:
        1. How directly the context addresses the query
        2. The completeness of information provided
        3. The accuracy and reliability of information

        Return only a number between 0.0 and 1.0:
        """

        try:
            response = self.llm_model.predict(relevance_prompt).strip()
            relevance_score = float(response)
            return max(0.0, min(1.0, relevance_score))
        except:
            return 0.5  # Default score if assessment fails

# Complete query enhancement system

class ComprehensiveQueryEnhancer:
    """Complete query enhancement pipeline for RAG systems."""

    def __init__(self, llm_model, embedding_model, tokenizer):
        # Initialize all enhancement components
        self._initialize_components(llm_model, embedding_model, tokenizer)

    def _initialize_components(self, llm_model, embedding_model, tokenizer):
        """Initialize all query enhancement components."""

        self.hyde_enhancer = HyDEQueryEnhancer(llm_model, embedding_model)
        self.query_expander = IntelligentQueryExpander(llm_model)
        self.multi_query_gen = MultiQueryGenerator(llm_model)
        self.context_optimizer = ContextWindowOptimizer(tokenizer)
        self.prompt_engineer = RAGPromptEngineer(llm_model)
        self.dynamic_adapter = DynamicPromptAdapter(llm_model)

    async def comprehensive_enhancement(self, query: str,
                                      vector_store,
                                      enhancement_config: Dict) -> Dict[str, Any]:
        """Run comprehensive query enhancement pipeline."""

        results = {'original_query': query}

        # Phase 1: Query Enhancement
        await self._apply_query_enhancements(query, results, enhancement_config)

        # Phase 2: Enhanced Retrieval
        enhanced_retrieval = await self._perform_enhanced_retrieval(
            results, vector_store, enhancement_config
        )
        results['enhanced_retrieval'] = enhanced_retrieval

        # Phase 3: Context and Prompt Optimization  
        final_results = await self._optimize_context_and_prompts(
            query, results, enhanced_retrieval
        )

        return final_results

    async def _apply_query_enhancements(self, query: str, results: Dict, 
                                      config: Dict):
        """Apply configured query enhancement techniques."""

        # HyDE Enhancement
        if config.get('use_hyde', True):
            results['hyde_enhancement'] = self.hyde_enhancer.enhance_query_with_hyde(query)

        # Query Expansion  
        if config.get('use_expansion', True):
            results['query_expansion'] = self.query_expander.expand_query(query)

        # Multi-Query Generation
        if config.get('use_multi_query', True):
            results['multi_query'] = self.multi_query_gen.generate_multi_query_set(query)

    async def _optimize_context_and_prompts(self, query: str, results: Dict,
                                          enhanced_retrieval: Dict) -> Dict:
        """Optimize context window and engineer prompts."""

        # Context Optimization
        context_result = self.context_optimizer.optimize_context_window(
            query, enhanced_retrieval['combined_results']
        )
        results['optimized_context'] = context_result

        # Dynamic Prompt Engineering
        prompt_result = self.dynamic_adapter.adapt_prompt_dynamically(
            query, context_result['optimized_context'], enhanced_retrieval
        )
        results['engineered_prompt'] = prompt_result

        return results