info/queuegate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
1307e4fd3a
queuegate/src/queuegate_proxy/app.py

1937 lines
76 KiB
Python
Raw Normal View History

update
2026-02-09 15:48:42 +00:00
# QueueGate Proxy v0.3.1
# - Queue + sticky routing over multiple upstream OpenAI-compatible endpoints
# - OpenAI-compatible endpoints: /v1/models, /v1/chat/completions
# - Tool-calling modes:
# * passthrough: forward native tool_calls; convert text toolcalls ([TOOL_CALLS]) into tool_calls for the client
# * execute: proxy executes known tools via TOOLSERVER_URL and continues until final answer
# - Extra endpoint for "spannende" clients that want to manage tools themselves: /v1/chat/completions_passthrough
Initial proxy MVP
2026-02-02 10:03:25 +00:00
import asyncio
update
2026-02-09 15:48:42 +00:00
import ast
Initial proxy MVP
2026-02-02 10:03:25 +00:00
import hashlib
update
2026-02-09 15:48:42 +00:00
import html
Initial proxy MVP
2026-02-02 10:03:25 +00:00
import json
update
2026-02-09 15:48:42 +00:00
import logging
Initial proxy MVP
2026-02-02 10:03:25 +00:00
import os
update
2026-02-09 15:48:42 +00:00
import re
Initial proxy MVP
2026-02-02 10:03:25 +00:00
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
import httpx
from fastapi import Body, FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
update
2026-02-09 15:48:42 +00:00
log = logging.getLogger("queuegate")
# -------------------------
# env helpers
# -------------------------
Initial proxy MVP
2026-02-02 10:03:25 +00:00
def env_bool(key: str, default: bool = False) -> bool:
v = os.getenv(key)
if v is None:
return default
return v.strip().lower() in {"1", "true", "yes", "on"}
def env_int(key: str, default: int) -> int:
v = os.getenv(key)
if v is None:
return default
try:
return int(v)
except ValueError:
return default
update
2026-02-09 15:48:42 +00:00
def env_str(key: str, default: str) -> str:
v = os.getenv(key)
return (v.strip() if v is not None else default)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
def now_ts() -> float:
return time.time()
def job_id() -> str:
return uuid.uuid4().hex
def sha1(text: str) -> str:
return hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest()
update
2026-02-09 15:48:42 +00:00
# -------------------------
# config
# -------------------------
Initial proxy MVP
2026-02-02 10:03:25 +00:00
@dataclass
class ProxyConfig:
upstreams: List[str]
update
2026-02-09 15:48:42 +00:00
models: List[str] = field(default_factory=lambda: ["default"])
owned_by: str = "queuegate"
Initial proxy MVP
2026-02-02 10:03:25 +00:00
sticky_header: str = "X-Chat-Id"
affinity_ttl_sec: int = 60
queue_notify_user: str = "auto" # auto|always|never
queue_notify_min_ms: int = 1200
read_timeout_sec: int = 3600
update
2026-02-09 15:48:42 +00:00
# toolserver
toolserver_url: Optional[str] = None
toolserver_prefix: str = "/openapi" # default path prefix for tool endpoints
toolserver_schema_ttl_sec: int = 300
# toolcall behavior
toolcall_mode: str = "execute" # execute|passthrough|suppress
unknown_tool_policy: str = "error" # error|passthrough|ignore
max_tool_iters: int = 6
# text toolcall detection (for models that print [TOOL_CALLS] in content)
text_toolcall_detect: bool = True
text_holdback_chars: int = 1024
text_ring_chars: int = 8192
# visible streaming holdback (small)
vis_holdback_chars: int = 32
# debug
toolcall_debug: bool = False
forward_reasoning: bool = True
# chat memory (RAG) via toolserver (memory_query/memory_upsert)
chat_memory_enable: bool = False
chat_memory_truncate_history: bool = True
chat_memory_keep_last: int = 4
chat_memory_query_k: int = 8
chat_memory_upsert: bool = True
chat_memory_inject_role: str = "system" # system|user
chat_memory_hint: bool = True
chat_memory_max_upsert_chars: int = 12000
chat_memory_for_agents: bool = False
Initial proxy MVP
2026-02-02 10:03:25 +00:00
def load_config() -> ProxyConfig:
ups = (os.getenv("LLM_UPSTREAMS") or "").strip()
if not ups:
raise RuntimeError("LLM_UPSTREAMS is required (comma-separated URLs)")
upstreams = [u.strip() for u in ups.split(",") if u.strip()]
update
2026-02-09 15:48:42 +00:00
models_env = (os.getenv("PROXY_MODELS") or os.getenv("LLM_MODELS") or "").strip()
models = [m.strip() for m in models_env.split(",") if m.strip()] if models_env else ["default"]
Initial proxy MVP
2026-02-02 10:03:25 +00:00
return ProxyConfig(
upstreams=upstreams,
update
2026-02-09 15:48:42 +00:00
models=models,
owned_by=env_str("PROXY_OWNED_BY", "queuegate") or "queuegate",
sticky_header=env_str("STICKY_HEADER", "X-Chat-Id") or "X-Chat-Id",
Initial proxy MVP
2026-02-02 10:03:25 +00:00
affinity_ttl_sec=env_int("AFFINITY_TTL_SEC", 60),
update
2026-02-09 15:48:42 +00:00
queue_notify_user=env_str("QUEUE_NOTIFY_USER", "auto").lower(),
Initial proxy MVP
2026-02-02 10:03:25 +00:00
queue_notify_min_ms=env_int("QUEUE_NOTIFY_MIN_MS", 1200),
read_timeout_sec=env_int("LLM_READ_TIMEOUT", 3600),
update
2026-02-09 15:48:42 +00:00
toolserver_url=(env_str("TOOLSERVER_URL", "") or None),
toolserver_prefix=env_str("TOOLSERVER_PREFIX", "/openapi") or "/openapi",
toolserver_schema_ttl_sec=env_int("TOOLSERVER_SCHEMA_TTL_SEC", 300),
toolcall_mode=env_str("TOOLCALL_MODE", "execute").lower(),
unknown_tool_policy=env_str("UNKNOWN_TOOL_POLICY", "error").lower(),
max_tool_iters=env_int("MAX_TOOL_ITERS", 6),
text_toolcall_detect=env_bool("TEXT_TOOLCALL_DETECT", True),
text_holdback_chars=env_int("TEXT_TOOLCALL_HOLDBACK_CHARS", 1024),
text_ring_chars=env_int("TEXT_TOOLCALL_RING_CHARS", 8192),
vis_holdback_chars=env_int("TEXT_VISIBLE_HOLDBACK_CHARS", 32),
toolcall_debug=env_bool("TOOLCALL_DEBUG", False),
forward_reasoning=env_bool("FORWARD_REASONING", True),
chat_memory_enable=env_bool("CHAT_MEMORY_ENABLE", False),
chat_memory_truncate_history=env_bool("CHAT_MEMORY_TRUNCATE_HISTORY", True),
chat_memory_keep_last=env_int("CHAT_MEMORY_KEEP_LAST", 4),
chat_memory_query_k=env_int("CHAT_MEMORY_QUERY_K", 8),
chat_memory_upsert=env_bool("CHAT_MEMORY_UPSERT", True),
chat_memory_inject_role=env_str("CHAT_MEMORY_INJECT_ROLE", "system").lower(),
chat_memory_hint=env_bool("CHAT_MEMORY_HINT", True),
chat_memory_max_upsert_chars=env_int("CHAT_MEMORY_MAX_UPSERT_CHARS", 12000),
chat_memory_for_agents=env_bool("CHAT_MEMORY_FOR_AGENTS", False),
Initial proxy MVP
2026-02-02 10:03:25 +00:00
)
update
2026-02-09 15:48:42 +00:00
# -------------------------
# SSE helpers
# -------------------------
def sse_pack(obj: Dict[str, Any]) -> bytes:
return ("data: " + json.dumps(obj, ensure_ascii=False) + "\n\n").encode("utf-8")
def sse_done() -> bytes:
return b"data: [DONE]\n\n"
def make_chunk(job_id_: str, model: str, delta: Dict[str, Any], finish_reason: Optional[str] = None) -> Dict[str, Any]:
return {
"id": job_id_,
"object": "chat.completion.chunk",
"created": int(now_ts()),
"model": model,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
# -------------------------
# tool extraction (native + text)
# -------------------------
THINK_TAG_RE = re.compile(r"</?think[^>]*>", re.IGNORECASE)
TOOL_TAG_RE = re.compile(r"\[\s*tool_calls\s*\]", re.IGNORECASE)
def strip_think_tags(txt: str) -> str:
return THINK_TAG_RE.sub("", txt or "")
def _delta_text(delta: Dict[str, Any]) -> Optional[str]:
"""
Some OpenAI-compatible backends stream reasoning/thinking in fields other than `content`.
We combine a small set of known keys so toolcall detection doesn't miss those.
"""
if not isinstance(delta, dict):
return None
parts: List[str] = []
# IMPORTANT: this is for *detection only*. Do NOT stream this merged text to the client,
# otherwise "reasoning" fields become visible in UIs like OpenWebUI.
for k in ("content", "reasoning_content", "reasoning", "thinking", "thought"):
v = delta.get(k)
if isinstance(v, str) and v:
parts.append(v)
if not parts:
return None
return "".join(parts)
def looks_like_tool_text(s: str) -> bool:
if not s:
return False
low = s.lower()
return ("[tool_calls]" in low) or ("tool_calls" in low) or ("\"name\"" in low and "\"arguments\"" in low)
def _toolname_candidates(tool_names: List[str]) -> List[str]:
# Small helper: prefer longer names first to avoid partial matches.
return sorted([t for t in tool_names if isinstance(t, str) and t], key=len, reverse=True)
def _find_functionstyle_call(text: str, tool_names: List[str]) -> Optional[Tuple[str, int]]:
"""Find `toolname{...}` or `toolname {...}` occurrences. Returns (toolname, brace_index)."""
if not text:
return None
low = text.lower()
for t in _toolname_candidates(tool_names):
tl = t.lower()
pos = low.find(tl)
while pos >= 0:
j = pos + len(tl)
k = j
while k < len(text) and text[k].isspace():
k += 1
if k < len(text) and text[k] in "[{":
return (t, k)
pos = low.find(tl, pos + 1)
return None
def extract_functionstyle_toolcall_from_text(content: str, tool_names: List[str]) -> Tuple[str, List[Dict[str, Any]]]:
"""Parse `repo_grep{...}` style calls (common when models leak toolcalls into text).
Returns (prefix_text, tool_calls). If none found: (content, []).
"""
if not isinstance(content, str) or not content or not tool_names:
return (content or ""), []
hit = _find_functionstyle_call(content, tool_names)
if not hit:
return content, []
tname, brace_idx = hit
js = _extract_balanced_json(content, brace_idx)
if not js:
return content, []
try:
args_obj = _parse_json_flexible(js)
except Exception:
return content, []
if not isinstance(args_obj, (dict, list)):
return content, []
args_str = json.dumps(args_obj if isinstance(args_obj, dict) else args_obj, ensure_ascii=False)
prefix = content[: content.lower().find(tname.lower())].rstrip()
tc = {
"id": uuid.uuid4().hex,
"type": "function",
"function": {"name": tname, "arguments": args_str if isinstance(args_str, str) else str(args_str)},
}
return prefix, [tc]
def salvage_args_from_text(text: str, tool_name: str) -> Optional[Dict[str, Any]]:
"""If native tool_calls had empty/broken args, try to recover from leaked text like `tool{...}`."""
if not text or not tool_name:
return None
low = text.lower()
tl = tool_name.lower()
pos = low.find(tl)
while pos >= 0:
j = pos + len(tl)
k = j
while k < len(text) and text[k].isspace():
k += 1
if k < len(text) and text[k] in "[{":
js = _extract_balanced_json(text, k)
if js:
try:
obj = _parse_json_flexible(js)
if isinstance(obj, dict):
return obj
except Exception:
pass
pos = low.find(tl, pos + 1)
return None
@dataclass
class InterceptResult:
tool_calls: List[Dict[str, Any]]
tail_text: str = ""
def _strip_code_fences(s: str) -> str:
s = (s or "").strip()
if not s.startswith("```"):
return s
parts = s.splitlines()
if parts and parts[0].lstrip().startswith("```"):
parts = parts[1:]
if parts and parts[-1].strip() == "```":
parts = parts[:-1]
return "\n".join(parts).strip()
def _cleanup_jsonish(text: str) -> str:
return re.sub(r",\s*([}\]])", r"\1", text)
def _parse_json_flexible(s: str) -> Any:
raw = html.unescape(_strip_code_fences(s)).strip()
if not raw:
raise ValueError("empty json")
for cand in (raw, _cleanup_jsonish(raw)):
try:
return json.loads(cand)
except Exception:
pass
for cand in (raw, _cleanup_jsonish(raw)):
try:
return ast.literal_eval(cand)
except Exception:
pass
raise ValueError("unable to parse json")
def _extract_balanced_json(s: str, start: int) -> Optional[str]:
if start < 0 or start >= len(s):
return None
open_ch = s[start]
close_ch = "]" if open_ch == "[" else "}" if open_ch == "{" else None
if close_ch is None:
return None
depth = 0
in_str = False
esc = False
for i in range(start, len(s)):
c = s[i]
if in_str:
if esc:
esc = False
elif c == "\\":
esc = True
elif c == '"':
in_str = False
continue
if c == '"':
in_str = True
continue
if c == open_ch:
depth += 1
continue
if c == close_ch:
depth -= 1
if depth == 0:
return s[start : i + 1]
continue
return None
def _normalize_tool_calls(raw: Any) -> List[Dict[str, Any]]:
"""
Normalize tool call payloads into OpenAI-compatible `tool_calls`.
Important: be conservative. If a call has no valid name, skip it (do NOT emit name="None").
Supports:
1) OpenAI style: [{"id":..,"type":"function","function":{"name":..,"arguments":"{...}"}}]
2) Text style: [{"name":..,"arguments":{...},"id":"..."}]
3) Wrapper: {"tool_calls":[...]}
"""
if isinstance(raw, dict) and isinstance(raw.get("tool_calls"), list):
raw = raw["tool_calls"]
if not isinstance(raw, list):
return []
out: List[Dict[str, Any]] = []
for item in raw:
if not isinstance(item, dict):
continue
# OpenAI style
if item.get("type") == "function" and isinstance(item.get("function"), dict):
fn = item["function"]
name = fn.get("name")
if not name or str(name).strip().lower() == "none":
continue
args = fn.get("arguments", "")
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
out.append({
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": str(name), "arguments": str(args)},
})
continue
# Text style / wrapper-ish
name = item.get("name") or (item.get("function") or {}).get("name")
if not name or str(name).strip().lower() == "none":
continue
args = item.get("arguments") or (item.get("function") or {}).get("arguments") or {}
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
out.append({
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": str(name), "arguments": str(args)},
})
return out
def extract_toolcalls_from_text(content: str) -> Tuple[str, List[Dict[str, Any]]]:
"""
Returns (prefix_text, tool_calls). If none found: (content_without_think, []).
Conservative parsing:
- Prefer explicit [TOOL_CALLS] tag.
- Otherwise, only parse JSON that is *near* the "tool_calls"/"name"+"arguments" hints
(avoid grabbing the first "{" in the message which might be unrelated code).
"""
if not isinstance(content, str) or not content:
return "", []
s = content
# 1) explicit [TOOL_CALLS] tag (case-insensitive)
m = TOOL_TAG_RE.search(s)
if m:
idx = m.start()
prefix = s[:idx].strip()
rest = s[m.end():].lstrip()
# find first JSON bracket after tag
jstart = None
for i, ch in enumerate(rest):
if ch in "[{":
jstart = i
break
if jstart is not None:
js = _extract_balanced_json(rest, jstart)
if js:
try:
parsed = _parse_json_flexible(js)
tcs = _normalize_tool_calls(parsed)
if tcs:
return prefix, tcs
except Exception:
pass
return prefix, []
low = s.lower()
# 2) OpenAI-ish wrapper in text: {"tool_calls":[...]} or "tool_calls": [...]
pos_tc = low.find("tool_calls")
if pos_tc >= 0:
# look for the JSON bracket *after* the keyword
for start_ch in ("{", "["):
i0 = s.find(start_ch, pos_tc)
if i0 >= 0:
js = _extract_balanced_json(s, i0)
if js:
try:
parsed = _parse_json_flexible(js)
tcs = _normalize_tool_calls(parsed)
if tcs:
return "", tcs
except Exception:
pass
break
# 3) Heuristic: JSON list/object containing name+arguments
# Search the bracket close to the first '"name"' occurrence to avoid unrelated code blocks.
pos_name = s.find('"name"')
pos_args = s.find('"arguments"')
if pos_name >= 0 and pos_args >= 0:
anchor = min(pos_name, pos_args)
# Prefer list form
i0 = s.find("[", anchor)
if i0 >= 0:
js = _extract_balanced_json(s, i0)
if js:
try:
parsed = _parse_json_flexible(js)
tcs = _normalize_tool_calls(parsed)
if tcs:
return "", tcs
except Exception:
pass
# Then object form
i0 = s.find("{", anchor)
if i0 >= 0:
js = _extract_balanced_json(s, i0)
if js:
try:
parsed = _parse_json_flexible(js)
tcs = _normalize_tool_calls(parsed if isinstance(parsed, list) else [parsed])
if tcs:
return "", tcs
except Exception:
pass
return s, []
def normalize_native_tool_calls(delta_tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Best-effort normalization for native `delta.tool_calls` objects.
Some backends stream tool calls in partial fragments. This function is conservative:
it will SKIP calls without a valid name (never emit name="None").
"""
out: List[Dict[str, Any]] = []
for item in (delta_tool_calls or []):
if not isinstance(item, dict):
continue
# OpenAI style
if item.get("type") == "function" and isinstance(item.get("function"), dict):
fn = item.get("function") or {}
name = fn.get("name")
if not name or str(name).strip().lower() == "none":
continue
args = fn.get("arguments", "")
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
out.append({
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": str(name), "arguments": str(args or "")},
})
continue
# Other shapes
name = item.get("name") or (item.get("function") or {}).get("name")
if not name or str(name).strip().lower() == "none":
continue
args = item.get("arguments") or (item.get("function") or {}).get("arguments") or ""
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
out.append({
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": str(name), "arguments": str(args or "")},
})
return out
def accumulate_native_tool_calls(acc: Dict[str, Dict[str, Any]], delta_tool_calls: List[Dict[str, Any]]) -> None:
"""Accumulate partial native tool_call deltas by id/index.
OpenAI streaming may send tool call fragments across chunks (especially arguments).
We merge fragments and only later finalize into executable tool calls.
"""
for item in (delta_tool_calls or []):
if not isinstance(item, dict):
continue
key = item.get("id")
if not key:
key = str(item.get("index") if item.get("index") is not None else uuid.uuid4().hex)
cur = acc.get(str(key))
if not cur:
cur = {
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": None, "arguments": ""},
}
acc[str(key)] = cur
fn = item.get("function")
if isinstance(fn, dict):
nm = fn.get("name")
if nm:
cur["function"]["name"] = nm
arg = fn.get("arguments")
if arg is not None:
if isinstance(arg, dict):
arg = json.dumps(arg, ensure_ascii=False)
cur["function"]["arguments"] = (cur["function"].get("arguments") or "") + str(arg)
continue
nm = item.get("name")
if nm:
cur["function"]["name"] = nm
arg = item.get("arguments")
if arg is not None:
if isinstance(arg, dict):
arg = json.dumps(arg, ensure_ascii=False)
cur["function"]["arguments"] = (cur["function"].get("arguments") or "") + str(arg)
def finalize_native_tool_calls(acc: Dict[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
out: List[Dict[str, Any]] = []
for _, item in (acc or {}).items():
if not isinstance(item, dict):
continue
fn = item.get("function") or {}
name = fn.get("name")
if not name or str(name).strip().lower() == "none":
continue
args = fn.get("arguments", "")
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
out.append({
"id": str(item.get("id") or uuid.uuid4().hex),
"type": "function",
"function": {"name": str(name), "arguments": str(args or "")},
})
return out
# -------------------------
# tool registry + execution
# -------------------------
class ToolRegistry:
def __init__(self, cfg: ProxyConfig):
self.cfg = cfg
self._tools: Dict[str, str] = {}
self._required: Dict[str, List[str]] = {}
self._prop_types: Dict[str, Dict[str, str]] = {}
self._last_fetch: float = 0.0
self._lock = asyncio.Lock()
def has(self, name: str) -> bool:
return name in self._tools
def path_for(self, name: str) -> str:
pfx = self.cfg.toolserver_prefix.rstrip("/") or "/openapi"
return self._tools.get(name) or f"{pfx}/{name}"
def required_for(self, name: str) -> List[str]:
return list(self._required.get(name) or [])
def prop_types_for(self, name: str) -> Dict[str, str]:
return dict(self._prop_types.get(name) or {})
async def refresh_if_needed(self, client: httpx.AsyncClient) -> None:
if not self.cfg.toolserver_url:
return
ttl = max(5, int(self.cfg.toolserver_schema_ttl_sec))
if self._tools and (now_ts() - self._last_fetch) < ttl:
return
async with self._lock:
if self._tools and (now_ts() - self._last_fetch) < ttl:
return
url = self.cfg.toolserver_url.rstrip("/") + "/openapi.json"
r = await client.get(url)
r.raise_for_status()
spec = r.json()
paths = spec.get("paths") or {}
components = (spec.get("components") or {}).get("schemas") or {}
def deref(schema: dict) -> dict:
if not isinstance(schema, dict):
return {}
ref = schema.get("$ref")
if isinstance(ref, str) and ref.startswith("#/components/schemas/"):
name = ref.split("#/components/schemas/", 1)[1]
if name in components and isinstance(components[name], dict):
return components[name]
return schema
def merge_schema(schema: dict) -> tuple[list[str], dict[str, str]]:
schema = deref(schema)
req: list[str] = []
props: dict[str, str] = {}
if not isinstance(schema, dict):
return req, props
# handle allOf (common with pydantic)
if isinstance(schema.get("allOf"), list):
for sub in schema.get("allOf"):
r2, p2 = merge_schema(sub if isinstance(sub, dict) else {})
for k in r2:
if k not in req:
req.append(k)
props.update(p2)
return req, props
rlist = schema.get("required")
if isinstance(rlist, list):
for k in rlist:
if isinstance(k, str) and k not in req:
req.append(k)
pmap = schema.get("properties")
if isinstance(pmap, dict):
for k, v in pmap.items():
if not isinstance(k, str) or not isinstance(v, dict):
continue
vv = deref(v)
t = vv.get("type")
if isinstance(t, str):
props[k] = t
elif isinstance(vv.get("enum"), list):
props[k] = "enum"
else:
props[k] = "object"
return req, props
tools: Dict[str, str] = {}
required_map: Dict[str, List[str]] = {}
prop_types_map: Dict[str, Dict[str, str]] = {}
prefix = self.cfg.toolserver_prefix.rstrip("/")
for p, methods in paths.items():
if not isinstance(p, str):
continue
if not p.startswith(prefix + "/"):
continue
if not isinstance(methods, dict):
continue
post = methods.get("post") or methods.get("POST")
if not isinstance(post, dict):
continue
name = p.split(prefix + "/", 1)[1].strip("/")
if not name:
continue
tools[name] = p
# extract requestBody schema
req: List[str] = []
props: Dict[str, str] = {}
rb = post.get("requestBody") or {}
if isinstance(rb, dict):
content = rb.get("content") or {}
if isinstance(content, dict) and content:
# prefer application/json
entry = content.get("application/json") or content.get("application/*+json")
if not isinstance(entry, dict):
entry = next(iter(content.values())) if content else {}
if isinstance(entry, dict):
sch = entry.get("schema")
if isinstance(sch, dict):
req, props = merge_schema(sch)
if req:
required_map[name] = req
if props:
prop_types_map[name] = props
self._tools = tools
self._required = required_map
self._prop_types = prop_types_map
self._last_fetch = now_ts()
# -------------------------
# job/worker
# -------------------------
Initial proxy MVP
2026-02-02 10:03:25 +00:00
@dataclass
class Job:
job_id: str
created_ts: float
kind: str # user_chat | agent_call
stream: bool
body: Dict[str, Any]
headers: Dict[str, str]
thread_key: str
assigned_worker: int
status: str = "queued" # queued|running|done|error
error: Optional[str] = None
result: Optional[Dict[str, Any]] = None
done_fut: asyncio.Future = field(default_factory=asyncio.Future)
stream_q: asyncio.Queue = field(default_factory=asyncio.Queue)
saw_any_output: bool = False
class Worker:
def __init__(self, idx: int, upstream_url: str, client: httpx.AsyncClient):
self.idx = idx
self.upstream_url = upstream_url
self.client = client
self.user_q: asyncio.Queue[Job] = asyncio.Queue()
self.agent_q: asyncio.Queue[Job] = asyncio.Queue()
self.current_job_id: Optional[str] = None
self.task: Optional[asyncio.Task] = None
def pending_count(self) -> int:
return self.user_q.qsize() + self.agent_q.qsize() + (1 if self.current_job_id else 0)
async def start(self, state: "ProxyState") -> None:
self.task = asyncio.create_task(self._run(state), name=f"worker-{self.idx}")
update
2026-02-09 15:48:42 +00:00
async def _next_job(self) -> Tuple[Job, str]:
if not self.user_q.empty():
return await self.user_q.get(), "user"
if not self.agent_q.empty():
return await self.agent_q.get(), "agent"
t_user = asyncio.create_task(self.user_q.get())
t_agent = asyncio.create_task(self.agent_q.get())
done, pending = await asyncio.wait({t_user, t_agent}, return_when=asyncio.FIRST_COMPLETED)
for p in pending:
p.cancel()
if t_user in done:
return t_user.result(), "user"
return t_agent.result(), "agent"
Initial proxy MVP
2026-02-02 10:03:25 +00:00
async def _run(self, state: "ProxyState") -> None:
while True:
update
2026-02-09 15:48:42 +00:00
job, src = await self._next_job()
Initial proxy MVP
2026-02-02 10:03:25 +00:00
self.current_job_id = job.job_id
job.status = "running"
try:
if job.stream:
await state.handle_stream_job(job, self)
else:
await state.handle_non_stream_job(job, self)
except Exception as e:
job.status = "error"
job.error = str(e)
if not job.done_fut.done():
job.done_fut.set_exception(e)
update
2026-02-09 15:48:42 +00:00
if job.stream:
try:
model = (job.body.get("model") or "unknown").strip() or "unknown"
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"role": "assistant", "content": f"(proxy error: {e})"})))
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {}, finish_reason="stop")))
await job.stream_q.put(sse_done())
except Exception:
pass
await job.stream_q.put(None)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
finally:
self.current_job_id = None
try:
update
2026-02-09 15:48:42 +00:00
(self.user_q if src == "user" else self.agent_q).task_done()
Initial proxy MVP
2026-02-02 10:03:25 +00:00
except Exception:
pass
update
2026-02-09 15:48:42 +00:00
# -------------------------
# routing helpers
# -------------------------
Initial proxy MVP
2026-02-02 10:03:25 +00:00
def get_header_any(headers: Dict[str, str], names: List[str]) -> Optional[str]:
for n in names:
v = headers.get(n)
if v:
return v
return None
def infer_thread_key(cfg: ProxyConfig, headers: Dict[str, str], body: Dict[str, Any]) -> str:
update
2026-02-09 15:48:42 +00:00
v = get_header_any(headers, [cfg.sticky_header, "X-OpenWebUI-Chat-Id",
"X-OpenWebUI-Conversation-Id",
"X-OpenWebUI-Thread-Id", "X-Chat-Id", "X-Conversation-Id"])
Initial proxy MVP
2026-02-02 10:03:25 +00:00
if v:
return v
seed = {
"model": (body.get("model") or "").strip(),
"messages": (body.get("messages") or [])[:2],
"user": body.get("user"),
}
try:
txt = json.dumps(seed, sort_keys=True, ensure_ascii=False)
except Exception:
txt = str(seed)
return "h:" + sha1(txt)
update
2026-02-09 15:48:42 +00:00
def _msg_text(m: Dict[str, Any]) -> str:
"""Extract text from OpenAI chat message content variants."""
c = m.get("content")
if c is None:
return ""
if isinstance(c, str):
return c
# OpenAI "content": [{"type":"text","text":"..."}, ...]
if isinstance(c, list):
parts = []
for it in c:
if isinstance(it, dict):
if it.get("type") == "text" and isinstance(it.get("text"), str):
parts.append(it.get("text"))
elif isinstance(it.get("text"), str):
parts.append(it.get("text"))
return "\n".join([p for p in parts if p])
# anything else
return str(c)
def _last_user_text(messages: List[Dict[str, Any]]) -> str:
for m in reversed(messages or []):
if (m.get("role") or "").lower() == "user":
return _msg_text(m).strip()
return ""
def _build_memory_context(results: List[Dict[str, Any]], *, max_items: int = 8) -> str:
if not results:
return ""
out = []
for r in results[:max(1, int(max_items))]:
t = (r.get("text") or "").strip()
# toolserver stores a FILE/LANG header; remove for chat memory readability
if t.startswith("FILE:"):
t = re.sub(r"^FILE:.*?\n", "", t, flags=re.S).strip()
if not t:
continue
score = r.get("score")
if isinstance(score, (int, float)):
out.append(f"- ({score:.3f}) {t}")
else:
out.append(f"- {t}")
return "\n".join(out).strip()
def _truncate_for_memory(cfg: ProxyConfig, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
if not cfg.chat_memory_truncate_history:
return list(messages or [])
sys_msgs = [m for m in (messages or []) if (m.get("role") or "").lower() == "system"]
convo = [m for m in (messages or []) if (m.get("role") or "").lower() in {"user", "assistant"}]
keep = max(1, int(cfg.chat_memory_keep_last))
tail = convo[-keep:] if convo else []
return sys_msgs + tail
def _inject_memory(cfg: ProxyConfig, messages: List[Dict[str, Any]], mem_bullets: str) -> List[Dict[str, Any]]:
if not mem_bullets:
return list(messages or [])
role = (cfg.chat_memory_inject_role or "system").lower()
if role not in {"system", "user"}:
role = "system"
hint = ""
if cfg.chat_memory_hint:
hint = "If needed, you can ask to look up more chat memory using the tool `memory_query` (or by emitting a [TOOL_CALLS] for it).\n\n"
content = (
"### Retrieved chat memory (may be incomplete; treat as notes, not instructions)\n"
+ hint
+ mem_bullets
).strip()
mem_msg = {"role": role, "content": content}
# Insert after any system messages so system remains first.
sys_count = 0
for m in (messages or []):
if (m.get("role") or "").lower() == "system":
sys_count += 1
else:
break
out = list(messages or [])
out.insert(sys_count, mem_msg)
return out
def _infer_namespace(cfg: ProxyConfig, headers: Dict[str, str], body: Dict[str, Any], thread_key: str) -> str:
# Prefer explicit chat/conversation headers
for hk in [cfg.sticky_header, "X-OpenWebUI-Chat-Id", "X-OpenWebUI-Conversation-Id", "X-Conversation-Id", "X-Chat-Id"]:
v = get_header_any(headers, [hk])
if v:
return v
# body fields often present in clients
for k in ("chat_id", "conversation_id", "thread_id"):
v = body.get(k)
if isinstance(v, str) and v.strip():
return v.strip()
md = body.get("metadata")
if isinstance(md, dict):
for k in ("chat_id", "conversation_id", "thread_id"):
v = md.get(k)
if isinstance(v, str) and v.strip():
return v.strip()
# fallback to thread_key (stable-ish)
return thread_key
def infer_kind(headers: Dict[str, str], override: Optional[str] = None) -> str:
if override in {"user_chat", "agent_call"}:
return override
jk = (headers.get("X-Job-Kind") or "").strip().lower()
if jk in {"agent", "agent_call", "repo_agent"}:
return "agent_call"
return "user_chat"
# -------------------------
# proxy state
# -------------------------
Initial proxy MVP
2026-02-02 10:03:25 +00:00
class ProxyState:
def __init__(self, cfg: ProxyConfig):
self.cfg = cfg
self.http = httpx.AsyncClient(timeout=httpx.Timeout(cfg.read_timeout_sec))
self.workers: List[Worker] = [Worker(i, u, self.http) for i, u in enumerate(cfg.upstreams)]
self.affinity: Dict[str, Tuple[int, float]] = {}
self.jobs: Dict[str, Job] = {}
update
2026-02-09 15:48:42 +00:00
self.toolreg = ToolRegistry(cfg)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
async def start(self) -> None:
for w in self.workers:
await w.start(self)
async def close(self) -> None:
await self.http.aclose()
def pick_worker(self, thread_key: str) -> int:
now = now_ts()
sticky = self.affinity.get(thread_key)
if sticky:
widx, last = sticky
if now - last <= self.cfg.affinity_ttl_sec and 0 <= widx < len(self.workers):
self.affinity[thread_key] = (widx, now)
return widx
best = 0
update
2026-02-09 15:48:42 +00:00
best_load: Optional[int] = None
Initial proxy MVP
2026-02-02 10:03:25 +00:00
for w in self.workers:
load = w.pending_count()
if best_load is None or load < best_load:
best_load = load
best = w.idx
self.affinity[thread_key] = (best, now)
return best
def enqueue(self, job: Job) -> None:
w = self.workers[job.assigned_worker]
update
2026-02-09 15:48:42 +00:00
(w.user_q if job.kind == "user_chat" else w.agent_q).put_nowait(job)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
self.jobs[job.job_id] = job
update
2026-02-09 15:48:42 +00:00
def queue_position(self, job: Job) -> int:
w = self.workers[job.assigned_worker]
return (w.user_q.qsize() if job.kind == "user_chat" else w.agent_q.qsize())
# ---- tool execution
async def _tool_call(self, tool_name: str, args: Dict[str, Any]) -> str:
if not self.cfg.toolserver_url:
raise RuntimeError("TOOLSERVER_URL is required for execute mode")
args = args or {}
await self.toolreg.refresh_if_needed(self.http)
# Validate required args if we can infer them from toolserver OpenAPI
req = self.toolreg.required_for(tool_name)
if req:
def _missing(v: Any) -> bool:
if v is None:
return True
if isinstance(v, str) and v.strip() == "":
return True
return False
missing = [k for k in req if (k not in args) or _missing(args.get(k))]
if missing:
# lightweight example values
props = self.toolreg.prop_types_for(tool_name)
ex: Dict[str, Any] = {}
for k in req:
t = (props.get(k) or "").lower()
if k in ("repo", "repo_url"):
ex[k] = "http://HOST:PORT/ORG/REPO.git"
elif k in ("question", "query"):
ex[k] = "<your question/query>"
elif k == "branch":
ex[k] = "main"
elif k in ("k", "max_hits", "n_results") or t in ("integer", "int"):
ex[k] = 10
else:
ex[k] = "<value>" if t in ("string", "") else 1
return json.dumps({
"error": "missing_required_args",
"tool": tool_name,
"missing": missing,
"required": req,
"example": ex,
"note": "Provide a JSON object with at least the required fields.",
}, ensure_ascii=False)
url = self.cfg.toolserver_url.rstrip("/") + self.toolreg.path_for(tool_name)
try:
r = await self.http.post(url, json=args)
r.raise_for_status()
except httpx.HTTPStatusError as e:
# Pass a helpful error back to the model so it can retry correctly
status = int(e.response.status_code)
txt = (e.response.text or "").strip()
out = {
"error": "tool_http_error",
"tool": tool_name,
"status": status,
"response": txt[:2000],
}
if req:
out["required"] = req
return json.dumps(out, ensure_ascii=False)
except Exception as e:
out = {"error": "tool_call_failed", "tool": tool_name, "detail": str(e)}
if req:
out["required"] = req
return json.dumps(out, ensure_ascii=False)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
try:
return json.dumps(r.json(), ensure_ascii=False)
except Exception:
return r.text
def _tool_args_help(self, tool_name: str, detail: str = "") -> str:
"""Build a compact, model-friendly hint for how to call a tool.
This is used when:
- arguments could not be parsed from tool call text, or
- the model emitted an empty arguments object.
"""
req = self.toolreg.required_for(tool_name)
props = self.toolreg.prop_types_for(tool_name)
# lightweight example values
ex: Dict[str, Any] = {}
for k in req:
t = (props.get(k) or "").lower()
if k in ("repo", "repo_url"):
ex[k] = "http://HOST:PORT/ORG/REPO.git"
elif k in ("question", "query"):
ex[k] = "<your question/query>"
elif k == "branch":
ex[k] = "main"
elif k in ("k", "max_hits", "n_results") or t in ("integer", "int"):
ex[k] = 10
else:
ex[k] = "<value>" if t in ("string", "") else 1
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
return json.dumps(
{
"error": "tool_args_parse_error",
"tool": tool_name,
"detail": (detail or "").strip(),
"required": req,
"example": ex,
"note": "Provide a valid JSON object for arguments (not empty).",
},
ensure_ascii=False,
)
async def _tool_call_json(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
raw = await self._tool_call(tool_name, args)
try:
obj = json.loads(raw) if isinstance(raw, str) else raw
return obj if isinstance(obj, dict) else {"raw": obj}
except Exception:
return {"raw": raw}
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
async def _maybe_apply_chat_memory(self, job: Job, body: Dict[str, Any], messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], Optional[str], str]:
"""Optionally inject retrieved chat memory and optionally truncate history."""
cfg = self.cfg
if not cfg.chat_memory_enable:
return messages, None, _last_user_text(messages)
if not cfg.toolserver_url:
return messages, None, _last_user_text(messages)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
kind_ok = (job.kind == "user_chat") or (cfg.chat_memory_for_agents and job.kind == "agent_call")
if not kind_ok:
return messages, None, _last_user_text(messages)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
user_q = _last_user_text(messages)
if not user_q:
return messages, None, ""
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
ns = _infer_namespace(cfg, job.headers, body, job.thread_key)
try:
res = await self._tool_call_json("memory_query", {"namespace": ns, "query": user_q, "k": int(cfg.chat_memory_query_k)})
items = res.get("results") if isinstance(res, dict) else None
items = items if isinstance(items, list) else []
bullets = _build_memory_context(items, max_items=int(cfg.chat_memory_query_k))
if not bullets:
return messages, ns, user_q
trimmed = _truncate_for_memory(cfg, messages)
injected = _inject_memory(cfg, trimmed, bullets)
return injected, ns, user_q
except Exception as e:
if cfg.toolcall_debug:
log.info("chat memory query failed: %s", e)
return messages, ns, user_q
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
async def _maybe_upsert_chat_memory(self, namespace: Optional[str], user_text: str, assistant_text: str, *, source: str = "queuegate") -> None:
cfg = self.cfg
if not cfg.chat_memory_enable or not cfg.chat_memory_upsert:
return
if not cfg.toolserver_url or not namespace:
return
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
def clip(s: str) -> str:
s = (s or "").strip()
mx = int(cfg.chat_memory_max_upsert_chars or 0)
if mx > 0 and len(s) > mx:
return s[:mx] + "…"
return s
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
u = clip(user_text)
a = clip(assistant_text)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
try:
if u:
await self._tool_call_json("memory_upsert", {"namespace": namespace, "text": u, "role": "user", "source": source, "ts_unix": int(time.time())})
if a:
await self._tool_call_json("memory_upsert", {"namespace": namespace, "text": a, "role": "assistant", "source": source, "ts_unix": int(time.time())})
except Exception as e:
if cfg.toolcall_debug:
log.info("chat memory upsert failed: %s", e)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
# ---- non-stream
Initial proxy MVP
2026-02-02 10:03:25 +00:00
async def handle_non_stream_job(self, job: Job, worker: Worker) -> None:
update
2026-02-09 15:48:42 +00:00
cfg = self.cfg
mode = (job.headers.get("X-Tool-Mode") or cfg.toolcall_mode).lower()
body0 = dict(job.body)
body0["stream"] = False
messages = list(body0.get("messages") or [])
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
# chat-memory injection (optional)
messages, mem_ns, mem_user = await self._maybe_apply_chat_memory(job, body0, messages)
body0["messages"] = messages
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
if mode != "execute":
r = await self.http.post(worker.upstream_url, json=body0)
r.raise_for_status()
data = r.json()
Initial proxy MVP
2026-02-02 10:03:25 +00:00
try:
update
2026-02-09 15:48:42 +00:00
msg = (data.get("choices") or [{}])[0].get("message") or {}
Initial proxy MVP
2026-02-02 10:03:25 +00:00
content = msg.get("content")
update
2026-02-09 15:48:42 +00:00
if cfg.text_toolcall_detect and isinstance(content, str) and looks_like_tool_text(content):
prefix, tcs = extract_toolcalls_from_text(content)
if tcs:
msg["tool_calls"] = tcs
msg["content"] = prefix or ""
(data.get("choices") or [{}])[0]["finish_reason"] = "tool_calls"
else:
msg["content"] = content
Initial proxy MVP
2026-02-02 10:03:25 +00:00
except Exception:
pass
update
2026-02-09 15:48:42 +00:00
# upsert last turn into chat memory (optional)
try:
_msg = (data.get("choices") or [{}])[0].get("message") or {}
_assistant_text = _msg.get("content") if isinstance(_msg.get("content"), str) else ""
except Exception:
_assistant_text = ""
await self._maybe_upsert_chat_memory(mem_ns, mem_user, _assistant_text, source="openwebui")
job.result = data
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result(data)
return
# execute loop
for _ in range(max(1, cfg.max_tool_iters)):
body = dict(body0)
body["messages"] = messages
r = await self.http.post(worker.upstream_url, json=body)
r.raise_for_status()
data = r.json()
choice0 = (data.get("choices") or [{}])[0]
msg = choice0.get("message") or {}
tool_calls = msg.get("tool_calls") if isinstance(msg.get("tool_calls"), list) else []
content = msg.get("content")
if (not tool_calls) and cfg.text_toolcall_detect and isinstance(content, str) and looks_like_tool_text(content):
prefix, tcs = extract_toolcalls_from_text(content)
if tcs:
tool_calls = tcs
msg["tool_calls"] = tcs
msg["content"] = prefix or ""
choice0["finish_reason"] = "tool_calls"
else:
msg["content"] = content
# Function-style toolcalls in text: repo_grep{...}
if (not tool_calls) and cfg.text_toolcall_detect and isinstance(content, str):
await self.toolreg.refresh_if_needed(self.http)
prefix2, tcs2 = extract_functionstyle_toolcall_from_text(content, list(self.toolreg._tools.keys()))
if tcs2:
tool_calls = tcs2
msg["tool_calls"] = tcs2
msg["content"] = prefix2 or ""
choice0["finish_reason"] = "tool_calls"
if not tool_calls:
# upsert last turn into chat memory (optional)
_assistant_text = msg.get("content") if isinstance(msg.get("content"), str) else ""
await self._maybe_upsert_chat_memory(mem_ns, mem_user, _assistant_text, source="openwebui")
job.result = data
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result(data)
return
await self.toolreg.refresh_if_needed(self.http)
messages.append({"role": "assistant", "content": msg.get("content") or "", "tool_calls": tool_calls})
for tc in tool_calls:
fn = (tc.get("function") or {})
tname = fn.get("name")
arg_s = fn.get("arguments") or "{}"
args: Optional[Dict[str, Any]] = None
if isinstance(arg_s, dict):
args = arg_s
if isinstance(args, dict) and not args:
_salv = salvage_args_from_text(content or "", tname or "")
if isinstance(_salv, dict) and _salv:
args = _salv
if isinstance(args, dict) and not args:
args = None
elif isinstance(arg_s, str):
try:
parsed = _parse_json_flexible(arg_s)
if isinstance(parsed, dict):
args = parsed
if isinstance(args, dict) and not args:
_salv = salvage_args_from_text(content or "", tname or "")
if isinstance(_salv, dict) and _salv:
args = _salv
if isinstance(args, dict) and not args:
args = None
except Exception:
args = salvage_args_from_text(content or "", tname or "")
if not tname:
continue
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
# Never call tools with unknown/empty args unless we are sure.
if args is None:
if cfg.toolcall_debug:
log.info("toolcall parse failed (non-stream) name=%s arg_s=%r", tname, arg_s)
# Feed a structured hint back to the model so it can retry with valid JSON.
messages.append({
"role": "tool",
"tool_call_id": tc.get("id"),
"name": tname,
"content": self._tool_args_help(tname, detail="could not parse tool arguments"),
})
continue
if not self.toolreg.has(tname):
pol = cfg.unknown_tool_policy
if pol == "ignore":
messages.append({"role": "tool", "tool_call_id": tc.get("id"), "name": tname, "content": "(unknown tool ignored)"})
continue
if pol == "passthrough":
job.result = data
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result(data)
return
raise RuntimeError(f"Unknown tool: {tname}")
if cfg.toolcall_debug:
log.info("toolcall execute (non-stream) name=%s args_keys=%s", tname, sorted(list(args.keys())) if isinstance(args, dict) else "?")
out = await self._tool_call(tname, args)
messages.append({"role": "tool", "tool_call_id": tc.get("id"), "name": tname, "content": out})
raise RuntimeError("Max tool iterations exceeded")
# ---- stream
Initial proxy MVP
2026-02-02 10:03:25 +00:00
async def handle_stream_job(self, job: Job, worker: Worker) -> None:
update
2026-02-09 15:48:42 +00:00
cfg = self.cfg
mode = (job.headers.get("X-Tool-Mode") or cfg.toolcall_mode).lower()
base = dict(job.body)
base["stream"] = True
model = (base.get("model") or "unknown").strip() or "unknown"
messages = list(base.get("messages") or [])
# chat-memory injection (optional)
messages, mem_ns, mem_user = await self._maybe_apply_chat_memory(job, base, messages)
base["messages"] = messages
assistant_capture: List[str] = []
if mode != "execute":
# stream (passthrough/suppress/etc). Optionally capture assistant visible content for chat-memory upsert.
cap = assistant_capture if (cfg.chat_memory_enable and cfg.chat_memory_upsert) else None
await self._stream_single_passthrough(job, worker, base, model, mode, assistant_capture=cap)
if cap is not None:
await self._maybe_upsert_chat_memory(mem_ns, mem_user, "".join(assistant_capture), source="openwebui")
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result({"status": "streamed"})
return
# execute-mode streaming loop
for _ in range(max(1, cfg.max_tool_iters)):
intercept = await self._stream_until_tool_or_done(job, worker, base, model, messages, assistant_capture)
tool_calls = intercept.tool_calls
if tool_calls and (not job.saw_any_output):
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"role": "assistant", "content": "(running tools...)"})))
job.saw_any_output = True
if not tool_calls:
await self._maybe_upsert_chat_memory(mem_ns, mem_user, "".join(assistant_capture), source="openwebui")
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result({"status": "streamed"})
return
await self.toolreg.refresh_if_needed(self.http)
messages.append({"role": "assistant", "content": "", "tool_calls": tool_calls})
for tc in tool_calls:
fn = (tc.get("function") or {})
tname = fn.get("name")
arg_s = fn.get("arguments") or "{}"
args: Optional[Dict[str, Any]] = None
if isinstance(arg_s, dict):
args = arg_s
if isinstance(args, dict) and not args:
_salv = salvage_args_from_text(intercept.tail_text or "", tname or "")
if isinstance(_salv, dict) and _salv:
args = _salv
if isinstance(args, dict) and not args:
args = None
elif isinstance(arg_s, str):
try:
parsed = _parse_json_flexible(arg_s)
if isinstance(parsed, dict):
args = parsed
if isinstance(args, dict) and not args:
_salv = salvage_args_from_text(intercept.tail_text or "", tname or "")
if isinstance(_salv, dict) and _salv:
args = _salv
if isinstance(args, dict) and not args:
args = None
except Exception:
args = salvage_args_from_text(intercept.tail_text or "", tname or "")
if not tname:
continue
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
if args is None:
if cfg.toolcall_debug:
log.info("toolcall parse failed (stream) name=%s arg_s=%r", tname, arg_s)
# Feed an explicit tool-error back into the conversation so the model can retry cleanly.
messages.append({
"role": "tool",
"tool_call_id": tc.get("id"),
"name": tname,
"content": self._tool_args_help(tname, detail="could not parse tool arguments"),
})
continue
if not self.toolreg.has(tname):
pol = cfg.unknown_tool_policy
if pol == "ignore":
messages.append({"role": "tool", "tool_call_id": tc.get("id"), "name": tname, "content": "(unknown tool ignored)"})
continue
if pol == "passthrough":
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"tool_calls": tool_calls}, finish_reason="tool_calls")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
job.status = "done"
if not job.done_fut.done():
job.done_fut.set_result({"status": "streamed"})
return
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"role": "assistant", "content": f"(unknown tool: {tname})"})))
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {}, finish_reason="stop")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
raise RuntimeError(f"Unknown tool: {tname}")
if cfg.toolcall_debug:
log.info("toolcall execute (stream) name=%s args_keys=%s", tname, sorted(list(args.keys())) if isinstance(args, dict) else "?")
out = await self._tool_call(tname, args)
messages.append({"role": "tool", "tool_call_id": tc.get("id"), "name": tname, "content": out})
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"role": "assistant", "content": "(max tool iterations exceeded)"})))
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {}, finish_reason="stop")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
raise RuntimeError("Max tool iterations exceeded")
async def _stream_single_passthrough(self, job: Job, worker: Worker, body: Dict[str, Any], model: str, mode: str, *, assistant_capture: Optional[List[str]] = None) -> None:
cfg = self.cfg
# Pure passthrough: forward upstream SSE unmodified (best for clients with their own tool logic).
if mode == "passthrough":
async with self.http.stream("POST", worker.upstream_url, json=body) as r:
r.raise_for_status()
async for line in r.aiter_lines():
if not line or not line.startswith("data:"):
continue
await job.stream_q.put((line + "\n\n").encode("utf-8"))
job.saw_any_output = True
payload = line[len("data:"):].strip()
if assistant_capture is not None and payload and payload != "[DONE]":
try:
obj = json.loads(payload)
choice0 = (obj.get("choices") or [{}])[0]
delta = choice0.get("delta") or {}
c0 = delta.get("content")
if isinstance(c0, str) and c0:
assistant_capture.append(c0)
except Exception:
pass
if payload == "[DONE]":
break
await job.stream_q.put(None)
return
hold_n = max(32, int(cfg.text_holdback_chars))
hold = ""
ring = ""
text_capturing = False
text_buf = ""
Initial proxy MVP
2026-02-02 10:03:25 +00:00
async with self.http.stream("POST", worker.upstream_url, json=body) as r:
r.raise_for_status()
async for line in r.aiter_lines():
update
2026-02-09 15:48:42 +00:00
if not line or not line.startswith("data:"):
Initial proxy MVP
2026-02-02 10:03:25 +00:00
continue
payload = line[len("data:"):].strip()
if payload == "[DONE]":
break
try:
obj = json.loads(payload)
except Exception:
await job.stream_q.put((line + "\n\n").encode("utf-8"))
job.saw_any_output = True
continue
choice0 = (obj.get("choices") or [{}])[0]
delta = choice0.get("delta") or {}
update
2026-02-09 15:48:42 +00:00
# native tool_calls
if isinstance(delta.get("tool_calls"), list):
if mode == "passthrough":
await job.stream_q.put((line + "\n\n").encode("utf-8"))
job.saw_any_output = True
# suppress: drop
Initial proxy MVP
2026-02-02 10:03:25 +00:00
continue
update
2026-02-09 15:48:42 +00:00
c = _delta_text(delta)
if cfg.text_toolcall_detect and isinstance(c, str):
ring = (ring + c)[-cfg.text_ring_chars:]
candidate = hold + c
if text_capturing:
text_buf += c
pref, tcs = extract_toolcalls_from_text(text_buf)
if tcs:
if pref:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": pref})))
if assistant_capture is not None and isinstance(pref, str) and pref:
assistant_capture.append(pref)
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"tool_calls": tcs}, finish_reason="tool_calls")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
return
Initial proxy MVP
2026-02-02 10:03:25 +00:00
continue
update
2026-02-09 15:48:42 +00:00
if looks_like_tool_text(candidate):
pref, tcs = extract_toolcalls_from_text(candidate)
if tcs:
if pref:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": pref})))
if assistant_capture is not None and isinstance(pref, str) and pref:
assistant_capture.append(pref)
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"tool_calls": tcs}, finish_reason="tool_calls")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
return
text_capturing = True
text_buf = candidate
hold = ""
continue
if len(candidate) > hold_n:
out = candidate[:-hold_n]
if out:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": out})))
if assistant_capture is not None and isinstance(out, str) and out:
assistant_capture.append(out)
job.saw_any_output = True
hold = candidate[-hold_n:]
else:
hold = candidate
continue
Initial proxy MVP
2026-02-02 10:03:25 +00:00
await job.stream_q.put((line + "\n\n").encode("utf-8"))
job.saw_any_output = True
update
2026-02-09 15:48:42 +00:00
# final parse
if cfg.text_toolcall_detect and (looks_like_tool_text(ring) or looks_like_tool_text(hold) or text_capturing):
probe = text_buf if text_capturing else (ring if looks_like_tool_text(ring) else hold)
pref, tcs = extract_toolcalls_from_text(probe)
if tcs:
if pref:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": pref})))
if assistant_capture is not None and isinstance(pref, str) and pref:
assistant_capture.append(pref)
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"tool_calls": tcs}, finish_reason="tool_calls")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
return
if hold:
out = hold
if out:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": out})))
if assistant_capture is not None and isinstance(out, str) and out:
assistant_capture.append(out)
job.saw_any_output = True
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {}, finish_reason="stop")))
await job.stream_q.put(sse_done())
await job.stream_q.put(None)
async def _stream_until_tool_or_done(self, job: Job, worker: Worker, base: Dict[str, Any], model: str, messages: List[Dict[str, Any]], capture: Optional[List[str]] = None) -> InterceptResult:
"""Execute-mode streaming.
- Streams only `delta.content` to the client (so backends that emit separate `reasoning_*` fields don't leak those into the UI).
- Intercepts both native tool_calls and common leaked text patterns like `[TOOL_CALLS]...` and `repo_grep{...}`.
- When a toolcall is detected, it *keeps the HTTP stream open* and returns the tool_calls to the caller.
"""
cfg = self.cfg
await self.toolreg.refresh_if_needed(self.http)
tool_names = list(self.toolreg._tools.keys())
body = dict(base)
body["messages"] = messages
hold_det_n = max(64, int(cfg.text_holdback_chars))
hold_vis_n = max(0, int(getattr(cfg, 'vis_holdback_chars', 32)))
hold_vis = "" # only assistant-visible content
hold_det = "" # detection tail
ring_det = "" # detection ring
text_capturing = False
text_buf = ""
found_tool_calls: List[Dict[str, Any]] = []
last_finish: Optional[str] = None
native_acc: Dict[str, Dict[str, Any]] = {}
native_seen = False
async with self.http.stream("POST", worker.upstream_url, json=body) as r:
r.raise_for_status()
async for line in r.aiter_lines():
if not line or not line.startswith("data:"):
continue
payload = line[len("data:"):].strip()
if payload == "[DONE]":
break
try:
obj = json.loads(payload)
except Exception:
continue
choice0 = (obj.get("choices") or [{}])[0]
delta = choice0.get("delta") or {}
fin = choice0.get("finish_reason")
if fin is not None:
last_finish = fin
# Native tool_calls (can be streamed in fragments)
if isinstance(delta.get("tool_calls"), list):
native_seen = True
accumulate_native_tool_calls(native_acc, delta.get("tool_calls") or [])
if cfg.toolcall_debug:
log.info("native tool_calls fragment seen (stream) fin=%s", fin)
if fin == "tool_calls":
break
continue
# Detection text (includes `content` + possible `reasoning_*` keys)
det = _delta_text(delta)
if isinstance(det, str) and det:
ring_det = (ring_det + det)[-cfg.text_ring_chars:]
det_candidate = hold_det + det
# If native tool_calls started, do NOT forward any content; just keep reading for complete args.
if native_seen:
hold_det = det_candidate[-hold_det_n:] if len(det_candidate) > hold_det_n else det_candidate
continue
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
if cfg.text_toolcall_detect:
if text_capturing:
text_buf += det
pref, tcs = extract_toolcalls_from_text(text_buf)
if not tcs:
pref, tcs = extract_functionstyle_toolcall_from_text(text_buf, tool_names)
if tcs:
found_tool_calls = tcs
if cfg.toolcall_debug:
log.info("text toolcall detected (stream) names=%s", [((tc.get('function') or {}).get('name')) for tc in tcs])
break
continue
# Primary: explicit [TOOL_CALLS] markers
if TOOL_TAG_RE.search(det_candidate):
pref, tcs = extract_toolcalls_from_text(det_candidate)
if not tcs:
pref, tcs = extract_functionstyle_toolcall_from_text(det_candidate, tool_names)
if tcs:
found_tool_calls = tcs
if cfg.toolcall_debug:
log.info("text toolcall detected (stream) names=%s", [((tc.get('function') or {}).get('name')) for tc in tcs])
break
# start capturing from here to avoid leaking partial tool JSON
text_capturing = True
text_buf = det_candidate
hold_vis = ""
hold_det = ""
continue
# Secondary: `tool{...}` patterns without markers
pref2, tcs2 = extract_functionstyle_toolcall_from_text(det_candidate, tool_names)
if tcs2:
found_tool_calls = tcs2
if cfg.toolcall_debug:
log.info("function-style toolcall detected (stream) name=%s", ((tcs2[0].get('function') or {}).get('name')))
break
hold_det = det_candidate[-hold_det_n:] if len(det_candidate) > hold_det_n else det_candidate
# Stream reasoning/thinking fields for UIs like OpenWebUI
if cfg.forward_reasoning and not text_capturing:
for rk in ("reasoning_content", "reasoning", "thinking", "thought"):
rv = delta.get(rk)
if isinstance(rv, str) and rv:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {rk: rv})))
job.saw_any_output = True
# Stream visible assistant content (never after native tool_calls started)
if native_seen or text_capturing:
continue
vis = delta.get("content")
if isinstance(vis, str) and vis:
if hold_vis_n <= 0:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": vis})))
if capture is not None:
capture.append(vis)
job.saw_any_output = True
else:
vis_candidate = hold_vis + vis
if len(vis_candidate) > hold_vis_n:
out = vis_candidate[:-hold_vis_n]
if out:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": out})))
if capture is not None:
capture.append(out)
job.saw_any_output = True
hold_vis = vis_candidate[-hold_vis_n:]
else:
# Keep accumulating until we can flush, then emit the tail at the end.
hold_vis = vis_candidate
# finalize native toolcalls if any
if native_acc and not found_tool_calls:
found_tool_calls = finalize_native_tool_calls(native_acc)
if found_tool_calls:
tail = text_buf if text_capturing else (ring_det + hold_det)
return InterceptResult(tool_calls=found_tool_calls, tail_text=tail)
if hold_vis:
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {"content": hold_vis})))
if capture is not None:
capture.append(hold_vis)
job.saw_any_output = True
await job.stream_q.put(sse_pack(make_chunk(job.job_id, model, {}, finish_reason=last_finish or "stop")))
Initial proxy MVP
2026-02-02 10:03:25 +00:00
await job.stream_q.put(sse_done())
update
2026-02-09 15:48:42 +00:00
await job.stream_q.put(None)
return InterceptResult(tool_calls=[], tail_text=(ring_det + hold_det))
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
# -------------------------
# FastAPI app
# -------------------------
app = FastAPI(title="QueueGate Proxy", version="0.3.1")
Initial proxy MVP
2026-02-02 10:03:25 +00:00
CFG: Optional[ProxyConfig] = None
STATE: Optional[ProxyState] = None
@app.on_event("startup")
async def _startup() -> None:
global CFG, STATE
CFG = load_config()
update
2026-02-09 15:48:42 +00:00
lvl = getattr(logging, (os.getenv('LOG_LEVEL') or 'INFO').upper(), logging.INFO)
logging.basicConfig(level=lvl, format='%(asctime)s %(levelname)s:%(name)s:%(message)s')
logging.getLogger('queuegate').setLevel(lvl)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
STATE = ProxyState(CFG)
await STATE.start()
@app.on_event("shutdown")
async def _shutdown() -> None:
global STATE
if STATE:
await STATE.close()
update
2026-02-09 15:48:42 +00:00
def _require_state() -> ProxyState:
Initial proxy MVP
2026-02-02 10:03:25 +00:00
if not STATE:
raise HTTPException(status_code=503, detail="not ready")
update
2026-02-09 15:48:42 +00:00
return STATE
def _copy_headers(req: Request) -> Dict[str, str]:
keep = {
"X-Job-Kind",
"X-Chat-Id",
"X-Conversation-Id",
"X-OpenWebUI-Chat-Id",
"X-OpenWebUI-Conversation-Id",
"X-OpenWebUI-Thread-Id",
"X-Tool-Mode",
}
out: Dict[str, str] = {}
for k, v in req.headers.items():
if k in keep:
out[k] = v
elif CFG and k.lower() == CFG.sticky_header.lower():
out[k] = v
return out
@app.get("/healthz")
async def healthz() -> Dict[str, Any]:
st = _require_state()
Initial proxy MVP
2026-02-02 10:03:25 +00:00
return {
"ok": True,
update
2026-02-09 15:48:42 +00:00
"upstreams": len(st.workers),
Initial proxy MVP
2026-02-02 10:03:25 +00:00
"workers": [
{
"id": w.idx,
"pending": w.pending_count(),
"busy": bool(w.current_job_id),
"upstream": w.upstream_url,
}
update
2026-02-09 15:48:42 +00:00
for w in st.workers
Initial proxy MVP
2026-02-02 10:03:25 +00:00
],
}
update
2026-02-09 15:48:42 +00:00
@app.get("/v1/models")
async def list_models() -> Dict[str, Any]:
if not CFG:
Initial proxy MVP
2026-02-02 10:03:25 +00:00
raise HTTPException(status_code=503, detail="not ready")
update
2026-02-09 15:48:42 +00:00
created = int(now_ts())
return {
"object": "list",
"data": [
{"id": mid, "object": "model", "created": created, "owned_by": CFG.owned_by}
for mid in (CFG.models or ["default"])
],
}
@app.get("/v1/jobs/{jid}")
async def job_status(jid: str) -> Dict[str, Any]:
st = _require_state()
job = st.jobs.get(jid)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
if not job:
raise HTTPException(status_code=404, detail="unknown job")
return {
"job_id": job.job_id,
"status": job.status,
"kind": job.kind,
"created_ts": job.created_ts,
"worker": job.assigned_worker,
update
2026-02-09 15:48:42 +00:00
"queue_position_est": st.queue_position(job) if job.status == "queued" else 0,
Initial proxy MVP
2026-02-02 10:03:25 +00:00
"error": job.error,
}
update
2026-02-09 15:48:42 +00:00
async def _enqueue_and_respond(body: Dict[str, Any], request: Request, *, kind_override: Optional[str] = None, tool_mode_override: Optional[str] = None) -> Any:
st = _require_state()
cfg = st.cfg
Initial proxy MVP
2026-02-02 10:03:25 +00:00
update
2026-02-09 15:48:42 +00:00
if not (body.get("model") or "").strip():
body["model"] = (cfg.models[0] if cfg.models else "default")
Initial proxy MVP
2026-02-02 10:03:25 +00:00
stream = bool(body.get("stream", False))
headers = _copy_headers(request)
update
2026-02-09 15:48:42 +00:00
if tool_mode_override:
headers["X-Tool-Mode"] = tool_mode_override
kind = infer_kind(headers, override=kind_override)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
thread_key = infer_thread_key(cfg, headers, body)
update
2026-02-09 15:48:42 +00:00
worker_idx = st.pick_worker(thread_key)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
jid = job_id()
job = Job(
job_id=jid,
created_ts=now_ts(),
kind=kind,
stream=stream,
body=body,
headers=headers,
thread_key=thread_key,
assigned_worker=worker_idx,
)
update
2026-02-09 15:48:42 +00:00
st.enqueue(job)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
if not stream:
try:
data = await job.done_fut
except Exception as e:
raise HTTPException(status_code=502, detail=f"upstream error: {e}")
return JSONResponse(content=data)
async def gen() -> Any:
if job.kind == "user_chat" and cfg.queue_notify_user != "never":
t0 = now_ts()
while job.status == "queued" and (now_ts() - t0) * 1000 < cfg.queue_notify_min_ms:
await asyncio.sleep(0.05)
if job.status == "queued" and cfg.queue_notify_user in {"auto", "always"}:
update
2026-02-09 15:48:42 +00:00
pos = st.queue_position(job)
Initial proxy MVP
2026-02-02 10:03:25 +00:00
model = (body.get("model") or "unknown").strip() or "unknown"
msg = f"⏳ In wachtrij (positie ~{pos})…"
yield sse_pack(make_chunk(job.job_id, model, {"role": "assistant", "content": msg}))
while True:
item = await job.stream_q.get()
if item is None:
break
yield item
return StreamingResponse(gen(), media_type="text/event-stream")
update
2026-02-09 15:48:42 +00:00
@app.post("/v1/chat/completions")
async def chat_completions(body: Dict[str, Any] = Body(...), request: Request = None):
# main endpoint: default tool behavior from TOOLCALL_MODE (default execute)
return await _enqueue_and_respond(body, request)
@app.post("/v1/chat/completions_passthrough")
async def chat_completions_passthrough(body: Dict[str, Any] = Body(...), request: Request = None):
# for clients that want to manage tools themselves
return await _enqueue_and_respond(body, request, tool_mode_override="passthrough")
@app.post("/v1/agent/chat/completions")
async def agent_chat_completions(body: Dict[str, Any] = Body(...), request: Request = None):
# agent queue + execute tools
return await _enqueue_and_respond(body, request, kind_override="agent_call", tool_mode_override="execute")
Reference in New Issue Copy Permalink