97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
"""LLM client: streaming and non-streaming chat completions, SSE framing.
|
|
|
|
Pure functions — no Flask imports. Uses ``requests`` for HTTP against an
|
|
OpenAI-compatible endpoint (``config.LLM_URL``). The ``_sse`` helper is
|
|
shared by deconstruct / chapter_writer so every SSE stream uses the same
|
|
wire format.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import time
|
|
|
|
import requests as http_requests
|
|
|
|
import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_KEEPALIVE_INTERVAL = 15 # seconds
|
|
|
|
|
|
def _sse(event, payload):
|
|
"""Return a single SSE frame: ``event: <event>\\ndata: <json>\\n\\n``."""
|
|
return f"event: {event}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
|
|
|
|
|
def stream_chat(messages, max_tokens=4096, model=None, timeout=None):
|
|
"""Generator yielding raw token strings from a streaming chat completion.
|
|
|
|
POSTs to ``{config.LLM_URL}/chat/completions`` with ``stream=True``.
|
|
Parses ``data:`` lines, handles ``[DONE]`` terminator, extracts
|
|
``choices[0].delta.content``. Raises ``RuntimeError`` on HTTP failure.
|
|
"""
|
|
url = f"{config.LLM_URL}/chat/completions"
|
|
body = {
|
|
"model": model or config.LLM_MODEL,
|
|
"messages": messages,
|
|
"stream": True,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
resp = http_requests.post(url, json=body, stream=True, timeout=timeout or config.LLM_TIMEOUT)
|
|
resp.raise_for_status()
|
|
|
|
for line in resp.iter_lines():
|
|
if not line:
|
|
continue
|
|
line = line.decode("utf-8") if isinstance(line, bytes) else line
|
|
if not line.startswith("data: "):
|
|
continue
|
|
chunk_data = line[6:]
|
|
if chunk_data.strip() == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(chunk_data)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
|
token_text = delta.get("content", "")
|
|
if token_text:
|
|
yield token_text
|
|
|
|
|
|
def complete_chat(messages, max_tokens=2048, model=None, timeout=None):
|
|
"""Non-streaming call returning the assistant content string."""
|
|
url = f"{config.LLM_URL}/chat/completions"
|
|
body = {
|
|
"model": model or config.LLM_MODEL,
|
|
"messages": messages,
|
|
"stream": False,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
resp = http_requests.post(url, json=body, timeout=timeout or config.LLM_TIMEOUT)
|
|
resp.raise_for_status()
|
|
return resp.json().get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
|
|
|
|
def make_token_stream(messages, on_token, max_tokens=4096):
|
|
"""Generator wrapping :func:`stream_chat` with SSE framing + keepalive.
|
|
|
|
Yields ``_sse('token', {'text': t})`` for each token. Calls
|
|
``on_token(t)`` per token so callers can accumulate the full text.
|
|
Emits ``: keepalive\\n\\n`` comments if >15 s silence between chunks.
|
|
On exception yields ``_sse('error', {'error': ...})`` and stops.
|
|
"""
|
|
last_emit = time.time()
|
|
try:
|
|
for token in stream_chat(messages, max_tokens=max_tokens):
|
|
now = time.time()
|
|
if now - last_emit > _KEEPALIVE_INTERVAL:
|
|
yield ": keepalive\n\n"
|
|
yield _sse("token", {"text": token})
|
|
on_token(token)
|
|
last_emit = time.time()
|
|
except Exception as exc:
|
|
logger.exception("make_token_stream: LLM error")
|
|
yield _sse("error", {"error": str(exc)})
|