fix: don't call exit stack close in stream iterator as it will be cal…

…led by finally from on_complete anyway
abetlen · Oct 15, 2024 · 9ae9d80 · 9ae9d80
1 parent a69634b
commit 9ae9d80
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -159,7 +159,7 @@ async def get_event_publisher(
     request: Request,
     inner_send_chan: MemoryObjectSendStream[typing.Any],
     iterator: Iterator[typing.Any],
-    on_complete: typing.Optional[typing.Callable[[], None]] = None,
+    on_complete: typing.Optional[typing.Callable[[], typing.Awaitable[None]]] = None,
 ):
     server_settings = next(get_server_settings())
     interrupt_requests = (
@@ -182,7 +182,7 @@ async def get_event_publisher(
                 raise e
         finally:
             if on_complete:
-                on_complete()
+                await on_complete()
 
 
 def _logit_bias_tokens_to_input_ids(
@@ -326,7 +326,6 @@ async def create_completion(
         def iterator() -> Iterator[llama_cpp.CreateCompletionStreamResponse]:
             yield first_response
             yield from iterator_or_completion
-            exit_stack.aclose()
 
         send_chan, recv_chan = anyio.create_memory_object_stream(10)
         return EventSourceResponse(
@@ -518,7 +517,6 @@ async def create_chat_completion(
         def iterator() -> Iterator[llama_cpp.ChatCompletionChunk]:
             yield first_response
             yield from iterator_or_completion
-            exit_stack.aclose()
 
         send_chan, recv_chan = anyio.create_memory_object_stream(10)
         return EventSourceResponse(