d209d8ac0b
Lingma streams answers via chat/answer + chat/finish notifications and
never sends a JSON-RPC response for chat/ask. The old code awaited
rpc.request("chat/ask") and swallowed the TimeoutError, so every chat
was forced to wait the full rpc_timeout (default 30s) before draining
the stream queue - even though the first token was already present in
the queue within ~2s.
Effect:
- non-stream TTFB dropped from ~30s to actual upstream latency (~2-3s).
- stream first-chunk dropped from ~30s to upstream first-token latency.
- consume_stream idle timeout decoupled from rpc_timeout so shortening
rpc_timeout no longer starves long completions.
Switch chat/ask to rpc.notify (fire-and-forget) and rely entirely on the
existing chat/answer + chat/finish handlers for result delivery.
Made-with: Cursor
620 lines
23 KiB
Python
620 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import AsyncIterator, Callable, Optional
|
|
|
|
import websockets
|
|
|
|
from .logging_config import get_logger
|
|
|
|
|
|
logger = get_logger("lingma_gateway.client")
|
|
|
|
|
|
# Some callers live on Python 3.10 where asyncio.TimeoutError is a distinct class,
|
|
# while 3.11+ unifies it with the builtin TimeoutError. Always catch both.
|
|
TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (
|
|
asyncio.TimeoutError,
|
|
TimeoutError,
|
|
)
|
|
|
|
|
|
def _is_port_open(host: str, port: int, timeout_sec: float = 0.5) -> bool:
|
|
try:
|
|
with socket.create_connection((host, port), timeout=timeout_sec):
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _read_info_file(info_path: Path):
|
|
if not info_path.exists():
|
|
return None, None
|
|
txt = info_path.read_text(encoding="utf-8", errors="ignore").strip()
|
|
if not txt:
|
|
return None, None
|
|
lines = txt.splitlines()
|
|
if len(lines) < 2:
|
|
return None, None
|
|
try:
|
|
return int(lines[0].strip()), int(lines[1].strip())
|
|
except ValueError:
|
|
return None, None
|
|
|
|
|
|
def _wait_info_any(info_paths: list[Path], timeout_sec: int):
|
|
start = time.time()
|
|
while time.time() - start < timeout_sec:
|
|
for p in info_paths:
|
|
port, pid = _read_info_file(p)
|
|
if port and pid:
|
|
return port, pid, p
|
|
time.sleep(0.2)
|
|
raise TimeoutError(".info not ready")
|
|
|
|
|
|
def _encode_lsp_frame(payload_obj: dict) -> bytes:
|
|
body = json.dumps(payload_obj, ensure_ascii=False).encode("utf-8")
|
|
header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
|
|
return header + body
|
|
|
|
|
|
def _parse_lsp_frames(buf: bytes):
|
|
frames = []
|
|
while True:
|
|
header_end = buf.find(b"\r\n\r\n")
|
|
if header_end < 0:
|
|
break
|
|
header = buf[:header_end]
|
|
body_start = header_end + 4
|
|
content_length = None
|
|
for line in header.split(b"\r\n"):
|
|
if line.lower().startswith(b"content-length:"):
|
|
content_length = int(line.split(b":", 1)[1].strip())
|
|
break
|
|
if content_length is None:
|
|
buf = buf[body_start:]
|
|
continue
|
|
if len(buf) < body_start + content_length:
|
|
break
|
|
body = buf[body_start : body_start + content_length]
|
|
frames.append(body.decode("utf-8", errors="ignore"))
|
|
buf = buf[body_start + content_length :]
|
|
return frames, buf
|
|
|
|
|
|
class LspWsRpcClient:
|
|
def __init__(self, ws, on_disconnect: Optional[Callable[[BaseException], None]] = None):
|
|
self.ws = ws
|
|
self._id = 1
|
|
self._pending: dict[int, asyncio.Future] = {}
|
|
self._send_lock = asyncio.Lock()
|
|
self._reader_task: asyncio.Task | None = None
|
|
self._rx_buffer = b""
|
|
self._chat_streams: dict[str, dict] = {}
|
|
self._on_disconnect = on_disconnect
|
|
self._closed = False
|
|
|
|
async def start(self):
|
|
self._reader_task = asyncio.create_task(self._reader_loop())
|
|
|
|
async def close(self):
|
|
self._closed = True
|
|
if self._reader_task:
|
|
self._reader_task.cancel()
|
|
with contextlib.suppress(Exception):
|
|
await self._reader_task
|
|
# Abort any pending futures so callers fail fast instead of hanging.
|
|
for fut in self._pending.values():
|
|
if not fut.done():
|
|
fut.set_exception(ConnectionError("lingma client closed"))
|
|
self._pending.clear()
|
|
# Signal open streams to terminate.
|
|
for stream in self._chat_streams.values():
|
|
if not stream["done"].is_set():
|
|
stream["done"].set()
|
|
stream["chunks"].put_nowait(None)
|
|
self._chat_streams.clear()
|
|
|
|
async def _send(self, payload: dict):
|
|
async with self._send_lock:
|
|
await self.ws.send(_encode_lsp_frame(payload))
|
|
|
|
async def _reader_loop(self):
|
|
try:
|
|
while True:
|
|
raw = await self.ws.recv()
|
|
chunk = raw if isinstance(raw, bytes) else raw.encode("utf-8", errors="ignore")
|
|
self._rx_buffer += chunk
|
|
bodies, self._rx_buffer = _parse_lsp_frames(self._rx_buffer)
|
|
for body in bodies:
|
|
try:
|
|
msg = json.loads(body)
|
|
except Exception:
|
|
continue
|
|
|
|
if "method" in msg and "result" not in msg and "error" not in msg:
|
|
await self._handle_server_message(msg)
|
|
continue
|
|
|
|
rid = msg.get("id")
|
|
if rid is None:
|
|
continue
|
|
fut = self._pending.pop(rid, None)
|
|
if fut and not fut.done():
|
|
fut.set_result(msg)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as exc:
|
|
if not self._closed:
|
|
logger.warning("lingma reader loop terminated: %s", exc)
|
|
# Propagate failure to anyone waiting on an RPC.
|
|
for fut in self._pending.values():
|
|
if not fut.done():
|
|
fut.set_exception(exc)
|
|
self._pending.clear()
|
|
# Also unblock any in-flight chat streams so consumers exit.
|
|
for stream in self._chat_streams.values():
|
|
if not stream["done"].is_set():
|
|
stream["done"].set()
|
|
stream["chunks"].put_nowait(None)
|
|
if not self._closed and self._on_disconnect is not None:
|
|
try:
|
|
self._on_disconnect(exc)
|
|
except Exception:
|
|
logger.exception("on_disconnect callback failed")
|
|
|
|
async def _handle_server_message(self, msg: dict):
|
|
method = msg.get("method")
|
|
params = msg.get("params") or {}
|
|
|
|
if method == "chat/answer":
|
|
req_id = params.get("requestId")
|
|
stream = self._chat_streams.get(req_id)
|
|
if stream is not None:
|
|
text = params.get("text") or params.get("content") or ""
|
|
if text:
|
|
stream["parts"].append(text)
|
|
if stream["first_chunk_at"] is None:
|
|
stream["first_chunk_at"] = time.monotonic()
|
|
stream["chunks"].put_nowait(text)
|
|
|
|
if method == "chat/finish":
|
|
req_id = params.get("requestId")
|
|
stream = self._chat_streams.get(req_id)
|
|
if stream is not None and not stream["done"].is_set():
|
|
stream["finish"] = params
|
|
stream["finish_at"] = time.monotonic()
|
|
stream["done"].set()
|
|
stream["chunks"].put_nowait(None)
|
|
|
|
if "id" in msg:
|
|
await self._send({"jsonrpc": "2.0", "id": msg.get("id"), "result": {}})
|
|
|
|
async def request(self, method, params=None, timeout=20):
|
|
rid = self._id
|
|
self._id += 1
|
|
payload = {"jsonrpc": "2.0", "id": rid, "method": method, "params": params or {}}
|
|
fut = asyncio.get_running_loop().create_future()
|
|
self._pending[rid] = fut
|
|
await self._send(payload)
|
|
try:
|
|
msg = await asyncio.wait_for(fut, timeout=timeout)
|
|
except TIMEOUT_EXCEPTIONS:
|
|
self._pending.pop(rid, None)
|
|
raise TimeoutError(f"RPC timeout: {method}")
|
|
if "error" in msg:
|
|
raise RuntimeError(f"RPC {method} error: {msg['error']}")
|
|
return msg.get("result")
|
|
|
|
async def notify(self, method, params=None):
|
|
await self._send({"jsonrpc": "2.0", "method": method, "params": params or {}})
|
|
|
|
def create_stream(self, request_id: str):
|
|
self._chat_streams[request_id] = {
|
|
"parts": [],
|
|
"chunks": asyncio.Queue(),
|
|
"done": asyncio.Event(),
|
|
"finish": None,
|
|
"started_at": time.monotonic(),
|
|
"first_chunk_at": None,
|
|
"finish_at": None,
|
|
}
|
|
|
|
def pop_stream(self, request_id: str) -> None:
|
|
stream = self._chat_streams.pop(request_id, None)
|
|
if stream is None:
|
|
return
|
|
# Drain queue so no stray future gets stuck if the consumer bailed early.
|
|
if not stream["done"].is_set():
|
|
stream["done"].set()
|
|
with contextlib.suppress(Exception):
|
|
stream["chunks"].put_nowait(None)
|
|
|
|
async def consume_stream(self, request_id: str, timeout: float) -> AsyncIterator[str]:
|
|
stream = self._chat_streams.get(request_id)
|
|
if stream is None:
|
|
return
|
|
start = time.monotonic()
|
|
while True:
|
|
remain = timeout - (time.monotonic() - start)
|
|
if remain <= 0:
|
|
raise TimeoutError("chat stream timeout")
|
|
chunk = await asyncio.wait_for(stream["chunks"].get(), timeout=remain)
|
|
if chunk is None:
|
|
break
|
|
yield chunk
|
|
|
|
def get_stream_result(self, request_id: str) -> dict:
|
|
stream = self._chat_streams.get(request_id) or {}
|
|
first_ms = None
|
|
total_ms = None
|
|
if stream.get("first_chunk_at") is not None:
|
|
first_ms = int((stream["first_chunk_at"] - stream["started_at"]) * 1000)
|
|
if stream.get("finish_at") is not None:
|
|
total_ms = int((stream["finish_at"] - stream["started_at"]) * 1000)
|
|
return {
|
|
"text": "".join(stream.get("parts") or []),
|
|
"finish": stream.get("finish") or {},
|
|
"firstTokenLatencyMs": first_ms,
|
|
"totalLatencyMs": total_ms,
|
|
}
|
|
|
|
|
|
class LingmaGatewayClient:
|
|
"""Owns the Lingma subprocess and the LSP-over-WS connection.
|
|
|
|
Adds a small state machine + reconnect loop so the gateway can survive Lingma
|
|
restarts and slow cold starts without bringing down the FastAPI app.
|
|
"""
|
|
|
|
STATE_STOPPED = "stopped"
|
|
STATE_STARTING = "starting"
|
|
STATE_READY = "ready"
|
|
STATE_RECONNECTING = "reconnecting"
|
|
STATE_FAILED = "failed"
|
|
STATE_CLOSED = "closed"
|
|
|
|
def __init__(
|
|
self,
|
|
lingma_bin: str,
|
|
work_dir: str,
|
|
socket_port: int,
|
|
startup_timeout: int,
|
|
rpc_timeout: int,
|
|
default_model: str,
|
|
default_ask_mode: str,
|
|
*,
|
|
name: str = "lingma",
|
|
extra_info_paths: list[Path] | None = None,
|
|
):
|
|
self.name = name
|
|
self.lingma_bin = Path(lingma_bin)
|
|
self.work_dir = Path(work_dir)
|
|
self.socket_port = socket_port
|
|
self.startup_timeout = startup_timeout
|
|
self.rpc_timeout = rpc_timeout
|
|
self.default_model = default_model
|
|
self.default_ask_mode = default_ask_mode
|
|
# Each pool instance should only look at its own workDir .info to avoid
|
|
# cross-instance clobbering via the shared ~/.lingma/.info path.
|
|
if extra_info_paths is None:
|
|
extra_info_paths = [Path.home() / ".lingma" / ".info"]
|
|
self._extra_info_paths = list(extra_info_paths)
|
|
self._rpc: LspWsRpcClient | None = None
|
|
self._ws = None
|
|
self._state = self.STATE_STOPPED
|
|
self._state_lock = asyncio.Lock()
|
|
self._ready_event = asyncio.Event()
|
|
self._reconnect_task: asyncio.Task | None = None
|
|
self._last_error: str = ""
|
|
|
|
# ------------------------------------------------------------------ state
|
|
|
|
@property
|
|
def state(self) -> str:
|
|
return self._state
|
|
|
|
@property
|
|
def last_error(self) -> str:
|
|
return self._last_error
|
|
|
|
def _set_state(self, state: str, err: str = "") -> None:
|
|
if state != self._state:
|
|
logger.info("lingma client state %s -> %s", self._state, state, extra={"ctx_new_state": state})
|
|
self._state = state
|
|
if err:
|
|
self._last_error = err
|
|
if state == self.STATE_READY:
|
|
self._ready_event.set()
|
|
else:
|
|
self._ready_event.clear()
|
|
|
|
# -------------------------------------------------------------- lifecycle
|
|
|
|
async def start(self) -> None:
|
|
"""Initial start. Failure is non-fatal: ensure_ready() will retry later."""
|
|
try:
|
|
await self._connect(initial=True)
|
|
except Exception as exc:
|
|
self._set_state(self.STATE_FAILED, err=str(exc))
|
|
logger.exception("initial lingma start failed; will retry on demand")
|
|
|
|
async def close(self) -> None:
|
|
self._set_state(self.STATE_CLOSED)
|
|
if self._reconnect_task and not self._reconnect_task.done():
|
|
self._reconnect_task.cancel()
|
|
with contextlib.suppress(Exception):
|
|
await self._reconnect_task
|
|
if self._rpc:
|
|
await self._rpc.close()
|
|
if self._ws:
|
|
with contextlib.suppress(Exception):
|
|
await self._ws.close()
|
|
|
|
async def ensure_ready(self, timeout: float | None = None) -> None:
|
|
"""Block until the RPC connection is usable, (re)connecting on demand."""
|
|
if self._state == self.STATE_CLOSED:
|
|
raise RuntimeError("lingma client is closed")
|
|
if self._state == self.STATE_READY and self._ws is not None:
|
|
return
|
|
|
|
async with self._state_lock:
|
|
if self._state == self.STATE_READY and self._ws is not None:
|
|
return
|
|
if self._state in (self.STATE_STOPPED, self.STATE_FAILED):
|
|
try:
|
|
await self._connect(initial=False)
|
|
return
|
|
except Exception as exc:
|
|
self._set_state(self.STATE_FAILED, err=str(exc))
|
|
raise
|
|
|
|
wait_timeout = timeout if timeout is not None else max(
|
|
30.0, float(self.startup_timeout) + 10.0
|
|
)
|
|
try:
|
|
await asyncio.wait_for(self._ready_event.wait(), timeout=wait_timeout)
|
|
except TIMEOUT_EXCEPTIONS:
|
|
raise RuntimeError(f"lingma not ready (state={self._state}, err={self._last_error})")
|
|
|
|
# --------------------------------------------------------------- connect
|
|
|
|
async def _connect(self, *, initial: bool) -> None:
|
|
self._set_state(self.STATE_STARTING)
|
|
|
|
if not self.lingma_bin.exists():
|
|
raise FileNotFoundError(f"Lingma not found: {self.lingma_bin}")
|
|
|
|
info_paths = [self.work_dir / ".info", *self._extra_info_paths]
|
|
|
|
# socket_port <= 0 is the pool-friendly "always spawn and read .info" mode.
|
|
port_prewarmed = self.socket_port > 0 and _is_port_open(
|
|
"127.0.0.1", self.socket_port
|
|
)
|
|
if not port_prewarmed:
|
|
self.work_dir.mkdir(parents=True, exist_ok=True)
|
|
# Remove stale info files from host-mounted workspace before boot.
|
|
for p in info_paths:
|
|
with contextlib.suppress(Exception):
|
|
if p.exists():
|
|
p.unlink()
|
|
logger.info(
|
|
"[%s] spawning lingma: %s start --workDir %s",
|
|
self.name,
|
|
self.lingma_bin,
|
|
self.work_dir,
|
|
)
|
|
subprocess.Popen(
|
|
[str(self.lingma_bin), "start", "--workDir", str(self.work_dir)],
|
|
cwd=str(self.lingma_bin.parent),
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
start_new_session=True,
|
|
)
|
|
info, _, _ = _wait_info_any(info_paths, timeout_sec=self.startup_timeout)
|
|
self.socket_port = info
|
|
|
|
deadline = time.time() + self.startup_timeout
|
|
while time.time() < deadline:
|
|
if _is_port_open("127.0.0.1", self.socket_port, timeout_sec=0.3):
|
|
break
|
|
await asyncio.sleep(0.2)
|
|
else:
|
|
raise TimeoutError(f"Lingma socket not open on port {self.socket_port}")
|
|
|
|
# Close any stale ws/rpc before creating fresh ones (reconnect path).
|
|
if self._rpc is not None:
|
|
with contextlib.suppress(Exception):
|
|
await self._rpc.close()
|
|
self._rpc = None
|
|
if self._ws is not None:
|
|
with contextlib.suppress(Exception):
|
|
await self._ws.close()
|
|
self._ws = None
|
|
|
|
ws_url = f"ws://127.0.0.1:{self.socket_port}"
|
|
self._ws = await websockets.connect(ws_url, max_size=10 * 1024 * 1024)
|
|
self._rpc = LspWsRpcClient(self._ws, on_disconnect=self._on_disconnect)
|
|
await self._rpc.start()
|
|
await self._rpc.request(
|
|
"initialize",
|
|
{
|
|
"processId": os.getpid(),
|
|
"clientInfo": {"name": "lingma-openai-gateway", "version": "0.1.0"},
|
|
"capabilities": {},
|
|
"workspaceFolders": [],
|
|
"rootUri": None,
|
|
},
|
|
timeout=self.rpc_timeout,
|
|
)
|
|
await self._rpc.notify("initialized", {})
|
|
self._set_state(self.STATE_READY)
|
|
logger.info(
|
|
"[%s] lingma ready on port %d (initial=%s)",
|
|
self.name,
|
|
self.socket_port,
|
|
initial,
|
|
)
|
|
|
|
def _on_disconnect(self, exc: BaseException) -> None:
|
|
if self._state == self.STATE_CLOSED:
|
|
return
|
|
self._set_state(self.STATE_RECONNECTING, err=str(exc))
|
|
if self._reconnect_task and not self._reconnect_task.done():
|
|
return
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return
|
|
self._reconnect_task = loop.create_task(self._reconnect_loop())
|
|
|
|
async def _reconnect_loop(self) -> None:
|
|
backoff = 1.0
|
|
max_backoff = 30.0
|
|
max_attempts = 20
|
|
for attempt in range(1, max_attempts + 1):
|
|
if self._state == self.STATE_CLOSED:
|
|
return
|
|
await asyncio.sleep(backoff)
|
|
try:
|
|
async with self._state_lock:
|
|
await self._connect(initial=False)
|
|
logger.info("lingma reconnected after %d attempt(s)", attempt)
|
|
return
|
|
except Exception as exc:
|
|
self._last_error = str(exc)
|
|
logger.warning("lingma reconnect attempt %d failed: %s", attempt, exc)
|
|
backoff = min(backoff * 2, max_backoff)
|
|
self._set_state(self.STATE_FAILED, err="reconnect exhausted")
|
|
|
|
# ------------------------------------------------------------------ RPC
|
|
|
|
@property
|
|
def rpc(self) -> LspWsRpcClient:
|
|
if self._rpc is None:
|
|
raise RuntimeError(f"Lingma RPC not initialized (state={self._state})")
|
|
return self._rpc
|
|
|
|
async def auth_status(self):
|
|
await self.ensure_ready()
|
|
return await self.rpc.request("auth/status", {}, timeout=self.rpc_timeout)
|
|
|
|
async def query_models(self):
|
|
await self.ensure_ready()
|
|
return await self.rpc.request("config/queryModels", {}, timeout=self.rpc_timeout)
|
|
|
|
async def get_endpoint(self):
|
|
await self.ensure_ready()
|
|
return await self.rpc.request("config/getEndpoint", {}, timeout=self.rpc_timeout)
|
|
|
|
async def update_endpoint(self, endpoint: str):
|
|
await self.ensure_ready()
|
|
return await self.rpc.request(
|
|
"config/updateEndpoint", {"endpoint": endpoint}, timeout=self.rpc_timeout
|
|
)
|
|
|
|
async def generate_login_url(self):
|
|
await self.ensure_ready()
|
|
result = await self.rpc.request("login/generateUrl", {}, timeout=self.rpc_timeout)
|
|
if isinstance(result, str):
|
|
return result, {"raw": result}
|
|
if isinstance(result, dict):
|
|
for key in ("loginUrl", "url", "login_url"):
|
|
if isinstance(result.get(key), str):
|
|
return result[key], result
|
|
return "", result
|
|
return "", {"raw": result}
|
|
|
|
# ------------------------------------------------------------------ chat
|
|
|
|
def _build_payload(self, prompt: str, model_key: str, ask_mode: str, session_id: str, request_id: str):
|
|
session_type = "developer" if ask_mode == "agent" else "chat"
|
|
return {
|
|
"requestId": request_id,
|
|
"sessionId": session_id,
|
|
"sessionType": session_type,
|
|
"chatTask": "FREE_INPUT",
|
|
"mode": ask_mode,
|
|
"stream": True,
|
|
"source": 1,
|
|
"isReply": False,
|
|
"taskDefinitionType": "system",
|
|
"content": prompt,
|
|
"text": prompt,
|
|
"message": prompt,
|
|
"questionText": prompt,
|
|
"extra": {
|
|
"modelConfig": {"key": model_key},
|
|
"workspacePath": str(Path.cwd()),
|
|
},
|
|
"pluginPayloadConfig": {
|
|
"isEnableAskAgent": ask_mode == "agent",
|
|
"isEnableAutoMemory": True,
|
|
},
|
|
"chatContext": {
|
|
"text": prompt,
|
|
"features": [],
|
|
"preferredLanguage": "zh-CN",
|
|
"localeLang": "zh-CN",
|
|
},
|
|
}
|
|
|
|
async def _kick_chat_ask(self, payload: dict) -> None:
|
|
"""Fire chat/ask as a notification.
|
|
|
|
Lingma streams answers back via `chat/answer` + `chat/finish` and never
|
|
returns a JSON-RPC `result` for `chat/ask`. Waiting for one wasted
|
|
`rpc_timeout` seconds before the first byte could leave the gateway —
|
|
matching our previous 30s TTFB bug. `notify` sidesteps that entirely
|
|
by not registering a pending future.
|
|
"""
|
|
await self.rpc.notify("chat/ask", payload)
|
|
|
|
async def chat_complete(self, prompt: str, model_key: str, ask_mode: str) -> dict:
|
|
await self.ensure_ready()
|
|
request_id = str(uuid.uuid4())
|
|
session_id = str(uuid.uuid4())
|
|
payload = self._build_payload(prompt, model_key, ask_mode, session_id, request_id)
|
|
self.rpc.create_stream(request_id)
|
|
try:
|
|
await self._kick_chat_ask(payload)
|
|
# Consume until chat/finish closes the stream or the upstream idles.
|
|
async for _ in self.rpc.consume_stream(
|
|
request_id, timeout=max(60.0, self.rpc_timeout + 30.0)
|
|
):
|
|
pass
|
|
result = self.rpc.get_stream_result(request_id)
|
|
finally:
|
|
self.rpc.pop_stream(request_id)
|
|
finish = result.get("finish") or {}
|
|
result["requestId"] = request_id
|
|
result["sessionId"] = finish.get("sessionId") or session_id
|
|
result["model"] = model_key
|
|
result["mode"] = ask_mode
|
|
return result
|
|
|
|
async def chat_stream(self, prompt: str, model_key: str, ask_mode: str) -> AsyncIterator[str]:
|
|
await self.ensure_ready()
|
|
request_id = str(uuid.uuid4())
|
|
session_id = str(uuid.uuid4())
|
|
payload = self._build_payload(prompt, model_key, ask_mode, session_id, request_id)
|
|
self.rpc.create_stream(request_id)
|
|
try:
|
|
await self._kick_chat_ask(payload)
|
|
async for chunk in self.rpc.consume_stream(
|
|
request_id, timeout=max(60.0, self.rpc_timeout + 60.0)
|
|
):
|
|
yield chunk
|
|
finally:
|
|
# Runs on normal completion, exception, or consumer GeneratorExit (client disconnect).
|
|
self.rpc.pop_stream(request_id)
|