Chat Streaming endpoint bare-minimum PoC

Author: TamiTakamiya

ansible_ai_connect/ai/api/streaming_chat.py

@@ -0,0 +1,40 @@
+import logging
+
+import aiohttp
+from django.apps import apps
+from django.http import StreamingHttpResponse
+from rest_framework.decorators import parser_classes
+from rest_framework.parsers import JSONParser
+from rest_framework.views import APIView
+
+from ansible_ai_connect.ai.api.model_pipelines.pipelines import ModelPipelineChatBot
+
+logger = logging.getLogger(__name__)
+
+
+class StreamingChat(APIView):
+
+    def __init__(self):
+        self.llm = apps.get_app_config("ai").get_model_pipeline(ModelPipelineChatBot)
+
+    async def call_chatservice(self, request):
+        async with aiohttp.ClientSession(raise_for_status=True) as session:
+            headers = {
+                "Content-Type": "application/json",
+                "Accept": "application/json,text/event-stream",
+            }
+            async with session.post(
+                self.llm.config.inference_url + "/v1/streaming_query",
+                json=request.data,
+                headers=headers,
+            ) as r:
+                async for chunk in r.content:
+                    logger.debug(chunk)
+                    yield chunk
+
+    @parser_classes([JSONParser])
+    def post(self, request):
+        return StreamingHttpResponse(
+            self.call_chatservice(request),
+            content_type="text/event-stream",
+        )

ansible_ai_connect/ai/api/versions/v1/ai/urls.py

@@ -24,4 +24,5 @@
     path("generations/role/", views.GenerationRole.as_view(), name="generations/role"),
     path("feedback/", views.Feedback.as_view(), name="feedback"),
     path("chat/", views.Chat.as_view(), name="chat"),
+    path("streaming_chat/", views.StreamingChat.as_view(), name="streaming_chat"),
 ]

ansible_ai_connect/ai/api/versions/v1/ai/views.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from ansible_ai_connect.ai.api.streaming_chat import StreamingChat
 from ansible_ai_connect.ai.api.views import (
     Chat,
     Completions,
@@ -30,4 +31,5 @@
     "GenerationRole",
     "Feedback",
     "Chat",
+    "StreamingChat",
 ]

ansible_ai_connect/main/settings/base.py

