Add report templates module and tests
This commit is contained in:
parent
33d159f00c
commit
053f4a99da
|
|
@ -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}']
|
||||
))
|
||||
|
|
@ -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())
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue