Option to enable chat streaming (#1516)

* Option to enable chat streaming

* Add stream to HttpConfiguration

* Refactor based on a review comment

Author: TamiTakamiya

ansible_ai_connect/ai/api/model_pipelines/http/configuration.py

@@ -41,11 +41,14 @@ def __init__(
         timeout: Optional[int],
         enable_health_check: Optional[bool],
         verify_ssl: bool,
+        stream: bool = False,
     ):
         super().__init__(inference_url, model_id, timeout, enable_health_check)
         self.verify_ssl = verify_ssl
+        self.stream = stream
 
     verify_ssl: bool
+    stream: bool
 
 
 @Register(api_type="http")
@@ -60,10 +63,12 @@ def __init__(self, **kwargs):
                 timeout=kwargs["timeout"],
                 enable_health_check=kwargs["enable_health_check"],
                 verify_ssl=kwargs["verify_ssl"],
+                stream=kwargs["stream"],
             ),
         )
 
 
 @Register(api_type="http")
 class HttpConfigurationSerializer(BaseConfigSerializer):
     verify_ssl = serializers.BooleanField(required=False, default=True)
+    stream = serializers.BooleanField(required=False, default=False)

ansible_ai_connect/main/settings/legacy.py

@@ -189,6 +189,7 @@ def load_from_env_vars():
             "inference_url": chatbot_service_url or "http://localhost:8000",
             "model_id": chatbot_service_model_id or "granite3-8b",
             "verify_ssl": model_service_verify_ssl,
+            "stream": False,
         },
     }
 

ansible_ai_connect/main/templates/chatbot/index.html

@@ -21,5 +21,6 @@
     <div id="user_name" hidden>{{user_name}}</div>
     <div id="bot_name" hidden>{{bot_name}}</div>
     <div id="debug" hidden>{{debug}}</div>
+    <div id="stream" hidden>{{stream}}</div>
 {% endblock content %}
 </html>

ansible_ai_connect/main/tests/test_views.py

@@ -18,13 +18,15 @@
 from http import HTTPStatus
 from textwrap import dedent
 
+from django.apps import apps
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import AnonymousUser, Group
 from django.http import HttpResponseRedirect
 from django.test import RequestFactory, TestCase, modify_settings, override_settings
 from django.urls import reverse
 from rest_framework.test import APITransactionTestCase
 
+from ansible_ai_connect.ai.api.model_pipelines.pipelines import ModelPipelineChatBot
 from ansible_ai_connect.main.settings.base import SOCIAL_AUTH_OIDC_KEY
 from ansible_ai_connect.main.views import LoginView
 from ansible_ai_connect.test_utils import (
@@ -339,10 +341,21 @@ def test_chatbot_view_with_rh_user(self):
         self.assertContains(r, TestChatbotView.CHATBOT_PAGE_TITLE)
         self.assertContains(r, self.rh_user.username)
         self.assertContains(r, '<div id="debug" hidden>false</div>')
+        self.assertContains(r, '<div id="stream" hidden>false</div>')
 
     @override_settings(CHATBOT_DEBUG_UI=True)
     def test_chatbot_view_with_debug_ui(self):
         self.client.force_login(user=self.rh_user)
         r = self.client.get(reverse("chatbot"), {"debug": "true"})
         self.assertEqual(r.status_code, HTTPStatus.OK)
         self.assertContains(r, '<div id="debug" hidden>true</div>')
+
+    def test_chatbot_view_with_streaming_enabled(self):
+        llm: ModelPipelineChatBot = apps.get_app_config("ai").get_model_pipeline(
+            ModelPipelineChatBot
+        )
+        llm.config.stream = True
+        self.client.force_login(user=self.rh_user)
+        r = self.client.get(reverse("chatbot"), {"stream": "true"})
+        self.assertEqual(r.status_code, HTTPStatus.OK)
+        self.assertContains(r, '<div id="stream" hidden>true</div>')

ansible_ai_connect/main/views.py

@@ -121,12 +121,21 @@ class ChatbotView(ProtectedTemplateView):
         IsRHInternalUser | IsTestUser,
     ]
 
+    llm: ModelPipelineChatBot
+    chatbot_enabled: bool
+
+    def __init__(self):
+        super().__init__()
+        self.llm = apps.get_app_config("ai").get_model_pipeline(ModelPipelineChatBot)
+        self.chatbot_enabled = (
+            self.llm.config.inference_url
+            and self.llm.config.model_id
+            and settings.CHATBOT_DEFAULT_PROVIDER
+        )
+
     def get(self, request):
         # Open the chatbot page when the chatbot service is configured.
-        llm: ModelPipelineChatBot = apps.get_app_config("ai").get_model_pipeline(
-            ModelPipelineChatBot
-        )
-        if llm.config.inference_url and llm.config.model_id and settings.CHATBOT_DEFAULT_PROVIDER:
+        if self.chatbot_enabled:
             return super().get(request)
 
         # Otherwise, redirect to the home page.
@@ -139,6 +148,7 @@ def get_context_data(self, **kwargs):
         if user and user.is_authenticated:
             context["user_name"] = user.username
         context["debug"] = "true" if settings.CHATBOT_DEBUG_UI else "false"
+        context["stream"] = "true" if self.llm.config.stream else "false"
 
         return context