medalion-core 0.1.12__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. medalion_core-0.1.12/.gitignore +16 -0
  2. medalion_core-0.1.12/PKG-INFO +48 -0
  3. medalion_core-0.1.12/README.md +14 -0
  4. medalion_core-0.1.12/pyproject.toml +53 -0
  5. medalion_core-0.1.12/src/medalion/__init__.py +16 -0
  6. medalion_core-0.1.12/src/medalion/api/__init__.py +0 -0
  7. medalion_core-0.1.12/src/medalion/api/app.py +376 -0
  8. medalion_core-0.1.12/src/medalion/api/studio.py +69 -0
  9. medalion_core-0.1.12/src/medalion/api/supervisor.py +208 -0
  10. medalion_core-0.1.12/src/medalion/api/usage.py +48 -0
  11. medalion_core-0.1.12/src/medalion/cli.py +906 -0
  12. medalion_core-0.1.12/src/medalion/core/__init__.py +0 -0
  13. medalion_core-0.1.12/src/medalion/core/artifact.py +145 -0
  14. medalion_core-0.1.12/src/medalion/core/bench.py +69 -0
  15. medalion_core-0.1.12/src/medalion/core/binaries.py +39 -0
  16. medalion_core-0.1.12/src/medalion/core/catalog.py +95 -0
  17. medalion_core-0.1.12/src/medalion/core/config.py +132 -0
  18. medalion_core-0.1.12/src/medalion/core/license.py +313 -0
  19. medalion_core-0.1.12/src/medalion/core/model.py +104 -0
  20. medalion_core-0.1.12/src/medalion/core/preflight.py +164 -0
  21. medalion_core-0.1.12/src/medalion/core/registry.py +87 -0
  22. medalion_core-0.1.12/src/medalion/core/render.py +231 -0
  23. medalion_core-0.1.12/src/medalion/core/runtime.py +120 -0
  24. medalion_core-0.1.12/src/medalion/core/task.py +150 -0
  25. medalion_core-0.1.12/src/medalion/core/worker.py +431 -0
  26. medalion_core-0.1.12/src/medalion/engines/__init__.py +0 -0
  27. medalion_core-0.1.12/src/medalion/engines/bench_data.py +21 -0
  28. medalion_core-0.1.12/src/medalion/engines/llama_server.py +238 -0
  29. medalion_core-0.1.12/src/medalion/models/__init__.py +5 -0
  30. medalion_core-0.1.12/src/medalion/tasks/__init__.py +8 -0
  31. medalion_core-0.1.12/src/medalion/tasks/asr.py +125 -0
  32. medalion_core-0.1.12/src/medalion/tasks/chat.py +47 -0
  33. medalion_core-0.1.12/src/medalion/tasks/document.py +59 -0
  34. medalion_core-0.1.12/src/medalion/tasks/ner.py +82 -0
