fix(chat-ui): avoid double LLM call and stale history on errors

fix

Author: Pouyanpi

nemoguardrails/server/api.py

@@ -67,7 +67,7 @@ def __init__(self, *args, **kwargs):
         # Initialize custom attributes
         self.default_config_id: Optional[str] = None
         self.rails_config_path: str = ""
-        self.disable_chat_ui: bool = False
+        self.disable_chat_ui: bool = os.getenv("NEMO_GUARDRAILS_DISABLE_CHAT_UI", "false").lower() == "true"
         self.auto_reload: bool = False
         self.stop_signal: bool = False
         self.single_config_mode: bool = False
@@ -186,8 +186,6 @@ async def lifespan(app: GuardrailsApp):
 app.single_config_mode = False
 app.single_config_id = None
 
-app.disable_chat_ui = os.getenv("NEMO_GUARDRAILS_DISABLE_CHAT_UI", "false").lower() == "true"
-
 
 @app.get(
     "/v1/rails/configs",
@@ -725,7 +723,7 @@ class GuardrailsConfigurationError(Exception):
 
     @app.get("/")
     async def root_redirect():
-        return RedirectResponse(url="/chat")
+        return RedirectResponse(url="chat")
 
 else:
     if not app.disable_chat_ui and mount_chainlit is None:

nemoguardrails/server/app.py

@@ -115,25 +115,28 @@ async def on_message(message: cl.Message):
 
     try:
         llm_rails = await _get_rails([config_id])
-    except ValueError:
+    except Exception:
         log.exception("Failed to load rails config '%s'", config_id)
+        messages.pop()
+        cl.user_session.set("messages", messages)
         await cl.Message(content=f"Error loading guardrails configuration '{config_id}'. Check server logs.").send()
         return
 
     response_msg = cl.Message(content="")
     await response_msg.send()
 
     full_response = ""
+    streaming_unsupported = False
     try:
         try:
             async for chunk in llm_rails.stream_async(messages=messages):
                 if isinstance(chunk, str) and chunk:
                     full_response += chunk
                     await response_msg.stream_token(chunk)
         except StreamingNotSupportedError:
-            full_response = ""
+            streaming_unsupported = True
 
-        if not full_response:
+        if streaming_unsupported:
             result = await llm_rails.generate_async(messages=messages)
             full_response = result.get("content", str(result)) if isinstance(result, dict) else str(result)
             response_msg.content = full_response
@@ -144,5 +147,7 @@ async def on_message(message: cl.Message):
 
     except Exception:
         log.exception("Error generating response for config '%s'", config_id)
+        messages.pop()
+        cl.user_session.set("messages", messages)
         response_msg.content = f"An error occurred for configuration '{config_id}'. Check server logs."
         await response_msg.update()