@@ -55,6 +55,7 @@
 # Application definition
 
 INSTALLED_APPS = [
+    "daphne",
     "django.contrib.admin",
     "django.contrib.auth",
     "django.contrib.contenttypes",
@@ -294,7 +295,7 @@ def is_ssl_enabled(value: str) -> bool:
         },
     },
     "handlers": {
-        "console": {"class": "logging.StreamHandler", "formatter": "simple", "level": "INFO"},
+        "console": {"class": "logging.StreamHandler", "formatter": "simple", "level": "DEBUG"},
     },
     "loggers": {
         "django": {
@@ -334,6 +335,11 @@ def is_ssl_enabled(value: str) -> bool:
             "level": "INFO",
             "propagate": False,
         },
+        "ansible_ai_connect.ai.api.streaming_chat": {
+            "handlers": ["console"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
     },
     "root": {
         "handlers": ["console"],
@@ -358,6 +364,7 @@ def is_ssl_enabled(value: str) -> bool:
 ]
 
 WSGI_APPLICATION = "ansible_ai_connect.main.wsgi.application"
+ASGI_APPLICATION = "ansible_ai_connect.main.asgi.application"
 
 # Database
 # https://docs.djangoproject.com/en/4.1/ref/settings/#databases

ansible_ai_connect_chatbot/src/useChatbot/useChatbot.ts

@@ -369,13 +369,12 @@ export const useChatbot = () => {
             method: "POST",
             headers: {
               "Content-Type": "application/json",
+              Accept: "application/json,text/event-stream",
               "X-CSRFToken": csrfToken!,
             },
             body: JSON.stringify(chatRequest),
             async onopen(resp: any) {
-              if (resp.ok && resp.status === 200) {
-                setIsLoading(false);
-              } else if (
+              if (
                 resp.status >= 400 &&
                 resp.status < 500 &&
                 resp.status !== 429
@@ -394,6 +393,7 @@ export const useChatbot = () => {
                   setConversationId(message.data.conversation_id);
                 }
               } else if (message.event === "token") {
+                setIsLoading(false);
                 appendMessageChunk(message.data.token);
               } else if (message.event === "end") {
                 if (message.data.referenced_documents.length > 0) {

requirements-aarch64.txt

@@ -7,7 +7,9 @@
 aiohappyeyeballs==2.3.5
     # via aiohttp
 aiohttp==3.10.11
-    # via langchain
+    # via
+    #   -r requirements.in
+    #   langchain
 aiosignal==1.3.1
     # via aiohttp
 annotated-types==0.6.0
@@ -29,14 +31,22 @@ anyio==4.6.2.post1
 argparse==1.4.0
     # via uwsgi-readiness-check
 asgiref==3.8.1
-    # via django
+    # via
+    #   daphne
+    #   django
 asttokens==2.4.1
     # via stack-data
 attrs==23.2.0
     # via
     #   aiohttp
     #   jsonschema
     #   referencing
+    #   service-identity
+    #   twisted
+autobahn==24.4.2
+    # via daphne
+automat==24.8.1
+    # via twisted
 backcall==0.2.0
     # via ipython
 backoff==2.2.1
@@ -68,13 +78,19 @@ charset-normalizer==3.3.2
     # via requests
 click==8.1.7
     # via black
+constantly==23.10.4
+    # via twisted
 cryptography==43.0.1
     # via
     #   -r requirements.in
     #   ansible-core
+    #   autobahn
     #   jwcrypto
     #   pyopenssl
+    #   service-identity
     #   social-auth-core
+daphne==4.1.2
+    # via -r requirements.in
 decorator==5.1.1
     # via ipython
 defusedxml==0.8.0rc2
@@ -167,13 +183,21 @@ httpx==0.27.2
     # via
     #   langsmith
     #   ollama
+hyperlink==21.0.0
+    # via
+    #   autobahn
+    #   twisted
 idna==3.7
     # via
     #   -r requirements.in
     #   anyio
     #   httpx
+    #   hyperlink
     #   requests
+    #   twisted
     #   yarl
+incremental==24.7.2
+    # via twisted
 inflection==0.5.1
     # via drf-spectacular
 ipython==8.10.0
@@ -307,10 +331,12 @@ pyasn1==0.6.0
     #   oauth2client
     #   pyasn1-modules
     #   rsa
+    #   service-identity
 pyasn1-modules==0.4.0
     # via
     #   google-auth
     #   oauth2client
+    #   service-identity
 pycparser==2.21
     # via cffi
 pydantic==2.9.2
@@ -336,6 +362,7 @@ pyopenssl==24.2.1
     # via
     #   -r requirements.in
     #   pydrive2
+    #   twisted
 pyparsing==3.1.2
     # via httplib2
 pyrfc3339==1.1
@@ -410,6 +437,8 @@ segment-analytics-python==2.2.2
     # via -r requirements.in
 semver==3.0.2
     # via launchdarkly-server-sdk
+service-identity==24.2.0
+    # via twisted
 six==1.16.0
     # via
     #   asttokens
@@ -455,6 +484,10 @@ traitlets==5.14.3
     # via
     #   ipython
     #   matplotlib-inline
+twisted[tls]==24.11.0
+    # via daphne
+txaio==23.1.1
+    # via autobahn
 typing-extensions==4.11.0
     # via
     #   django-test-migrations
@@ -464,6 +497,7 @@ typing-extensions==4.11.0
     #   pydantic
     #   pydantic-core
     #   sqlalchemy
+    #   twisted
 uritemplate==4.1.1
     # via
     #   drf-spectacular
@@ -490,3 +524,8 @@ yamllint==1.35.1
     # via ansible-lint
 yarl==1.17.2
     # via aiohttp
+zope-interface==7.2
+    # via twisted
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools

requirements-x86_64.txt

@@ -7,7 +7,9 @@
 aiohappyeyeballs==2.3.5
     # via aiohttp
 aiohttp==3.10.11
-    # via langchain
+    # via
+    #   -r requirements.in
+    #   langchain
 aiosignal==1.3.1
     # via aiohttp
 annotated-types==0.6.0
@@ -29,14 +31,22 @@ anyio==4.6.2.post1
 argparse==1.4.0
     # via uwsgi-readiness-check
 asgiref==3.8.1
-    # via django
+    # via
+    #   daphne
+    #   django
 asttokens==2.4.1
     # via stack-data
 attrs==23.2.0
     # via
     #   aiohttp
     #   jsonschema
     #   referencing
+    #   service-identity
+    #   twisted
+autobahn==24.4.2
+    # via daphne
+automat==24.8.1
+    # via twisted
 backcall==0.2.0
     # via ipython
 backoff==2.2.1
@@ -68,13 +78,19 @@ charset-normalizer==3.3.2
     # via requests
 click==8.1.7
     # via black
+constantly==23.10.4
+    # via twisted
 cryptography==43.0.1
     # via
     #   -r requirements.in
     #   ansible-core
+    #   autobahn
     #   jwcrypto
     #   pyopenssl
+    #   service-identity
     #   social-auth-core
+daphne==4.1.2
+    # via -r requirements.in
 decorator==5.1.1
     # via ipython
 defusedxml==0.8.0rc2
@@ -167,13 +183,21 @@ httpx==0.27.2
     # via
     #   langsmith
     #   ollama
+hyperlink==21.0.0
+    # via
+    #   autobahn
+    #   twisted
 idna==3.7
     # via
     #   -r requirements.in
     #   anyio
     #   httpx
+    #   hyperlink
     #   requests
+    #   twisted
     #   yarl
+incremental==24.7.2
+    # via twisted
 inflection==0.5.1
     # via drf-spectacular
 ipython==8.10.0
@@ -307,10 +331,12 @@ pyasn1==0.6.0
     #   oauth2client
     #   pyasn1-modules
     #   rsa
+    #   service-identity
 pyasn1-modules==0.4.0
     # via
     #   google-auth
     #   oauth2client
+    #   service-identity
 pycparser==2.21
     # via cffi
 pydantic==2.9.2
@@ -336,6 +362,7 @@ pyopenssl==24.2.1
     # via
     #   -r requirements.in
     #   pydrive2
+    #   twisted
 pyparsing==3.1.2
     # via httplib2
 pyrfc3339==1.1
@@ -410,6 +437,8 @@ segment-analytics-python==2.2.2
     # via -r requirements.in
 semver==3.0.2
     # via launchdarkly-server-sdk
+service-identity==24.2.0
+    # via twisted
 six==1.16.0
     # via
     #   asttokens
@@ -455,6 +484,10 @@ traitlets==5.14.3
     # via
     #   ipython
     #   matplotlib-inline
+twisted[tls]==24.11.0
+    # via daphne
+txaio==23.1.1
+    # via autobahn
 typing-extensions==4.11.0
     # via
     #   django-test-migrations
@@ -464,6 +497,7 @@ typing-extensions==4.11.0
     #   pydantic
     #   pydantic-core
     #   sqlalchemy
+    #   twisted
 uritemplate==4.1.1
     # via
     #   drf-spectacular
@@ -490,3 +524,8 @@ yamllint==1.35.1
     # via ansible-lint
 yarl==1.17.2
     # via aiohttp
+zope-interface==7.2
+    # via twisted
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools