Update app.py via AI Editor

app.py

@@ -5,6 +5,7 @@ import uuid
 import shutil
 import json
 import tempfile
+import glob
 from flask import Flask, request as flask_request, make_response
 import dash
 from dash import dcc, html, Input, Output, State, callback_context
@@ -25,6 +26,10 @@ os.makedirs(SESSION_DIR_BASE, exist_ok=True)
 
 openai.api_key = os.environ.get("OPENAI_API_KEY")
 
+EMBEDDING_INDEX = {}
+EMBEDDING_TEXTS = {}
+EMBEDDING_MODEL = "text-embedding-ada-002"
+
 def get_session_id():
     sid = flask_request.cookies.get("asktricare_session_id")
     if not sid:
@@ -75,6 +80,39 @@ def load_system_prompt():
         logger.error(f"Failed to load system prompt: {e}")
         return "You are Ask Tricare, a helpful assistant for TRICARE health benefits. Respond conversationally, and cite relevant sources when possible. If you do not know, say so."
 
+def embed_docs_folder():
+    global EMBEDDING_INDEX, EMBEDDING_TEXTS
+    docs_folder = os.path.join(os.getcwd(), "docs")
+    if not os.path.isdir(docs_folder):
+        logger.warning(f"Docs folder '{docs_folder}' does not exist. Skipping embedding.")
+        return
+    doc_files = []
+    for ext in ("*.txt", "*.md", "*.pdf"):
+        doc_files.extend(glob.glob(os.path.join(docs_folder, ext)))
+    for doc_path in doc_files:
+        fname = os.path.basename(doc_path)
+        if fname in EMBEDDING_INDEX:
+            continue
+        try:
+            with open(doc_path, "r", encoding="utf-8", errors="ignore") as f:
+                text = f.read()
+            if not text.strip():
+                continue
+            # OpenAI recommends chunking long texts; here, we take the first 1000 tokens per doc as a simple approach
+            chunk = text[:4000]
+            response = openai.Embedding.create(
+                input=[chunk],
+                model=EMBEDDING_MODEL
+            )
+            embedding = response['data'][0]['embedding']
+            EMBEDDING_INDEX[fname] = embedding
+            EMBEDDING_TEXTS[fname] = chunk
+            logger.info(f"Embedded doc: {fname}")
+        except Exception as e:
+            logger.error(f"Embedding failed for {fname}: {e}")
+
+embed_docs_folder()
+
 app = dash.Dash(
     __name__,
     server=app_flask,
@@ -201,18 +239,20 @@ def assign_session_id(_):
     Output("stream-interval", "disabled"),
     Output("stream-interval", "n_intervals"),
     Output("chat-name-input", "style"),
+    Output("chat-window", "children"),
+    Output("stream-interval", "disabled"),
     Input("session-id", "data"),
     Input("send-btn", "n_clicks"),
     Input("file-upload", "contents"),
     Input("new-chat-btn", "n_clicks"),
+    Input("stream-interval", "n_intervals"),
     State("file-upload", "filename"),
     State("user-input", "value"),
     State("chat-name-input", "value"),
-    State("stream-interval", "n_intervals"),
     State("chat-name-input", "style"),
     prevent_initial_call=False
 )
-def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, file_names, user_input, chat_name, stream_n, chat_name_style):
+def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, stream_n, file_names, user_input, chat_name, chat_name_style):
     trigger = callback_context.triggered[0]['prop_id'].split('.')[0] if callback_context.triggered else ""
     if not session_id:
         session_id = get_session_id()
@@ -303,27 +343,32 @@ def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, file_
             # If chat name input box is not yet visible, show it
             if show_chat_name.get("display", "none") == "none":
                 show_chat_name["display"] = "block"
+                upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])]
+                chat_history_items = [
+                    html.Li(
+                        html.Span(
+                            chat["name"],
+                            style={"fontSize": "0.92rem"}
+                        )
+                    ) for chat in state.get("chat_histories", [])[-6:]
+                ]
+                chat_cards = [
+                    chat_message_card(msg['content'], is_user=(msg['role'] == "user"))
+                    for msg in state.get("messages", [])
+                ] + (
+                    [chat_message_card(state["stream_buffer"], is_user=False)]
+                    if state.get("streaming", False) and state.get("stream_buffer", "") else []
+                )
                 return (
-                    [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])],
-                    [
-                        html.Li(
-                            html.Span(
-                                chat["name"],
-                                style={"fontSize": "0.92rem"}
-                            )
-                        ) for chat in state.get("chat_histories", [])[-6:]
-                    ],
-                    [
-                        chat_message_card(msg['content'], is_user=(msg['role'] == "user"))
-                        for msg in state.get("messages", [])
-                    ] + (
-                        [chat_message_card(state["stream_buffer"], is_user=False)]
-                        if state.get("streaming", False) and state.get("stream_buffer", "") else []
-                    ),
+                    upload_cards,
+                    chat_history_items,
+                    chat_cards,
                     "",
                     not state.get("streaming", False),
                     0,
