mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 18:15:33 +00:00
fix: better v1 compat and keepalive on long thinking
This commit is contained in:
1 parent
ca5f80995d
commit
1e4e2ff24a
7 files changed
+252
-51
No files matched your search
@@ -10,6 +10,7 @@ from application.api.answer.routes.base import answer_ns, BaseAnswerResource
|
||||
|
||||
from application.api.answer.services.persistence_policy import resolve_persistence
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.streaming.sse_keepalive import with_sse_keepalive
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -113,24 +114,26 @@ class StreamResource(Resource, BaseAnswerResource):
|
||||
if error := self.check_usage(processor.agent_config):
|
||||
return error
|
||||
return Response(
|
||||
self.complete_stream(
|
||||
question="",
|
||||
agent=agent,
|
||||
conversation_id=processor.conversation_id,
|
||||
user_api_key=processor.agent_config.get("user_api_key"),
|
||||
decoded_token=processor.decoded_token,
|
||||
agent_id=processor.agent_id,
|
||||
model_id=processor.model_id,
|
||||
model_user_id=processor.model_user_id,
|
||||
_continuation={
|
||||
"messages": messages,
|
||||
"tools_dict": tools_dict,
|
||||
"pending_tool_calls": pending_tool_calls,
|
||||
"tool_actions": tool_actions,
|
||||
"reserved_message_id": processor.reserved_message_id,
|
||||
"request_id": processor.request_id,
|
||||
"reasoning_content": reasoning_content,
|
||||
},
|
||||
with_sse_keepalive(
|
||||
self.complete_stream(
|
||||
question="",
|
||||
agent=agent,
|
||||
conversation_id=processor.conversation_id,
|
||||
user_api_key=processor.agent_config.get("user_api_key"),
|
||||
decoded_token=processor.decoded_token,
|
||||
agent_id=processor.agent_id,
|
||||
model_id=processor.model_id,
|
||||
model_user_id=processor.model_user_id,
|
||||
_continuation={
|
||||
"messages": messages,
|
||||
"tools_dict": tools_dict,
|
||||
"pending_tool_calls": pending_tool_calls,
|
||||
"tool_actions": tool_actions,
|
||||
"reserved_message_id": processor.reserved_message_id,
|
||||
"request_id": processor.request_id,
|
||||
"reasoning_content": reasoning_content,
|
||||
},
|
||||
),
|
||||
),
|
||||
mimetype="text/event-stream",
|
||||
)
|
||||
@@ -151,22 +154,24 @@ class StreamResource(Resource, BaseAnswerResource):
|
||||
persist_flag=data.get("persist"),
|
||||
)
|
||||
return Response(
|
||||
self.complete_stream(
|
||||
question=data["question"],
|
||||
agent=agent,
|
||||
conversation_id=processor.conversation_id,
|
||||
user_api_key=processor.agent_config.get("user_api_key"),
|
||||
decoded_token=processor.decoded_token,
|
||||
isNoneDoc=data.get("isNoneDoc"),
|
||||
index=data.get("index"),
|
||||
should_persist=should_persist,
|
||||
visibility=visibility,
|
||||
attachment_ids=data.get("attachments", []),
|
||||
agent_id=processor.agent_id,
|
||||
is_shared_usage=processor.is_shared_usage,
|
||||
shared_token=processor.shared_token,
|
||||
model_id=processor.model_id,
|
||||
model_user_id=processor.model_user_id,
|
||||
with_sse_keepalive(
|
||||
self.complete_stream(
|
||||
question=data["question"],
|
||||
agent=agent,
|
||||
conversation_id=processor.conversation_id,
|
||||
user_api_key=processor.agent_config.get("user_api_key"),
|
||||
decoded_token=processor.decoded_token,
|
||||
isNoneDoc=data.get("isNoneDoc"),
|
||||
index=data.get("index"),
|
||||
should_persist=should_persist,
|
||||
visibility=visibility,
|
||||
attachment_ids=data.get("attachments", []),
|
||||
agent_id=processor.agent_id,
|
||||
is_shared_usage=processor.is_shared_usage,
|
||||
shared_token=processor.shared_token,
|
||||
model_id=processor.model_id,
|
||||
model_user_id=processor.model_user_id,
|
||||
),
|
||||
),
|
||||
mimetype="text/event-stream",
|
||||
)
|
||||
|
||||
@@ -39,6 +39,7 @@ from application.api.v1.translator import (
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.conversations import ConversationsRepository
|
||||
from application.storage.db.session import db_readonly
|
||||
from application.streaming.sse_keepalive import with_sse_keepalive
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -353,19 +354,21 @@ def chat_completions():
|
||||
# b2b client is non-streaming), so a streaming request never
|
||||
# claims a key. This is a known, accepted limitation.
|
||||
return Response(
|
||||
_stream_response(
|
||||
helper,
|
||||
question,
|
||||
agent,
|
||||
processor,
|
||||
model_name,
|
||||
continuation,
|
||||
should_persist,
|
||||
visibility,
|
||||
strip_reasoning_leak,
|
||||
bool((data.get("stream_options") or {}).get("include_usage")),
|
||||
client_session,
|
||||
finalize_stateless_tool_pause,
|
||||
with_sse_keepalive(
|
||||
_stream_response(
|
||||
helper,
|
||||
question,
|
||||
agent,
|
||||
processor,
|
||||
model_name,
|
||||
continuation,
|
||||
should_persist,
|
||||
visibility,
|
||||
strip_reasoning_leak,
|
||||
bool((data.get("stream_options") or {}).get("include_usage")),
|
||||
client_session,
|
||||
finalize_stateless_tool_pause,
|
||||
),
|
||||
),
|
||||
mimetype="text/event-stream",
|
||||
headers={
|
||||
|
||||
@@ -2,7 +2,7 @@ import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import field_validator
|
||||
from pydantic import Field, field_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
@@ -249,6 +249,13 @@ class Settings(BaseSettings):
|
||||
# reasoning items instead.
|
||||
OPENAI_RESPONSES_STORE: bool = False
|
||||
|
||||
# Reasoning-summary verbosity requested on the Responses path when a
|
||||
# model has a reasoning_effort configured ("auto" | "concise" |
|
||||
# "detailed"). Some models/gateways reject "detailed", so the default
|
||||
# stays "auto"; stream liveness is handled by the wire-level SSE
|
||||
# keepalive (see SSE_KEEPALIVE_SECONDS) rather than summary verbosity.
|
||||
OPENAI_REASONING_SUMMARY: str = "auto"
|
||||
|
||||
# OpenAI-compatible clients can identify a logical chat with session
|
||||
# headers even though chat-completions itself has no conversation field.
|
||||
V1_SESSION_TTL_SECONDS: int = 24 * 60 * 60
|
||||
@@ -288,8 +295,10 @@ class Settings(BaseSettings):
|
||||
GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS: int = 30
|
||||
WSGI_THREADPOOL_WORKERS: int = 96
|
||||
# SSE keepalive comment cadence. Must sit under Cloudflare's 100s idle
|
||||
# close and iOS Safari's ~60s — 15s gives generous headroom.
|
||||
SSE_KEEPALIVE_SECONDS: int = 15
|
||||
# close and iOS Safari's ~60s — 15s gives generous headroom. ge=1
|
||||
# because a zero/negative interval would busy-loop or crash the
|
||||
# keepalive wrapper's queue timeout.
|
||||
SSE_KEEPALIVE_SECONDS: int = Field(default=15, ge=1)
|
||||
# Cap on simultaneous SSE connections per user. Each connection holds
|
||||
# one WSGI thread (32 per gunicorn worker) and one Redis pub/sub
|
||||
# connection. 8 covers normal multi-tab use without letting one user
|
||||
|
||||
@@ -819,7 +819,8 @@ class OpenAILLM(BaseLLM):
|
||||
else None
|
||||
)
|
||||
if effort:
|
||||
params["reasoning"] = {"effort": effort, "summary": "auto"}
|
||||
summary = settings.OPENAI_REASONING_SUMMARY or "auto"
|
||||
params["reasoning"] = {"effort": effort, "summary": summary}
|
||||
|
||||
if response_format:
|
||||
fmt = self._responses_text_format(response_format)
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Wire-level keepalive wrapper for synchronous SSE streaming routes."""
|
||||
|
||||
import contextvars
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
from typing import Generator, Optional
|
||||
|
||||
from application.core.settings import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def with_sse_keepalive(
|
||||
inner: Generator[str, None, None],
|
||||
interval_seconds: Optional[float] = None,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Yield from ``inner``, emitting ``: keepalive`` SSE comments whenever it
|
||||
stays quiet for ``interval_seconds`` (defaults to
|
||||
``settings.SSE_KEEPALIVE_SECONDS``).
|
||||
|
||||
Long tool-argument generations, summary-less reasoning stretches, and
|
||||
server-side tool runs keep the origin busy for minutes without producing
|
||||
a single chunk; proxies in front of the API then cut the byte-silent
|
||||
connection (Cloudflare read-times-out at ~100s) while the origin
|
||||
finishes into a dead socket. SSE comment frames keep bytes flowing and
|
||||
are ignored by every SSE/OpenAI-compatible client — the same
|
||||
``: keepalive`` convention used by
|
||||
``application/streaming/async_event_replay.py``.
|
||||
|
||||
``inner`` is pumped from a daemon thread so this wrapper can time out on
|
||||
silence. Callers must pass generators that don't rely on Flask's request
|
||||
context (both chat stream routes return bare generators without
|
||||
``stream_with_context``). The pump runs in a copy of the caller's
|
||||
contextvars context so request-scoped log bindings survive the thread
|
||||
hop. The queue is unbounded and the thread is a daemon: if the consumer
|
||||
disconnects mid-stream, the pump drains the upstream generator to
|
||||
completion — matching the production ASGI adapter's pre-existing
|
||||
disconnect behavior, where the WSGI iterable is drained before close().
|
||||
"""
|
||||
if interval_seconds is None:
|
||||
interval_seconds = float(settings.SSE_KEEPALIVE_SECONDS)
|
||||
frames: queue.Queue = queue.Queue()
|
||||
|
||||
def _pump() -> None:
|
||||
try:
|
||||
for item in inner:
|
||||
frames.put(("item", item))
|
||||
except BaseException as exc:
|
||||
# Logged here because after a client disconnect the consumer
|
||||
# loop is gone and nothing ever dequeues (or re-raises) this.
|
||||
logger.exception("SSE keepalive pump: upstream stream failed")
|
||||
frames.put(("error", exc))
|
||||
else:
|
||||
frames.put(("end", None))
|
||||
|
||||
ctx = contextvars.copy_context()
|
||||
threading.Thread(
|
||||
target=ctx.run, args=(_pump,), daemon=True, name="sse-keepalive-pump"
|
||||
).start()
|
||||
while True:
|
||||
try:
|
||||
kind, payload = frames.get(timeout=interval_seconds)
|
||||
except queue.Empty:
|
||||
yield ": keepalive\n\n"
|
||||
continue
|
||||
if kind == "item":
|
||||
yield payload
|
||||
elif kind == "error":
|
||||
raise payload
|
||||
else:
|
||||
return
|
||||
@@ -28,6 +28,7 @@ def _make_llm(monkeypatch, capabilities=None, store_responses=False):
|
||||
OPENAI_BASE_URL="",
|
||||
AZURE_DEPLOYMENT_NAME="dep",
|
||||
OPENAI_RESPONSES_STORE=store_responses,
|
||||
OPENAI_REASONING_SUMMARY="auto",
|
||||
),
|
||||
)
|
||||
from application.llm.openai import OpenAILLM
|
||||
@@ -388,12 +389,30 @@ def test_build_responses_params_stateless(monkeypatch):
|
||||
assert params["model"] == "gpt-5.5"
|
||||
assert params["stream"] is True
|
||||
assert params["max_output_tokens"] == 256
|
||||
# Summary verbosity comes from settings.OPENAI_REASONING_SUMMARY
|
||||
# ("auto" unless overridden).
|
||||
assert params["reasoning"] == {"effort": "high", "summary": "auto"}
|
||||
assert params["store"] is False
|
||||
assert params["include"] == ["reasoning.encrypted_content"]
|
||||
assert "previous_response_id" not in params
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_build_responses_params_summary_override(monkeypatch):
|
||||
llm = _make_llm(monkeypatch, _responses_caps(reasoning_effort="high"))
|
||||
from application.llm import openai as openai_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
openai_mod.settings, "OPENAI_REASONING_SUMMARY", "detailed", raising=False
|
||||
)
|
||||
params = llm._build_responses_params(
|
||||
"gpt-5.5", [{"role": "user", "content": []}], tools=None,
|
||||
response_format=None, previous_response_id=None, stream=True,
|
||||
kwargs={},
|
||||
)
|
||||
assert params["reasoning"] == {"effort": "high", "summary": "detailed"}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_chat_path_drops_tool_controls_when_tools_are_unavailable(monkeypatch):
|
||||
caps = ModelCapabilities(
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Wire-level keepalive for synchronous SSE streaming routes.
|
||||
|
||||
``with_sse_keepalive`` pumps the stream from a daemon thread and emits
|
||||
``: keepalive`` SSE comment frames whenever the inner generator stays
|
||||
quiet — long tool-argument generations, summary-less reasoning stretches,
|
||||
and server-side tool runs otherwise leave the connection byte-silent until
|
||||
a fronting proxy (Cloudflare, ~100s idle) cuts it while the origin is
|
||||
still working. Used by both ``/stream`` and ``/v1/chat/completions``.
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from application.streaming.sse_keepalive import with_sse_keepalive
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_keepalive_comments_emitted_while_inner_is_quiet():
|
||||
release = threading.Event()
|
||||
|
||||
def slow_inner():
|
||||
release.wait(5)
|
||||
yield "data: one\n\n"
|
||||
|
||||
gen = with_sse_keepalive(slow_inner(), 0.02)
|
||||
# The inner generator is gated on ``release``, so the first frame must
|
||||
# be a keepalive — no timing race on loaded runners.
|
||||
assert next(gen) == ": keepalive\n\n"
|
||||
release.set()
|
||||
rest = list(gen)
|
||||
assert rest[-1] == "data: one\n\n"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_items_pass_through_in_order_without_spurious_keepalives():
|
||||
out = list(with_sse_keepalive(iter(["a", "b", "c"]), 5.0))
|
||||
assert out == ["a", "b", "c"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_inner_exception_propagates_to_consumer():
|
||||
def failing_inner():
|
||||
yield "a"
|
||||
raise RuntimeError("boom")
|
||||
|
||||
frames = with_sse_keepalive(failing_inner(), 5.0)
|
||||
assert next(frames) == "a"
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
next(frames)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_empty_inner_ends_cleanly():
|
||||
assert list(with_sse_keepalive(iter([]), 5.0)) == []
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_close_returns_promptly_while_pump_drains_inner():
|
||||
"""Consumer close() (client disconnect) ends the wrapper without raising
|
||||
or blocking; the pump drains the inner generator to completion."""
|
||||
release = threading.Event()
|
||||
drained = threading.Event()
|
||||
|
||||
def inner():
|
||||
yield "a"
|
||||
release.wait(5)
|
||||
yield "b"
|
||||
drained.set()
|
||||
|
||||
gen = with_sse_keepalive(inner(), 5.0)
|
||||
assert next(gen) == "a"
|
||||
gen.close()
|
||||
release.set()
|
||||
assert drained.wait(2), "pump should drain the inner generator"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_pump_runs_in_callers_contextvars_context():
|
||||
"""Request-scoped log bindings (ContextVars) must survive the pump's
|
||||
thread hop so generation logs keep their correlation ids."""
|
||||
var = contextvars.ContextVar("sse_keepalive_test_var", default="unset")
|
||||
var.set("bound")
|
||||
seen = {}
|
||||
|
||||
def inner():
|
||||
seen["value"] = var.get()
|
||||
yield "data: x\n\n"
|
||||
|
||||
assert list(with_sse_keepalive(inner(), 5.0)) == ["data: x\n\n"]
|
||||
assert seen["value"] == "bound"
|
||||
Reference in new issue
Block a user