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/__init__.py +107 -0
- proxytool/__main__.py +14 -0
- proxytool/cli.py +634 -0
- proxytool/conduit.py +88 -0
- proxytool/config.py +496 -0
- proxytool/loopguard.py +92 -0
- proxytool/main.py +572 -0
- proxytool/routing.py +113 -0
- proxytool/settings.py +65 -0
- proxytool/trace.py +81 -0
- proxytool-0.4.0.dist-info/METADATA +557 -0
- proxytool-0.4.0.dist-info/RECORD +14 -0
- proxytool-0.4.0.dist-info/WHEEL +4 -0
- proxytool-0.4.0.dist-info/entry_points.txt +3 -0
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
|