ira/report/report_templates.py

135 lines
7.2 KiB
Python

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}']
))