Update report_gen.py

report_gen.py

@@ -1,162 +1,115 @@
 import io
-import re
 from reportlab.lib import colors
 from reportlab.lib.pagesizes import A4
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Flowable
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
 from reportlab.lib.units import mm
 from reportlab.graphics.shapes import Drawing, Rect, String, Line
 
-class SliderFlowable(Flowable):
-    def __init__(self, name, value, min_val, max_val, is_percentage=False):
-        Flowable.__init__(self)
+class SliderFlowable(Drawing):
+    def __init__(self, name, value, min_val=1, max_val=10, width=400, height=50):
+        Drawing.__init__(self, width, height)
         self.name = name
         self.value = value
         self.min_val = min_val
         self.max_val = max_val
-        self.is_percentage = is_percentage
-        self.width = 400
-        self.height = 80
-
-    def draw(self):
-        drawing = Drawing(self.width, self.height)
-
+        
         # Draw slider bar
-        bar = Rect(50, 30, 300, 20, fillColor=colors.HexColor("#f7fbfd"), strokeColor=colors.HexColor("#9999ff"))
-        drawing.add(bar)
-
+        self.add(Rect(50, 15, 300, 20, fillColor=colors.lightgrey, strokeColor=colors.grey))
+        
         # Draw slider value
-        if self.max_val == self.min_val:
-            value_width = 50  # or some default width
-        else:
-            value_width = 50 + ((self.value - self.min_val) / (self.max_val - self.min_val) * 300)
-        value_bar = Rect(50, 30, value_width - 50, 20, fillColor=colors.HexColor("#9999ff"), strokeColor=None)
-        drawing.add(value_bar)
-
-        # Add slider name
-        drawing.add(String(0, 60, self.name, fontSize=12, fillColor=colors.HexColor("#26004d")))
-
-        # Add range labels
-        min_str = f"{self.min_val:.1f}%" if self.is_percentage else f"{self.min_val:.1f}"
-        max_str = f"{self.max_val:.1f}%" if self.is_percentage else f"{self.max_val:.1f}"
-        drawing.add(String(40, 10, min_str, fontSize=10, fillColor=colors.HexColor("#26004d")))
-        drawing.add(String(340, 10, max_str, fontSize=10, fillColor=colors.HexColor("#26004d")))
-
-        # Add value label
-        value_str = f"{self.value:.1f}%" if self.is_percentage else f"{self.value:.1f}"
-        drawing.add(String(value_width - 20, 55, value_str, fontSize=10, fillColor=colors.HexColor("#26004d")))
-
-        # Add value marker
-        drawing.add(Line(value_width, 25, value_width, 55, strokeColor=colors.HexColor("#26004d"), strokeWidth=2))
-
-        drawing.drawOn(self.canv, 0, 0)
+        value_width = 50 + ((self.value - self.min_val) / (self.max_val - self.min_val) * 300)
+        self.add(Rect(50, 15, value_width - 50, 20, fillColor=colors.blue, strokeColor=None))
+        
+        # Add slider name and value
+        self.add(String(0, 40, self.name, fontSize=10))
+        self.add(String(value_width - 10, 40, str(self.value), fontSize=10))
 
 def create_styles():
     styles = getSampleStyleSheet()
-    styles['Title'].fontName = 'Helvetica-Bold'
-    styles['Title'].fontSize = 18
-    styles['Title'].spaceAfter = 16
-    styles['Title'].textColor = colors.HexColor("#26004d")
-
-    styles['Heading1'].fontName = 'Helvetica-Bold'
-    styles['Heading1'].fontSize = 16
-    styles['Heading1'].spaceAfter = 10
-    styles['Heading1'].textColor = colors.HexColor("#3b0b75")
-
-    styles['Heading2'].fontName = 'Helvetica'
-    styles['Heading2'].fontSize = 14
-    styles['Heading2'].spaceAfter = 12
-    styles['Heading2'].textColor = colors.HexColor("#52176a")
-
-    styles['BodyText'].fontName = 'Helvetica'
-    styles['BodyText'].fontSize = 12
-    styles['BodyText'].spaceAfter = 12
-    styles['BodyText'].textColor = colors.HexColor("#26004d")
-
     styles.add(ParagraphStyle(
-        name='Answer',
+        name='SectionTitle',
+        parent=styles['Heading1'],
+        fontSize=16,
+        textColor=colors.HexColor("#3b0b75"),
+        spaceAfter=10
+    ))
+    styles.add(ParagraphStyle(
+        name='SubsectionTitle',
+        parent=styles['Heading2'],
+        fontSize=14,
+        textColor=colors.HexColor("#52176a"),
+        spaceAfter=8
+    ))
+    styles.add(ParagraphStyle(
+        name='BodyText',
+        parent=styles['Normal'],
+        fontSize=10,
+        textColor=colors.HexColor("#26004d"),
+        spaceAfter=6
+    ))
+    styles.add(ParagraphStyle(
+        name='AnswerText',
         parent=styles['BodyText'],
         backColor=colors.HexColor("#f0f2fd"),
-        borderColor=colors.HexColor("#9999ff"),
-        borderWidth=0.5,
-        borderPadding=(5, 5, 5, 5),
+        borderPadding=(4, 4, 4, 4),
         spaceAfter=10
     ))
-
     return styles
 
