src/llm/templates.py

"""LLM prompt templates and message formatting utilities."""

from typing import List, Dict, Any, Union
from dataclasses import dataclass
from langchain.schema import SystemMessage, HumanMessage


@dataclass
class PromptTemplate:
    """Template for managing prompts with variables."""
    
    system_prompt: str
    user_prompt_template: str
    
    def format(self, **kwargs) -> tuple:
        """Format the template with provided variables."""
        formatted_user = self.user_prompt_template.format(**kwargs)
        return self.system_prompt, formatted_user


# Default system prompt for audit Q&A
DEFAULT_AUDIT_SYSTEM_PROMPT = """
You are AuditQ&A, an AI Assistant for audit reports. Answer questions directly and factually based on the provided context.

Guidelines:
- Answer directly and concisely (2-3 sentences maximum)
- Use specific facts and numbers from the context
- Cite sources using [Doc i] format
- Be factual, not opinionated
- Avoid phrases like "From my point of view", "I think", "It seems"

Examples:

Query: "What challenges arise from contradictory PDM implementation guidelines?"
Context: [Retrieved documents about PDM guidelines contradictions]
Answer: "Contradictory PDM implementation guidelines cause challenges during implementation, as entities receive numerous and often conflicting directives from different authorities. For example, guidelines on transfer of funds to PDM SACCOs differ between the PDM Secretariat and PSST, and there are conflicting directives on fund diversion from various authorities."

Query: "What was the supplementary funding obtained for the wage budget?"
Context: [Retrieved documents about wage budget funding]
Answer: "The supplementary funding obtained for the wage budget was UGX.2,208,040,656."

Now answer the following question based on the provided context:
"""

# Default user prompt template
DEFAULT_USER_PROMPT_TEMPLATE = """Passages:
{context}
-----------------------
Question: {question} - Explained to audit expert
Answer in english with the passages citations:
"""


def create_audit_prompt(context_list: List[str], query: str) -> List[Dict[str, str]]:
    """
    Create audit Q&A prompt messages from context and query.
    
    Args:
        context_list: List of context passages
        query: User query
        
    Returns:
        List of message dictionaries for LLM
    """
    # Join context passages with numbering
    numbered_context = []
    for i, passage in enumerate(context_list, 1):
        numbered_context.append(f"Doc {i}: {passage}")
    
    context_str = "\n\n".join(numbered_context)
    
    # Format user prompt
    user_prompt = DEFAULT_USER_PROMPT_TEMPLATE.format(
        context=context_str,
        question=query
    )
    
    # Return as message format
    messages = [
        {"role": "system", "content": DEFAULT_AUDIT_SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt}
    ]
    
    return messages


def get_message_template(
    provider_type: str, 
    system_prompt: str, 
    user_prompt: str
) -> List[Union[Dict[str, str], SystemMessage, HumanMessage]]:
    """
    Get message template based on LLM provider type.
    
    Args:
        provider_type: Type of LLM provider
        system_prompt: System prompt content
        user_prompt: User prompt content
        
    Returns:
        List of messages in the appropriate format for the provider
    """
    provider_type = provider_type.upper()
    
    if provider_type in ['NVIDIA', 'INF_PROVIDERS', 'MISTRAL', 'OPENAI', 'OPENROUTER']:
        # Dictionary format for API-based providers
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    elif provider_type in ['DEDICATED', 'SERVERLESS', 'OLLAMA']:
        # LangChain message objects for local/dedicated providers
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt)
        ]
    else:
        # Default to dictionary format
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    
    return messages


def create_custom_prompt_template(
    system_prompt: str,
    user_template: str
) -> PromptTemplate:
    """
    Create a custom prompt template.
    
    Args:
        system_prompt: System prompt content
        user_template: User prompt template with placeholders
        
    Returns:
        PromptTemplate instance
    """
    return PromptTemplate(
        system_prompt=system_prompt,
        user_prompt_template=user_template
    )


def create_evaluation_prompt(context_list: List[str], query: str, expected_answer: str) -> List[Dict[str, str]]:
    """
    Create prompt for evaluation purposes with expected answer.
    
    Args:
        context_list: List of context passages
        query: User query
        expected_answer: Expected/ground truth answer
        
    Returns:
        List of message dictionaries for evaluation
    """
    # Join context passages
    context_str = "\n\n".join([f"Doc {i}: {passage}" for i, passage in enumerate(context_list, 1)])
    
    evaluation_system_prompt = """
    You are an evaluation assistant. Given context passages, a question, and an expected answer,
    evaluate how well the provided context supports answering the question accurately.
    
    Provide your evaluation focusing on:
    1. Relevance of the context to the question
    2. Completeness of information needed to answer
    3. Quality and accuracy of supporting details
    """
    
    user_prompt = f"""Context Passages:
{context_str}

Question: {query}
Expected Answer: {expected_answer}

Evaluation:"""
    
    return [
        {"role": "system", "content": evaluation_system_prompt},
        {"role": "user", "content": user_prompt}
    ]


def get_prompt_variants() -> Dict[str, PromptTemplate]:
    """
    Get different prompt template variants for testing.
    
    Returns:
        Dictionary of named prompt templates
    """
    variants = {
        "standard": create_custom_prompt_template(
            DEFAULT_AUDIT_SYSTEM_PROMPT,
            DEFAULT_USER_PROMPT_TEMPLATE
        ),
        
        "concise": create_custom_prompt_template(
            """You are an audit report AI assistant. Provide clear, concise answers based on the given context passages. Always cite sources using [Doc i] format.""",
            """Context:\n{context}\n\nQuestion: {question}\nAnswer:"""
        ),
        
        "detailed": create_custom_prompt_template(
            DEFAULT_AUDIT_SYSTEM_PROMPT + """\n\nAdditional Instructions:
            - Provide detailed explanations with specific examples
            - Include relevant numbers, dates, and financial figures when available
            - Structure your response with clear headings when appropriate
            - Explain the significance of findings in the context of governance and accountability""",
            DEFAULT_USER_PROMPT_TEMPLATE
        )
    }
    
    return variants


# Backward compatibility function
def format_context_with_citations(context_list: List[str]) -> str:
    """
    Format context list with document citations.
    
    Args:
        context_list: List of context passages
        
    Returns:
        Formatted context string with citations
    """
    formatted_passages = []
    for i, passage in enumerate(context_list, 1):
        formatted_passages.append(f"Doc {i}: {passage}")
    
    return "\n\n".join(formatted_passages)