feat: mount heartflow backend at /heartflow with prefetch and websocket support

Dockerfile

@@ -12,9 +12,9 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && rm -rf /var/lib/apt/lists/*
 
 RUN useradd -m -u 1000 user && \
-    mkdir -p /home/user/hf-cache /tmp/p2 && \
-    chown -R user:user /home/user/hf-cache /tmp/p2 && \
-    chmod 777 /tmp/p2
+    mkdir -p /home/user/hf-cache /tmp/p1 /tmp/p2 && \
+    chown -R user:user /home/user/hf-cache /tmp/p1 /tmp/p2 && \
+    chmod 777 /tmp/p1 /tmp/p2
 
 USER user
 ENV HOME=/home/user \

registry.py

@@ -1,3 +1,4 @@
 PROJECT_REGISTRY = {
     "/semantic": "semantic_backend.router:router",
+    "/heartflow": "heartflow_backend.router:router",
 }

requirements.in

@@ -3,3 +3,4 @@ uvicorn[standard]
 pydantic
 starlette
 git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@v1.3#subdirectory=backend
+git+https://github.com/rachit-mangawa/HEARTFLOW_OS.git@v1.0#subdirectory=backend

requirements.txt

@@ -3,6 +3,8 @@
 --index-url https://pypi.org/simple
 --extra-index-url https://download.pytorch.org/whl/cpu
 
+aiosqlite==0.22.1
+    # via heartflow-backend
 annotated-doc==0.0.5
     # via
     #   fastapi
@@ -19,10 +21,13 @@ click==8.5.0
     #   huggingface-hub
     #   uvicorn
 cloudpickle==3.1.2
-    # via joblib
+    # via
+    #   joblib
+    #   shap
 fastapi==0.162.0
     # via
     #   -r requirements.in
+    #   heartflow-backend
     #   semantic-backend
 filelock==3.32.3
     # via
@@ -36,6 +41,8 @@ h11==0.16.0
     # via
     #   httpcore2
     #   uvicorn
+heartflow-backend @ git+https://github.com/rachit-mangawa/HEARTFLOW_OS.git@c81f1ffadf10f23df2767498230f54b14d39084f#subdirectory=backend
+    # via -r requirements.in
 hf-xet==1.7.0
     # via huggingface-hub
 httpcore2==2.3.0
@@ -57,8 +64,13 @@ jinja2==3.1.6
     # via torch
 joblib==1.6.0
     # via
+    #   heartflow-backend
     #   scikit-learn
     #   semantic-backend
+llvmlite==0.50.0
+    # via
+    #   numba
+    #   shap
 markdown-it-py==4.2.0
     # via rich
 markupsafe==3.0.3
@@ -71,26 +83,36 @@ narwhals==2.27.1
     # via scikit-learn
 networkx==3.6.1
     # via torch
+numba==0.68.0
+    # via shap
 numpy==2.4.6
     # via
+    #   heartflow-backend
+    #   numba
     #   pandas
     #   scikit-learn
     #   scipy
     #   semantic-backend
     #   sentence-transformers
+    #   shap
     #   transformers
 opentelemetry-api==1.45.1
     # via fastapi
 packaging==24.1
     # via
     #   huggingface-hub
+    #   shap
     #   transformers
 pandas==3.1.0
-    # via semantic-backend
+    # via
+    #   heartflow-backend
+    #   semantic-backend
+    #   shap
 pydantic==2.14.1
     # via
     #   -r requirements.in
     #   fastapi
+    #   heartflow-backend
 pydantic-core==2.50.1
     # via pydantic
 pygments==2.21.0
@@ -100,7 +122,9 @@ python-dateutil==3.0.0
 python-dotenv==1.2.4
     # via uvicorn
 python-multipart==0.0.32
-    # via semantic-backend
+    # via
+    #   heartflow-backend
+    #   semantic-backend
 pyyaml==6.0.3
     # via
     #   huggingface-hub
@@ -114,22 +138,30 @@ safetensors==0.8.0
     # via transformers
 scikit-learn==1.9.1
     # via
+    #   heartflow-backend
     #   semantic-backend
     #   sentence-transformers
+    #   shap
 scipy==1.17.1
     # via
+    #   heartflow-backend
     #   scikit-learn
     #   sentence-transformers
+    #   shap
 semantic-backend @ git+https://github.com/rachit-mangawa/Semantic-Comment-Analyze.git@eb73c60a2d94fafa9d224d75d1ae8602a0fe372f#subdirectory=backend
     # via -r requirements.in
 sentence-transformers==6.2.0
     # via semantic-backend
 setuptools==78.1.0
     # via torch
+shap==0.51.0
+    # via heartflow-backend
 shellingham==1.5.4
     # via typer
 six==1.17.0
     # via python-dateutil
+slicer==0.0.8
+    # via shap
 starlette==1.10.0
     # via
     #   -r requirements.in
@@ -148,6 +180,7 @@ tqdm==4.66.5
     # via
     #   huggingface-hub
     #   sentence-transformers
+    #   shap
     #   transformers
 transformers==5.20.2
     # via
@@ -168,6 +201,7 @@ typing-extensions==4.16.0
     #   pydantic
     #   pydantic-core
     #   sentence-transformers
+    #   shap
     #   starlette
     #   torch
     #   typing-inspection
@@ -184,4 +218,6 @@ uvloop==0.23.0
 watchfiles==1.3.0
     # via uvicorn
 websockets==17.3
-    # via uvicorn
+    # via
+    #   heartflow-backend
+    #   uvicorn

tests/test_host.py

@@ -284,8 +284,9 @@ def test_real_registry():
     import importlib
     import registry
     importlib.reload(registry)
-    assert list(registry.PROJECT_REGISTRY.keys()) == ["/semantic"]
+    assert set(registry.PROJECT_REGISTRY.keys()) == {"/semantic", "/heartflow"}
     assert registry.PROJECT_REGISTRY["/semantic"] == "semantic_backend.router:router"
+    assert registry.PROJECT_REGISTRY["/heartflow"] == "heartflow_backend.router:router"
 
     if "main" in sys.modules:
         del sys.modules["main"]
@@ -297,8 +298,10 @@ def test_real_registry():
         data = res.json()
         assert data["status"] == "healthy"
         assert "/semantic" in data["loaded"]
+        assert "/heartflow" in data["loaded"]
         assert not data["failed"]
         assert data["warm_status"]["/semantic"] == "warmed"
+        assert data["warm_status"]["/heartflow"] == "warmed"
         assert data["total_warmup_seconds"] >= 0
 
         res_root = client.get("/")
@@ -315,6 +318,13 @@ def test_real_registry():
         assert sem_data["models_loaded"]["classifier_head"] is True
         assert sem_data["models_loaded"]["sentiment"] is True
 
+        # Check Heartflow /health endpoint reports models_loaded
+        res_hf = client.get("/heartflow/health")
+        assert res_hf.status_code == 200
+        hf_data = res_hf.json()
+        assert hf_data["status"] == "Healthy"
+        assert hf_data["models_loaded"] is True
+
 
 def test_status_page_and_content_negotiation():
     if "main" in sys.modules:
@@ -351,3 +361,42 @@ def test_status_page_and_content_negotiation():
         assert "application/json" in res_api.headers.get("content-type", "")
         assert res_api.json()["status"] in ["healthy", "degraded", "unhealthy"]
 
+
+
+def test_host_websocket_through_middleware(monkeypatch):
+    from fastapi import WebSocket
+
+    ws_pkg = types.ModuleType("mock_ws_pkg")
+    ws_router_mod = types.ModuleType("mock_ws_pkg.router")
+    r = APIRouter()
+
+    @r.websocket("/ws")
+    async def dummy_ws(websocket: WebSocket):
+        await websocket.accept()
+        for i in range(3):
+            await websocket.send_text(f"frame-{i}")
+        await websocket.close()
+
+    ws_router_mod.router = r
+    ws_pkg.router = ws_router_mod
+    ws_pkg.warmup = lambda: True
+    sys.modules["mock_ws_pkg"] = ws_pkg
+    sys.modules["mock_ws_pkg.router"] = ws_router_mod
+
+    monkeypatch.setenv("DISABLE_PROJECTS", "")
+    monkeypatch.setenv("MAX_BODY_SIZE_BYTES", "1024")
+    monkeypatch.setenv("REQUEST_TIMEOUT_SECONDS", "5.0")
+
+    import registry
+    registry.PROJECT_REGISTRY = {
+        "/mockws": "mock_ws_pkg.router:router",
+    }
+
+    if "main" in sys.modules:
+        del sys.modules["main"]
+
+    import main
+    with TestClient(main.app) as client:
+        with client.websocket_connect("/mockws/ws") as ws:
+            frames = [ws.receive_text() for _ in range(3)]
+            assert frames == ["frame-0", "frame-1", "frame-2"]