app.py

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.")