jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

API: Fix disconnect handling on streaming responses

Browse source 
Starlette's StreamingResponse has an issue where it yields after
a request has disconnected. A bugfix to starlette will fix this
issue, but FastAPI uses starlette <= 0.36 which isn't ideal.

Therefore, switch back to sse-starlette which handles these disconnects
correctly.

Also don't try yielding after the request is disconnected. Just return
out of the generator instead.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri 2024-03-10 17:31:47 -04:00
parent 6b4f100db2
commit d45e847c7a
2 changed files with 21 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

7
common/utils.py
View file

@ -29,7 +29,7 @@ def get_generator_error(message: str, exc_info: bool = True):
generator_error = handle_request_error(message)
return get_sse_packet(generator_error.model_dump_json())
return generator_error.model_dump_json()
def handle_request_error(message: str, exc_info: bool = True):
@ -50,11 +50,6 @@ def handle_request_error(message: str, exc_info: bool = True):
return request_error
def get_sse_packet(json_data: str):
"""Get an SSE packet."""
return f"data: {json_data}\n\n"
def unwrap(wrapped, default=None):
"""Unwrap function for Optionals."""
if wrapped is None: