Update app.py

app.py

@@ -3,70 +3,80 @@ import io
 import os
 import threading
 import time
-from typing import List, Dict
+from typing import List, Tuple
 
 import dash
 import dash_bootstrap_components as dbc
-from dash import html, dcc, callback, Input, Output, State
+from dash import html, dcc, Input, Output, State, ctx
 import google.generativeai as genai
 from docx import Document
+from PyPDF2 import PdfReader
 
 # Initialize Dash app
 app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
 
-# Gemini API setup
-genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
+# Configure Gemini AI
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 model = genai.GenerativeModel('gemini-pro')
 
 def process_document(contents: str, filename: str) -> str:
     content_type, content_string = contents.split(',')
     decoded = base64.b64decode(content_string)
     
-    if filename.endswith('.docx'):
+    if filename.endswith('.pdf'):
+        pdf = PdfReader(io.BytesIO(decoded))
+        text = ""
+        for page in pdf.pages:
+            text += page.extract_text()
+    elif filename.endswith('.docx'):
         doc = Document(io.BytesIO(decoded))
-        text = '\n'.join([para.text for para in doc.paragraphs])
-    else:  # Assume it's a .doc file
-        text = decoded.decode('utf-8')
+        text = "\n".join([para.text for para in doc.paragraphs])
+    else:
+        return "Unsupported file format. Please upload a PDF or DOCX file."
     
     return text
 
 def generate_outline(text: str) -> str:
     prompt = f"""
-    Analyze the following Performance Work Statement (PWS), focusing on sections L&M:
+    Analyze the following Project Work Statement (PWS) and create an outline 
+    focusing on sections L&M. Extract the main headers, subheaders, and specific 
+    requirements in each section. Summarize the key points:
+
     {text}
-    
-    Generate an outline that follows the PWS headers and sub-headers, summarizing specific requirements in each section.
-    Format the outline with clear hierarchy and brief summaries.
+
+    Provide the outline in a structured format.
     """
     response = model.generate_content(prompt)
     return response.text
 
-def generate_pink_team_response(outline: str) -> str:
+def generate_pink_team_document(outline: str) -> str:
     prompt = f"""
-    Based on the following outline of a Performance Work Statement (PWS):
+    Based on the following outline of a Project Work Statement (PWS):
+
     {outline}
-    
-    Draft a detailed response as if MicroHealth is responding to the PWS. Follow these guidelines:
-    1. Write in Wikipedia style with active voice.
-    2. Describe in detail how MicroHealth will innovate to address each requirement.
-    3. Explain the industry best practices applied.
+
+    Create a detailed response document as if MicroHealth is responding to this PWS. 
+    Follow these guidelines:
+    1. Use Wikipedia style writing with active voice.
+    2. For each requirement, describe in detail how MicroHealth will innovate to address it.
+    3. Explain the industry best practices that will be applied.
     4. Provide measurable outcomes for the customer.
     5. Limit the use of bullet points and write predominantly in paragraph format.
     6. Ensure a logical flow of steps taken by MicroHealth for each requirement.
-    
-    Your response should be comprehensive and professional, showcasing MicroHealth's expertise and approach.
+
+    Generate a comprehensive response that showcases MicroHealth's expertise and approach.
     """
     response = model.generate_content(prompt)
     return response.text
 
 # Layout
 app.layout = dbc.Container([
-    html.H1("MicroHealth PWS Analysis Tool", className="my-4"),
+    html.H1("MicroHealth PWS Analysis and Response Generator", className="my-4"),
     dbc.Tabs([
         dbc.Tab(label="Shred", tab_id="shred", children=[
             dcc.Upload(
                 id='upload-document',
-                children=html.Div(['Drag and Drop or ', html.A('Select a File')]),
+                children=html.Div(['Drag and Drop or ', html.A('Select Files')]),
                 style={
                     'width': '100%',
                     'height': '60px',
@@ -79,44 +89,65 @@ app.layout = dbc.Container([
                 },
                 multiple=False
             ),
-            html.Div(id='output-document-upload'),
-            dbc.Spinner(html.Div(id='output-outline')),
+            dbc.Spinner(html.Div(id='shred-output')),
+            dbc.Button("Download Outline", id="download-shred", className="mt-3"),
+            dcc.Download(id="download-shred-doc")
         ]),
-        dbc.Tab(label="Pink", tab_id="pink-team", children=[
-            dbc.Button("Generate Pink Team Response", id="generate-pink-team", color="primary", className="mt-3"),
-            dbc.Spinner(html.Div(id='output-pink-team')),
+        dbc.Tab(label="Pink", tab_id="pink", children=[
+            dbc.Button("Generate Pink Team Document", id="generate-pink", className="mt-3"),
+            dbc.Spinner(html.Div(id='pink-output')),
+            dbc.Button("Download Pink Team Document", id="download-pink", className="mt-3"),
+            dcc.Download(id="download-pink-doc")
         ]),
     ], id="tabs", active_tab="shred"),
 ])
 
 @app.callback(
-    Output('output-document-upload', 'children'),
-    Output('output-outline', 'children'),
+    Output('shred-output', 'children'),
     Input('upload-document', 'contents'),
-    State('upload-document', 'filename'),
-    prevent_initial_call=True
+    State('upload-document', 'filename')
 )
-def update_output(contents, filename):
+def update_shred_output(contents, filename):
     if contents is None:
-        return "No file uploaded.", ""
+        return "Upload a document to begin."
     
     text = process_document(contents, filename)
     outline = generate_outline(text)
-    
-    return f"Processed file: {filename}", dcc.Markdown(outline)
+    return dcc.Markdown(outline)
 
 @app.callback(
-    Output('output-pink-team', 'children'),
-    Input('generate-pink-team', 'n_clicks'),
-    State('output-outline', 'children'),
-    prevent_initial_call=True
+    Output('pink-output', 'children'),
+    Input('generate-pink', 'n_clicks'),
+    State('shred-output', 'children')
 )
-def update_pink_team(n_clicks, outline):
-    if outline is None:
-        return "Please generate an outline first."
+def update_pink_output(n_clicks, shred_output):
+    if n_clicks is None or shred_output is None:
+        return "Generate an outline in the Shred tab first."
     
-    pink_team_response = generate_pink_team_response(outline['props']['children'])
-    return dcc.Markdown(pink_team_response)
+    pink_doc = generate_pink_team_document(shred_output)
+    return dcc.Markdown(pink_doc)
+
+@app.callback(
+    Output("download-shred-doc", "data"),
+    Input("download-shred", "n_clicks"),
+    State('shred-output', 'children'),
+    prevent_initial_call=True,
+)
+def download_shred(n_clicks, shred_output):
+    if shred_output is None:
+        return dash.no_update
+    return dict(content=shred_output, filename="shred_outline.md")
+
+@app.callback(
+    Output("download-pink-doc", "data"),
+    Input("download-pink", "n_clicks"),
+    State('pink-output', 'children'),
+    prevent_initial_call=True,
+)
+def download_pink(n_clicks, pink_output):
+    if pink_output is None:
+        return dash.no_update
+    return dict(content=pink_output, filename="pink_team_document.md")
 
 if __name__ == '__main__':
     print("Starting the Dash application...")