proxytool 0.4.0__py3-none-any.whl

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.
proxytool/main.py ADDED
@@ -0,0 +1,572 @@
1
+ """FastAPI application: proxytool sits between the agent client and the upstream
2
+ model server (llama-server / mlx_lm.server / llama-swap / any OpenAI dialect).
3
+
4
+ - POST /v1/chat/completions : sanitize the outbound request, repair leaked
5
+ tool calls in the inbound response (both streaming SSE and non-streaming).
6
+ - Everything else: byte pass-through.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ import re
13
+ import time
14
+ from collections.abc import AsyncGenerator
15
+ from typing import Any
16
+
17
+ import httpx
18
+ from fastapi import FastAPI, Request
19
+ from fastapi.responses import JSONResponse, Response, StreamingResponse
20
+
21
+ import proxytool as _pkg
22
+ from proxytool import trace
23
+ from proxytool.conduit import UpstreamConduit
24
+ from proxytool.loopguard import ABORT_TEXT, NUDGE_TEXT, LoopGuard
25
+ from proxytool.settings import get_settings
26
+ from toolpatch import StreamHealer, extract_tool_calls, normalize_chat_request
27
+ from toolpatch.logging import get_logger
28
+ from toolpatch.request import Backend
29
+
30
+ _log = get_logger("proxytool.main")
31
+
32
+
33
+ def _can_emit(obj: object) -> bool:
34
+ """True only for a configured stdlib logger (no-op loggers drop lines)."""
35
+ import logging as _std_logging
36
+
37
+ return obj is not None and isinstance(obj, _std_logging.Logger)
38
+
39
+
40
+ settings = get_settings()
41
+ # When the app module is imported without the CLI (e.g. `uvicorn
42
+ # proxytool.main:app`), no one has configured app logging yet; apply the
43
+ # env-driven default exactly once. A CLI configure call has already set the
44
+ # "proxytool" logger level, so this is a no-op then.
45
+ if _pkg._logging.getLogger("proxytool").level == _pkg._logging.NOTSET:
46
+ _pkg.configure_logging()
47
+ trace.attach_logger(_log)
48
+ trace.enable(settings.trace)
49
+ conduit = UpstreamConduit(
50
+ settings.upstream_url,
51
+ timeout=settings.upstream_timeout,
52
+ api_key=settings.upstream_api_key,
53
+ )
54
+ loop_guard = LoopGuard(
55
+ nudge_after=settings.loop_nudge,
56
+ abort_after=settings.loop_abort,
57
+ )
58
+ app = FastAPI(title="proxytool")
59
+
60
+ if _can_emit(_log):
61
+ _log.debug("proxytool initialized: upstream=%s",
62
+ settings.upstream_url)
63
+
64
+ PASS_HEADERS_OUT = ("content-type", "x-request-id")
65
+
66
+ _UPSTREAM_MODELS: list[str] | None = None # lazily discovered, cached
67
+
68
+ # ---- optional multi-backend routing -----------------------------------------
69
+ # When PTP_UPSTREAM_<NAME> vars are present, requests are routed per-model
70
+ # through BackendRouter; each routed backend gets its own conduit. With no
71
+ # PTP_UPSTREAM_* vars, everything flows through the single `conduit` above
72
+ # (PTP_UPSTREAM_URL), preserving single-backend behavior byte-for-byte.
73
+
74
+ from proxytool.routing import BackendRouter, backends_from_env # noqa: E402
75
+
76
+ _ROUTER: BackendRouter | None = None
77
+ _CONDUITS: dict[str, UpstreamConduit] = {}
78
+
79
+
80
+ def get_router() -> BackendRouter | None:
81
+ """Build the router from env once (PTP_UPSTREAM_*); None = single-backend."""
82
+ global _ROUTER
83
+ if _ROUTER is None:
84
+ declared = backends_from_env()
85
+ if len(declared) <= 1:
86
+ _ROUTER = BackendRouter() # empty, stays None forever in single mode
87
+ else:
88
+ router = BackendRouter()
89
+ for name, backend in declared.items():
90
+ router.register(backend)
91
+ _CONDUITS[name] = UpstreamConduit(
92
+ backend.url,
93
+ timeout=settings.upstream_timeout,
94
+ api_key=settings.upstream_api_key,
95
+ )
96
+ _ROUTER = router
97
+ trace.log("PT.ROUTER", {n: b.url for n, b in declared.items()})
98
+ return _ROUTER if _ROUTER.backends else None
99
+
100
+
101
+ def conduit_for(backend_name: str | None) -> UpstreamConduit:
102
+ """Conduit for the routed backend, or the default single conduit."""
103
+ if backend_name and backend_name in _CONDUITS:
104
+ return _CONDUITS[backend_name]
105
+ return conduit
106
+
107
+
108
+ def _request_tool_names(payload: dict[str, Any]) -> set[str] | None:
109
+ tools = payload.get("tools")
110
+ if not isinstance(tools, list) or not tools:
111
+ return None
112
+ names = set()
113
+ for t in tools:
114
+ fn = t.get("function") or {}
115
+ name = fn.get("name")
116
+ if isinstance(name, str):
117
+ names.add(name)
118
+ return names or None
119
+
120
+
121
+ class _SSEParser:
122
+ """Stateful SSE frame parser for chunked upstream streams.
123
+
124
+ HTTP chunks are *not* guaranteed to end on an SSE frame boundary: a
125
+ `data:` JSON payload can be split across two ``aiter_bytes()`` chunks,
126
+ and a multibyte UTF-8 character can be split mid-sequence. Bytes are
127
+ accumulated until a complete ``\n``-terminated line exists and only
128
+ whole lines are decoded, so both cases resolve before ``json.loads``
129
+ ever sees a payload (”data events carry the none-None frame dict).
130
+ """
131
+
132
+ def __init__(self) -> None:
133
+ self._buf = b""
134
+
135
+ def feed(self, chunk: bytes) -> list[tuple[str, dict[str, Any] | None]]:
136
+ out: list[tuple[str, dict[str, Any] | None]] = []
137
+ self._buf += chunk
138
+ while b"\n" in self._buf:
139
+ line, self._buf = self._buf.split(b"\n", 1)
140
+ out.append(_sse_line_event(line))
141
+ return out
142
+
143
+ def finish(self) -> list[tuple[str, dict[str, Any] | None]]:
144
+ """Drain a trailing unterminated line (tolerates sloppy servers)."""
145
+ if not self._buf:
146
+ return []
147
+ line, self._buf = self._buf, b""
148
+ return [_sse_line_event(line)]
149
+
150
+
151
+ def _sse_line_event(line_bytes: bytes) -> tuple[str, dict[str, Any] | None]:
152
+ """Parse one complete SSE line into a (kind, data-or-None) tuple."""
153
+ line = line_bytes.decode("utf-8", errors="replace").rstrip("\r")
154
+ body = None
155
+ if line.startswith("data: "):
156
+ body = line[6:]
157
+ elif line.startswith("data:"):
158
+ body = line[5:].lstrip()
159
+ if body is None:
160
+ return ("raw" if line.strip() else "comment", None)
161
+ if body == "[DONE]":
162
+ return ("done", None)
163
+ try:
164
+ return ("data", json.loads(body))
165
+ except json.JSONDecodeError:
166
+ return ("raw", None)
167
+
168
+
169
+ async def _resolve_model(model: str) -> str:
170
+ """Map a client-facing model id to one the upstream actually serves.
171
+
172
+ Order: PTP_MODEL_MAP → upstream /v1/models membership → single loaded
173
+ model. mlx_lm.server resolves unknown model ids against HF (401/404
174
+ "Repository Not Found"), so an unmapped alias must never be forwarded
175
+ raw when the real id is discoverable.
176
+
177
+ In routed mode only the explicit map applies: backends have separate id
178
+ spaces, and probing one backend's /v1/models would rewrite an id that
179
+ belongs to another (`qwen-mlx-omlx` → `qwen-mlx` on the swap group).
180
+ """
181
+ mapped = settings.model_map.get(model)
182
+ if mapped:
183
+ return mapped
184
+ if get_router() is not None:
185
+ return model
186
+ global _UPSTREAM_MODELS
187
+ if _UPSTREAM_MODELS is None:
188
+ try:
189
+ upstream = await conduit.passthrough("GET", "v1/models", {}, b"")
190
+ data = upstream.json().get("data", [])
191
+ _UPSTREAM_MODELS = [m.get("id", "") for m in data if isinstance(m, dict)]
192
+ except Exception:
193
+ _UPSTREAM_MODELS = [] # unreachable/mock — keep alias as-is
194
+ trace.log("PT.MODELS", _UPSTREAM_MODELS)
195
+ known = [m for m in _UPSTREAM_MODELS if m]
196
+ if model in known:
197
+ return model
198
+ best = _pick_by_alias(model, known)
199
+ return best or model
200
+
201
+
202
+ def _pick_by_alias(alias: str, known: list[str]) -> str | None:
203
+ """Best upstream id for an unmapped alias, or None to pass through.
204
+
205
+ Scores each upstream id by token overlap with the alias (e.g.
206
+ ``qwen-mlx`` → tokens {qwen, mlx}). Ties prefer the artifact the server
207
+ was actually launched with — a filesystem path beats an HF repo id,
208
+ which matters when the cache contains several Qwen mlx snapshots.
209
+ """
210
+ tokens = {t for t in re.split(r"[-/_ .]", alias.lower()) if t}
211
+ if not tokens:
212
+ return None
213
+ scored: list[tuple[int, int, str]] = []
214
+ for mid in known:
215
+ low = mid.lower()
216
+ overlap = sum(1 for t in tokens if t in re.split(r"[-/_ .]", low))
217
+ if overlap == 0:
218
+ continue
219
+ path_like = 1 if (mid.startswith("/") or "snapshots/" in mid) else 0
220
+ scored.append((overlap, path_like, mid))
221
+ if not scored:
222
+ return None
223
+ scored.sort(reverse=True)
224
+ top_overlap, top_path, _ = scored[0]
225
+ winners = [s for s in scored if (s[0], s[1]) == (top_overlap, top_path)]
226
+ if len(winners) != 1:
227
+ return None # genuinely ambiguous — don't guess
228
+ return winners[0][2]
229
+
230
+
231
+ @app.get("/health")
232
+ async def health() -> dict[str, str]:
233
+ return {"status": "ok", "upstream": settings.upstream_url}
234
+
235
+
236
+ def _resolve_backend(payload: dict[str, Any]) -> Backend:
237
+ """Dialect auto-detection, decided per request from the payload itself.
238
+
239
+ Routing stays llama-swap-style (the request's ``model`` field selects the
240
+ upstream); the only dialect split left is request-side: mlx servers don't
241
+ support ``tool_choice``, gguf/llama.cpp dialects do. A model id mentioning
242
+ mlx, or an explicit mlx-community/HF repo id, implies mlx; anything else
243
+ is treated as gguf-style. No env knob — proxies that need a different
244
+ dropdown use the repair hints in the toolpatch library instead.
245
+ """
246
+ model = str(payload.get("model") or "")
247
+ # Check the alias the client sent *and* the real id PTP_MODEL_MAP points
248
+ # it at: a client may call an MLX model "qwen", leaving the only MLX hint
249
+ # in the mapping, and the mapped id is what the backend actually runs.
250
+ candidates = (model, settings.model_map.get(model, ""))
251
+ if any(m and "mlx" in m.lower() for m in candidates):
252
+ return "mlx"
253
+ return "gguf"
254
+
255
+
256
+ @app.post("/v1/chat/completions")
257
+ async def chat_completions(request: Request):
258
+ body = await request.body()
259
+ try:
260
+ payload = json.loads(body)
261
+ except json.JSONDecodeError:
262
+ return JSONResponse({"error": "invalid JSON body"}, status_code=400)
263
+
264
+ trace.log("CLIENT->PT", payload)
265
+ # Route to a multi-backend upstream by the client's model id (no-op in
266
+ # single-backend mode). Routing happens before sanitization because the
267
+ # sanitized payload's model may have been remapped by PTP_MODEL_MAP.
268
+ chat_conduit = conduit
269
+ router = get_router()
270
+ if router is not None:
271
+ client_model = str(payload.get("model") or "")
272
+ decision = router.route(client_model)
273
+ chat_conduit = conduit_for(decision.backend.name)
274
+ trace.log("PT.ROUTE", {"model": client_model,
275
+ "backend": decision.backend.name,
276
+ "matched_by": decision.matched_by})
277
+ backend = _resolve_backend(payload)
278
+ payload, notes = normalize_chat_request(payload, backend=backend)
279
+ if backend == "mlx":
280
+ payload.pop("tool_choice", None)
281
+ notes.append("tool_choice stripped (mlx backend)")
282
+ # Remap the client-facing model id to what the upstream actually serves.
283
+ # mlx_lm.server resolves unknown model ids against HF ("Repository Not
284
+ # Found"); llama-server ignores names, but remapping helps logs too.
285
+ model = str(payload.get("model") or "qwen")
286
+ payload["model"] = await _resolve_model(model)
287
+ if notes:
288
+ trace.log("PT.NOTES", notes)
289
+ trace.log("PT->UPSTREAM", payload)
290
+
291
+ # LoopGuard: identical consecutive prompts make llama-server replay the
292
+ # same continuation forever. Nudge the model, then cut the loop off.
293
+ decision = loop_guard.check(LoopGuard.request_key(payload))
294
+ trace.log("PT.GUARD", {"action": decision.action, "count": decision.count})
295
+ if decision.action == "abort":
296
+ return _loop_guard_abort(decision.count, stream=bool(payload.get("stream")))
297
+ if decision.action == "nudge":
298
+ payload["messages"] = [
299
+ *payload["messages"],
300
+ {"role": "user", "content": NUDGE_TEXT},
301
+ ]
302
+ trace.log("PT->UPSTREAM", payload)
303
+
304
+ if payload.get("stream"):
305
+ return await _streaming(
306
+ request=request, payload=payload,
307
+ client_headers={k.lower(): v for k, v in request.headers.items()},
308
+ conduit=chat_conduit,
309
+ )
310
+ return await _non_streaming(payload=payload, client_headers=
311
+ {k.lower(): v for k, v in request.headers.items()}, conduit=chat_conduit)
312
+
313
+
314
+ async def _non_streaming(payload: dict[str, Any],
315
+ client_headers: dict[str, str] | None = None,
316
+ conduit: UpstreamConduit | None = None):
317
+ conduit = conduit or globals()["conduit"]
318
+ assert conduit is not None
319
+ upstream = await conduit.forward_json(
320
+ "/v1/chat/completions", payload, headers=client_headers
321
+ )
322
+ if upstream.status_code >= 400:
323
+ trace.log("UPSTREAM->PT", {"error": upstream.text[:500]})
324
+ return JSONResponse(
325
+ {"error": {"message": f"upstream error (HTTP {upstream.status_code}): "
326
+ f"{upstream.text[:500]}",
327
+ "type": "upstream_error",
328
+ "code": upstream.status_code}},
329
+ status_code=upstream.status_code,
330
+ )
331
+ try:
332
+ data = upstream.json()
333
+ except (json.JSONDecodeError, UnicodeDecodeError):
334
+ return _as_response(upstream, upstream.content)
335
+
336
+ trace.log("UPSTREAM->PT", data)
337
+ tool_names = _request_tool_names(payload)
338
+ _repair_response(data, tool_names)
339
+ trace.log("PT->CLIENT", data)
340
+ return JSONResponse(content=data, status_code=upstream.status_code)
341
+
342
+
343
+ def _repair_response(data: dict[str, Any], tool_names: set[str] | None) -> bool:
344
+ """Rewrite leaked tool calls inside data in place. True if repaired."""
345
+ repaired = False
346
+ for choice in data.get("choices", []) or []:
347
+ message = choice.get("message") or {}
348
+ content = message.get("content")
349
+ if not isinstance(content, str) or not content:
350
+ continue
351
+ ex = extract_tool_calls(content, tool_names, allow_unterminated_hermes=True)
352
+ if ex is None:
353
+ continue
354
+ message["content"] = ex.content
355
+ existing = message.get("tool_calls") or []
356
+ message["tool_calls"] = existing + ex.tool_calls
357
+ choice["message"] = message
358
+ repaired = True
359
+ return repaired
360
+
361
+
362
+ async def _streaming(request: Request, payload: dict[str, Any],
363
+ client_headers: dict[str, str] | None = None,
364
+ conduit: UpstreamConduit | None = None):
365
+ conduit = conduit or globals()["conduit"]
366
+ assert conduit is not None
367
+ tool_names = _request_tool_names(payload)
368
+ upstream = await conduit.forward_stream(
369
+ "/v1/chat/completions", payload, headers=client_headers
370
+ )
371
+ if upstream.status_code >= 400:
372
+ body = await upstream.aread()
373
+ await upstream.aclose()
374
+ return _upstream_error_sse(upstream.status_code, body)
375
+
376
+ async def gen() -> AsyncGenerator[bytes]:
377
+ healer = StreamHealer(tool_names=tool_names)
378
+ sse = _SSEParser() # reassembles frames split across HTTP chunks
379
+ meta: dict[str, Any] = {}
380
+ usage: dict[str, Any] | None = None
381
+ upstream_finish: str | None = None
382
+ terminated = False # a finish_reason has been forwarded downstream
383
+ first_role: str | None = None # upstream role delta, forwarded once
384
+
385
+ def _ev_json(ev: dict[str, Any],
386
+ finish_reason: str | None = None,
387
+ usage_: dict[str, Any] | None = None) -> bytes:
388
+ nonlocal first_role
389
+ r, first_role = first_role, None
390
+ return _event_json(ev, meta, finish_reason, usage_, role=r)
391
+ try:
392
+ async for chunk in upstream.aiter_bytes():
393
+ for kind, data in sse.feed(chunk):
394
+ if kind in ("raw", "done"):
395
+ continue
396
+ if data is None:
397
+ continue
398
+ if not meta:
399
+ # Forward id/created/model so clients correlating
400
+ # chunks to the request don't see a different object.
401
+ meta = {
402
+ k: data[k]
403
+ for k in ("id", "created", "model", "system_fingerprint")
404
+ if k in data
405
+ }
406
+ if data.get("usage"):
407
+ usage = data["usage"]
408
+ choice = (data.get("choices") or [{}])[0]
409
+ if choice.get("finish_reason"):
410
+ upstream_finish = choice["finish_reason"]
411
+ delta = choice.get("delta") or {}
412
+ if first_role is None and delta.get("role"):
413
+ first_role = delta["role"]
414
+ text = delta.get("content")
415
+ if not text:
416
+ continue
417
+ for ev in healer.feed(text):
418
+ if ev["type"] == "tool_calls" and ev.get("stop"):
419
+ yield _ev_json(ev, "tool_calls", usage)
420
+ terminated = True
421
+ else:
422
+ yield _ev_json(ev)
423
+ for ev in healer.finish():
424
+ if ev["type"] == "tool_calls" and ev.get("stop"):
425
+ yield _ev_json(ev, "tool_calls", usage)
426
+ terminated = True
427
+ else:
428
+ yield _ev_json(ev)
429
+ if not terminated:
430
+ # Repaired/replayed streams still need OpenAI-protocol
431
+ # termination: without a finish_reason chunk, opencode treats
432
+ # every response as an unterminated turn and re-prompts.
433
+ yield _ev_json({"type": "control"}, upstream_finish or "stop", usage)
434
+ yield b"data: [DONE]\n\n"
435
+ finally:
436
+ await upstream.aclose()
437
+
438
+ return StreamingResponse(
439
+ gen(), media_type="text/event-stream", status_code=upstream.status_code
440
+ )
441
+
442
+
443
+ def _event_json(
444
+ ev: dict[str, Any],
445
+ meta: dict[str, Any],
446
+ finish_reason: str | None = None,
447
+ usage: dict[str, Any] | None = None,
448
+ role: str | None = None,
449
+ ) -> bytes:
450
+ """Render one StreamHealer event as an OpenAI-compatible SSE chunk."""
451
+ if ev["type"] == "tool_calls":
452
+ delta: dict[str, Any] = {"tool_calls": ev["tool_calls"]}
453
+ elif ev["type"] == "content":
454
+ delta = {"content": ev.get("text", "")}
455
+ else:
456
+ delta = {}
457
+ # OpenAI streams start with a role-bearing first delta; clients use it to
458
+ # initialize the assistant message, so the first emitted chunk carries it.
459
+ if role is not None:
460
+ delta = {"role": role, **delta}
461
+ out: dict[str, Any] = {
462
+ **meta,
463
+ "choices": [
464
+ {
465
+ "delta": delta,
466
+ "index": 0,
467
+ "finish_reason": finish_reason,
468
+ }
469
+ ],
470
+ }
471
+ if usage is not None:
472
+ out["usage"] = usage
473
+ return f"data: {json.dumps(out, separators=(',', ':'))}\n\n".encode()
474
+
475
+
476
+ def _upstream_error_sse(status: int, body: bytes) -> StreamingResponse:
477
+ """Convert an upstream error into a single clean OpenAI-style SSE stream.
478
+
479
+ Without this, error pages / SSE keepalives from the upstream leak raw into
480
+ the client (opencode renders the garbled text or hangs waiting for deltas).
481
+ """
482
+ detail = body.decode("utf-8", errors="replace").strip()[:500]
483
+ msg = f"upstream error (HTTP {status}): {detail}"
484
+ chunk = {
485
+ "choices": [
486
+ {
487
+ "delta": {"content": msg},
488
+ "index": 0,
489
+ "finish_reason": "stop",
490
+ }
491
+ ]
492
+ }
493
+ payload = (
494
+ b"data: " + json.dumps(chunk, separators=(",", ":")).encode() + b"\n\ndata: [DONE]\n\n"
495
+ )
496
+ return StreamingResponse(
497
+ iter([payload]),
498
+ media_type="text/event-stream",
499
+ status_code=200,
500
+ )
501
+
502
+
503
+ def _loop_guard_abort(count: int, stream: bool) -> Response | StreamingResponse:
504
+ """Return a clean, short completion explaining the cut-off."""
505
+ text = ABORT_TEXT.format(count=count)
506
+ created = int(time.time())
507
+ if stream:
508
+ first = {
509
+ "id": "loopguard-abort",
510
+ "created": created,
511
+ "model": "loopguard",
512
+ "choices": [
513
+ {"delta": {"role": "assistant", "content": text},
514
+ "index": 0, "finish_reason": None}
515
+ ],
516
+ }
517
+ last = {
518
+ "id": "loopguard-abort",
519
+ "created": created,
520
+ "model": "loopguard",
521
+ "choices": [
522
+ {"delta": {}, "index": 0, "finish_reason": "stop"}
523
+ ],
524
+ }
525
+ sse = "".join(
526
+ f"data: {json.dumps(c, separators=(',', ':'))}\n\n" for c in (first, last)
527
+ ) + "data: [DONE]\n\n"
528
+ return StreamingResponse(
529
+ iter([sse.encode()]), media_type="text/event-stream", status_code=200
530
+ )
531
+ return JSONResponse(
532
+ {
533
+ "id": "loopguard-abort",
534
+ "created": created,
535
+ "model": "loopguard",
536
+ "choices": [
537
+ {
538
+ "message": {"role": "assistant", "content": text},
539
+ "index": 0,
540
+ "finish_reason": "stop",
541
+ }
542
+ ],
543
+ "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
544
+ },
545
+ status_code=200,
546
+ )
547
+
548
+
549
+ @app.get("/v1/models")
550
+ @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
551
+ async def passthrough(request: Request, path: str = "v1/models"):
552
+ content = await request.body()
553
+ headers = {k.lower(): v for k, v in request.headers.items() if k.lower() != "host"}
554
+ try:
555
+ upstream = await conduit.passthrough(request.method, path, headers, content)
556
+ except httpx.ConnectError:
557
+ return JSONResponse({"error": f"upstream unreachable: {settings.upstream_url}"},
558
+ status_code=502)
559
+ return _as_response(upstream, upstream.content)
560
+
561
+
562
+ def _as_response(upstream: httpx.Response, content: bytes) -> Response:
563
+ headers = {}
564
+ for h in PASS_HEADERS_OUT:
565
+ if h in upstream.headers:
566
+ headers[h] = upstream.headers[h]
567
+ return Response(
568
+ content=content,
569
+ status_code=upstream.status_code,
570
+ headers=headers,
571
+ media_type=upstream.headers.get("content-type"),
572
+ )
proxytool/routing.py ADDED
@@ -0,0 +1,113 @@
1
+ """Backend routing table: which upstream does a request hit.
2
+
3
+ Backends can be declared on the CLI via repeatable ``--upstream NAME=URL``
4
+ flags or through ``PTP_UPSTREAM_<NAME>`` / ``PTP_UPSTREAM_DEFAULT`` env vars
5
+ (environment wins unless the flag is passed explicitly). Each backend carries
6
+ an optional comma-separated model list — ``--upstream gguf=http://...:8181
7
+ --upstream-models gguf=qwen-gguf,qwen`` — and the client-facing ``model`` id is
8
+ matched against it first, then against ``PTP_DEFAULT_BACKEND`` / the flag
9
+ default. Exactly one registered backend receives everything.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import os
15
+ from dataclasses import dataclass, field
16
+
17
+
18
+ @dataclass(frozen=True)
19
+ class Backend:
20
+ name: str
21
+ url: str
22
+ models: tuple[str, ...] = () # client-facing model ids this backend serves
23
+ default: bool = False
24
+
25
+
26
+ @dataclass(frozen=True)
27
+ class RouteDecision:
28
+ backend: Backend
29
+ matched_by: str # "model-list" | "default" | "only-backend"
30
+
31
+
32
+ @dataclass
33
+ class BackendRouter:
34
+ backends: dict[str, Backend] = field(default_factory=dict)
35
+
36
+ def register(self, backend: Backend) -> None:
37
+ if backend.name in self.backends:
38
+ raise ValueError(f"duplicate backend name: {backend.name}")
39
+ self.backends[backend.name] = backend
40
+
41
+ @property
42
+ def default(self) -> Backend | None:
43
+ defaults = [b for b in self.backends.values() if b.default]
44
+ if len(defaults) == 1:
45
+ return defaults[0]
46
+ if defaults:
47
+ raise ValueError("multiple backends flagged default")
48
+ return next(iter(self.backends.values()), None)
49
+
50
+ def route(self, model: str | None) -> RouteDecision:
51
+ model = (model or "").strip().lower()
52
+ candidates = [b for b in self.backends.values() if b.models]
53
+ for b in candidates:
54
+ if model in {m.lower() for m in b.models}:
55
+ return RouteDecision(b, "model-list")
56
+ d = self.default
57
+ if d is None:
58
+ raise RuntimeError("no backends registered")
59
+ if len(self.backends) == 1:
60
+ return RouteDecision(d, "only-backend")
61
+ return RouteDecision(d, "default")
62
+
63
+
64
+ # Var suffixes under PTP_UPSTREAM_* that are single-upstream *settings*, not
65
+ # named routes: PTP_UPSTREAM_URL / _DEFAULT / _API_KEY / _TIMEOUT.
66
+ # Everything else in that namespace is a route name (PTP_UPSTREAM_<NAME>=URL,
67
+ # PTP_UPSTREAM_<NAME>_MODELS=m1,m2).
68
+ RESERVED_ROUTE_SUFFIXES = ("URL", "DEFAULT", "API_KEY", "TIMEOUT")
69
+
70
+
71
+ def backends_from_env() -> dict[str, Backend]:
72
+ """Build backends from PTP_UPSTREAM_<NAME> / _MODELS / _DEFAULT env vars."""
73
+ out: dict[str, Backend] = {}
74
+ for var, val in os.environ.items():
75
+ if not var.startswith("PTP_UPSTREAM_"):
76
+ continue
77
+ rest = var[len("PTP_UPSTREAM_"):]
78
+ if rest.endswith("_MODELS"):
79
+ continue
80
+ if rest in RESERVED_ROUTE_SUFFIXES:
81
+ continue
82
+ name = rest.lower().replace("-", "_")
83
+ models_var = f"PTP_UPSTREAM_{rest}_MODELS"
84
+ models = tuple(
85
+ m.strip() for m in os.environ.get(models_var, "").split(",") if m.strip()
86
+ )
87
+ out[name] = Backend(
88
+ name=name,
89
+ url=val.rstrip("/"),
90
+ models=models,
91
+ default=os.environ.get("PTP_DEFAULT_BACKEND", "").lower() == name,
92
+ )
93
+ return out
94
+
95
+
96
+ def backends_from_args(upstreams: list[str], models: dict[str, str],
97
+ default_backend: str | None) -> dict[str, Backend]:
98
+ """Build backends from repeatable --upstream NAME=URL and --upstream-models NAME=a,b."""
99
+ out: dict[str, Backend] = {}
100
+ for spec in upstreams:
101
+ name, _, url = spec.partition("=")
102
+ name = name.strip().lower()
103
+ if not name or not url:
104
+ raise ValueError(f"--upstream expects NAME=URL, got: {spec!r}")
105
+ out[name] = Backend(
106
+ name=name,
107
+ url=url.rstrip("/"),
108
+ models=tuple(
109
+ m.strip() for m in models.get(name, "").split(",") if m.strip()
110
+ ),
111
+ default=(default_backend or "").lower() == name,
112
+ )
113
+ return out