diff --git a/report/report_templates.py b/report/report_templates.py new file mode 100644 index 0000000..b9b526b --- /dev/null +++ b/report/report_templates.py @@ -0,0 +1,134 @@ +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}'] + )) diff --git a/tests/report/test_brief_report.py b/tests/report/test_brief_report.py new file mode 100644 index 0000000..69bd569 --- /dev/null +++ b/tests/report/test_brief_report.py @@ -0,0 +1,36 @@ +import sys +import os +import asyncio + +sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from report.report_synthesis import ReportSynthesizer +from report.report_templates import QueryType, DetailLevel + +async def main(): + # Initialize synthesizer + synthesizer = ReportSynthesizer() + + # Test data + query = "What is the capital of France?" + chunks = [ + { + 'content': 'Paris is the capital of France.', + 'source': 'Wikipedia', + 'url': 'https://en.wikipedia.org/wiki/Paris' + } + ] + + # Generate brief report + report = await synthesizer.synthesize_report( + query_type=QueryType.FACTUAL.value, + detail_level=DetailLevel.BRIEF.value, + query=query, + chunks=chunks + ) + + print('Generated Report:\n') + print(report) + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/report/test_report_templates.py b/tests/report/test_report_templates.py new file mode 100644 index 0000000..9894eaa --- /dev/null +++ b/tests/report/test_report_templates.py @@ -0,0 +1,24 @@ +import unittest +from report.report_templates import ReportTemplateManager, QueryType, DetailLevel + +class TestReportTemplates(unittest.TestCase): + def setUp(self): + self.manager = ReportTemplateManager() + self.manager.initialize_default_templates() + + def test_template_retrieval(self): + template = self.manager.get_template(QueryType.FACTUAL, DetailLevel.BRIEF) + self.assertIsNotNone(template) + self.assertEqual(template.detail_level, DetailLevel.BRIEF) + self.assertEqual(template.query_type, QueryType.FACTUAL) + + def test_template_validation(self): + template = self.manager.get_template(QueryType.FACTUAL, DetailLevel.BRIEF) + self.assertTrue(template.validate()) + + def test_all_templates_available(self): + templates = self.manager.get_available_templates() + self.assertEqual(len(templates), 12) # 3 query types * 4 detail levels + +if __name__ == '__main__': + unittest.main()