mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/lmproto: await streaming AsyncResponse close
Support both sync and async close/aclose so stream cleanup does not leave un-awaited coroutines.
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import inspect
|
||||||
from typing import TYPE_CHECKING, Literal, overload
|
from typing import TYPE_CHECKING, Literal, overload
|
||||||
|
|
||||||
import msgspec
|
import msgspec
|
||||||
@@ -24,6 +25,18 @@ if TYPE_CHECKING:
|
|||||||
from niquests.models import AsyncResponse
|
from niquests.models import AsyncResponse
|
||||||
|
|
||||||
|
|
||||||
|
async def _close_async_response(resp: AsyncResponse) -> None:
|
||||||
|
"""Close a streaming response; support sync or async close()."""
|
||||||
|
for name in ("aclose", "close"):
|
||||||
|
method = getattr(resp, name, None)
|
||||||
|
if not callable(method):
|
||||||
|
continue
|
||||||
|
result = method()
|
||||||
|
if inspect.isawaitable(result):
|
||||||
|
await result
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
class BaseOpenAIClient:
|
class BaseOpenAIClient:
|
||||||
session: AsyncSession
|
session: AsyncSession
|
||||||
encoder: msgspec.json.Encoder
|
encoder: msgspec.json.Encoder
|
||||||
@@ -47,13 +60,7 @@ class BaseOpenAIClient:
|
|||||||
if line.startswith(b"data: "):
|
if line.startswith(b"data: "):
|
||||||
yield self.chunk_decoder.decode(line[6:])
|
yield self.chunk_decoder.decode(line[6:])
|
||||||
finally:
|
finally:
|
||||||
aclose = getattr(resp, "aclose", None)
|
await _close_async_response(resp)
|
||||||
if callable(aclose):
|
|
||||||
await aclose() # pyright: ignore[reportGeneralTypeIssues]
|
|
||||||
else:
|
|
||||||
close = getattr(resp, "close", None)
|
|
||||||
if callable(close):
|
|
||||||
_ = close()
|
|
||||||
|
|
||||||
|
|
||||||
class OpenAIClient(BaseOpenAIClient):
|
class OpenAIClient(BaseOpenAIClient):
|
||||||
|
|||||||
Reference in New Issue
Block a user