@@ -0,0 +1,16 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.pyc
4
+ .pytest_cache/
5
+ .ruff_cache/
6
+ studio/node_modules/
7
+ studio/.next/
8
+ studio/next-env.d.ts
9
+ *.log
10
+ .DS_Store
11
+ dist/
12
+ services/distribution/keys/
13
+ services/distribution/.venv/
14
+ services/distribution/models/
15
+ packages/medalion-studio/src/medalion/studio/dist/
16
+ studio/out/
@@ -0,0 +1,48 @@
1
+ Metadata-Version: 2.5
2
+ Name: medalion-core
3
+ Version: 0.1.12
4
+ Summary: Medalion Platform — runtime, task contracts, API gateway, CLI and licensing
5
+ Project-URL: Homepage, https://medalion.tech
6
+ Project-URL: Documentation, https://medalion.tech
7
+ Project-URL: Source, https://github.com/Medalion-Tech/medalion-platform
8
+ Author-email: Medalion <kontakt@medalion.tech>
9
+ License: Proprietary
10
+ Keywords: asr,fhir,llm,medical,nlp,on-premise
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Environment :: Console
13
+ Classifier: Framework :: FastAPI
14
+ Classifier: Intended Audience :: Healthcare Industry
15
+ Classifier: License :: Other/Proprietary License
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Python: <3.13,>=3.12
19
+ Requires-Dist: cryptography>=42
20
+ Requires-Dist: fastapi>=0.115
21
+ Requires-Dist: httpx>=0.27
22
+ Requires-Dist: huggingface-hub>=0.25
23
+ Requires-Dist: orjson>=3.10
24
+ Requires-Dist: pydantic-settings>=2.4
25
+ Requires-Dist: pydantic>=2.8
26
+ Requires-Dist: python-multipart>=0.0.9
27
+ Requires-Dist: pyyaml>=6
28
+ Requires-Dist: ruamel-yaml>=0.18
29
+ Requires-Dist: structlog>=24
30
+ Requires-Dist: typer>=0.12
31
+ Requires-Dist: uvicorn[standard]>=0.30
32
+ Requires-Dist: websockets>=13
33
+ Description-Content-Type: text/markdown
34
+
35
+ # Medalion Platform — core
36
+
37
+ The runtime of the Medalion Platform: task contracts, the API gateway that supervises one worker
38
+ process per model, the CLI (`medalion init/update/serve/status/doctor`) and licensing.
39
+
40
+ Core on its own gives you the API, the Studio UI (`medalion-studio`) and a llama.cpp chat engine.
41
+ The clinical modules — speech to text, clinical NER, anonymisation, documents to FHIR, workflows —
42
+ are licensed and installed from Medalion's own index:
43
+
44
+ ```bash
45
+ curl -fsSL https://get.medalion.tech | sh
46
+ ```
47
+
48
+ Documentation and licensing: <https://medalion.tech>.
@@ -0,0 +1,14 @@
1
+ # Medalion Platform — core
2
+
3
+ The runtime of the Medalion Platform: task contracts, the API gateway that supervises one worker
4
+ process per model, the CLI (`medalion init/update/serve/status/doctor`) and licensing.
5
+
6
+ Core on its own gives you the API, the Studio UI (`medalion-studio`) and a llama.cpp chat engine.
7
+ The clinical modules — speech to text, clinical NER, anonymisation, documents to FHIR, workflows —
8
+ are licensed and installed from Medalion's own index:
9
+
10
+ ```bash
11
+ curl -fsSL https://get.medalion.tech | sh
12
+ ```
13
+
14
+ Documentation and licensing: <https://medalion.tech>.
@@ -0,0 +1,53 @@
1
+ [project]
2
+ name = "medalion-core"
3
+ version = "0.1.12"
4
+ description = "Medalion Platform — runtime, task contracts, API gateway, CLI and licensing"
5
+ readme = "README.md"
6
+ requires-python = ">=3.12,<3.13"
7
+ license = { text = "Proprietary" }
8
+ authors = [{ name = "Medalion", email = "kontakt@medalion.tech" }]
9
+ keywords = ["medical", "nlp", "asr", "fhir", "llm", "on-premise"]
10
+ classifiers = [
11
+ "Development Status :: 4 - Beta",
12
+ "Environment :: Console",
13
+ "Framework :: FastAPI",
14
+ "Intended Audience :: Healthcare Industry",
15
+ "License :: Other/Proprietary License",
16
+ "Programming Language :: Python :: 3.12",
17
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
18
+ ]
19
+
20
+ dependencies = [
21
+ "fastapi>=0.115",
22
+ "uvicorn[standard]>=0.30",
23
+ "pydantic>=2.8",
24
+ "pydantic-settings>=2.4",
25
+ "httpx>=0.27",
26
+ "websockets>=13",
27
+ "python-multipart>=0.0.9",
28
+ "structlog>=24",
29
+ "typer>=0.12",
30
+ "pyyaml>=6",
31
+ "ruamel.yaml>=0.18",
32
+ "cryptography>=42",
33
+ "huggingface_hub>=0.25",
34
+ "orjson>=3.10",
35
+ ]
36
+
37
+ [project.urls]
38
+ Homepage = "https://medalion.tech"
39
+ Documentation = "https://medalion.tech"
40
+ Source = "https://github.com/Medalion-Tech/medalion-platform"
41
+
42
+ [project.scripts]
43
+ medalion = "medalion.cli:app"
44
+
45
+ [project.entry-points."medalion.models"]
46
+ llama-server = "medalion.engines.llama_server:LlamaServer"
47
+
48
+ [build-system]
49
+ requires = ["hatchling"]
50
+ build-backend = "hatchling.build"
51
+
52
+ [tool.hatch.build.targets.wheel]
53
+ packages = ["src/medalion"]
@@ -0,0 +1,16 @@
1
+ """Medalion Platform — one runtime, every model a plugin."""
2
+
3
+ # `medalion` spans several distributions (medalion-core, medalion-asr, ...); merge
4
+ # their `medalion/` directories into one package, also for editable installs.
5
+ from pkgutil import extend_path
6
+
7
+ __path__ = extend_path(__path__, __name__)
8
+
9
+ from medalion.core.artifact import artifact # noqa: E402
10
+ from medalion.core.bench import BenchmarkResult # noqa: E402
11
+ from medalion.core.model import Model # noqa: E402
12
+ from medalion.core.runtime import Runtime # noqa: E402
13
+ from medalion.core.task import Op, Task, Unit, op # noqa: E402
14
+
15
+ __all__ = ["Model", "artifact", "Runtime", "Task", "Op", "Unit", "op", "BenchmarkResult"]
16
+ __version__ = "0.1.12"
File without changes
@@ -0,0 +1,376 @@
1
+ """Control-plane API: license, registry, router/proxy, usage, benchmarks. No user auth in the POC."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import inspect
7
+ import json
8
+ import os
9
+ import time
10
+ import uuid
11
+ from contextlib import asynccontextmanager
12
+ from typing import Any
13
+
14
+ import httpx
15
+ import structlog
16
+ import websockets
17
+ from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
18
+ from fastapi.middleware.cors import CORSMiddleware
19
+ from fastapi.responses import JSONResponse, Response, StreamingResponse
20
+
21
+ from medalion import __version__
22
+ from medalion.api.supervisor import Supervisor, WorkerHandle
23
+ from medalion.api.usage import UsageStore
24
+ from medalion.core.config import Config, load_config
25
+ from medalion.core.license import LicenseManager
26
+ from medalion.core.registry import scan_plugins
27
+ from medalion.core.task import Op, all_tasks, get_task
28
+
29
+ log = structlog.get_logger("medalion.api")
30
+
31
+ HOP_HEADERS = {"content-length", "transfer-encoding", "connection", "host", "content-encoding"}
32
+
33
+
34
+ def create_app(config_path: str | None = None) -> FastAPI:
35
+ cfg: Config = load_config(config_path)
36
+ api_url = f"http://{cfg.api.host}:{cfg.api.port}"
37
+ lic = LicenseManager(cfg, version=__version__)
38
+ ent = lic.start().entitlements
39
+ sup = Supervisor(cfg, config_path or os.environ.get("MEDALION_CONFIG", "medalion.yaml"), api_url, entitlements=ent)
40
+ usage = UsageStore(cfg.data_dir / "usage.db")
41
+
42
+ @asynccontextmanager
43
+ async def lifespan(app: FastAPI):
44
+ sup.start()
45
+ app.state.client = httpx.AsyncClient(timeout=httpx.Timeout(900.0, connect=5.0))
46
+ heartbeat = asyncio.get_event_loop().create_task(lic.heartbeat_loop())
47
+ yield
48
+ heartbeat.cancel()
49
+ await app.state.client.aclose()
50
+ await sup.stop()
51
+
52
+ app = FastAPI(
53
+ title="Medalion Platform API",
54
+ version=__version__,
55
+ description="One runtime, every model a plugin. Routes under /v1/{task}/{op} are derived from task contracts.",
56
+ lifespan=lifespan,
57
+ )
58
+ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], expose_headers=["*"])
59
+ app.state.cfg, app.state.sup, app.state.usage, app.state.license = cfg, sup, usage, lic
60
+
61
+ modules: list[str] = []
62
+
63
+ # ---------------------------------------------------------------- meta
64
+ @app.get("/v1/info", tags=["platform"])
65
+ def info():
66
+ return {"version": app.version, "deployment": cfg.deployment, "hardware": cfg.hardware, "modules": modules,
67
+ "license": lic.state.status, "studio": app.state.studio}
68
+
69
+ @app.get("/v1/license", tags=["platform"])
70
+ def license_status():
71
+ return lic.state.to_dict()
72
+
73
+ @app.get("/v1/capabilities", tags=["platform"])
74
+ def capabilities():
75
+ """What this deployment can do: licensed products, mounted modules, tasks with a usable model."""
76
+ plugins = scan_plugins(lic.state.entitlements)
77
+ return {
78
+ "products": sorted(lic.state.entitlements.products),
79
+ "modules": modules,
80
+ "tasks": {t.id: any(p.licensed and p.task == t.id for p in plugins) for t in all_tasks()},
81
+ "plugins": [{"id": p.id, "task": p.task, "product": p.product, "licensed": p.licensed} for p in plugins],
82
+ "license": lic.state.to_dict(),
83
+ }
84
+
85
+ @app.get("/healthz", tags=["platform"])
86
+ def healthz():
87
+ return {"ok": True, "deployment": cfg.deployment, "hardware": cfg.hardware, "models": {k: h.status for k, h in sup.workers.items()}}
88
+
89
+ @app.get("/v1/models", tags=["platform"])
90
+ def list_models():
91
+ return {"models": sup.snapshot(), "defaults": cfg.defaults}
92
+
93
+ @app.get("/v1/models/{model_id}", tags=["platform"])
94
+ async def get_model(model_id: str):
95
+ h = _handle(sup, model_id)
96
+ d = next(m for m in sup.snapshot() if m["model"] == model_id)
97
+ if h.ready:
98
+ try:
99
+ r = await app.state.client.get(f"{h.base_url}/manifest")
100
+ d["manifest"] = r.json()
101
+ except Exception:
102
+ pass
103
+ return d
104
+
105
+ @app.get("/v1/tasks", tags=["platform"])
106
+ def list_tasks():
107
+ out = []
108
+ for t in all_tasks():
109
+ m = t.manifest()
110
+ m["models"] = [{"id": h.spec.id, "title": h.info.title, "status": h.status} for h in sup.for_task(t.id)]
111
+ m["default_model"] = cfg.defaults.get(t.id) or (m["models"][0]["id"] if m["models"] else None)
112
+ out.append(m)
113
+ return {"tasks": out}
114
+
115
+ @app.get("/v1/usage", tags=["platform"])
116
+ def get_usage(limit: int = 50):
117
+ return {"summary": usage.summary(), "recent": usage.recent(limit)}
118
+
119
+ # ---------------------------------------------------------- benchmarks
120
+ @app.get("/v1/benchmarks", tags=["benchmarks"])
121
+ async def list_benchmarks():
122
+ out = []
123
+ for h in sup.workers.values():
124
+ item: dict[str, Any] = {"model": h.spec.id, "task": h.info.task, "has_benchmark": h.info.has_benchmark, "running": False, "result": None, "error": None}
125
+ if h.ready:
126
+ try:
127
+ r = await app.state.client.get(f"{h.base_url}/v1/benchmark", timeout=10)
128
+ item.update(r.json())
129
+ except Exception as e: # noqa: BLE001
130
+ item["error"] = str(e)
131
+ out.append(item)
132
+ return {"benchmarks": out}
133
+
134
+ @app.get("/v1/models/{model_id}/benchmark", tags=["benchmarks"])
135
+ async def get_benchmark(model_id: str):
136
+ h = _handle(sup, model_id)
137
+ r = await app.state.client.get(f"{h.base_url}/v1/benchmark", timeout=10)
138
+ return JSONResponse(r.json(), status_code=r.status_code)
139
+
140
+ @app.post("/v1/models/{model_id}/benchmark", tags=["benchmarks"])
141
+ async def run_benchmark(model_id: str, limit: int | None = None):
142
+ h = _handle(sup, model_id)
143
+ r = await app.state.client.post(f"{h.base_url}/v1/benchmark", params={"limit": limit} if limit else None, timeout=10)
144
+ return JSONResponse(r.json(), status_code=r.status_code)
145
+
146
+ # ---------------------------------------------------- derived op routes
147
+ for t in all_tasks():
148
+ for op in t.ops:
149
+ _mount_proxy(app, sup, usage, t.id, op)
150
+
151
+ # ------------------------------------------------------ optional modules
152
+ if cfg.workflows.enabled and not ent.allows("workflows"):
153
+ log.warning("workflows_not_licensed")
154
+ elif cfg.workflows.enabled:
155
+ try:
156
+ from medalion.workflows import build_router
157
+
158
+ app.include_router(build_router(app))
159
+ modules.append("workflows")
160
+ except ImportError as e:
161
+ log.info("workflows_unavailable", error=str(e))
162
+
163
+ # ------------------------------------------------------------- the UI
164
+ from medalion.api.studio import mount_studio
165
+
166
+ app.state.studio = mount_studio(app)
167
+ if not app.state.studio:
168
+ log.info("studio_not_installed") # API-only deployment, or a dev checkout without MEDALION_STUDIO_DIR
169
+
170
+ # note: the derived route for task `chat` / op `completions` is already /v1/chat/completions (OpenAI-compatible)
171
+ return app
172
+
173
+
174
+ # ------------------------------------------------------------------ helpers
175
+ def _handle(sup: Supervisor, model_id: str) -> WorkerHandle:
176
+ h = sup.handle(model_id)
177
+ if not h:
178
+ raise HTTPException(404, f"unknown model {model_id!r}")
179
+ return h
180
+
181
+
182
+ def _mount_proxy(app: FastAPI, sup: Supervisor, usage: UsageStore, task_id: str, op: Op) -> None:
183
+ path = f"/v1/{task_id}/{op.path}"
184
+
185
+ if op.protocol == "websocket":
186
+
187
+ @app.websocket(path)
188
+ async def ws_proxy(ws: WebSocket):
189
+ await _proxy_ws(sup, usage, task_id, op, ws)
190
+
191
+ return
192
+
193
+ if op.request is not None and op.accepts == "json":
194
+ ReqModel = op.request
195
+
196
+ async def typed_proxy(request, body): # signature set below so the OpenAPI shows the op's request model
197
+ return await _proxy_http(app, sup, usage, task_id, op, request)
198
+
199
+ typed_proxy.__signature__ = inspect.Signature( # type: ignore[attr-defined]
200
+ [inspect.Parameter("request", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Request),
201
+ inspect.Parameter("body", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=ReqModel)]
202
+ )
203
+ app.post(path, tags=[task_id], summary=op.summary, name=f"{task_id}_{op.name}", response_model=op.response)(typed_proxy)
204
+
205
+ else:
206
+
207
+ @app.post(path, tags=[task_id], summary=op.summary, name=f"{task_id}_{op.name}", response_model=op.response)
208
+ async def raw_proxy(request: Request):
209
+ return await _proxy_http(app, sup, usage, task_id, op, request)
210
+
211
+
212
+ async def call_op(
213
+ app: FastAPI,
214
+ task_id: str,
215
+ op_name: str,
216
+ *,
217
+ json_body: dict | None = None,
218
+ data: dict | None = None,
219
+ files: dict | None = None,
220
+ model: str | None = None,
221
+ stream: bool = False,
222
+ request_id: str | None = None,
223
+ ) -> httpx.Response:
224
+ """Invoke a model op from inside the API process (used by the workflows module).
225
+
226
+ Shares model selection, readiness checks and usage metering with the HTTP proxy;
227
+ calls carry ``request_id`` so workflow steps are traceable in /v1/usage.
228
+ The returned response has ``extensions["medalion_model"]`` set to the model used.
229
+ Raises LookupError (unknown model/none for task) or RuntimeError (not ready).
230
+ """
231
+ sup: Supervisor = app.state.sup
232
+ usage: UsageStore = app.state.usage
233
+ client: httpx.AsyncClient = app.state.client
234
+ op = get_task(task_id).op(op_name)
235
+ h = sup.pick(task_id, model)
236
+ if not h.ready:
237
+ raise RuntimeError(f"model {h.spec.id!r} is {h.status}" + (f": {h.error}" if h.error else ""))
238
+ rid = request_id or uuid.uuid4().hex[:12]
239
+ t0 = time.time()
240
+ req = client.build_request(
241
+ "POST", f"{h.base_url}/v1/{task_id}/{op.path}",
242
+ json=json_body, data=data, files=files, headers={"x-request-id": rid},
243
+ )
244
+ upstream = await client.send(req, stream=stream)
245
+ if not stream:
246
+ await upstream.aread()
247
+ await upstream.aclose()
248
+ u = upstream.headers.get("x-medalion-usage", "")
249
+ try:
250
+ amount = float(u.split("=", 1)[1]) if "=" in u else 1.0
251
+ except ValueError:
252
+ amount = 1.0
253
+ usage.record(task=task_id, op=op.name, model=h.spec.id, unit=op.unit.name, amount=amount,
254
+ latency_ms=(time.time() - t0) * 1000, status=upstream.status_code, request_id=rid)
255
+ upstream.extensions["medalion_model"] = h.spec.id
256
+ return upstream
257
+
258
+
259
+ async def _pick_model(sup: Supervisor, task_id: str, request: Request, body: bytes | None) -> WorkerHandle:
260
+ model_id = request.query_params.get("model")
261
+ if not model_id and body and request.headers.get("content-type", "").startswith("application/json"):
262
+ try:
263
+ model_id = json.loads(body).get("model") or None
264
+ except Exception:
265
+ model_id = None
266
+ if not model_id and request.headers.get("content-type", "").startswith("multipart/form-data"):
267
+ # cheap sniff for a `model` field without parsing the whole form twice
268
+ pass
269
+ try:
270
+ return sup.pick(task_id, model_id)
271
+ except LookupError as e:
272
+ raise HTTPException(404, str(e))
273
+
274
+
275
+ async def _proxy_http(app: FastAPI, sup: Supervisor, usage: UsageStore, task_id: str, op: Op, request: Request) -> Response:
276
+ client: httpx.AsyncClient = app.state.client
277
+ body = await request.body()
278
+ h = await _pick_model(sup, task_id, request, body)
279
+ if not h.ready:
280
+ raise HTTPException(503, f"model {h.spec.id!r} is {h.status}" + (f": {h.error}" if h.error else ""))
281
+ rid = request.headers.get("x-request-id") or uuid.uuid4().hex[:12]
282
+ headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_HEADERS}
283
+ headers["x-request-id"] = rid
284
+ url = f"{h.base_url}/v1/{task_id}/{op.path}"
285
+ t0 = time.time()
286
+
287
+ req = client.build_request("POST", url, content=body, headers=headers, params=dict(request.query_params))
288
+ try:
289
+ upstream = await client.send(req, stream=True)
290
+ except httpx.HTTPError as e:
291
+ raise HTTPException(502, f"worker unreachable: {e}")
292
+
293
+ resp_headers = {k: v for k, v in upstream.headers.items() if k.lower() not in HOP_HEADERS}
294
+ resp_headers["x-request-id"] = rid
295
+ resp_headers["x-medalion-model"] = h.spec.id
296
+
297
+ def _record(status: int, amount: float | None = None):
298
+ unit_amount = amount
299
+ if unit_amount is None:
300
+ u = upstream.headers.get("x-medalion-usage", "")
301
+ try:
302
+ unit_amount = float(u.split("=", 1)[1]) if "=" in u else 1.0
303
+ except ValueError:
304
+ unit_amount = 1.0
305
+ usage.record(task=task_id, op=op.name, model=h.spec.id, unit=op.unit.name, amount=unit_amount,
306
+ latency_ms=(time.time() - t0) * 1000, status=status, request_id=rid)
307
+
308
+ if upstream.headers.get("content-type", "").startswith("text/event-stream"):
309
+ async def gen():
310
+ try:
311
+ async for chunk in upstream.aiter_raw():
312
+ yield chunk
313
+ finally:
314
+ await upstream.aclose()
315
+ _record(upstream.status_code)
316
+
317
+ return StreamingResponse(gen(), status_code=upstream.status_code, headers=resp_headers, media_type="text/event-stream")
318
+
319
+ content = await upstream.aread()
320
+ await upstream.aclose()
321
+ _record(upstream.status_code)
322
+ return Response(content=content, status_code=upstream.status_code, headers=resp_headers, media_type=upstream.headers.get("content-type"))
323
+
324
+
325
+ async def _proxy_ws(sup: Supervisor, usage: UsageStore, task_id: str, op: Op, ws: WebSocket) -> None:
326
+ await ws.accept()
327
+ model_id = ws.query_params.get("model")
328
+ try:
329
+ h = sup.pick(task_id, model_id)
330
+ except LookupError as e:
331
+ await ws.send_text(json.dumps({"type": "error", "error": {"code": "not_found", "message": str(e)}}))
332
+ await ws.close(code=1008)
333
+ return
334
+ if not h.ready:
335
+ await ws.send_text(json.dumps({"type": "error", "error": {"code": "unavailable", "message": f"model {h.spec.id} is {h.status}"}}))
336
+ await ws.close(code=1013)
337
+ return
338
+ rid = uuid.uuid4().hex[:12]
339
+ t0 = time.time()
340
+ upstream_url = f"{h.ws_url}/v1/{task_id}/{op.path}"
341
+ try:
342
+ async with websockets.connect(upstream_url, max_size=16 * 1024 * 1024) as up:
343
+
344
+ async def c2s():
345
+ try:
346
+ while True:
347
+ msg = await ws.receive()
348
+ if msg["type"] == "websocket.disconnect":
349
+ await up.close()
350
+ return
351
+ if "text" in msg and msg["text"] is not None:
352
+ await up.send(msg["text"])
353
+ elif "bytes" in msg and msg["bytes"] is not None:
354
+ await up.send(msg["bytes"])
355
+ except (WebSocketDisconnect, websockets.ConnectionClosed):
356
+ pass
357
+
358
+ async def s2c():
359
+ try:
360
+ async for m in up:
361
+ if isinstance(m, bytes):
362
+ await ws.send_bytes(m)
363
+ else:
364
+ await ws.send_text(m)
365
+ except (websockets.ConnectionClosed, WebSocketDisconnect, RuntimeError):
366
+ pass
367
+
368
+ await asyncio.wait([asyncio.create_task(c2s()), asyncio.create_task(s2c())], return_when=asyncio.FIRST_COMPLETED)
369
+ except Exception as e: # noqa: BLE001
370
+ log.warning("ws.proxy_failed", error=str(e))
371
+ finally:
372
+ usage.record(task=task_id, op=op.name, model=h.spec.id, unit=op.unit.name, amount=1.0, latency_ms=(time.time() - t0) * 1000, status=101, request_id=rid)
373
+ try:
374
+ await ws.close()
375
+ except Exception:
376
+ pass
@@ -0,0 +1,69 @@
1
+ """Serve the Studio from the platform API.
2
+
3
+ The `medalion-studio` wheel carries a static export (no Node on the customer's machine).
4
+ Mounting it under the API means one origin, one port and one service: the UI talks to
5
+ `/v1/...` on itself, so nothing has to be configured after `medalion init`.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import os
11
+ from pathlib import Path
12
+
13
+ import structlog
14
+ from fastapi import FastAPI
15
+ from fastapi.responses import FileResponse
16
+ from fastapi.staticfiles import StaticFiles
17
+
18
+ log = structlog.get_logger("medalion.api")
19
+
20
+ # The export renders one page per dynamic route; both sides agree on this id
21
+ # (studio/src/lib/routes.ts: PLACEHOLDER_WORKFLOW_ID).
22
+ PLACEHOLDER = "_id"
23
+
24
+
25
+ def studio_root() -> Path | None:
26
+ """MEDALION_STUDIO_DIR (a dev checkout's `studio/out`) or the installed wheel's assets."""
27
+ env = os.environ.get("MEDALION_STUDIO_DIR")
28
+ if env:
29
+ p = Path(env).expanduser()
30
+ return p if (p / "index.html").is_file() else None
31
+ try:
32
+ from medalion.studio import root
33
+ except ImportError:
34
+ return None
35
+ return root()
36
+
37
+
38
+ def mount_studio(app: FastAPI, root: Path | None = None) -> bool:
39
+ """Mount the UI at `/`. Returns False when the Studio package is not installed."""
40
+ root = root or studio_root()
41
+ if root is None:
42
+ return False
43
+
44
+ @app.get("/workflows/{workflow_id}", include_in_schema=False)
45
+ @app.get("/workflows/{workflow_id}/", include_in_schema=False)
46
+ def workflow_page(workflow_id: str) -> FileResponse:
47
+ return _placeholder(root, workflow_id, "index.html")
48
+
49
+ @app.get("/workflows/{workflow_id}/edit", include_in_schema=False)
50
+ @app.get("/workflows/{workflow_id}/edit/", include_in_schema=False)
51
+ def workflow_editor_page(workflow_id: str) -> FileResponse:
52
+ return _placeholder(root, workflow_id, "edit/index.html")
53
+
54
+ # last route: everything not claimed by /v1, /docs, … is a static file or a directory index
55
+ app.mount("/", StaticFiles(directory=root, html=True), name="studio")
56
+ log.info("studio.mounted", path=str(root))
57
+ return True
58
+
59
+
60
+ def _placeholder(root: Path, workflow_id: str, rel: str) -> FileResponse:
61
+ """Every /workflows/<id> URL is served by the one exported placeholder page."""
62
+ from fastapi import HTTPException
63
+
64
+ if workflow_id in ("index.html", ""): # `/workflows/` itself is a real page
65
+ raise HTTPException(404)
66
+ page = root / "workflows" / PLACEHOLDER / rel
67
+ if not page.is_file():
68
+ raise HTTPException(404)
69
+ return FileResponse(page, media_type="text/html")