Add labelling studio

main.py

@@ -2,7 +2,7 @@ import os
 import time
 import requests
 from fastapi import FastAPI, UploadFile, File, Form, Query, Request, Body
-from fastapi.responses import Response
+from fastapi.responses import Response, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.concurrency import run_in_threadpool # <--- FIX: Added Threadpool
 import uvicorn
@@ -178,6 +178,156 @@ async def proxy_compatibility(request: Request, password: str = Form(""), userna
         return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "results": []}
 
 
+@app.get("/proxy_experts")
+async def proxy_experts(username: str = Query(""), password: str = Query("")):
+    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
+    target_url = f"{base_url}/experts"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+    def make_get():
+        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)
+
+    try:
+        response = await run_in_threadpool(make_get)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return proxy_error_payload("Proxy Error (Hugging Face)", e, experts=[])
+    except Exception as e:
+        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "experts": []}
+
+
+@app.post("/proxy_finetune")
+async def proxy_finetune(payload: dict = Body(...)):
+    username = str(payload.get("username", ""))
+    expert_id = str(payload.get("expert_id", "")).strip()
+    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
+    target_url = f"{base_url}/experts/{expert_id}/finetune"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+    def make_post():
+        return request_with_hf_backoff(requests.post, target_url, max_retries=0, headers=headers, json=payload, timeout=60)
+
+    try:
+        response = await run_in_threadpool(make_post)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return proxy_error_payload("Proxy Error (Hugging Face)", e)
+    except Exception as e:
+        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}
+
+
+@app.get("/proxy_train_job")
+async def proxy_train_job(job_id: str = Query(...), username: str = Query(""), password: str = Query("")):
+    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
+    target_url = f"{base_url}/train_jobs/status/{job_id}"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+    def make_get():
+        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)
+
+    try:
+        response = await run_in_threadpool(make_get)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return proxy_error_payload("Proxy Error (Hugging Face)", e)
+    except Exception as e:
+        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}
+
+
+@app.get("/proxy_train_jobs")
+async def proxy_train_jobs(username: str = Query(""), password: str = Query("")):
+    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
+    target_url = f"{base_url}/train_jobs"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+    def make_get():
+        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)
+
+    try:
+        response = await run_in_threadpool(make_get)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return proxy_error_payload("Proxy Error (Hugging Face)", e, jobs=[])
+    except Exception as e:
+        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "jobs": []}
+
+
+DATASET_PROXY_TIMEOUT_SECONDS = 180
+
+
+@app.api_route("/proxy_datasets{rest:path}", methods=["GET", "POST", "PUT", "DELETE"])
+async def proxy_datasets(rest: str, request: Request):
+    """Generic forwarder for all Labeling Studio dataset endpoints.
+
+    Mirrors /proxy_datasets{rest} -> {base}/datasets{rest}, preserving the query
+    string, body (multipart / JSON), and returning either JSON or raw binary
+    (image/zip) responses. Routes dev vs prod by username (query, form, or JSON).
+    """
+    username = request.query_params.get("username", "")
+    content_type = request.headers.get("content-type", "")
+
+    files = []
+    data = {}
+    json_body = None
+    if request.method in ("POST", "PUT", "DELETE"):
+        if "multipart/form-data" in content_type:
+            form = await request.form()
+            for key, value in form.multi_items():
+                if hasattr(value, "filename"):
+                    files.append((key, (value.filename, await value.read(), value.content_type)))
+                else:
+                    data[key] = str(value)
+            username = username or data.get("username", "")
+        elif "application/json" in content_type:
+            try:
+                json_body = await request.json()
+            except Exception:
+                json_body = None
+            if isinstance(json_body, dict):
+                username = username or str(json_body.get("username", ""))
+
+    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
+    target_url = f"{base_url}/datasets{rest}"
+    if request.url.query:
+        target_url += f"?{request.url.query}"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+
+    def make_request():
+        return request_with_hf_backoff(
+            getattr(requests, request.method.lower()),
+            target_url,
+            max_retries=0,
+            headers=headers,
+            files=files or None,
+            data=data or None,
+            json=json_body if json_body is not None else None,
+            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
+        )
+
+    try:
+        response = await run_in_threadpool(make_request)
+    except requests.exceptions.RequestException as e:
+        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", e))
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Error (Internal): {str(e)}"})
+
+    resp_ct = response.headers.get("content-type", "application/octet-stream")
+    passthrough_headers = {}
+    disposition = response.headers.get("content-disposition")
+    if disposition:
+        passthrough_headers["content-disposition"] = disposition
+    return Response(
+        content=response.content,
+        status_code=response.status_code,
+        media_type=resp_ct,
+        headers=passthrough_headers,
+    )
+
+
 @app.post("/proxy_batch_stage")
 async def proxy_batch_stage(payload: dict = Body(...)):
     username = str(payload.get("username", ""))