Update app.py

app.py

@@ -1,9 +1,20 @@
-# app.py
 import gradio as gr
 
 from rag_data import THREAD_OPTIONS
-from rag_sessions import start_session, reset_session, get_session
-from rag_retrieval import rewrite_query, retrieve_chunks, build_answer, log_trace
+from rag_sessions import (
+    start_session,
+    reset_session,
+    get_session,
+    update_entity_memory,
+)
+from rag_retrieval import (
+    rewrite_query,
+    retrieve_chunks,
+    build_answer,
+    log_trace,
+    extract_entities_for_turn,
+)
+from rag_timeline import build_timeline
 
 
 def init_session_ui(thread_id: str):
@@ -21,17 +32,29 @@ def chat_ui(user_text: str, session_id: str, search_outside_thread: bool):
     if session is None:
         return "Session not found. Please start again.", "", ""
 
+    # 1) Rewrite query using thread + entity memory
     rewrite = rewrite_query(user_text, session)
+
+    # 2) Retrieve chunks
     retrieved = retrieve_chunks(rewrite, session, search_outside_thread)
+
+    # 3) Extract entities from this turn + retrieved evidence, update memory
+    new_entities = extract_entities_for_turn(user_text, retrieved)
+    if new_entities:
+        update_entity_memory(session_id, new_entities)
+
+    # 4) Build grounded answer
     answer, citations = build_answer(user_text, rewrite, retrieved)
 
-    # Update simple memory
+    # 5) Update simple turn memory
     session["recent_turns"].append({"user": user_text, "answer": answer})
     if len(session["recent_turns"]) > 5:
         session["recent_turns"] = session["recent_turns"][-5:]
 
+    # 6) Log trace for this turn
     log_trace(session_id, user_text, rewrite, retrieved, answer, citations)
 
+    # 7) Debug: show retrieved chunk ids + scores
     debug_retrieved = "\n".join(
         [
             f"{r['chunk_id']} (msg={r['message_id']}, "
@@ -50,6 +73,16 @@ def reset_session_ui(session_id: str):
     return "", "Session reset."
 
 
+def timeline_ui(session_id: str):
+    if not session_id:
+        return "Please start a session by selecting a thread."
+    session = get_session(session_id)
+    if session is None:
+        return "Session not found. Please start again."
+    tid = session["thread_id"]
+    return build_timeline(tid)
+
+
 with gr.Blocks() as demo:
     gr.Markdown("# 📧 Email Thread RAG Assistant\nAsk questions about a selected Enron email thread.")
 
@@ -77,8 +110,11 @@ with gr.Blocks() as demo:
         search_toggle = gr.Checkbox(label="Search outside selected thread", value=False)
         ask_btn = gr.Button("Ask")
         reset_btn = gr.Button("Reset Session")
+        timeline_btn = gr.Button("Show Timeline")
 
     answer_box = gr.Markdown(label="Answer")
+    timeline_box = gr.Markdown(label="Thread timeline")
+
     with gr.Accordion("Debug info", open=False):
         rewrite_box = gr.Textbox(label="Rewritten query", interactive=False)
         retrieved_box = gr.Textbox(label="Retrieved chunks", interactive=False)
@@ -95,5 +131,11 @@ with gr.Blocks() as demo:
         outputs=[session_state, status_box],
     )
 
+    timeline_btn.click(
+        fn=timeline_ui,
+        inputs=[session_state],
+        outputs=[timeline_box],
+    )
+
 if __name__ == "__main__":
     demo.launch()