mirror of
https://github.com/langgenius/dify.git
synced 2026-10-16 15:35:32 +08:00
GareArc
13ba97512e
# Conflicts: # api/controllers/openapi/auth/requirements.py # api/tests/unit_tests/controllers/openapi/test_auth_matrix.py # api/tests/unit_tests/services/enterprise/test_enterprise_service.py
103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
"""Server-built next steps (`hints`) for the /openapi/v1 surface.
|
|
|
|
A hint is a normal response field, and the handler that owns the target op composes it.
|
|
This module holds only the two mechanics a handler cannot do inline: the next page of a
|
|
list (the contract layer knows the op id and raw query) and adding hints to one kind of
|
|
SSE event in a stream the handler only holds as a generator. Neither knows any op id;
|
|
the caller passes it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import Callable, Generator, Iterable, Mapping
|
|
from typing import Any, Final, Protocol, runtime_checkable
|
|
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from controllers.openapi._models import Hint, Hinted, PageQuery
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# The core's SSE framing and its event key; it exports no name for either.
|
|
_DATA_PREFIX: Final = "data: "
|
|
_EVENT_FIELD: Final = "event"
|
|
|
|
HintBuilder = Callable[[Mapping[str, Any]], list[Hint]]
|
|
|
|
|
|
@runtime_checkable
|
|
class _ClosableStream(Protocol):
|
|
def close(self) -> None: ...
|
|
|
|
|
|
def page_after(*, page: int, limit: int, has_more: bool) -> dict[str, Any] | None:
|
|
"""The page-number input of the next page, or None on the last one."""
|
|
return PageQuery(page=page + 1, limit=limit).model_dump() if has_more else None
|
|
|
|
|
|
def cursor_after(*, last_id: str | None, has_more: bool) -> dict[str, Any] | None:
|
|
"""The cursor input of the next page, or None on the last one."""
|
|
return {"last_id": last_id} if has_more and last_id else None
|
|
|
|
|
|
def next_page_hint(
|
|
*, op: str, path_args: Mapping[str, Any], query: BaseModel | None, page: Mapping[str, Any] | None
|
|
) -> list[Hint]:
|
|
"""The call's own input with `page` (from `page_after` or `cursor_after`) on top; none on the last page."""
|
|
if page is None:
|
|
return []
|
|
params: dict[str, Any] = dict(path_args)
|
|
if query is not None:
|
|
params |= query.model_dump(exclude_none=True)
|
|
return [Hint(summary="Next page", op=op, input=params | dict(page))]
|
|
|
|
|
|
def _wanted_event(chunk: str, event: str) -> dict[str, Any] | None:
|
|
"""The parsed event when `chunk` is a `data:` frame of that kind, else None.
|
|
|
|
Only chunks that can be the wanted event are parsed — every other chunk of a long run
|
|
passes on a substring test. A frame that is not JSON, or not an object, is not the
|
|
wanted event either; the stream owns its bytes and this layer never breaks it.
|
|
"""
|
|
if not chunk.startswith(_DATA_PREFIX) or event not in chunk:
|
|
return None
|
|
try:
|
|
parsed = json.loads(chunk[len(_DATA_PREFIX) :])
|
|
except ValueError:
|
|
return None
|
|
if not isinstance(parsed, dict) or parsed.get(_EVENT_FIELD) != event:
|
|
return None
|
|
return parsed
|
|
|
|
|
|
def attach_stream_hints(events: Iterable[str], *, event: str, build: HintBuilder) -> Generator[str]:
|
|
"""Yield the source SSE chunks, adding a top-level `hints` list to every `event` that `build` hints.
|
|
|
|
`event: ping` chunks, other event kinds, events `build` returns nothing for, and events
|
|
`build` cannot read are yielded as the original string so the wire bytes stay identical.
|
|
Closing this generator closes the source (the run stream is a `RateLimitGenerator` that
|
|
releases its slot on close).
|
|
"""
|
|
try:
|
|
for chunk in events:
|
|
parsed = _wanted_event(chunk, event)
|
|
if parsed is None:
|
|
yield chunk
|
|
continue
|
|
try:
|
|
hints = build(parsed)
|
|
except ValidationError:
|
|
logger.warning("%s event without the fields its hint needs; passed through.", event, exc_info=True)
|
|
yield chunk
|
|
continue
|
|
if not hints:
|
|
yield chunk
|
|
continue
|
|
parsed |= Hinted(hints=hints).model_dump(exclude_none=True)
|
|
yield f"{_DATA_PREFIX}{json.dumps(parsed, ensure_ascii=False)}\n\n"
|
|
finally:
|
|
if isinstance(events, _ClosableStream):
|
|
events.close()
|