Update app.py via AI Editor

app.py

@@ -12,49 +12,43 @@ import threading
 import tempfile
 import shutil
 import logging
-from flask import request
+from flask import request, make_response, g
 from dash.exceptions import PreventUpdate
 import PyPDF2
 import docx
 import chardet
 
-# --- Logging Setup ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("maiko_matrix_app")
 
-# --- App and Session Setup ---
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
-server = app.server
-
-openai.api_key = os.environ.get('OPENAI_API_KEY')
-
-# Session management globals
 SESSION_DATA = {}
 SESSION_LOCKS = {}
 
-def get_session_id():
-    if 'session-id' in request.cookies:
-        return request.cookies['session-id']
-    # Generate new session id (uuid4) if not present
-    session_id = str(uuid.uuid4())
-    logger.info(f"Generated new session id: {session_id}")
-    return session_id
+def get_session_id(set_cookie=False):
+    # Try to get session id from browser cookies
+    sid = request.cookies.get('session-id')
+    if sid and sid in SESSION_DATA:
+        return sid
+    # If not present, generate a new one
+    sid = str(uuid.uuid4())
+    if set_cookie:
+        g.set_cookie_sid = sid
+    return sid
 
 def get_session_data():
-    session_id = get_session_id()
-    if session_id not in SESSION_DATA:
-        SESSION_DATA[session_id] = {
-            'uploaded_files': {},  # filename -> temp file path
-            'file_texts': {},      # filename -> extracted text
+    sid = get_session_id()
+    if sid not in SESSION_DATA:
+        SESSION_DATA[sid] = {
+            'uploaded_files': {},
+            'file_texts': {},
             'current_matrix': None,
             'matrix_type': None,
-            'temp_dir': tempfile.mkdtemp(prefix=f"maiko_{session_id}_"),
+            'temp_dir': tempfile.mkdtemp(prefix=f"maiko_{sid}_"),
         }
-        SESSION_LOCKS[session_id] = threading.Lock()
-    return SESSION_DATA[session_id], SESSION_LOCKS[session_id]
+        SESSION_LOCKS[sid] = threading.Lock()
+    return SESSION_DATA[sid], SESSION_LOCKS[sid]
 
 def cleanup_session_tempdirs():
-    # Called optionally on server shutdown for cleanup
     for sess in SESSION_DATA.values():
         try:
             shutil.rmtree(sess['temp_dir'])
@@ -81,7 +75,102 @@ matrix_types = {
     "SWOT Matrix": "Create a matrix analyzing Strengths, Weaknesses, Opportunities, and Threats."
 }
 
-# --- UI Layout ---
+app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+server = app.server
+
+openai.api_key = os.environ.get('OPENAI_API_KEY')
+
+@server.before_request
+def ensure_session_cookie():
+    # Set session cookie if not present
+    sid = request.cookies.get('session-id')
+    if not sid or sid not in SESSION_DATA:
+        sid_new = str(uuid.uuid4())
+        g.set_cookie_sid = sid_new
+        get_session_id(set_cookie=True)
+    else:
+        g.set_cookie_sid = None
+
+@server.after_request
+def set_session_cookie(response):
+    sid = getattr(g, 'set_cookie_sid', None)
+    if sid:
+        response.set_cookie('session-id', sid, max_age=60*60*48, httponly=True)
+    return response
+
+def parse_file_content(file_path, filename):
+    try:
+        with open(file_path, "rb") as f:
+            decoded = f.read()
+        if filename.endswith('.pdf'):
+            with io.BytesIO(decoded) as pdf_file:
+                reader = PyPDF2.PdfReader(pdf_file)
+                return ' '.join([page.extract_text() or "" for page in reader.pages])
+        elif filename.endswith('.docx'):
+            with io.BytesIO(decoded) as docx_file:
+                doc = docx.Document(docx_file)
+                return ' '.join([para.text for para in doc.paragraphs])
+        elif filename.endswith('.txt') or filename.endswith('.rtf'):
+            encoding = chardet.detect(decoded)['encoding']
+            return decoded.decode(encoding)
+        else:
+            return "Unsupported file format"
+    except Exception as e:
+        logger.exception(f"Error processing file {filename}: {str(e)}")
+        return "Error processing file"
+
+def truncate_filename(filename, max_length=24):
+    if len(filename) <= max_length:
+        return filename
+    else:
+        return filename[:max_length - 3] + '...'
+
+def get_file_cards(file_dict):
+    cards = []
+    for name in file_dict:
+        cards.append(
+            dbc.Card(
+                dbc.CardBody(
+                    dbc.Row([
+                        dbc.Col(
+                            html.Span(
+                                truncate_filename(name),
+                                title=name,
+                                style={
+                                    'display': 'inline-block',
+                                    'overflow': 'hidden',
+                                    'textOverflow': 'ellipsis',
+                                    'whiteSpace': 'nowrap',
+                                    'maxWidth': '90%',
+                                    'verticalAlign': 'middle',
+                                }
+                            ),
+                            width='auto',
+                            style={'display': 'flex', 'alignItems': 'center', 'padding': '0'}
+                        ),
+                        dbc.Col(
+                            dbc.Button(
+                                "Delete",
+                                id={'type': 'remove-file', 'index': name},
+                                color="danger",
+                                size="sm",
+                                style={'marginLeft': 'auto', 'float': 'right'}
+                            ),
+                            width='auto',
+                            style={'display': 'flex', 'alignItems': 'center', 'justifyContent': 'flex-end', 'padding': '0'}
+                        ),
+                    ],
+                    justify="between",
+                    align="center",
+                    style={"margin": "0", "padding": "0"}
+                    ),
+                    style={'padding': '6px 8px', 'margin': '0', 'display': 'flex', 'alignItems': 'center', 'background': 'none', 'boxShadow': 'none'}
+                ),
+                style={'border': 'none', 'boxShadow': 'none', 'background': 'none', 'marginBottom': '2px'}
+            )
+        )
+    return cards
+
 app.layout = dbc.Container([
     dbc.Row([
         dbc.Col([
@@ -166,83 +255,6 @@ app.layout = dbc.Container([
     ])
 ], fluid=True, style={'padding': '0'})
 
-
-# --- File Parsing ---
-def parse_file_content(file_path, filename):
-    try:
-        with open(file_path, "rb") as f:
-            decoded = f.read()
-        if filename.endswith('.pdf'):
-            with io.BytesIO(decoded) as pdf_file:
-                reader = PyPDF2.PdfReader(pdf_file)
-                return ' '.join([page.extract_text() or "" for page in reader.pages])
-        elif filename.endswith('.docx'):
-            with io.BytesIO(decoded) as docx_file:
-                doc = docx.Document(docx_file)
-                return ' '.join([para.text for para in doc.paragraphs])
-        elif filename.endswith('.txt') or filename.endswith('.rtf'):
-            encoding = chardet.detect(decoded)['encoding']
-            return decoded.decode(encoding)
-        else:
-            return "Unsupported file format"
-    except Exception as e:
-        logger.exception(f"Error processing file {filename}: {str(e)}")
-        return "Error processing file"
-
-def truncate_filename(filename, max_length=24):
-    if len(filename) <= max_length:
-        return filename
-    else:
-        return filename[:max_length - 3] + '...'
-
-def get_file_cards(file_dict):
-    cards = []
-    for name in file_dict:
-        cards.append(
-            dbc.Card(
-                dbc.CardBody(
-                    dbc.Row([
-                        dbc.Col(
-                            html.Span(
-                                truncate_filename(name),
-                                title=name,
-                                style={
-                                    'display': 'inline-block',
-                                    'overflow': 'hidden',
-                                    'textOverflow': 'ellipsis',
-                                    'whiteSpace': 'nowrap',
-                                    'maxWidth': '90%',
-                                    'verticalAlign': 'middle',
-                                }
-                            ),
-                            width='auto',
-                            style={'display': 'flex', 'alignItems': 'center', 'padding': '0'}
-                        ),
-                        dbc.Col(
-                            dbc.Button(
-                                "Delete",
-                                id={'type': 'remove-file', 'index': name},
-                                color="danger",
-                                size="sm",
-                                style={'marginLeft': 'auto', 'float': 'right'}
-                            ),
-                            width='auto',
-                            style={'display': 'flex', 'alignItems': 'center', 'justifyContent': 'flex-end', 'padding': '0'}
-                        ),
-                    ],
-                    justify="between",
-                    align="center",
-                    style={"margin": "0", "padding": "0"}
-                    ),
-                    style={'padding': '6px 8px', 'margin': '0', 'display': 'flex', 'alignItems': 'center', 'background': 'none', 'boxShadow': 'none'}
-                ),
-                style={'border': 'none', 'boxShadow': 'none', 'background': 'none', 'marginBottom': '2px'}
-            )
-        )
-    return cards
-
-# --- Callbacks ---
-
 @app.callback(
     Output('file-list', 'children'),
     Input('upload-files', 'contents'),
@@ -269,12 +281,11 @@ def update_output(list_of_contents, list_of_names):
 @app.callback(
     Output('file-list', 'children', allow_duplicate=True),
     Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
-    State('file-list', 'children'),
     prevent_initial_call=True
 )
-def remove_file(n_clicks, existing_cards):
-    session_data, lock = get_session_data()
+def remove_file(n_clicks):
     ctx = dash.callback_context
+    session_data, lock = get_session_data()
     if not ctx.triggered:
         raise PreventUpdate
     triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]
@@ -352,7 +363,7 @@ Now, generate the {matrix_type}:
         ]
     )
     matrix_text = response.choices[0].message.content.strip()
-    logger.info(f"Raw matrix text from GPT: {matrix_text[:200]}...")  # log only first 200 chars
+    logger.info(f"Raw matrix text from GPT: {matrix_text[:200]}...")
     lines = [line.strip() for line in matrix_text.split('\n') if '|' in line]
     data = [line.split('|') for line in lines]
     data = [[cell.strip() for cell in row] for row in data]