-def create_page_template(canvas, doc):
-    canvas.saveState()
-    canvas.setFillColor(colors.HexColor("#e6ebfb"))
-    canvas.rect(0, 0, doc.pagesize[0], doc.pagesize[1], fill=1)
-    canvas.setFillColor(colors.HexColor("#26004d"))
-    canvas.setFont('Helvetica', 9)
-    canvas.drawString(30, 20, f"Page {doc.page}")
-    canvas.restoreState()
-
 def generate_pdf(pages, answers):
     buffer = io.BytesIO()
     doc = SimpleDocTemplate(buffer, pagesize=A4, rightMargin=20*mm, leftMargin=20*mm, topMargin=20*mm, bottomMargin=20*mm)
     styles = create_styles()
-    story = [Paragraph("AI Trust and Opacity Evaluation", styles['Title'])]
+    story = [Paragraph("Critical Infrastructure - LLM Considerations Report", styles['Title'])]
 
-    for page in pages[:-1]:  # Skip the last page
-        if 'input_key' in page and page['input_key'] is not None:
-            story.append(Paragraph(page['title'], styles['Heading1']))
+    for page in pages[:-1]:  # Skip the last page (report generation page)
+        if page['type'] == 'description':
+            story.append(Paragraph(page['title'], styles['SectionTitle']))
             story.append(Paragraph(page['content'], styles['BodyText']))
-
+        
+        elif page['type'] == 'quiz':
+            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
+            story.append(Paragraph(page['description'], styles['BodyText']))
+            
             answer = answers.get(page['input_key'], "")
-            if isinstance(answer, list):
-                answer = ', '.join(answer)
-
-            if page.get('input_type') == 'combined':
-                option = answers.get(page['input_key'], "")
-                conclusion = answers.get(f"{page['input_key']}_conclusion", "")
-                story.append(Paragraph(f"<b>Option selected:</b> {option}", styles['Answer']))
-                story.append(Paragraph(f"<b>Conclusion:</b> {conclusion}", styles['Answer']))
-            else:
-                story.append(Paragraph(f"<b>Answer:</b> {answer}", styles['Answer']))
-
-            # Add quantitative criteria visualization if applicable
-            if page['input_key'] in ['technological_literacy', 'cognitive_mismatch']:
-                story.extend(process_quantitative_criteria(answer, styles))
-
-            story.append(Spacer(1, 12))
-
-    doc.build(story, onFirstPage=create_page_template, onLaterPages=create_page_template)
+            story.append(Paragraph(f"<b>Analysis:</b>", styles['BodyText']))
+            story.append(Paragraph(answer, styles['AnswerText']))
+            
+            risk_level = answers.get(page['slider_key'], 1)
+            story.append(SliderFlowable(f"Risk Level: {risk_level}", risk_level))
+        
+        elif page['type'] == 'inference':
+            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
+            story.append(Paragraph(page['prompt_description'], styles['BodyText']))
+            
+            user_input = answers.get(page['input_key'], "")
+            inference_output = answers.get(page['output_key'], "")
+            
+            story.append(Paragraph("<b>User Input:</b>", styles['BodyText']))
+            story.append(Paragraph(user_input, styles['AnswerText']))
+            story.append(Paragraph("<b>Inference Output:</b>", styles['BodyText']))
+            story.append(Paragraph(inference_output, styles['AnswerText']))
+        
+        elif page['type'] == 'architecture':
+            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
+            
+            # We can't include the actual architecture diagram in the PDF,
+            # so we'll just mention that it was reviewed
+            story.append(Paragraph("Architecture diagram was reviewed.", styles['BodyText']))
+            
+            analysis = answers.get(page['input_key'], "")
+            story.append(Paragraph("<b>Security Analysis:</b>", styles['BodyText']))
+            story.append(Paragraph(analysis, styles['AnswerText']))
+        
+        story.append(Spacer(1, 12))
+
+    doc.build(story)
     buffer.seek(0)
     return buffer
 
-def process_quantitative_criteria(answer, styles):
-    story = []
-    lines = answer.split('\n')
-    for line in lines:
-        parsed = parse_quantitative_criteria(line)
-        if parsed:
-            name, value, min_val, max_val, is_percentage = parsed
-            if is_percentage:
-                slider = SliderFlowable(name, value*100, min_val*100, max_val*100, is_percentage=True)
-            else:
-                slider = SliderFlowable(name, value, min_val, max_val, is_percentage=False)
-            story.append(slider)
-    return story
-
-def parse_quantitative_criteria(input_string):
-    match = re.match(r'(.+):\s*([-+]?(?:\d*\.*\d+)(?:%)?)(?:\s*\[([-+]?(?:\d*\.*\d+)(?:%)?)\s*-\s*([-+]?(?:\d*\.*\d+)(?:%)?)?\])?', input_string)
-    if match:
-        name, value, min_val, max_val = match.groups()
-        name = name.strip()
-
-        # Handle percentage inputs
-        is_percentage = '%' in value or '%' in min_val or '%' in max_val
-        value = float(value.rstrip('%'))
-        min_val = float(min_val.rstrip('%') if min_val else 0)
-        max_val = float(max_val.rstrip('%') if max_val else 100)
-
-        if is_percentage:
-            value /= 100
-            min_val /= 100
-            max_val /= 100
-
-        return name, value, min_val, max_val, is_percentage
-    return None
+def create_page_template(canvas, doc):
+    canvas.saveState()
+    canvas.setFont('Helvetica', 9)
+    canvas.drawString(20*mm, 10*mm, f"Page {doc.page}")
+    canvas.restoreState()