wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve Jupyter workspace bootstrap

Browse Source
This commit is contained in:
wuyang6
2026-05-12 22:33:15 +08:00
parent 295d5d1ea1
commit 86f49134a4
5 changed files with 858 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+156 -4
View File
@@ -23,10 +23,11 @@ from pathlib import Path
from threading import Lock, RLock
from typing import Any
from urllib import error, request
from urllib.parse import quote, unquote, urlparse
from uuid import uuid4
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
@@ -1004,6 +1005,66 @@ def create_app(state: AgentState) -> FastAPI:
safe_session_id,
)
@app.get('/api/jupyter/files')
async def list_jupyter_files(
session_id: str,
account_id: str | None = None,
) -> dict[str, Any]:
safe_session_id = _safe_session_id(session_id)
if not safe_session_id:
raise HTTPException(status_code=400, detail='session_id is required')
runtime = state.jupyter_runtime_manager.get(
state._account_key(account_id),
safe_session_id,
)
if runtime is None:
return {
'connected': False,
'session_id': safe_session_id,
'input': [],
'output': [],
}
try:
return {
'connected': True,
'session_id': safe_session_id,
'workspace_cwd': runtime.binding.workspace_cwd,
'input': runtime.list_files('input', kind='input'),
'output': runtime.list_files('output', kind='output'),
}
except JupyterRuntimeError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
@app.get('/api/jupyter/file')
async def download_jupyter_file(
session_id: str,
path: str,
account_id: str | None = None,
) -> Response:
safe_session_id = _safe_session_id(session_id)
if not safe_session_id:
raise HTTPException(status_code=400, detail='session_id is required')
runtime = state.jupyter_runtime_manager.get(
state._account_key(account_id),
safe_session_id,
)
if runtime is None:
raise HTTPException(status_code=404, detail='Jupyter runtime is not connected')
try:
data, filename = runtime.read_file_bytes(path)
except JupyterRuntimeError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
return Response(
data,
media_type=_content_type_for_filename(filename),
headers={
'content-disposition': (
f'attachment; filename="{_ascii_download_filename(filename)}"; '
f"filename*=UTF-8''{_url_quote_filename(filename)}"
)
},
)
@app.post('/api/jupyter/bind-session')
async def bind_jupyter_session(
payload: JupyterWorkspaceBindRequest,
@@ -1020,6 +1081,7 @@ def create_app(state: AgentState) -> FastAPI:
base_url=payload.base_url,
password=payload.password,
workspace_root=payload.workspace_root,
project_root=state.config_for(payload.account_id).cwd,
)
except JupyterRuntimeError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -1092,6 +1154,13 @@ def create_app(state: AgentState) -> FastAPI:
except (ValueError, subprocess.TimeoutExpired) as exc:
raise HTTPException(status_code=400, detail=str(exc))
result['skills'] = await list_skills(payload.account_id)
account_key = state._account_key(payload.account_id)
config = state.config_for(payload.account_id)
result['remote_runtime_sync'] = await asyncio.to_thread(
state.jupyter_runtime_manager.sync_account,
account_key,
config.cwd,
)
return result
@app.get('/api/models')
@@ -1127,7 +1196,11 @@ def create_app(state: AgentState) -> FastAPI:
@app.post('/api/files/online-doc')
async def create_feishu_online_doc(payload: FeishuOnlineDocRequest) -> dict[str, Any]:
file_path = _resolve_account_file(state, payload.account_id, payload.path)
file_path, online_doc_key = _resolve_feishu_source_file(
state,
payload.account_id,
payload.path,
)
suffix = file_path.suffix.lower()
if suffix in FEISHU_UNSUPPORTED_DOCUMENT_SUFFIXES:
raise HTTPException(
@@ -1186,6 +1259,7 @@ def create_app(state: AgentState) -> FastAPI:
state,
payload.account_id,
file_path=file_path,
map_key=online_doc_key,
title=title,
url=url,
kind=kind,
@@ -2165,6 +2239,28 @@ def _join_url(base_url: str, suffix: str) -> str:
return f'{base_url.rstrip("/")}/{suffix.lstrip("/")}'
def _content_type_for_filename(filename: str) -> str:
suffix = Path(filename).suffix.lower()
if suffix == '.json':
return 'application/json; charset=utf-8'
if suffix in {'.jsonl', '.txt', '.md', '.csv'}:
return 'text/plain; charset=utf-8'
if suffix == '.xlsx':
return 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
if suffix == '.docx':
return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
return 'application/octet-stream'
def _ascii_download_filename(filename: str) -> str:
safe = re.sub(r'[^A-Za-z0-9._-]+', '_', filename).strip('._')
return safe[:120] or 'download'
def _url_quote_filename(filename: str) -> str:
return quote(filename, safe='')
def _annotate_transcript_elapsed(payload: dict[str, Any], elapsed_ms: int) -> None:
transcript = payload.get('transcript')
if not isinstance(transcript, list):
@@ -2651,6 +2747,7 @@ def _record_feishu_online_doc(
account_id: str | None,
*,
file_path: Path,
map_key: str | None = None,
title: str,
url: str,
kind: str = 'doc',
@@ -2667,17 +2764,19 @@ def _record_feishu_online_doc(
files = payload.get('files') if isinstance(payload, dict) else None
if not isinstance(files, dict):
files = {}
previous = files.get(str(file_path))
key = map_key or str(file_path)
previous = files.get(key)
created_at = (
previous.get('created_at')
if isinstance(previous, dict) and isinstance(previous.get('created_at'), int)
else now
)
files[str(file_path)] = {
files[key] = {
'url': clean_url,
'title': title,
'kind': kind,
'file_path': str(file_path),
'source_path': key,
'created_at': created_at,
'updated_at': now,
}
@@ -2687,6 +2786,59 @@ def _record_feishu_online_doc(
)
def _resolve_feishu_source_file(
state: AgentState,
account_id: str | None,
raw_path: str,
) -> tuple[Path, str | None]:
if raw_path.startswith('jupyter://'):
return _materialize_jupyter_file_for_feishu(state, account_id, raw_path), raw_path
return _resolve_account_file(state, account_id, raw_path), None
def _materialize_jupyter_file_for_feishu(
state: AgentState,
account_id: str | None,
raw_uri: str,
) -> Path:
session_id, remote_path = _parse_jupyter_file_uri(raw_uri)
runtime = state.jupyter_runtime_manager.get(
state._account_key(account_id),
session_id,
)
if runtime is None:
raise HTTPException(status_code=404, detail='Jupyter runtime is not connected')
try:
data, filename = runtime.read_file_bytes(remote_path)
except JupyterRuntimeError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
target_dir = (
state.account_paths(account_id)['base']
/ 'integrations'
/ 'feishu'
/ 'remote-files'
/ session_id
)
target_dir.mkdir(parents=True, exist_ok=True)
target = target_dir / _safe_uploaded_filename(filename)
target.write_bytes(data)
return target
def _parse_jupyter_file_uri(raw_uri: str) -> tuple[str, str]:
parsed = urlparse(raw_uri)
session_id = _safe_session_id(parsed.netloc)
remote_path = unquote(parsed.path or '')
if not session_id or not remote_path.startswith('/'):
raise HTTPException(status_code=400, detail='Invalid Jupyter file URI')
return session_id, remote_path
def _safe_uploaded_filename(filename: str) -> str:
cleaned = re.sub(r'[\x00-\x1f/\\]+', '_', filename).strip(' ._')
return cleaned[:180] or 'remote-file'
def _resolve_account_file(state: AgentState, account_id: str | None, raw_path: str) -> Path:
path = Path(raw_path).expanduser().resolve()
account_base = state.account_paths(account_id)['base'].resolve()