ira/report/report_templates.py

from dataclasses import dataclass
from typing import Dict, Optional, List
from enum import Enum

class QueryType(Enum):
    FACTUAL = 'factual'
    EXPLORATORY = 'exploratory'
    COMPARATIVE = 'comparative'

class DetailLevel(Enum):
    BRIEF = 'brief'
    STANDARD = 'standard'
    DETAILED = 'detailed'
    COMPREHENSIVE = 'comprehensive'

@dataclass
class ReportTemplate:
    template: str
    detail_level: DetailLevel
    query_type: QueryType
    model: Optional[str] = None
    required_sections: Optional[List[str]] = None

    def validate(self) -> bool:
        """Validate that the template contains all required sections"""
        if not self.required_sections:
            return True
        return all(section in self.template for section in self.required_sections)

class ReportTemplateManager:
    def __init__(self):
        self.templates: Dict[str, ReportTemplate] = {}

    def add_template(self, template: ReportTemplate):
        if not template.validate():
            raise ValueError(f"Template missing required sections: {template.required_sections}")
        key = f"{template.query_type.value}_{template.detail_level.value}"
        self.templates[key] = template

    def get_template(self, query_type: QueryType, detail_level: DetailLevel) -> ReportTemplate:
        key = f"{query_type.value}_{detail_level.value}"
        return self.templates.get(key)

    def get_available_templates(self) -> List[str]:
        return list(self.templates.keys())

    def initialize_default_templates(self):
        # Brief templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Summary\n{summary}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{summary}', '{key_findings}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{key_findings}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Key Findings\n{key_findings}",
            detail_level=DetailLevel.BRIEF,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{key_findings}']
        ))

        # Standard templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Introduction\n{introduction}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{introduction}', '{key_findings}', '{analysis}', '{conclusion}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{methodology}', '{key_findings}', '{analysis}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.STANDARD,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}']
        ))

        # Detailed templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Introduction\n{introduction}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{introduction}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Research Questions\n{research_questions}\n\n## Literature Review\n{literature_review}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{research_questions}', '{literature_review}', '{methodology}', '{key_findings}', '{analysis}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}",
            detail_level=DetailLevel.DETAILED,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}']
        ))

        # Comprehensive templates
        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Introduction\n{introduction}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.FACTUAL,
            required_sections=['{title}', '{exec_summary}', '{introduction}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Research Questions\n{research_questions}\n\n## Literature Review\n{literature_review}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.EXPLORATORY,
            required_sections=['{title}', '{exec_summary}', '{research_questions}', '{literature_review}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))

        self.add_template(ReportTemplate(
            template="# {title}\n\n## Executive Summary\n{exec_summary}\n\n## Comparison Criteria\n{comparison_criteria}\n\n## Methodology\n{methodology}\n\n## Key Findings\n{key_findings}\n\n## Analysis\n{analysis}\n\n## Conclusion\n{conclusion}\n\n## References\n{references}\n\n## Appendices\n{appendices}",
            detail_level=DetailLevel.COMPREHENSIVE,
            query_type=QueryType.COMPARATIVE,
            required_sections=['{title}', '{exec_summary}', '{comparison_criteria}', '{methodology}', '{key_findings}', '{analysis}', '{conclusion}', '{references}', '{appendices}']
        ))