fix: added missing exit_stack.close() to /v1/chat/completions (abetlen#1796)

Ian321 · web-flow · commit 073b7e421e41 · 2024-12-06T05:02:34.000-05:00
* fix: added missing exit_stack.close() to /v1/chat/completions

* fix: added missing exit_stack.close() to /v1/completions
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -314,10 +314,14 @@ async def create_completion(
         else:
             kwargs["logits_processor"].extend(_min_tokens_logits_processor)
 
-    iterator_or_completion: Union[
-        llama_cpp.CreateCompletionResponse,
-        Iterator[llama_cpp.CreateCompletionStreamResponse],
-    ] = await run_in_threadpool(llama, **kwargs)
+    try:
+        iterator_or_completion: Union[
+            llama_cpp.CreateCompletionResponse,
+            Iterator[llama_cpp.CreateCompletionStreamResponse],
+        ] = await run_in_threadpool(llama, **kwargs)
+    except Exception as err:
+        exit_stack.close()
+        raise err
 
     if isinstance(iterator_or_completion, Iterator):
         # EAFP: It's easier to ask for forgiveness than permission
@@ -344,6 +348,7 @@ def iterator() -> Iterator[llama_cpp.CreateCompletionStreamResponse]:
             ping_message_factory=_ping_message_factory,
         )
     else:
+        exit_stack.close()
         return iterator_or_completion
 
 
@@ -508,9 +513,13 @@ async def create_chat_completion(
         else:
             kwargs["logits_processor"].extend(_min_tokens_logits_processor)
 
-    iterator_or_completion: Union[
-        llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk]
-    ] = await run_in_threadpool(llama.create_chat_completion, **kwargs)
+    try:
+        iterator_or_completion: Union[
+            llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk]
+        ] = await run_in_threadpool(llama.create_chat_completion, **kwargs)
+    except Exception as err:
+        exit_stack.close()
+        raise err
 
     if isinstance(iterator_or_completion, Iterator):
         # EAFP: It's easier to ask for forgiveness than permission
diff --git a/llama_cpp/server/errors.py b/llama_cpp/server/errors.py
@@ -134,8 +134,6 @@ def error_message_wrapper(
         ] = None,
     ) -> Tuple[int, ErrorResponse]:
         """Wraps error message in OpenAI style error response"""
-        print(f"Exception: {str(error)}", file=sys.stderr)
-        traceback.print_exc(file=sys.stderr)
         if body is not None and isinstance(
             body,
             (
@@ -149,6 +147,10 @@ def error_message_wrapper(
                 if match is not None:
                     return callback(body, match)
 
+        # Only print the trace on unexpected exceptions
+        print(f"Exception: {str(error)}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
+
         # Wrap other errors as internal server error
         return 500, ErrorResponse(
             message=str(error),
