Updated output for modelfit_chat

main.py

@@ -283,6 +283,111 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
     # Fallback: return short text envelope if non-JSON
     return JSONResponse(status_code=r.status_code, content={"text": r.text[:1000]})
 
+
+async def _forward_and_get_body(path: str, method: str = "GET", json_body=None, user_token: str | None = None) -> tuple[int, dict | list | None]:
+    """Same as _forward but returns (status_code, parsed_json_body) so caller can build a custom response."""
+    url = f"{UPSTREAM_URL}{path}"
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "Accept": "application/json",
+    }
+    if user_token:
+        headers["X-API-Key"] = user_token
+    ping_url = f"{UPSTREAM_URL}/secure-ping/"
+    ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    if user_token:
+        ping_headers["X-API-Key"] = user_token
+    timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+    r = None
+
+    async def do_request():
+        nonlocal r
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            r = await client.request(method, url, headers=headers, json=json_body)
+
+    request_task = asyncio.create_task(do_request())
+    wait_sec = PING_INTERVAL if PING_INTERVAL > 0 else 0.0
+    try:
+        while not request_task.done():
+            if wait_sec <= 0:
+                await request_task
+                break
+            ping_sleep = asyncio.create_task(asyncio.sleep(wait_sec))
+            done, pending = await asyncio.wait(
+                {request_task, ping_sleep},
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=UPSTREAM_TIMEOUT + 10,
+            )
+            for t in pending:
+                if t is not request_task:
+                    t.cancel()
+                    try:
+                        await t
+                    except asyncio.CancelledError:
+                        pass
+            if request_task in done:
+                break
+            async with httpx.AsyncClient(timeout=10.0) as c:
+                if SELF_URL:
+                    try:
+                        await c.get(f"{SELF_URL}/keep-alive")
+                    except Exception:
+                        pass
+                try:
+                    await c.get(ping_url, headers=ping_headers)
+                except Exception:
+                    pass
+        if not request_task.done():
+            request_task.cancel()
+            try:
+                await request_task
+            except asyncio.CancelledError:
+                pass
+        await request_task
+    except asyncio.CancelledError:
+        request_task.cancel()
+        try:
+            await request_task
+        except asyncio.CancelledError:
+            pass
+        raise
+    exc = request_task.exception()
+    if exc is not None:
+        raise exc
+    if r is None:
+        raise RuntimeError("Upstream request did not complete")
+
+    ct = r.headers.get("content-type", "")
+    if "application/json" in ct:
+        try:
+            return (r.status_code, r.json())
+        except Exception:
+            return (r.status_code, {"error": r.text[:500]})
+    return (r.status_code, {"text": r.text[:1000]})
+
+
+def _build_modelfit_chat_minimal(body: dict) -> dict:
+    """Extract formula, alpha, beta, and Performance subset (col 2, rows 3,4,5) for chat response."""
+    out = {"ok": True, "formula": body.get("formula")}
+    out["alpha"] = body.get("alpha")
+    out["beta"] = body.get("beta")
+    perf = body.get("Performance")
+    if isinstance(perf, list) and len(perf) >= 5:
+        keys = list(perf[0].keys()) if perf[0] else []
+        # Col 2 (1-based) = index 1; if only one column use index 0. Rows 3,4,5 = indices 2,3,4.
+        col_idx = 1 if len(keys) > 1 else 0
+        col2_name = keys[col_idx] if keys else None
+        if col2_name is not None:
+            row_names = ["R2 overall", "MSE", "MAE"]
+            values = [perf[i].get(col2_name) for i in [2, 3, 4]]
+            out["Performance_subset"] = {"metric_names": row_names, "column": col2_name, "values": values}
+        else:
+            out["Performance_subset"] = None
+    else:
+        out["Performance_subset"] = None
+    return out
+
+
 async def _forward_stream(path: str, files=None, data=None, user_token: str | None = None, method: str = "POST"):
     url = f"{UPSTREAM_URL}{path}"
     headers = {
@@ -590,8 +695,7 @@ async def modelfit_chat(
     """
     Fetch training data from data_url (HTTPS), infer schema from column order, then call /modelfit/ on the API.
     Column order: 1st = id_col, 2nd = time_col, 3rd..second-to-last = current_features, last = y.
-    All other parameters use defaults. On success returns a minimal JSON (ok, message) to stay under
-    chat platform response size limits; errors are returned in full.
+    On success returns a small informative JSON: formula, alpha, beta, Performance_subset (col 2, rows 3–5).
     """
     user_token = _extract_user_token(req)
     if not user_token:
@@ -619,13 +723,11 @@ async def modelfit_chat(
         "meanvar_test": False,
         "signif": 0.05,
     }
-    return await _forward(
-        "/modelfit/",
-        "POST",
-        json_body=payload,
-        user_token=user_token,
-        success_response_override={"ok": True, "message": "Model fitted. Call modelforecast_chat with your forecast data URL next."},
-    )
+    status, body = await _forward_and_get_body("/modelfit/", "POST", json_body=payload, user_token=user_token)
+    if status != 200 or not isinstance(body, dict):
+        return JSONResponse(status_code=status, content=body if isinstance(body, dict) else {"error": str(body)})
+    minimal = _build_modelfit_chat_minimal(body)
+    return JSONResponse(status_code=200, content=minimal)
 
 
 @app.post("/modelforecast_chat/")