ira/tests/report/test_all_detail_levels.py

import sys
import os
import asyncio
import argparse
from datetime import datetime

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 generate_report(query_type, detail_level, query, chunks):
    """Generate a report with the specified parameters."""
    synthesizer = ReportSynthesizer()
    
    print(f"\n{'='*80}")
    print(f"Generating {detail_level} report with {query_type} query type")
    print(f"{'='*80}")
    
    # Convert string values to enum objects
    query_type_enum = QueryType(query_type)
    detail_level_enum = DetailLevel(detail_level)
    
    report = await synthesizer.synthesize_report(
        query_type=query_type_enum.value,
        detail_level=detail_level_enum.value,
        query=query,
        chunks=chunks
    )
    
    # Save the report to a file
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"tests/report/{query_type}_{detail_level}_report_{timestamp}.md"
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(report)
    print(f"Report saved to: {filename}")
    
    # Print a snippet of the report
    report_preview = report[:500] + "..." if len(report) > 500 else report
    print(f"\nReport Preview:\n")
    print(report_preview)
    
    return report

async def main():
    parser = argparse.ArgumentParser(description='Test report generation with different detail levels')
    parser.add_argument('--query-type', choices=['factual', 'exploratory', 'comparative', 'code'], default='factual',
                        help='Query type to test (default: factual)')
    parser.add_argument('--detail-level', choices=['brief', 'standard', 'detailed', 'comprehensive'], default=None,
                        help='Detail level to test (default: test all)')
    args = parser.parse_args()
    
    # Test data
    queries = {
        'factual': "What is the capital of France?",
        'exploratory': "How do electric vehicles impact the environment?",
        'comparative': "Compare solar and wind energy technologies.",
        'code': "How to implement a binary search tree in Python?"
    }
    
    chunks = {
        'factual': [
            {
                'content': 'Paris is the capital of France. It is located in the north-central part of the country.',
                'source': 'Wikipedia',
                'url': 'https://en.wikipedia.org/wiki/Paris'
            }
        ],
        'exploratory': [
            {
                'content': 'Electric vehicles produce zero direct emissions, which improves air quality in urban areas.',
                'source': 'EPA',
                'url': 'https://www.epa.gov/greenvehicles/electric-vehicles'
            },
            {
                'content': 'The environmental impact of electric vehicles depends on how the electricity is generated. Renewable sources make EVs more environmentally friendly.',
                'source': 'Energy.gov',
                'url': 'https://www.energy.gov/eere/electricvehicles/electric-vehicle-benefits'
            }
        ],
        'comparative': [
            {
                'content': 'Solar energy is generated by converting sunlight into electricity using photovoltaic cells or concentrated solar power.',
                'source': 'National Renewable Energy Laboratory',
                'url': 'https://www.nrel.gov/research/re-solar.html'
            },
            {
                'content': 'Wind energy is generated by using wind turbines to create mechanical power that can be converted into electricity.',
                'source': 'Department of Energy',
                'url': 'https://www.energy.gov/eere/wind/how-do-wind-turbines-work'
            },
            {
                'content': 'Solar energy works best in sunny areas, while wind energy is more effective in windy regions. Both have different land use requirements.',
                'source': 'Renewable Energy World',
                'url': 'https://www.renewableenergyworld.com/solar/solar-vs-wind/'
            }
        ],
        'code': [
            {
                'content': 'A Binary Search Tree (BST) is a node-based binary tree data structure which has the following properties: The left subtree of a node contains only nodes with keys lesser than the node\'s key. The right subtree of a node contains only nodes with keys greater than the node\'s key.',
                'source': 'GeeksforGeeks',
                'url': 'https://www.geeksforgeeks.org/binary-search-tree-data-structure/'
            },
            {
                'content': '''
# Python program to implement a binary search tree

class Node:
    def __init__(self, key):
        self.left = None
        self.right = None
        self.val = key

# A utility function to insert a new node with the given key
def insert(root, key):
    if root is None:
        return Node(key)
    else:
        if root.val == key:
            return root
        elif root.val < key:
            root.right = insert(root.right, key)
        else:
            root.left = insert(root.left, key)
    return root

# A utility function to search a given key in BST
def search(root, key):
    # Base Cases: root is null or key is present at root
    if root is None or root.val == key:
        return root

    # Key is greater than root's key
    if root.val < key:
        return search(root.right, key)
    
    # Key is smaller than root's key
    return search(root.left, key)
''',
                'source': 'GitHub',
                'url': 'https://github.com/example/bst-implementation'
            },
            {
                'content': 'The time complexity of operations on a binary search tree is O(h) where h is the height of the tree. In the worst case, the height can be O(n) (when the tree becomes a linked list), but on average it is O(log n) for a balanced tree.',
                'source': 'Algorithm Textbook',
                'url': 'https://example.com/algorithms/bst'
            }
        ]
    }
    
    # Get the query type to test
    query_type = args.query_type
    query = queries[query_type]
    test_chunks = chunks[query_type]
    
    # Test all detail levels or just the specified one
    detail_levels = ['brief', 'standard', 'detailed', 'comprehensive'] if args.detail_level is None else [args.detail_level]
    
    for detail_level in detail_levels:
        await generate_report(query_type, detail_level, query, test_chunks)

if __name__ == "__main__":
    asyncio.run(main())