added mlflow tracing and context size logging. added heartbeat to keep the frontend alive while waiting for long llm calls

Author: sean-zhang-dbx

app.yaml

@@ -29,7 +29,7 @@ env:
   - name: MLFLOW_REGISTRY_URI
     value: "databricks-uc"
   - name: MLFLOW_EXPERIMENT_ID
-    value: ""  # OPTIONAL: Set to your MLflow experiment ID to enable tracing
+    value: "2624606869914616"
 
   # ---------------------------------------------------------------------------
   # LLM Model

backend/routers/create.py

@@ -1,6 +1,7 @@
 """
 /api/create — UC discovery + config validation + Genie Space creation wizard + agent chat.
 """
+import asyncio
 import json
 import logging
 from fastapi import APIRouter, HTTPException, Request
@@ -158,13 +159,26 @@ async def agent_chat(body: AgentChatRequest, request: Request):
     # that outlive the middleware's call_next).
     user_token = getattr(request.state, "user_token", "")
 
+    _KEEPALIVE_INTERVAL = 15  # seconds between SSE keepalive comments
+
     async def event_stream():
         if user_token:
             set_obo_user_token(user_token)
         try:
             yield _sse_event("session", {"session_id": session.session_id})
-            async for event in agent.chat(session, user_message):
-                yield _sse_event(event["event"], event["data"])
+
+            agent_iter = agent.chat(session, user_message).__aiter__()
+            while True:
+                try:
+                    event = await asyncio.wait_for(
+                        agent_iter.__anext__(), timeout=_KEEPALIVE_INTERVAL
+                    )
+                    yield _sse_event(event["event"], event["data"])
+                except asyncio.TimeoutError:
+                    yield ": keepalive\n\n"
+                except StopAsyncIteration:
+                    break
+
             await persist_session(session)
         finally:
             clear_obo_user_token()

backend/services/create_agent.py

@@ -86,6 +86,21 @@ async def chat(
         step = detect_step(session)
         step_idx = STEP_ORDER.index(step) if step in STEP_ORDER else 0
 
+        # Start a root MLflow trace for the entire agent turn.
+        # Wrapped in try/except so tracing failures never break the app.
+        _trace = None
+        try:
+            _trace = mlflow.start_trace(
+                name="agent_chat",
+                inputs={"user_message": user_message, "session_id": session.session_id, "step": step},
+            )
+            mlflow.update_current_trace(tags={
+                "session_id": session.session_id,
+                "workflow_step": step,
+            })
+        except Exception:
+            logger.debug("MLflow start_trace failed, continuing without tracing", exc_info=True)
+
         yield {"event": "step", "data": {
             "step": step,
             "label": STEP_LABELS.get(step, step),
@@ -301,9 +316,20 @@ async def chat(
                 step, tools_used, error_msg,
             )
 
+        try:
+            if _trace is not None:
+                status = "ERROR" if error_msg else "OK"
+                mlflow.end_trace(
+                    request_id=_trace.request_id,
+                    outputs={"tools_used": tools_used, "error": error_msg},
+                    status=status,
+                )
+        except Exception:
+            logger.debug("MLflow end_trace failed", exc_info=True)
+
         yield {"event": "done", "data": {}}
 
-    _TOOL_RESULT_CHAR_LIMIT = 6000
+    _TOOL_RESULT_CHAR_LIMIT = 3000
     _COMPRESSIBLE_TOOLS = frozenset({
         "describe_table", "profile_columns", "profile_table_usage",
         "assess_data_quality", "test_sql",

frontend/src/lib/api.ts

@@ -687,7 +687,9 @@ export function streamAgentChat(
     })
     .catch((error) => {
       if (error.name !== "AbortError") {
-        callbacks.onError(error.message || "Connection failed")
+        callbacks.onError(error.message === "network error"
+          ? "Connection interrupted — your progress is saved. Send another message to continue."
+          : (error.message || "Connection failed"))
         callbacks.onDone()
       }
     })