wuyang
105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import AsyncIterator
|
|
from typing import Annotated
|
|
|
|
import httpx
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from fastapi.responses import StreamingResponse
|
|
from open_webui.utils.auth import get_verified_user
|
|
from open_webui.utils.headers import include_user_info_headers
|
|
|
|
router = APIRouter()
|
|
|
|
GATEWAY_URL = os.getenv("K1412_GATEWAY_URL", "http://gateway:8001").rstrip("/")
|
|
GATEWAY_KEY = os.getenv("K1412_GATEWAY_KEY", "")
|
|
REQUEST_TIMEOUT = httpx.Timeout(120, connect=10)
|
|
STREAM_TIMEOUT = httpx.Timeout(1800, connect=10)
|
|
VerifiedUser = Annotated[object, Depends(get_verified_user)]
|
|
|
|
|
|
def gateway_headers(user) -> dict[str, str]:
|
|
if not GATEWAY_KEY:
|
|
raise HTTPException(status_code=503, detail="Workspace service is not configured")
|
|
return include_user_info_headers(
|
|
{"Authorization": f"Bearer {GATEWAY_KEY}"},
|
|
user,
|
|
)
|
|
|
|
|
|
async def gateway_error(response: httpx.Response) -> HTTPException:
|
|
try:
|
|
payload = response.json()
|
|
detail = payload.get("detail", payload)
|
|
except Exception:
|
|
detail = (await response.aread()).decode("utf-8", errors="replace")[:2000]
|
|
return HTTPException(status_code=response.status_code, detail=detail)
|
|
|
|
|
|
@router.get("/files")
|
|
async def list_workspace_files(
|
|
user: VerifiedUser,
|
|
path: str = Query(default=".", max_length=4096),
|
|
):
|
|
async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
|
|
response = await client.get(
|
|
f"{GATEWAY_URL}/v1/files",
|
|
params={"path": path},
|
|
headers=gateway_headers(user),
|
|
)
|
|
if response.status_code >= 400:
|
|
raise await gateway_error(response)
|
|
return response.json()
|
|
|
|
|
|
async def close_after_stream(response: httpx.Response, client: httpx.AsyncClient) -> AsyncIterator[bytes]:
|
|
try:
|
|
async for chunk in response.aiter_bytes():
|
|
yield chunk
|
|
finally:
|
|
await response.aclose()
|
|
await client.aclose()
|
|
|
|
|
|
async def stream_workspace_response(endpoint: str, path: str, user) -> StreamingResponse:
|
|
client = httpx.AsyncClient(timeout=STREAM_TIMEOUT)
|
|
request = client.build_request(
|
|
"GET",
|
|
f"{GATEWAY_URL}{endpoint}",
|
|
params={"path": path},
|
|
headers=gateway_headers(user),
|
|
)
|
|
response = await client.send(request, stream=True)
|
|
if response.status_code >= 400:
|
|
error = await gateway_error(response)
|
|
await response.aclose()
|
|
await client.aclose()
|
|
raise error
|
|
headers = {
|
|
key: value
|
|
for key, value in response.headers.items()
|
|
if key.lower() in {"content-disposition", "content-length"}
|
|
}
|
|
return StreamingResponse(
|
|
close_after_stream(response, client),
|
|
media_type=response.headers.get("content-type", "application/octet-stream"),
|
|
headers=headers,
|
|
)
|
|
|
|
|
|
@router.get("/download")
|
|
async def download_workspace_file(
|
|
user: VerifiedUser,
|
|
path: str = Query(min_length=1, max_length=4096),
|
|
):
|
|
return await stream_workspace_response("/v1/files/download", path, user)
|
|
|
|
|
|
@router.get("/archive")
|
|
async def archive_workspace_files(
|
|
user: VerifiedUser,
|
|
path: str = Query(default=".", max_length=4096),
|
|
):
|
|
return await stream_workspace_response("/v1/files/archive", path, user)
|