[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
e945f43652
commit
0f0e02603b
4 changed files with 278 additions and 139 deletions
|
|
@ -1486,7 +1486,10 @@ class LlamaCppBackend:
|
|||
pool = 10,
|
||||
)
|
||||
with client.stream(
|
||||
"POST", url, json = payload, timeout = prefill_timeout,
|
||||
"POST",
|
||||
url,
|
||||
json = payload,
|
||||
timeout = prefill_timeout,
|
||||
headers = headers,
|
||||
) as response:
|
||||
_response_ref[0] = response
|
||||
|
|
@ -1561,10 +1564,16 @@ class LlamaCppBackend:
|
|||
# can finish. Cancel during streaming is handled by the
|
||||
# watcher thread (closes the response on cancel_event).
|
||||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||||
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||||
_auth_headers = (
|
||||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||||
)
|
||||
with httpx.Client(timeout = stream_timeout) as client:
|
||||
with self._stream_with_retry(
|
||||
client, url, payload, cancel_event, headers = _auth_headers,
|
||||
client,
|
||||
url,
|
||||
payload,
|
||||
cancel_event,
|
||||
headers = _auth_headers,
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
error_body = response.read().decode()
|
||||
|
|
@ -1721,7 +1730,11 @@ class LlamaCppBackend:
|
|||
payload["stop"] = stop
|
||||
|
||||
try:
|
||||
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||||
_auth_headers = (
|
||||
{"Authorization": f"Bearer {self._api_key}"}
|
||||
if self._api_key
|
||||
else None
|
||||
)
|
||||
with httpx.Client(timeout = None) as client:
|
||||
resp = client.post(url, json = payload, headers = _auth_headers)
|
||||
if resp.status_code != 200:
|
||||
|
|
@ -1966,10 +1979,16 @@ class LlamaCppBackend:
|
|||
|
||||
try:
|
||||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||||
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||||
_auth_headers = (
|
||||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||||
)
|
||||
with httpx.Client(timeout = stream_timeout) as client:
|
||||
with self._stream_with_retry(
|
||||
client, url, stream_payload, cancel_event, headers = _auth_headers,
|
||||
client,
|
||||
url,
|
||||
stream_payload,
|
||||
cancel_event,
|
||||
headers = _auth_headers,
|
||||
) as response:
|
||||
if response.status_code != 200:
|
||||
error_body = response.read().decode()
|
||||
|
|
@ -2095,7 +2114,9 @@ class LlamaCppBackend:
|
|||
if not self.is_loaded:
|
||||
return None
|
||||
try:
|
||||
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
_auth_headers = (
|
||||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
)
|
||||
with httpx.Client(timeout = 10, headers = _auth_headers) as client:
|
||||
|
||||
def _detok(tid: int) -> str:
|
||||
|
|
@ -2214,8 +2235,12 @@ class LlamaCppBackend:
|
|||
if need_ids:
|
||||
payload["n_probs"] = 1
|
||||
|
||||
_auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client:
|
||||
_auth_headers = (
|
||||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
)
|
||||
with httpx.Client(
|
||||
timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
|
||||
) as client:
|
||||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(
|
||||
|
|
|
|||
|
|
@ -115,7 +115,10 @@ async def lifespan(app: FastAPI):
|
|||
|
||||
# Start dedicated streaming server in a daemon thread (Option B).
|
||||
# Accepts both UNSLOTH_FAST_SSE=1 and UNSLOTH_STREAM_SERVER=1.
|
||||
if os.getenv("UNSLOTH_FAST_SSE", "0") == "1" or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1":
|
||||
if (
|
||||
os.getenv("UNSLOTH_FAST_SSE", "0") == "1"
|
||||
or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1"
|
||||
):
|
||||
import threading as _threading
|
||||
|
||||
from streaming_server import start_streaming_server, find_free_port
|
||||
|
|
@ -123,9 +126,9 @@ async def lifespan(app: FastAPI):
|
|||
stream_port = find_free_port()
|
||||
app.state.stream_port = stream_port
|
||||
_stream_thread = _threading.Thread(
|
||||
target=start_streaming_server,
|
||||
args=(stream_port,),
|
||||
daemon=True,
|
||||
target = start_streaming_server,
|
||||
args = (stream_port,),
|
||||
daemon = True,
|
||||
)
|
||||
_stream_thread.start()
|
||||
print(f"[streaming_server] Started on 127.0.0.1:{stream_port}")
|
||||
|
|
|
|||
|
|
@ -1633,25 +1633,27 @@ async def consume_stream_token_endpoint(request: Request):
|
|||
body = await request.json()
|
||||
token = body.get("token")
|
||||
if not token:
|
||||
return JSONResponse({"valid": False}, status_code=400)
|
||||
return JSONResponse({"valid": False}, status_code = 400)
|
||||
|
||||
from stream_token_store import consume_stream_token
|
||||
|
||||
username = consume_stream_token(token)
|
||||
if not username:
|
||||
return JSONResponse({"valid": False}, status_code=401)
|
||||
return JSONResponse({"valid": False}, status_code = 401)
|
||||
|
||||
llama = get_llama_cpp_backend()
|
||||
return JSONResponse({
|
||||
"valid": True,
|
||||
"username": username,
|
||||
"llama_port": llama._port if llama.is_loaded else None,
|
||||
"llama_api_key": llama._api_key,
|
||||
"model_name": llama.model_identifier or "unknown",
|
||||
"supports_reasoning": llama.supports_reasoning,
|
||||
"supports_tools": llama.supports_tools,
|
||||
"is_vision": llama.is_vision,
|
||||
})
|
||||
return JSONResponse(
|
||||
{
|
||||
"valid": True,
|
||||
"username": username,
|
||||
"llama_port": llama._port if llama.is_loaded else None,
|
||||
"llama_api_key": llama._api_key,
|
||||
"model_name": llama.model_identifier or "unknown",
|
||||
"supports_reasoning": llama.supports_reasoning,
|
||||
"supports_tools": llama.supports_tools,
|
||||
"is_vision": llama.is_vision,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
|
|||
from stream_token_store import consume_stream_token
|
||||
|
||||
|
||||
stream_app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
|
||||
stream_app = FastAPI(docs_url = None, redoc_url = None, openapi_url = None)
|
||||
|
||||
# ── Shared helpers ────────────────────────────────────────────
|
||||
|
||||
|
|
@ -92,17 +92,17 @@ def _process_image(image_b64, llama_backend):
|
|||
if img.mode == "RGBA":
|
||||
img = img.convert("RGB")
|
||||
buf = _BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
img.save(buf, format = "PNG")
|
||||
return _b64.b64encode(buf.getvalue()).decode("ascii")
|
||||
elif image_b64 and not llama_backend.is_vision:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Image provided but current GGUF model does not support vision.",
|
||||
status_code = 400,
|
||||
detail = "Image provided but current GGUF model does not support vision.",
|
||||
)
|
||||
return image_b64
|
||||
|
||||
|
||||
def _build_llama_payload(llama_backend, openai_messages, payload, stream=True):
|
||||
def _build_llama_payload(llama_backend, openai_messages, payload, stream = True):
|
||||
"""Build the payload for llama-server /v1/chat/completions.
|
||||
|
||||
Only sends repeat_penalty when the client explicitly provides it.
|
||||
|
|
@ -125,7 +125,9 @@ def _build_llama_payload(llama_backend, openai_messages, payload, stream=True):
|
|||
if stream:
|
||||
llama_payload["stream_options"] = {"include_usage": True}
|
||||
if llama_backend.supports_reasoning and payload.get("enable_thinking") is not None:
|
||||
llama_payload["chat_template_kwargs"] = {"enable_thinking": payload["enable_thinking"]}
|
||||
llama_payload["chat_template_kwargs"] = {
|
||||
"enable_thinking": payload["enable_thinking"]
|
||||
}
|
||||
if payload.get("max_tokens") is not None:
|
||||
llama_payload["max_tokens"] = payload["max_tokens"]
|
||||
if payload.get("stop"):
|
||||
|
|
@ -139,12 +141,15 @@ def _build_llama_payload(llama_backend, openai_messages, payload, stream=True):
|
|||
@stream_app.options("/stream")
|
||||
async def stream_preflight():
|
||||
"""Handle CORS preflight for the /stream endpoint."""
|
||||
return JSONResponse(content={}, headers={
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "POST, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "Content-Type, X-Stream-Token",
|
||||
"Access-Control-Max-Age": "86400",
|
||||
})
|
||||
return JSONResponse(
|
||||
content = {},
|
||||
headers = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "POST, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "Content-Type, X-Stream-Token",
|
||||
"Access-Control-Max-Age": "86400",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ── Request validation ────────────────────────────────────────
|
||||
|
|
@ -154,21 +159,22 @@ async def _validate_request(request: Request):
|
|||
"""Validate token, parse body, get backend. Returns all needed context."""
|
||||
token = request.headers.get("X-Stream-Token")
|
||||
if not token:
|
||||
raise HTTPException(status_code=401, detail="Missing X-Stream-Token header")
|
||||
raise HTTPException(status_code = 401, detail = "Missing X-Stream-Token header")
|
||||
username = consume_stream_token(token)
|
||||
if username is None:
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired stream token")
|
||||
raise HTTPException(status_code = 401, detail = "Invalid or expired stream token")
|
||||
|
||||
body_bytes = await request.body()
|
||||
try:
|
||||
payload = json.loads(body_bytes)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON body")
|
||||
raise HTTPException(status_code = 400, detail = "Invalid JSON body")
|
||||
|
||||
from routes.inference import get_llama_cpp_backend
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
if not llama_backend.is_loaded:
|
||||
raise HTTPException(status_code=400, detail="No GGUF model loaded")
|
||||
raise HTTPException(status_code = 400, detail = "No GGUF model loaded")
|
||||
|
||||
messages = payload.get("messages", [])
|
||||
gguf_messages, image_b64 = _extract_content_parts(messages)
|
||||
|
|
@ -184,7 +190,15 @@ async def _validate_request(request: Request):
|
|||
created = int(time.time())
|
||||
model_name = llama_backend.model_identifier or "unknown"
|
||||
|
||||
return payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name
|
||||
return (
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
)
|
||||
|
||||
|
||||
# ── Path A: Direct async streaming (HOT PATH) ────────────────
|
||||
|
|
@ -200,8 +214,16 @@ _SSE_HEADERS = {
|
|||
_SEP = (",", ":")
|
||||
|
||||
|
||||
async def _handle_async_stream(request, payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name):
|
||||
async def _handle_async_stream(
|
||||
request,
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
):
|
||||
"""
|
||||
Stream directly from llama-server using httpx.AsyncClient.
|
||||
|
||||
|
|
@ -210,24 +232,34 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i
|
|||
with delta tokens natively, so no cumulative-to-delta conversion needed.
|
||||
"""
|
||||
openai_messages = llama_backend._build_openai_messages(gguf_messages, image_b64)
|
||||
llama_payload = _build_llama_payload(llama_backend, openai_messages, payload, stream=True)
|
||||
llama_payload = _build_llama_payload(
|
||||
llama_backend, openai_messages, payload, stream = True
|
||||
)
|
||||
|
||||
port = llama_backend._port
|
||||
api_key = llama_backend._api_key
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
url = f"http://127.0.0.1:{port}/v1/chat/completions"
|
||||
timeout = httpx.Timeout(connect=30, read=120.0, write=10, pool=10)
|
||||
timeout = httpx.Timeout(connect = 30, read = 120.0, write = 10, pool = 10)
|
||||
|
||||
async def sse_generator():
|
||||
try:
|
||||
# Role chunk
|
||||
role = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name,
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]}
|
||||
yield f"data: {json.dumps(role, separators=_SEP)}\n\n"
|
||||
role = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
|
||||
],
|
||||
}
|
||||
yield f"data: {json.dumps(role, separators = _SEP)}\n\n"
|
||||
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
async with client.stream("POST", url, json=llama_payload, headers=headers) as resp:
|
||||
async with httpx.AsyncClient(timeout = timeout) as client:
|
||||
async with client.stream(
|
||||
"POST", url, json = llama_payload, headers = headers
|
||||
) as resp:
|
||||
if resp.status_code != 200:
|
||||
error_body = await resp.aread()
|
||||
raise RuntimeError(
|
||||
|
|
@ -253,9 +285,9 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i
|
|||
if line == "data: [DONE]":
|
||||
if in_thinking:
|
||||
if has_content_tokens:
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
else:
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
stream_done = True
|
||||
break
|
||||
if not line.startswith("data: "):
|
||||
|
|
@ -284,8 +316,8 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i
|
|||
reasoning_text += reasoning
|
||||
if not in_thinking:
|
||||
in_thinking = True
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
|
||||
# Handle content tokens
|
||||
token = delta.get("content", "")
|
||||
|
|
@ -293,32 +325,41 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i
|
|||
has_content_tokens = True
|
||||
if in_thinking:
|
||||
in_thinking = False
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
|
||||
|
||||
if stream_done:
|
||||
break
|
||||
|
||||
# Final stop chunk
|
||||
final = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
|
||||
yield f"data: {json.dumps(final, separators=_SEP)}\n\n"
|
||||
final = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}
|
||||
yield f"data: {json.dumps(final, separators = _SEP)}\n\n"
|
||||
|
||||
# Usage chunk
|
||||
if stream_usage or stream_timings:
|
||||
usage_chunk = {
|
||||
"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name, "choices": [],
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": (stream_usage or {}).get("prompt_tokens", 0),
|
||||
"completion_tokens": (stream_usage or {}).get("completion_tokens", 0),
|
||||
"completion_tokens": (stream_usage or {}).get(
|
||||
"completion_tokens", 0
|
||||
),
|
||||
"total_tokens": (stream_usage or {}).get("total_tokens", 0),
|
||||
},
|
||||
}
|
||||
if stream_timings:
|
||||
usage_chunk["timings"] = stream_timings
|
||||
yield f"data: {json.dumps(usage_chunk, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps(usage_chunk, separators = _SEP)}\n\n"
|
||||
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
|
|
@ -327,14 +368,24 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i
|
|||
except Exception as e:
|
||||
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
|
||||
|
||||
return StreamingResponse(sse_generator(), media_type="text/event-stream", headers=_SSE_HEADERS)
|
||||
return StreamingResponse(
|
||||
sse_generator(), media_type = "text/event-stream", headers = _SSE_HEADERS
|
||||
)
|
||||
|
||||
|
||||
# ── Path B: Tool calling (asyncio.to_thread) ─────────────────
|
||||
|
||||
|
||||
async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name):
|
||||
async def _handle_tool_stream(
|
||||
request,
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
):
|
||||
"""Handle a tool-calling streaming request via asyncio.to_thread.
|
||||
|
||||
Tool execution is the bottleneck, not streaming, so the thread overhead
|
||||
|
|
@ -346,7 +397,9 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im
|
|||
|
||||
p_enabled_tools = payload.get("enabled_tools")
|
||||
if p_enabled_tools is not None:
|
||||
tools_to_use = [t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools]
|
||||
tools_to_use = [
|
||||
t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools
|
||||
]
|
||||
else:
|
||||
tools_to_use = ALL_TOOLS
|
||||
|
||||
|
|
@ -354,27 +407,33 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im
|
|||
|
||||
async def tool_sse():
|
||||
try:
|
||||
first = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name,
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]}
|
||||
yield f"data: {json.dumps(first, separators=_SEP)}\n\n"
|
||||
first = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
|
||||
],
|
||||
}
|
||||
yield f"data: {json.dumps(first, separators = _SEP)}\n\n"
|
||||
|
||||
gen = llama_backend.generate_chat_completion_with_tools(
|
||||
messages=gguf_messages,
|
||||
tools=tools_to_use,
|
||||
temperature=payload.get("temperature", 0.6),
|
||||
top_p=payload.get("top_p", 0.95),
|
||||
top_k=payload.get("top_k", 20),
|
||||
min_p=payload.get("min_p", 0.01),
|
||||
max_tokens=payload.get("max_tokens"),
|
||||
repetition_penalty=payload.get("repetition_penalty", 1.1),
|
||||
presence_penalty=payload.get("presence_penalty", 0.0),
|
||||
cancel_event=cancel_event,
|
||||
enable_thinking=payload.get("enable_thinking"),
|
||||
auto_heal_tool_calls=payload.get("auto_heal_tool_calls", True),
|
||||
max_tool_iterations=payload.get("max_tool_calls_per_message", 10),
|
||||
tool_call_timeout=payload.get("tool_call_timeout", 300),
|
||||
session_id=payload.get("session_id"),
|
||||
messages = gguf_messages,
|
||||
tools = tools_to_use,
|
||||
temperature = payload.get("temperature", 0.6),
|
||||
top_p = payload.get("top_p", 0.95),
|
||||
top_k = payload.get("top_k", 20),
|
||||
min_p = payload.get("min_p", 0.01),
|
||||
max_tokens = payload.get("max_tokens"),
|
||||
repetition_penalty = payload.get("repetition_penalty", 1.1),
|
||||
presence_penalty = payload.get("presence_penalty", 0.0),
|
||||
cancel_event = cancel_event,
|
||||
enable_thinking = payload.get("enable_thinking"),
|
||||
auto_heal_tool_calls = payload.get("auto_heal_tool_calls", True),
|
||||
max_tool_iterations = payload.get("max_tool_calls_per_message", 10),
|
||||
tool_call_timeout = payload.get("tool_call_timeout", 300),
|
||||
session_id = payload.get("session_id"),
|
||||
)
|
||||
|
||||
prev_text = ""
|
||||
|
|
@ -399,29 +458,50 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im
|
|||
_timings = event.get("timings")
|
||||
continue
|
||||
cumulative = event.get("text", "")
|
||||
new_text = cumulative[len(prev_text):]
|
||||
new_text = cumulative[len(prev_text) :]
|
||||
prev_text = cumulative
|
||||
if not new_text:
|
||||
continue
|
||||
chunk = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name,
|
||||
"choices": [{"index": 0, "delta": {"content": new_text}, "finish_reason": None}]}
|
||||
yield f"data: {json.dumps(chunk, separators=_SEP)}\n\n"
|
||||
chunk = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": new_text},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
yield f"data: {json.dumps(chunk, separators = _SEP)}\n\n"
|
||||
|
||||
final = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
|
||||
yield f"data: {json.dumps(final, separators=_SEP)}\n\n"
|
||||
final = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}
|
||||
yield f"data: {json.dumps(final, separators = _SEP)}\n\n"
|
||||
|
||||
if _usage or _timings:
|
||||
uc = {"id": completion_id, "object": "chat.completion.chunk",
|
||||
"created": created, "model": model_name, "choices": [],
|
||||
"usage": {"prompt_tokens": (_usage or {}).get("prompt_tokens", 0),
|
||||
"completion_tokens": (_usage or {}).get("completion_tokens", 0),
|
||||
"total_tokens": (_usage or {}).get("total_tokens", 0)}}
|
||||
uc = {
|
||||
"id": completion_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": (_usage or {}).get("prompt_tokens", 0),
|
||||
"completion_tokens": (_usage or {}).get("completion_tokens", 0),
|
||||
"total_tokens": (_usage or {}).get("total_tokens", 0),
|
||||
},
|
||||
}
|
||||
if _timings:
|
||||
uc["timings"] = _timings
|
||||
yield f"data: {json.dumps(uc, separators=_SEP)}\n\n"
|
||||
yield f"data: {json.dumps(uc, separators = _SEP)}\n\n"
|
||||
|
||||
yield "data: [DONE]\n\n"
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -430,31 +510,34 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im
|
|||
except Exception as e:
|
||||
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
|
||||
|
||||
return StreamingResponse(tool_sse(), media_type="text/event-stream", headers=_SSE_HEADERS)
|
||||
return StreamingResponse(
|
||||
tool_sse(), media_type = "text/event-stream", headers = _SSE_HEADERS
|
||||
)
|
||||
|
||||
|
||||
# ── Path C: Non-streaming ─────────────────────────────────────
|
||||
|
||||
|
||||
async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name):
|
||||
async def _handle_non_streaming(
|
||||
payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name
|
||||
):
|
||||
"""Handle a non-streaming request. Returns a JSON response."""
|
||||
cancel_event = threading.Event()
|
||||
|
||||
def _run_sync():
|
||||
gen = llama_backend.generate_chat_completion(
|
||||
messages=gguf_messages,
|
||||
image_b64=image_b64,
|
||||
temperature=payload.get("temperature", 0.6),
|
||||
top_p=payload.get("top_p", 0.95),
|
||||
top_k=payload.get("top_k", 20),
|
||||
min_p=payload.get("min_p", 0.01),
|
||||
max_tokens=payload.get("max_tokens"),
|
||||
repetition_penalty=payload.get("repetition_penalty", 1.0),
|
||||
presence_penalty=payload.get("presence_penalty", 0.0),
|
||||
stop=payload.get("stop"),
|
||||
cancel_event=cancel_event,
|
||||
enable_thinking=payload.get("enable_thinking"),
|
||||
messages = gguf_messages,
|
||||
image_b64 = image_b64,
|
||||
temperature = payload.get("temperature", 0.6),
|
||||
top_p = payload.get("top_p", 0.95),
|
||||
top_k = payload.get("top_k", 20),
|
||||
min_p = payload.get("min_p", 0.01),
|
||||
max_tokens = payload.get("max_tokens"),
|
||||
repetition_penalty = payload.get("repetition_penalty", 1.0),
|
||||
presence_penalty = payload.get("presence_penalty", 0.0),
|
||||
stop = payload.get("stop"),
|
||||
cancel_event = cancel_event,
|
||||
enable_thinking = payload.get("enable_thinking"),
|
||||
)
|
||||
text = ""
|
||||
usage = None
|
||||
|
|
@ -474,11 +557,13 @@ async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64
|
|||
"object": "chat.completion",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": text},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": text},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": (usage or {}).get("prompt_tokens", 0),
|
||||
"completion_tokens": (usage or {}).get("completion_tokens", 0),
|
||||
|
|
@ -489,8 +574,8 @@ async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64
|
|||
result["timings"] = timings
|
||||
|
||||
return JSONResponse(
|
||||
content=result,
|
||||
headers={"Access-Control-Allow-Origin": "*"},
|
||||
content = result,
|
||||
headers = {"Access-Control-Allow-Origin": "*"},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -507,29 +592,53 @@ async def stream_endpoint(request: Request):
|
|||
- Path B: tool calling via asyncio.to_thread (tool exec is the bottleneck)
|
||||
- Path C: non-streaming one-shot JSON response
|
||||
"""
|
||||
payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name = \
|
||||
await _validate_request(request)
|
||||
(
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
) = await _validate_request(request)
|
||||
|
||||
# Path C: Non-streaming
|
||||
stream = payload.get("stream", True)
|
||||
if not stream:
|
||||
return await _handle_non_streaming(
|
||||
payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name,
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
)
|
||||
|
||||
# Path B: Tool calling
|
||||
use_tools = payload.get("use_tools", False)
|
||||
if use_tools and llama_backend.supports_tools:
|
||||
return await _handle_tool_stream(
|
||||
request, payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name,
|
||||
request,
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
)
|
||||
|
||||
# Path A: Direct async streaming (hot path)
|
||||
return await _handle_async_stream(
|
||||
request, payload, llama_backend, gguf_messages, image_b64,
|
||||
completion_id, created, model_name,
|
||||
request,
|
||||
payload,
|
||||
llama_backend,
|
||||
gguf_messages,
|
||||
image_b64,
|
||||
completion_id,
|
||||
created,
|
||||
model_name,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -542,10 +651,10 @@ def start_streaming_server(port: int) -> None:
|
|||
|
||||
uvicorn.run(
|
||||
stream_app,
|
||||
host="127.0.0.1",
|
||||
port=port,
|
||||
log_level="warning",
|
||||
access_log=False,
|
||||
host = "127.0.0.1",
|
||||
port = port,
|
||||
log_level = "warning",
|
||||
access_log = False,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue