feat: add disable_ping_events flag (abetlen#1257)

khimaros · xhedit · commit c9a6e2396792 · 2024-04-30T12:50:55.000Z
for backward compatibility, this is false by default it can be set to true to disable EventSource pings which are not supported by some OpenAI clients. fixes abetlen#1256
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -87,6 +87,13 @@ def get_llama_proxy():
             llama_outer_lock.release()
 
 
+_ping_message_factory = None
+
+def set_ping_message_factory(factory):
+   global  _ping_message_factory
+   _ping_message_factory = factory
+
+
 def create_app(
     settings: Settings | None = None,
     server_settings: ServerSettings | None = None,
@@ -138,6 +145,9 @@ def create_app(
     assert model_settings is not None
     set_llama_proxy(model_settings=model_settings)
 
+    if server_settings.disable_ping_events:
+        set_ping_message_factory(lambda: bytes())
+
     return app
 
 
@@ -302,6 +312,7 @@ def iterator() -> Iterator[llama_cpp.CreateCompletionStreamResponse]:
                 iterator=iterator(),
             ),
             sep="\n",
+            ping_message_factory=_ping_message_factory,
         )
     else:
         return iterator_or_completion
@@ -470,6 +481,7 @@ def iterator() -> Iterator[llama_cpp.ChatCompletionChunk]:
                 iterator=iterator(),
             ),
             sep="\n",
+            ping_message_factory=_ping_message_factory,
         )
     else:
         return iterator_or_completion
diff --git a/llama_cpp/server/settings.py b/llama_cpp/server/settings.py
@@ -195,6 +195,10 @@ class ServerSettings(BaseSettings):
         default=True,
         description="Whether to interrupt requests when a new request is received.",
     )
+    disable_ping_events: bool = Field(
+        default=False,
+        description="Disable EventSource pings (may be needed for some clients).",
+    )
 
 
 class Settings(ServerSettings, ModelSettings):

Original file line number	Diff line number	Diff line change
`@@ -195,6 +195,10 @@ class ServerSettings(BaseSettings):`
`195`	`195`	`default=True,`
`196`	`196`	`description="Whether to interrupt requests when a new request is received.",`
`197`	`197`	`)`
	`198`	`+ disable_ping_events: bool = Field(`
	`199`	`+ default=False,`
	`200`	`+ description="Disable EventSource pings (may be needed for some clients).",`
	`201`	`+ )`
`198`	`202`
`199`	`203`
`200`	`204`	`class Settings(ServerSettings, ModelSettings):`