generation/prompt_builder.py

# DEPENDENCIES
from typing import List
from typing import Dict
from typing import Optional
from config.models import PromptType
from config.settings import get_settings
from config.models import ChunkWithScore
from config.logging_config import get_logger
from utils.error_handler import handle_errors
from generation.token_manager import TokenManager
from utils.error_handler import PromptBuildingError


# Setup Settings and Logging
settings = get_settings()
logger   = get_logger(__name__)


class PromptBuilder:
    """
    Intelligent prompt building for LLM generation: Constructs optimized prompts for different task types with proper context management and citation handling
    """
    def __init__(self, model_name: str = None):
        """
        Initialize prompt builder
        
        Arguments:
        ----------
            model_name { str } : Model name for token management
        """
        self.logger           = logger
        self.settings         = get_settings()
        self.model_name       = model_name or self.settings.OLLAMA_MODEL
        self.token_manager    = TokenManager(model_name)
        
        # Prompt templates for different tasks
        self.prompt_templates = {PromptType.QA             : {"system": self._get_qa_system_prompt(), "user": self._get_qa_user_template(), "max_context_ratio": 0.6},
                                 PromptType.SUMMARY        : {"system": self._get_summary_system_prompt(), "user": self._get_summary_user_template(), "max_context_ratio": 0.8,},
                                 PromptType.ANALYTICAL     : {"system": self._get_analytical_system_prompt(), "user": self._get_analytical_user_template(), "max_context_ratio": 0.7},
                                 PromptType.COMPARISON     : {"system": self._get_comparison_system_prompt(), "user": self._get_comparison_user_template(), "max_context_ratio": 0.5},
                                 PromptType.EXTRACTION     : {"system": self._get_extraction_system_prompt(), "user": self._get_extraction_user_template(), "max_context_ratio": 0.6},
                                 PromptType.CREATIVE       : {"system": self._get_creative_system_prompt(), "user": self._get_creative_user_template(), "max_context_ratio": 0.4},
                                 PromptType.CONVERSATIONAL : {"system": self._get_conversational_system_prompt(), "user": self._get_conversational_user_template(), "max_context_ratio": 0.5}
                                }
    

    def build_prompt(self, query: str, context: str, sources: List[ChunkWithScore], prompt_type: PromptType = PromptType.QA,  include_citations: bool = True,
                     max_completion_tokens: int = 1000) -> Dict[str, str]:
        """
        Build complete prompt for LLM generation
        
        Arguments:
        ----------
            query                      { str }   : User query
            
            context                    { str }   : Retrieved context
            
            sources                   { list }   : Source chunks
            
            prompt_type           { PromptType } : Type of prompt to build
            
            include_citations         { bool }   : Whether to include citation instructions
            
            max_completion_tokens     { int }    : Tokens to reserve for completion
        
        Returns:
        --------
                         { dict }                : Dictionary with 'system' and 'user' prompts
        """
        if not query or not context:
            raise PromptBuildingError("Query and context cannot be empty")
        
        try:
            # Get template for prompt type
            template          = self.prompt_templates.get(prompt_type, self.prompt_templates[PromptType.QA])
            
            # Optimize context to fit within token limits
            optimized_context = self._optimize_context(context               = context, 
                                                       template              = template, 
                                                       max_completion_tokens = max_completion_tokens,
                                                      )
            
            # Build system prompt
            system_prompt     = self._build_system_prompt(template          = template, 
                                                          include_citations = include_citations, 
                                                          prompt_type       = prompt_type,
                                                         )
            
            # Build user prompt
            user_prompt       = self._build_user_prompt(template    = template, 
                                                        query       = query, 
                                                        context     = optimized_context, 
                                                        sources     = sources, 
                                                        prompt_type = prompt_type,
                                                       )
            
            # Validate token usage
            self._validate_prompt_length(system_prompt         = system_prompt, 
                                         user_prompt           = user_prompt, 
                                         max_completion_tokens = max_completion_tokens,
                                        )
            
            self.logger.debug(f"Built {prompt_type.value} prompt: {self.token_manager.count_tokens(system_prompt + user_prompt)} tokens")
            
            return {"system"   : system_prompt,
                    "user"     : user_prompt,
                    "metadata" : {"prompt_type"    : prompt_type.value,
                                  "context_tokens" : self.token_manager.count_tokens(optimized_context),
                                  "total_tokens"   : self.token_manager.count_tokens(system_prompt + user_prompt),
                                  "sources_count"  : len(sources),
                                 }
                  }
        
        except Exception as e:
            self.logger.error(f"Prompt building failed: {repr(e)}")
            raise PromptBuildingError(f"Prompt building failed: {repr(e)}")
    

    def _optimize_context(self, context: str, template: Dict, max_completion_tokens: int) -> str:
        """
        Optimize context to fit within token limits
        
        Arguments:
        ----------
            context               { str }  : Context text to optimize

            template              { Dict } : Prompt template
            
            max_completion_tokens { int }  : Tokens to reserve for completion
        
        Returns:
        --------
                        { str }            : Optimized context
        """
        # Calculate tokens for system prompt
        system_tokens          = self.token_manager.count_tokens(template["system"])
        
        # Estimate user template tokens by removing placeholders
        user_template_clean    = template["user"]
        
        # Remove all known placeholders
        placeholders_to_remove = ["{query}", "{context}", "{sources_info}", "{focus}"]
        
        for placeholder in placeholders_to_remove:
            user_template_clean = user_template_clean.replace(placeholder, "")
        
        user_template_tokens = self.token_manager.count_tokens(user_template_clean)
        
        # Calculate available tokens for context - Reserve: system + user template + completion + buffer
        reserved_tokens      = system_tokens + user_template_tokens + max_completion_tokens + 100
        max_context_tokens   = self.token_manager.context_window - reserved_tokens
        
        # Ensure we don't exceed zero
        if (max_context_tokens <= 0):
            self.logger.warning(f"No tokens available for context. Reserved: {reserved_tokens}, Window: {self.token_manager.context_window}")
            return ""
        
        # Apply max context ratio from template config
        ratio_limit        = int(self.token_manager.context_window * template["max_context_ratio"])
        max_context_tokens = min(max_context_tokens, ratio_limit)
        
        # Log optimization details
        self.logger.debug(f"Context optimization: max_tokens={max_context_tokens}, ratio={template['max_context_ratio']}")
        
        # Truncate context to fit
        optimized_context  = self.token_manager.truncate_to_fit(context, max_context_tokens, strategy="end")
        
        # Calculate reduction percentage
        original_tokens    = self.token_manager.count_tokens(context)
        optimized_tokens   = self.token_manager.count_tokens(optimized_context)
        
        if (original_tokens > optimized_tokens):
            reduction_pct = ((original_tokens - optimized_tokens) / original_tokens) * 100
            
            self.logger.info(f"Context reduced by {reduction_pct:.1f}% ({original_tokens} → {optimized_tokens} tokens)")
        
        return optimized_context
    

    def _build_system_prompt(self, template: Dict, include_citations: bool, prompt_type: PromptType) -> str:
        """
        Build system prompt
        """
        system_prompt = template["system"]
        
        # Add citation instructions if needed
        if (include_citations and (prompt_type != PromptType.CREATIVE)):
            citation_instructions = self._get_citation_instructions()
            system_prompt        += "\n\n" + citation_instructions
        
        return system_prompt
    

    def _build_user_prompt(self, template: Dict, query: str, context: str, sources: List[ChunkWithScore], prompt_type: PromptType) -> str:
        """
        Build user prompt
        """
        # Format sources information
        sources_info = self._format_sources_info(sources) if sources else ""
        
        # Build user prompt using template
        user_prompt  = template["user"].format(query        = query,
                                               context      = context,
                                               sources_info = sources_info,
                                             )
        
        # Add task-specific formatting
        if (prompt_type == PromptType.COMPARISON):
            user_prompt = self._enhance_comparison_prompt(user_prompt = user_prompt, 
                                                          sources     = sources,
                                                         )

        elif (prompt_type == PromptType.ANALYTICAL):
            user_prompt = self._enhance_analytical_prompt(user_prompt = user_prompt, 
                                                          sources     = sources,
                                                         )
        
        return user_prompt
    

    def _format_sources_info(self, sources: List[ChunkWithScore]) -> str:
        """
        Format sources information for the prompt
        """
        if not sources:
            return ""
        
        sources_list = list()

        for i, source in enumerate(sources, 1):
            chunk       = source.chunk
            source_info = f"Source [{i}]:"
            
            if hasattr(chunk, 'metadata') and 'filename' in chunk.metadata:
                source_info += f" {chunk.metadata['filename']}"
            
            if chunk.page_number:
                source_info += f" (page {chunk.page_number})"
            
            if chunk.section_title:
                source_info += f" - {chunk.section_title}"
            
            sources_list.append(source_info)
        
        return "Available sources:\n" + "\n".join(sources_list)
    

    def _enhance_comparison_prompt(self, user_prompt: str, sources: List[ChunkWithScore]) -> str:
        """
        Enhance prompt for comparison tasks
        
        Arguments:
        ----------
            user_prompt { str }  : Base user prompt

            sources     { list } : Source chunks
        
        Returns:
        --------
                { str }          : Enhanced prompt
        """
        if (len(sources) < 2):
            return user_prompt
        
        enhancement  = "\n\nPlease compare and contrast the information from different sources. "
        enhancement += "Highlight agreements, disagreements, and complementary information. "
        enhancement += "If sources conflict, present both perspectives clearly."
        
        return user_prompt + enhancement
    

    def _enhance_analytical_prompt(self, user_prompt: str, sources: List[ChunkWithScore]) -> str:
        """
        Enhance prompt for analytical tasks - FIXED SIGNATURE
        
        Arguments:
        ----------
            user_prompt { str }  : Base user prompt

            sources     { list } : Source chunks
        
        Returns:
        --------
                { str }          : Enhanced prompt
        """
        enhancement  = "\n\nProvide analytical insights by:"
        enhancement += "\n1. Identifying patterns and relationships in the information"
        enhancement += "\n2. Analyzing implications and consequences"
        enhancement += "\n3. Evaluating the strength of evidence from different sources"
        enhancement += "\n4. Drawing well-reasoned conclusions"
        
        return user_prompt + enhancement
    

    def _validate_prompt_length(self, system_prompt: str, user_prompt: str, max_completion_tokens: int):
        """
        Validate that prompt fits within context window
        
        Arguments:
        ----------
            system_prompt         { str } : System prompt

            user_prompt           { str } : User prompt
            
            max_completion_tokens { int } : Tokens needed for completion
        
        Raises:
        -------
            PromptBuildingError           : If prompt exceeds context window
        """
        system_tokens  = self.token_manager.count_tokens(system_prompt)
        user_tokens    = self.token_manager.count_tokens(user_prompt)
        total_tokens   = system_tokens + user_tokens
        total_required = total_tokens + max_completion_tokens
        
        if (total_required > self.token_manager.context_window):
            error_msg = (f"Prompt exceeds context window:\n"
                         f"- System prompt: {system_tokens} tokens\n"
                         f"- User prompt: {user_tokens} tokens\n"
                         f"- Completion reserve: {max_completion_tokens} tokens\n"
                         f"- Total required: {total_required} tokens\n"
                         f"- Context window: {self.token_manager.context_window} tokens\n"
                         f"- Overflow: {total_required - self.token_manager.context_window} tokens"
                        )
            self.logger.error(error_msg)

            raise PromptBuildingError(error_msg)
        
        # Log successful validation
        utilization = (total_required / self.token_manager.context_window) * 100

        self.logger.debug(f"Prompt validation passed: {total_required}/{self.token_manager.context_window} tokens ({utilization:.1f}% utilization)")


    def _get_citation_instructions(self) -> str:
        """
        Get citation instructions for system prompt
        """
        return ("CITATION INSTRUCTIONS:\n"
                "1. Always cite your sources using [number] notation\n"
                "2. Use the citation number that corresponds to the source in the context\n"
                "3. Cite sources for all factual claims and specific information\n"
                "4. If information appears in multiple sources, cite the most relevant one\n"
                "5. Do not make up information not present in the provided context\n"
                "6. If the context doesn't contain the answer, explicitly state this"
               )
    

    # Template methods for different prompt types
    def _get_qa_system_prompt(self) -> str:
        return ("You are a precise and helpful AI assistant that answers questions based solely on the provided context.\n\n"
                "Core Principles:\n"
                "1. ONLY use information from the provided context\n"
                "2. Be concise but complete - don't omit important details\n"
                "3. Structure complex answers clearly\n"
                "4. If information is ambiguous or conflicting, acknowledge this\n"
                "5. Never make up or infer information not present in the context"
               )
    

    def _get_qa_user_template(self) -> str:
        return ("Context Information:\n{context}\n\n"
                "{sources_info}\n\n"
                "Question: {query}\n\n"
                "Instructions: Answer the question using ONLY the information provided in the context above. "
                "Cite your sources using [number] notation. If the context doesn't contain enough information "
                "to answer fully, state this clearly.\n\n"
                "Answer:"
               )
    

    def _get_summary_system_prompt(self) -> str:
        return ("You are a thorough AI assistant that provides comprehensive summaries based on the provided context.\n\n"
                "Summary Guidelines:\n"
                "1. Capture all key points and main ideas\n"
                "2. Maintain the original meaning and intent\n"
                "3. Organize information logically\n"
                "4. Be comprehensive but concise\n"
                "5. Highlight important findings and conclusions"
               )
    

    def _get_summary_user_template(self) -> str:
        return ("Content to summarize:\n{context}\n\n"
                "Please provide a comprehensive summary that captures all key points and main ideas. "
                "Organize the summary logically and ensure it reflects the original content accurately.\n\n"
                "Summary:"
               )
    

    def _get_analytical_system_prompt(self) -> str:
        return ("You are an analytical AI assistant that provides insights based on the provided context.\n\n"
                "Analytical Guidelines:\n"
                "1. Analyze patterns and connections in the information\n"
                "2. Compare different perspectives if multiple sources exist\n"
                "3. Highlight key findings and implications\n"
                "4. Identify gaps or limitations in the available information\n"
                "5. Provide well-reasoned analysis and conclusions"
               )

    
    def _get_analytical_user_template(self) -> str:
        return ("Context for analysis:\n{context}\n\n"
                "{sources_info}\n\n"
                "Analytical task: {query}\n\n"
                "Please provide a detailed analysis based on the context above. "
                "Identify patterns, relationships, and implications. "
                "Cite sources for all analytical claims.\n\n"
                "Analysis:"
               )

    
    def _get_comparison_system_prompt(self) -> str:
        return ("You are an AI assistant that compares information across multiple sources.\n\n"
                "Comparison Guidelines:\n"
                "1. Identify similarities and differences between sources\n"
                "2. Note if sources agree or disagree on specific points\n"
                "3. Highlight complementary information\n"
                "4. Present conflicting perspectives fairly\n"
                "5. Draw conclusions about the overall consensus or disagreement"
               )

    
    def _get_comparison_user_template(self) -> str:
        return ("Context from multiple sources:\n{context}\n\n"
                "{sources_info}\n\n"
                "Comparison task: {query}\n\n"
                "Please compare how different sources address this topic. "
                "Identify agreements, disagreements, and complementary information. "
                "Cite specific sources for each point of comparison.\n\n"
                "Comparison:"
               )

    
    def _get_extraction_system_prompt(self) -> str:
        return ("You are a precise AI assistant that extracts specific information from context.\n\n"
                "Extraction Guidelines:\n"
                "1. Extract only the requested information\n"
                "2. Be thorough and complete\n"
                "3. Maintain accuracy and precision\n"
                "4. Organize extracted information clearly\n"
                "5. Cite sources for all extracted information"
               )

    
    def _get_extraction_user_template(self) -> str:
        return ("Context:\n{context}\n\n"
                "{sources_info}\n\n"
                "Extraction task: {query}\n\n"
                "Please extract the requested information from the context above. "
                "Be thorough and precise. Cite sources for all extracted information.\n\n"
                "Extracted Information:"
               )

    
    def _get_creative_system_prompt(self) -> str:
        return ("You are a creative AI assistant that generates content based on provided context.\n\n"
                "Creative Guidelines:\n"
                "1. Use the context as inspiration and foundation\n"
                "2. Be creative and engaging\n"
                "3. Maintain coherence with the source material\n"
                "4. You may extrapolate and build upon the provided information\n"
                "5. Clearly distinguish between source-based content and creative additions"
               )

    
    def _get_creative_user_template(self) -> str:
        return ("Context and inspiration:\n{context}\n\n"
                "Creative task: {query}\n\n"
                "Please create content based on the context above. "
                "You may use the context as inspiration and build upon it creatively. "
                "If you add information beyond what's in the context, make this clear.\n\n"
                "Creative Response:"
               )

    
    def _get_conversational_system_prompt(self) -> str:
        return ("You are a helpful and engaging conversational AI assistant.\n\n"
                "Conversational Guidelines:\n"
                "1. Be natural and engaging in conversation\n"
                "2. Use the provided context to inform your responses\n"
                "3. Maintain a friendly and helpful tone\n"
                "4. Ask clarifying questions when needed\n"
                "5. Cite sources when providing specific information"
               )

    
    def _get_conversational_user_template(self) -> str:
        return ("Context for our conversation:\n{context}\n\n"
                "Current message: {query}\n\n"
                "Please respond naturally and helpfully based on the context above. "
                "If providing specific information, cite your sources.\n\n"
                "Response:"
               )


# Global prompt builder instance
_prompt_builder = None


def get_prompt_builder(model_name: str = None) -> PromptBuilder:
    """
    Get global prompt builder instance (singleton)
    
    Arguments:
    ----------
        model_name { str } : Model name for token management
    
    Returns:
    --------
        { PromptBuilder }  : PromptBuilder instance
    """
    global _prompt_builder
    
    if _prompt_builder is None or (model_name and _prompt_builder.model_name != model_name):
        _prompt_builder = PromptBuilder(model_name)
    
    return _prompt_builder


@handle_errors(error_type = PromptBuildingError, log_error = True, reraise = False)
def build_qa_prompt(query: str, context: str, sources: List[ChunkWithScore], **kwargs) -> Dict[str, str]:
    """
    Convenience function for building QA prompts
    
    Arguments:
    ----------
        query   { str }  : User query
        
        context { str }  : Retrieved context
        
        sources { list } : Source chunks
        
        **kwargs         : Additional prompt building arguments
    
    Returns:
    --------
             { dict }    : Dictionary with system and user prompts
    """
    builder = get_prompt_builder()
    
    return builder.build_prompt(query, context, sources, PromptType.QA, **kwargs)