Update app.py

app.py

@@ -50,6 +50,7 @@ def image_to_base64(image_path: str) -> str:
         mime = "image/png"
     return f"data:{mime};base64,{encoded_string}"
 
+
 # ================== User Guide Content ==================
 USER_GUIDE_SECTIONS = {
     "getting_started": """
@@ -389,6 +390,9 @@ with gr.Blocks(
     # ✅ 当前“最近一次回答”是否已经被点赞/点踩（只允许一次）
     feedback_used_state = gr.State(False)
 
+    # ✅ 性能输出
+    perf_state = gr.State({})
+
     # --- Header ---
     with gr.Row(elem_classes="header-container"):
         with gr.Column(scale=3):
@@ -532,6 +536,9 @@ with gr.Blocks(
                 type="tuples",
             )
 
+            # ✅ profiling output JSON (shows your TTFT / tokens/sec etc.)
+            perf_output = gr.JSON(label="Output", value={})
+
             # Rating bar (last answer)
             gr.Markdown("#### Rate Clare’s last answer")
             with gr.Row():
@@ -701,6 +708,8 @@ with gr.Blocks(
                 feedback_toggle_btn: gr.update(interactive=False),
                 feedback_text: gr.update(visible=False, value=""),
                 feedback_submit_btn: gr.update(interactive=False, visible=False),
+                perf_state: {},                         # ✅
+                perf_output: gr.update(value={}),        # ✅
             }
 
         info_html = f"""
@@ -738,6 +747,8 @@ with gr.Blocks(
             feedback_toggle_btn: gr.update(interactive=True),
             feedback_text: gr.update(visible=False, value=""),
             feedback_submit_btn: gr.update(interactive=True, visible=False),
+            perf_state: {},                         # ✅
+            perf_output: gr.update(value={}),        # ✅
         }
 
     login_confirm_btn.click(
@@ -768,6 +779,8 @@ with gr.Blocks(
             feedback_toggle_btn,
             feedback_text,
             feedback_submit_btn,
+            perf_state,         # ✅
+            perf_output,        # ✅
         ],
     )
 
@@ -802,6 +815,8 @@ with gr.Blocks(
             feedback_toggle_btn: gr.update(interactive=False),
             feedback_text: gr.update(visible=False, value=""),
             feedback_submit_btn: gr.update(interactive=False, visible=False),
+            perf_state: {},                          # ✅
+            perf_output: gr.update(value={}),         # ✅
         }
 
     logout_btn.click(
@@ -832,6 +847,8 @@ with gr.Blocks(
             feedback_toggle_btn,
             feedback_text,
             feedback_submit_btn,
+            perf_state,        # ✅
+            perf_output,       # ✅
         ],
     )
 
@@ -895,6 +912,14 @@ with gr.Blocks(
         user_id_val,
         feedback_used,
     ):
+        # perf container
+        t0 = time.perf_counter()
+        marks = {"start": 0}
+        segs = {}
+
+        def mark(name: str):
+            marks[name] = (time.perf_counter() - t0) * 1000.0
+
         # 未登录：不解锁按钮
         if not user_id_val:
             out_msg = (
@@ -907,6 +932,7 @@ with gr.Blocks(
                 weaknesses or [],
                 cognitive_state or {"confusion": 0, "mastery": 0},
             )
+            perf = {"marks_ms": marks, "segments_ms": segs, "total_ms": (time.perf_counter() - t0) * 1000.0}
             return (
                 "",
                 new_history,
@@ -918,9 +944,14 @@ with gr.Blocks(
                 feedback_used,
                 gr.update(interactive=False, value="👍 Helpful"),
                 gr.update(interactive=False, value="👎 Not helpful"),
+                perf,                      # ✅ perf_state
+                perf,                      # ✅ perf_output
             )
 
+        # language detect
         resolved_lang = detect_language(message or "", lang_pref)
+        mark("language_detect_done")
+        segs["language_detect_done"] = marks["language_detect_done"] - marks["start"]
 
         # 空输入：不改变按钮状态
         if not message or not message.strip():
@@ -929,6 +960,7 @@ with gr.Blocks(
                 weaknesses or [],
                 cognitive_state or {"confusion": 0, "mastery": 0},
             )
+            perf = {"marks_ms": marks, "segments_ms": segs, "total_ms": (time.perf_counter() - t0) * 1000.0}
             return (
                 "",
                 chat_history,
@@ -940,20 +972,35 @@ with gr.Blocks(
                 feedback_used,
                 gr.update(),
                 gr.update(),
+                perf,                      # ✅
+                perf,                      # ✅
             )
 
+        # weakness/cognitive updates
+        t_w0 = time.perf_counter()
         weaknesses = update_weaknesses_from_message(message, weaknesses or [])
+        mark("weakness_update_done")
+        segs["weakness_update_done"] = (time.perf_counter() - t_w0) * 1000.0
+
+        t_c0 = time.perf_counter()
         cognitive_state = update_cognitive_state_from_message(message, cognitive_state)
+        mark("cognitive_update_done")
+        segs["cognitive_update_done"] = (time.perf_counter() - t_c0) * 1000.0
 
+        # RAG
+        t_r0 = time.perf_counter()
         if is_academic_query(message):
             rag_context_text, rag_used_chunks = retrieve_relevant_chunks(
                 message, rag_chunks or []
             )
         else:
             rag_context_text, rag_used_chunks = "", []
+        mark("rag_retrieve_done")
+        segs["rag_retrieve_done"] = (time.perf_counter() - t_r0) * 1000.0
 
-        start_ts = time.time()
-        answer, new_history = chat_with_clare(
+        # LLM (chat_with_clare must return 3 values)
+        t_llm0 = time.perf_counter()
+        answer, new_history, llm_stats = chat_with_clare(
             message=message,
             history=chat_history,
             model_name=model_name_val,
@@ -965,9 +1012,31 @@ with gr.Blocks(
             cognitive_state=cognitive_state,
             rag_context=rag_context_text,
         )
-        end_ts = time.time()
-        latency_ms = (end_ts - start_ts) * 1000.0
+        mark("llm_done")
+        segs["llm_done"] = (time.perf_counter() - t_llm0) * 1000.0
+
+        # merge llm_stats into perf
+        perf = {
+            "marks_ms": marks,
+            "segments_ms": segs,
+            "total_ms": marks.get("llm_done", (time.perf_counter() - t0) * 1000.0),
+        }
 
+        # llm_stats may contain marks_ms/segments_ms/llm_profile
+        if isinstance(llm_stats, dict):
+            if "llm_profile" in llm_stats:
+                perf.update({"llm_profile": llm_stats.get("llm_profile", {})})
+            # merge marks/segments from llm_stats if present
+            ms2 = llm_stats.get("marks_ms") if isinstance(llm_stats.get("marks_ms"), dict) else {}
+            sg2 = llm_stats.get("segments_ms") if isinstance(llm_stats.get("segments_ms"), dict) else {}
+            for k, v in ms2.items():
+                if v is not None:
+                    perf["marks_ms"][k] = v
+            for k, v in sg2.items():
+                if v is not None:
+                    perf["segments_ms"][k] = v
+
+        # References formatting
         if is_academic_query(message) and rag_used_chunks:
             ref_text = format_references(rag_used_chunks)
         else:
@@ -980,6 +1049,7 @@ with gr.Blocks(
                 new_history[-1] = [last_user, last_assistant]
                 answer = last_assistant
 
+        # LangSmith event
         student_id = user_id_val or "ANON"
         experiment_id = "RESP_AI_W10"
         try:
@@ -988,13 +1058,18 @@ with gr.Blocks(
                     "experiment_id": experiment_id,
                     "student_id": student_id,
                     "event_type": "chat_turn",
-                    "timestamp": end_ts,
-                    "latency_ms": latency_ms,
+                    "timestamp": time.time(),
+                    "latency_ms": perf["total_ms"],
                     "question": message,
                     "answer": answer,
                     "model_name": model_name_val,
                     "language": resolved_lang,
                     "learning_mode": mode_val,
+                    # ✅ extra profiling fields if available
+                    "ttft_ms": perf.get("segments_ms", {}).get("llm_ttft_ms"),
+                    "tokens_per_sec": (perf.get("llm_profile", {}) or {}).get("tokens_per_sec_est"),
+                    "output_tokens_est": (perf.get("llm_profile", {}) or {}).get("output_tokens_est"),
+                    "streaming_used": (perf.get("llm_profile", {}) or {}).get("streaming_used"),
                 }
             )
         except Exception as e:
@@ -1014,6 +1089,8 @@ with gr.Blocks(
             False,
             gr.update(interactive=True, value="👍 Helpful"),
             gr.update(interactive=True, value="👎 Not helpful"),
+            perf,                 # ✅ perf_state
+            perf,                 # ✅ perf_output
         )
 
     user_input.submit(
@@ -1043,6 +1120,8 @@ with gr.Blocks(
             feedback_used_state,
             thumb_up_btn,
             thumb_down_btn,
+            perf_state,      # ✅
+            perf_output,     # ✅
         ],
     )
 
@@ -1059,8 +1138,16 @@ with gr.Blocks(
         doc_type_val,
         user_id_val,
     ):
+        t0 = time.perf_counter()
+        marks = {"start": 0}
+        segs = {}
+
+        def mark(name: str):
+            marks[name] = (time.perf_counter() - t0) * 1000.0
+
         if not user_id_val:
             gr.Info("Please log in first to start a micro-quiz.", title="Login required")
+            perf = {"marks_ms": marks, "segments_ms": segs, "total_ms": (time.perf_counter() - t0) * 1000.0}
             return (
                 chat_history,
                 weaknesses,
@@ -1070,6 +1157,8 @@ with gr.Blocks(
                     weaknesses or [],
                     cognitive_state or {"confusion": 0, "mastery": 0},
                 ),
+                perf,
+                perf,
             )
 
         quiz_instruction = (
@@ -1100,13 +1189,18 @@ with gr.Blocks(
         )
 
         resolved_lang = lang_pref
+        mark("language_detect_done")
+        segs["language_detect_done"] = marks["language_detect_done"]
 
-        start_ts = time.time()
+        t_r0 = time.perf_counter()
         quiz_ctx_text, _quiz_ctx_chunks = retrieve_relevant_chunks(
             "Module 10 quiz", rag_chunks or []
         )
+        mark("rag_retrieve_done")
+        segs["rag_retrieve_done"] = (time.perf_counter() - t_r0) * 1000.0
 
-        answer, new_history = chat_with_clare(
+        t_llm0 = time.perf_counter()
+        answer, new_history, llm_stats = chat_with_clare(
             message=quiz_instruction,
             history=chat_history,
             model_name=model_name_val,
@@ -1118,32 +1212,49 @@ with gr.Blocks(
             cognitive_state=cognitive_state,
             rag_context=quiz_ctx_text,
         )
-        end_ts = time.time()
-        latency_ms = (end_ts - start_ts) * 1000.0
+        mark("llm_done")
+        segs["llm_done"] = (time.perf_counter() - t_llm0) * 1000.0
+
+        perf = {
+            "marks_ms": marks,
+            "segments_ms": segs,
+            "total_ms": marks.get("llm_done", (time.perf_counter() - t0) * 1000.0),
+        }
+        if isinstance(llm_stats, dict):
+            if "llm_profile" in llm_stats:
+                perf.update({"llm_profile": llm_stats.get("llm_profile", {})})
+            ms2 = llm_stats.get("marks_ms") if isinstance(llm_stats.get("marks_ms"), dict) else {}
+            sg2 = llm_stats.get("segments_ms") if isinstance(llm_stats.get("segments_ms"), dict) else {}
+            for k, v in ms2.items():
+                if v is not None:
+                    perf["marks_ms"][k] = v
+            for k, v in sg2.items():
+                if v is not None:
+                    perf["segments_ms"][k] = v
 
         student_id = user_id_val or "ANON"
         experiment_id = "RESP_AI_W10"
-
         try:
             log_event(
                 {
                     "experiment_id": experiment_id,
                     "student_id": student_id,
                     "event_type": "micro_quiz_start",
-                    "timestamp": end_ts,
-                    "latency_ms": latency_ms,
+                    "timestamp": time.time(),
+                    "latency_ms": perf["total_ms"],
                     "question": quiz_instruction,
                     "answer": answer,
                     "model_name": model_name_val,
                     "language": resolved_lang,
                     "learning_mode": mode_val,
+                    "ttft_ms": perf.get("segments_ms", {}).get("llm_ttft_ms"),
                 }
             )
         except Exception as e:
             print("log_event error:", e)
 
         new_status = render_session_status(mode_val, weaknesses, cognitive_state)
-        return new_history, weaknesses, cognitive_state, new_status
+        return new_history, weaknesses, cognitive_state, new_status, perf, perf
 
     quiz_btn.click(
         start_micro_quiz,
@@ -1159,7 +1270,7 @@ with gr.Blocks(
             doc_type,
             user_id_state,
         ],
-        [chatbot, weakness_state, cognitive_state_state, session_status],
+        [chatbot, weakness_state, cognitive_state_state, session_status, perf_state, perf_output],
     )
 
     # ===== Feedback Handlers (thumb + detailed) =====
@@ -1389,6 +1500,8 @@ with gr.Blocks(
             False,
             gr.update(interactive=False, value="👍 Helpful"),
             gr.update(interactive=False, value="👎 Not helpful"),
+            {},
+            {},
         )
 
     clear_btn.click(
@@ -1406,6 +1519,8 @@ with gr.Blocks(
             feedback_used_state,
             thumb_up_btn,
             thumb_down_btn,
+            perf_state,
+            perf_output,
         ],
         queue=False,
     )