-                    show_chat_name
+                    show_chat_name,
+                    chat_cards,
+                    not state.get("streaming", False)
                 )
             # If input box is visible and has a value, save chat history
             else:
@@ -342,24 +387,82 @@ def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, file_
                     save_session_state(session_id)
                     logger.info(f"Session {session_id}: Saved chat history '{chat_title}'")
                 show_chat_name["display"] = "none"
+                upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])]
+                chat_history_items = [
+                    html.Li(
+                        html.Span(
+                            chat["name"],
+                            style={"fontSize": "0.92rem", "fontWeight": "bold"}
+                        )
+                    ) for chat in state.get("chat_histories", [])[-6:]
+                ]
                 return (
-                    [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])],
-                    [
-                        html.Li(
-                            html.Span(
-                                chat["name"],
-                                style={"fontSize": "0.92rem", "fontWeight": "bold"}
-                            )
-                        ) for chat in state.get("chat_histories", [])[-6:]
-                    ],
+                    upload_cards,
+                    chat_history_items,
                     [],
                     error,
                     not state.get("streaming", False),
                     0,
-                    show_chat_name
+                    show_chat_name,
+                    [],
+                    not state.get("streaming", False)
                 )
 
-        # Build Uploads, Chat History and Chat Window
+        # Handle polling for streaming
+        if trigger == "stream-interval":
+            chat_history = state.get("messages", [])
+            chat_cards = []
+            for msg in chat_history:
+                chat_cards.append(chat_message_card(msg['content'], is_user=(msg['role'] == "user")))
+            if state.get("streaming", False):
+                if state.get("stream_buffer", ""):
+                    chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
+                upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])]
+                chat_history_items = [
+                    html.Li(
+                        html.Span(
+                            chat["name"],
+                            style={"fontSize": "0.92rem", "fontWeight": "bold" if i == len(state.get("chat_histories", []))-1 else "normal"}
+                        )
+                    ) for i, chat in enumerate(state.get("chat_histories", [])[-6:])
+                ]
+                return (
+                    upload_cards,
+                    chat_history_items,
+                    chat_cards,
+                    "",
+                    False,
+                    stream_n+1,
+                    chat_name_style,
+                    chat_cards,
+                    False
+                )
+            else:
+                chat_cards = []
+                for msg in state.get("messages", []):
+                    chat_cards.append(chat_message_card(msg['content'], is_user=(msg['role'] == "user")))
+                upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])]
+                chat_history_items = [
+                    html.Li(
+                        html.Span(
+                            chat["name"],
+                            style={"fontSize": "0.92rem", "fontWeight": "bold" if i == len(state.get("chat_histories", []))-1 else "normal"}
+                        )
+                    ) for i, chat in enumerate(state.get("chat_histories", [])[-6:])
+                ]
+                return (
+                    upload_cards,
+                    chat_history_items,
+                    chat_cards,
+                    "",
+                    True,
+                    0,
+                    chat_name_style,
+                    chat_cards,
+                    True
+                )
+
+        # Default: Build Uploads, Chat History and Chat Window
         chat_history = state.get("messages", [])
         uploads = state.get("uploads", [])
         chat_histories = state.get("chat_histories", [])
@@ -378,30 +481,8 @@ def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, file_
         if state.get("streaming", False):
             if state.get("stream_buffer", ""):
                 chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
-            return upload_cards, chat_history_items, chat_cards, error, False, 0, show_chat_name
-        return upload_cards, chat_history_items, chat_cards, error, (not state.get("streaming", False)), 0, show_chat_name
-
-@app.callback(
-    Output("chat-window", "children"),
-    Output("stream-interval", "disabled"),
-    Input("stream-interval", "n_intervals"),
-    State("session-id", "data"),
-    prevent_initial_call=True
-)
-def poll_stream(n_intervals, session_id):
-    session_lock = get_session_lock(session_id)
-    with session_lock:
-        load_session_state(session_id)
-        state = get_session_state(session_id)
-        chat_history = state.get("messages", [])
-        chat_cards = []
-        for msg in chat_history:
-            chat_cards.append(chat_message_card(msg['content'], is_user=(msg['role'] == "user")))
-        if state.get("streaming", False):
-            if state.get("stream_buffer", ""):
-                chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
-            return chat_cards, False
-        return chat_cards, True
+            return upload_cards, chat_history_items, chat_cards, error, False, 0, chat_name_style, chat_cards, False
+        return upload_cards, chat_history_items, chat_cards, error, (not state.get("streaming", False)), 0, chat_name_style, chat_cards, True
 
 @app_flask.after_request
 def set_session_cookie(resp):