import base64 import io import os import pandas as pd from docx import Document from io import BytesIO import dash import dash_bootstrap_components as dbc from dash import html, dcc, Input, Output, State, callback_context import google.generativeai as genai from docx.shared import Pt from docx.enum.style import WD_STYLE_TYPE from PyPDF2 import PdfReader from io import StringIO # Initialize Dash app app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP]) # Configure Gemini AI genai.configure(api_key=os.environ["GEMINI_API_KEY"]) model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25') # Global variables uploaded_files = {} current_document = None document_type = None # Document types and their descriptions document_types = { "Shred": "Generate an outline of the Project Work Statement (PWS)", "Pink": "Create a Pink Team document based on the PWS outline", "P.Review": "Evaluate compliance of the Pink Team document", "Red": "Generate a Red Team document based on the P.Review", "R.Review": "Evaluate compliance of the Red Team document", "G.Review": "Perform a final compliance review", "LOE": "Generate a Level of Effort (LOE) breakdown" } app.layout = dbc.Container([ dbc.Row([ dbc.Col([ html.H4("Proposal Documents", className="mt-3 mb-4"), dcc.Upload( id='upload-document', children=html.Div([ 'Drag and Drop or ', html.A('Select Files') ]), style={ 'width': '100%', 'height': '60px', 'lineHeight': '60px', 'borderWidth': '1px', 'borderStyle': 'dashed', 'borderRadius': '5px', 'textAlign': 'center', 'margin': '10px 0' }, multiple=True ), html.Div(id='file-list'), html.Hr(), html.Div([ dbc.Button( doc_type, id=f'btn-{doc_type.lower().replace(" ", "-")}', color="link", className="mb-2 w-100 text-left custom-button", style={'overflow': 'hidden', 'text-overflow': 'ellipsis', 'white-space': 'nowrap'} ) for doc_type in document_types.keys() ]) ], width=3), dbc.Col([ html.Div(style={"height": "20px"}), # Added small gap dcc.Loading( id="loading-indicator", type="dot", children=[html.Div(id="loading-output")] ), html.Div(id='document-preview', className="border p-3 mb-3"), dbc.Button("Download Document", id="btn-download", color="success", className="mt-3"), dcc.Download(id="download-document"), html.Hr(), html.Div(style={"height": "20px"}), # Added small gap dcc.Loading( id="chat-loading", type="dot", children=[ dbc.Input(id="chat-input", type="text", placeholder="Chat with AI to update document...", className="mb-2"), dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3"), html.Div(id="chat-output") ] ) ], width=9) ]) ], fluid=True) def process_document(contents, filename): content_type, content_string = contents.split(',') decoded = base64.b64decode(content_string) try: if filename.lower().endswith('.docx'): doc = Document(BytesIO(decoded)) text = "\n".join([para.text for para in doc.paragraphs]) return text elif filename.lower().endswith('.pdf'): pdf = PdfReader(BytesIO(decoded)) text = "" for page in pdf.pages: text += page.extract_text() return text else: return f"Unsupported file format: {filename}. Please upload a PDF or DOCX file." except Exception as e: return f"Error processing document: {str(e)}" @app.callback( Output('file-list', 'children'), Input('upload-document', 'contents'), State('upload-document', 'filename'), State('file-list', 'children') ) def update_output(list_of_contents, list_of_names, existing_files): global uploaded_files if list_of_contents is not None: new_files = [] for i, (content, name) in enumerate(zip(list_of_contents, list_of_names)): file_content = process_document(content, name) uploaded_files[name] = file_content new_files.append(html.Div([ html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}), html.Span(name) ])) if existing_files is None: existing_files = [] return existing_files + new_files return existing_files @app.callback( Output('file-list', 'children', allow_duplicate=True), Input({'type': 'remove-file', 'index': dash.ALL}, 'n_clicks'), State('file-list', 'children'), prevent_initial_call=True ) def remove_file(n_clicks, existing_files): global uploaded_files ctx = dash.callback_context if not ctx.triggered: raise dash.exceptions.PreventUpdate removed_file = ctx.triggered[0]['prop_id'].split(',')[0].split(':')[-1].strip('}') uploaded_files.pop(removed_file, None) return [file for file in existing_files if file['props']['children'][1]['props']['children'] != removed_file] def generate_document(document_type, file_contents): prompt = f"""Generate a {document_type} based on the following project artifacts: {' '.join(file_contents)} Instructions: 1. Create the {document_type} as a detailed document. 2. Use proper formatting and structure. 3. Include all necessary sections and details. 4. Start the output immediately with the document content. Now, generate the {document_type}: """ response = model.generate_content(prompt) return response.text @app.callback( Output('document-preview', 'children'), Output('loading-output', 'children'), [Input(f'btn-{doc_type.lower().replace(" ", "-")}', 'n_clicks') for doc_type in document_types.keys()], prevent_initial_call=True ) def generate_document_preview(*args): global current_document, document_type ctx = dash.callback_context if not ctx.triggered: raise dash.exceptions.PreventUpdate button_id = ctx.triggered[0]['prop_id'].split('.')[0] document_type = button_id.replace('btn-', '').replace('-', ' ').title() if not uploaded_files: return html.Div("Please upload project artifacts before generating a document."), "" file_contents = list(uploaded_files.values()) try: current_document = generate_document(document_type, file_contents) return dcc.Markdown(current_document), f"{document_type} generated" except Exception as e: print(f"Error generating document: {str(e)}") return html.Div(f"Error generating document: {str(e)}"), "Error" @app.callback( Output('chat-output', 'children'), Output('document-preview', 'children', allow_duplicate=True), Input('btn-send-chat', 'n_clicks'), State('chat-input', 'value'), prevent_initial_call=True ) def update_document_via_chat(n_clicks, chat_input): global current_document, document_type if not chat_input or current_document is None: raise dash.exceptions.PreventUpdate prompt = f"""Update the following {document_type} based on this instruction: {chat_input} Current document: {current_document} Instructions: 1. Provide the updated document content. 2. Maintain proper formatting and structure. 3. Incorporate the requested changes seamlessly. Now, provide the updated {document_type}: """ response = model.generate_content(prompt) current_document = response.text return f"Document updated based on: {chat_input}", dcc.Markdown(current_document) @app.callback( Output("download-document", "data"), Input("btn-download", "n_clicks"), prevent_initial_call=True ) def download_document(n_clicks): global current_document, document_type if current_document is None: raise dash.exceptions.PreventUpdate # Create an in-memory Word document doc = Document() doc.add_paragraph(current_document) # Save the document to a BytesIO object output = BytesIO() doc.save(output) return dcc.send_bytes(output.getvalue(), f"{document_type}.docx") if __name__ == '__main__': print("Starting the Dash application...") app.run(debug=False, host='0.0.0.0', port=7860) print("Dash application has finished running.")