Repository: Quansight/ragna (python)

Issue:
Use JSONL for streaming rather then SSE
Looking at the builtin assistants, there are currently three ways of streaming:

1. Using SSE
   - [Anthropic](https://github.com/Quansight/ragna/blob/4a667f9aced1970f8eab8aa9e01aaf8c67b05db0/ragna/assistants/_anthropic.py#L38)
   - [OpenAI](https://github.com/Quansight/ragna/blob/4a667f9aced1970f8eab8aa9e01aaf8c67b05db0/ragna/assistants/_openai.py#L38)
2. Using streaming responses
   1. with one object per line (JSONL):
      - [Cohere](https://github.com/Quansight/ragna/blob/4a667f9aced1970f8eab8aa9e01aaf8c67b05db0/ragna/assistants/_cohere.py#L59-L60)
   2. with streaming JSON: 
      - [Google](https://github.com/Quansight/ragna/blob/4a667f9aced1970f8eab8aa9e01aaf8c67b05db0/ragna/assistants/_google.py#L64)

Out of those, I think 2.ii. is by far the worst as it is the hardest to parse. When we implemented streaming for Ragna, I went with 1., since that is what I was familiar with and seen other providers to as well. At that point, I was not aware of streaming JSONL (2.i.) as an alternative.

2.i. has the upside, that one does not need an extra library to handle as it is "required" for the other two cases (of course we could write the logic ourselves as well).

Due to its simplicity I would actually prefer JSONL streaming. We should make a decision here before we release streaming with `0.2.0`.

Pull request:
use JSONL streaming over SSE
Closes #356.

Patch:
diff --git a/README.md b/README.md
index a60c1b3..807701e 100644
--- a/README.md
+++ b/README.md
@@ -50,6 +50,7 @@ A RAG orchestration framework.
 
 Check out the Ragna documentation to get started:
 
+- [Overview video](https://youtu.be/xuNnjFS1XqQ?si=AssQq_1CF2o1jv3M)
 - [Installation steps](https://ragna.chat/en/stable/install/)
 - Tutorials
   - [Python API](https://ragna.chat/en/stable/tutorials/python-api/)
diff --git a/docs/examples/gallery_streaming.py b/docs/examples/gallery_streaming.py
index cdba77c..b0b5b3f 100644
--- a/docs/examples/gallery_streaming.py
+++ b/docs/examples/gallery_streaming.py
@@ -141,41 +141,38 @@ chat = (
 client.post(f"/chats/{chat['id']}/prepare").raise_for_status()
 
 # %%
-# Streaming the response is performed with
-# [server-sent events (SSE)](https://en.wikipedia.org/wiki/Server-sent_events).
+# Streaming the response is performed with [JSONL](https://jsonlines.org/). Each line
+# in the response is valid JSON and corresponds to one chunk.
 
-import httpx_sse
 import json
 
-chunks = []
-with httpx_sse.connect_sse(
-    client,
+
+with client.stream(
     "POST",
     f"/chats/{chat['id']}/answer",
     json={"prompt": "What is Ragna?", "stream": True},
-) as event_source:
-    for sse in event_source.iter_sse():
-        chunks.append(json.loads(sse.data))
+) as response:
+    chunks = [json.loads(data) for data in response.iter_lines()]
 
 # %%
-# The first event contains the full message object including the sources along the first
+# The first chunk contains the full message object including the sources along the first
 # chunk of the content.
 
 print(len(chunks))
 print(json.dumps(chunks[0], indent=2))
 
 # %%
-# Subsequent events no longer contain the sources.
+# Subsequent chunks no longer contain the sources.
 
 print(json.dumps(chunks[1], indent=2))
 
 # %%
-# Joining the chunks together results in the full message.
+# Joining the content of the chunks together results in the full message.
 
 print("".join(chunk["content"] for chunk in chunks))
 
 # %%
-# Before we close the tutorial, let's stop the REST API and have a look at what would
+# Before we close the example, let's stop the REST API and have a look at what would
 # have printed in the terminal if we had started it the regular way.
 
 rest_api.stop()
diff --git a/docs/index.md b/docs/index.md
index fbeb289..e2d365b 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -2,14 +2,14 @@
 title: Home
 ---
 
-# ![Ragna logo](assets/brand/logo-lockup-horizontal/logo-lockup-horizontal.png)
-
-Ragna is an open source RAG orchestration framework.
+# Ragna — Open source RAG orchestration framework
 
 With an intuitive API for quick experimentation and built-in tools for creating
 production-ready application, you can quickly leverage Large Language Models (LLMs) for
 your work.
 
+<iframe width="560" height="315" src="https://www.youtube.com/embed/xuNnjFS1XqQ?si=zKMVXSSjQWZvYvkf" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen></iframe>
+
 !!! note
 
     Ragna is a framework for experimenting with LLMs. It does not install one by default. To use Ragna you either need to provide API keys to a cloud LLM API, e.g. OpenAI, Anthropic, etc. or run your own LLM.
diff --git a/docs/references/faq.md b/docs/references/faq.md
index c89376a..d6ebfee 100644
--- a/docs/references/faq.md
+++ b/docs/references/faq.md
@@ -4,8 +4,9 @@
 
 !!! tip "TL;DR"
 
-    Ragna is the only tool out there that specializes in orchestrating RAG uses cases with arbitrary components,
-    as well as offers a Python API, a REST API, and a web UI for that.
+    Ragna is the only tool out there that specializes in orchestrating RAG use cases
+    with arbitrary components, as well as offering a Python API, a REST API, and a web
+    UI for that.
 
 !!! note
 
diff --git a/pyproject.toml b/pyproject.toml
index 574813c..705ead0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,6 @@ dependencies = [
     "emoji",
     "fastapi",
     "httpx",
-    "httpx-sse",
     "importlib_metadata>=4.6; python_version<'3.10'",
     "packaging",
     "panel==1.3.8",
@@ -38,7 +37,6 @@ dependencies = [
     "questionary",
     "rich",
     "sqlalchemy>=2",
-    "sse-starlette",
     "starlette",
     "tomlkit",
     "typer",
@@ -49,13 +47,14 @@ dynamic = ["version"]
 [project.urls]
 Homepage = "https://ragna.chat"
 Documentation = "https://ragna.chat"
-Changelog = "https://ragna.chat/en/latest/references/changelog/"
+Changelog = "https://ragna.chat/en/stable/references/release-notes/"
 Repository = "https://github.com/Quansight/ragna"
 
 [project.optional-dependencies]
 # to update the array below, run scripts/update_optional_dependencies.py
 all = [
     "chromadb>=0.4.13",
+    "httpx_sse",
     "ijson",
     "lancedb>=0.2",
     "pyarrow",
diff --git a/ragna/assistants/_anthropic.py b/ragna/assistants/_anthropic.py
index de09e53..f5f4c53 100644
--- a/ragna/assistants/_anthropic.py
+++ b/ragna/assistants/_anthropic.py
@@ -1,9 +1,7 @@
 import json
 from typing import AsyncIterator, cast
 
-import httpx_sse
-
-from ragna.core import RagnaException, Source
+from ragna.core import PackageRequirement, RagnaException, Requirement, Source
 
 from ._api import ApiAssistant
 
@@ -12,6 +10,10 @@ class AnthropicApiAssistant(ApiAssistant):
     _API_KEY_ENV_VAR = "ANTHROPIC_API_KEY"
     _MODEL: str
 
+    @classmethod
+    def _extra_requirements(cls) -> list[Requirement]:
+        return [PackageRequirement("httpx_sse")]
+
     @classmethod
     def display_name(cls) -> str:
         return f"Anthropic/{cls._MODEL}"
@@ -29,6 +31,8 @@ class AnthropicApiAssistant(ApiAssistant):
     async def _call_api(
         self, prompt: str, sources: list[Source], *, max_new_tokens: int
     ) -> AsyncIterator[str]:
+        import httpx_sse
+
         # See https://docs.anthropic.com/claude/reference/streaming
         async with httpx_sse.aconnect_sse(
             self._client,
@@ -68,6 +72,10 @@ class ClaudeInstant(AnthropicApiAssistant):
     !!! info "Required environment variables"
 
         - `ANTHROPIC_API_KEY`
+
+    !!! info "Required packages"
+
+        - `httpx_sse`
     """
 
     _MODEL = "claude-instant-1"
@@ -79,6 +87,10 @@ class Claude(AnthropicApiAssistant):
     !!! info "Required environment variables"
 
         - `ANTHROPIC_API_KEY`
+
+    !!! info "Required packages"
+
+        - `httpx_sse`
     """
 
     _MODEL = "claude-2"
diff --git a/ragna/assistants/_api.py b/ragna/assistants/_api.py
index 1a5bd3a..01041c2 100644
--- a/ragna/assistants/_api.py
+++ b/ragna/assistants/_api.py
@@ -16,7 +16,11 @@ class ApiAssistant(Assistant):
 
     @classmethod
     def requirements(cls) -> list[Requirement]:
-        return [EnvVarRequirement(cls._API_KEY_ENV_VAR)]
+        return [EnvVarRequirement(cls._API_KEY_ENV_VAR), *cls._extra_requirements()]
+
+    @classmethod
+    def _extra_requirements(cls) -> list[Requirement]:
+        return []
 
     def __init__(self) -> None:
         self._client = httpx.AsyncClient(
diff --git a/ragna/assistants/_google.py b/ragna/assistants/_google.py
index 4de69ce..afbb829 100644
--- a/ragna/assistants/_google.py
+++ b/ragna/assistants/_google.py
@@ -30,11 +30,8 @@ class GoogleApiAssistant(ApiAssistant):
     _MODEL: str
 
     @classmethod
-    def requirements(cls) -> list[Requirement]:
-        return [
-            *super().requirements(),
-            PackageRequirement("ijson"),
-        ]
+    def _extra_requirements(cls) -> list[Requirement]:
+        return [PackageRequirement("ijson")]
 
     @classmethod
     def display_name(cls) -> str:
diff --git a/ragna/assistants/_openai.py b/ragna/assistants/_openai.py
index 095ae9c..a9dad3e 100644
--- a/ragna/assistants/_openai.py
+++ b/ragna/assistants/_openai.py
@@ -1,9 +1,7 @@
 import json
 from typing import AsyncIterator, cast
 
-import httpx_sse
-
-from ragna.core import Source
+from ragna.core import PackageRequirement, Requirement, Source
 
 from ._api import ApiAssistant
 
@@ -12,6 +10,10 @@ class OpenaiApiAssistant(ApiAssistant):
     _API_KEY_ENV_VAR = "OPENAI_API_KEY"
     _MODEL: str
 
+    @classmethod
+    def _extra_requirements(cls) -> list[Requirement]:
+        return [PackageRequirement("httpx_sse")]
+
     @classmethod
     def display_name(cls) -> str:
         return f"OpenAI/{cls._MODEL}"
@@ -28,6 +30,8 @@ class OpenaiApiAssistant(ApiAssistant):
     async def _call_api(
         self, prompt: str, sources: list[Source], *, max_new_tokens: int
     ) -> AsyncIterator[str]:
+        import httpx_sse
+
         # See https://platform.openai.com/docs/api-reference/chat/create
         # and https://platform.openai.com/docs/api-reference/chat/streaming
         async with httpx_sse.aconnect_sse(
@@ -72,6 +76,10 @@ class Gpt35Turbo16k(OpenaiApiAssistant):
     !!! info "Required environment variables"
 
         - `OPENAI_API_KEY`
+
+    !!! info "Required packages"
+
+        - `httpx_sse`
     """
 
     _MODEL = "gpt-3.5-turbo-16k"
@@ -83,6 +91,10 @@ class Gpt4(OpenaiApiAssistant):
     !!! info "Required environment variables"
 
         - `OPENAI_API_KEY`
+
+    !!! info "Required packages"
+
+        - `httpx_sse`
     """
 
     _MODEL = "gpt-4"
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index 2236c01..9c25456 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -3,7 +3,6 @@ import uuid
 from typing import Annotated, Any, AsyncIterator, Iterator, Type, cast
 
 import aiofiles
-import sse_starlette
 from fastapi import (
     Body,
     Depends,
@@ -15,7 +14,8 @@ from fastapi import (
     status,
 )
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
 
 import ragna
 import ragna.core
@@ -273,7 +273,7 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
 
         if stream:
 
-            async def message_chunks() -> AsyncIterator[sse_starlette.ServerSentEvent]:
+            async def message_chunks() -> AsyncIterator[BaseModel]:
                 core_answer_stream = aiter(core_answer)
                 content_chunk = await anext(core_answer_stream)
 
@@ -285,7 +285,7 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
                         for source in core_answer.sources
                     ],
                 )
-                yield sse_starlette.ServerSentEvent(answer.model_dump_json())
+                yield answer
 
                 # Avoid sending the sources multiple times
                 answer_chunk = answer.model_copy(update=dict(sources=None))
@@ -293,15 +293,19 @@ def app(*, config: Config, ignore_unavailable_components: bool) -> FastAPI:
                 async for content_chunk in core_answer_stream:
                     content_chunks.append(content_chunk)
                     answer_chunk.content = content_chunk
-                    yield sse_starlette.ServerSentEvent(answer_chunk.model_dump_json())
+                    yield answer_chunk
 
                 with get_session() as session:
                     answer.content = "".join(content_chunks)
                     chat.messages.append(answer)
                     database.update_chat(session, user=user, chat=chat)
 
-            return sse_starlette.EventSourceResponse(  # type: ignore[return-value]
-                message_chunks()
+            async def to_jsonl(models: AsyncIterator[Any]) -> AsyncIterator[str]:
+                async for model in models:
+                    yield f"{model.model_dump_json()}\n"
+
+            return StreamingResponse(  # type: ignore[return-value]
+                to_jsonl(message_chunks())
             )
         else:
             answer = schemas.Message.from_core(core_answer)
diff --git a/ragna/deploy/_ui/api_wrapper.py b/ragna/deploy/_ui/api_wrapper.py
index 4d7ae1b..f96375d 100644
--- a/ragna/deploy/_ui/api_wrapper.py
+++ b/ragna/deploy/_ui/api_wrapper.py
@@ -3,7 +3,6 @@ from datetime import datetime
 
 import emoji
 import httpx
-import httpx_sse
 import param
 
 
@@ -64,14 +63,13 @@ class ApiWrapper(param.Parameterized):
         return json_data
 
     async def answer(self, chat_id, prompt):
-        async with httpx_sse.aconnect_sse(
-            self.client,
+        async with self.client.stream(
             "POST",
             f"/chats/{chat_id}/answer",
             json={"prompt": prompt, "stream": True},
-        ) as event_source:
-            async for sse in event_source.aiter_sse():
-                yield self.improve_message(json.loads(sse.data))
+        ) as response:
+            async for data in response.aiter_lines():
+                yield self.improve_message(json.loads(data))
 
     async def get_components(self):
         return (await self.client.get("/components")).raise_for_status().json()
diff --git a/ragna/deploy/_ui/central_view.py b/ragna/deploy/_ui/central_view.py
index 78b2d93..b1614d4 100644
--- a/ragna/deploy/_ui/central_view.py
+++ b/ragna/deploy/_ui/central_view.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 from typing import Callable, Literal, Optional, cast
 
 import panel as pn
@@ -162,6 +163,7 @@ class RagnaChatMessage(pn.chat.ChatMessage):
             show_user=False,
             show_copy_icon=False,
             css_classes=[f"message-{role}"],
+            avatar_lookup=functools.partial(self._ragna_avatar_lookup, role=role),
         )
 
     def _copy_and_source_view_buttons(self) -> pn.Row:
@@ -185,10 +187,14 @@ class RagnaChatMessage(pn.chat.ChatMessage):
             ),
         )
 
-    def avatar_lookup(self, user: str) -> str:
-        if self.role == "system":
+    # This cannot be a bound method, because it creates a reference cycle when trying
+    # to access the repr of the message. See
+    # https://github.com/Quansight/ragna/issues/359 for details.
+    @staticmethod
+    def _ragna_avatar_lookup(user: str, *, role: str) -> str:
+        if role == "system":
             return "imgs/ragna_logo.svg"
-        elif self.role == "user":
+        elif role == "user":
             return "👤"
 
         try:
diff --git a/requirements-docker.lock b/requirements-docker.lock
index 34a2128..bb63315 100644
--- a/requirements-docker.lock
+++ b/requirements-docker.lock
@@ -11,7 +11,6 @@ annotated-types==0.6.0
 anyio==4.2.0
     # via
     #   httpx
-    #   sse-starlette
     #   starlette
     #   watchfiles
 asgiref==3.7.2
@@ -71,7 +70,6 @@ fastapi==0.109.0
     # via
     #   Ragna (pyproject.toml)
     #   chromadb
-    #   sse-starlette
 filelock==3.13.1
     # via huggingface-hub
 flatbuffers==23.5.26
@@ -342,13 +340,10 @@ sniffio==1.3.0
     #   httpx
 sqlalchemy==2.0.25
     # via Ragna (pyproject.toml)
-sse-starlette==1.8.2
-    # via Ragna (pyproject.toml)
 starlette==0.35.1
     # via
     #   Ragna (pyproject.toml)
     #   fastapi
-    #   sse-starlette
 sympy==1.12
     # via onnxruntime
 tenacity==8.2.3
@@ -395,7 +390,6 @@ uvicorn==0.26.0
     # via
     #   Ragna (pyproject.toml)
     #   chromadb
-    #   sse-starlette
 uvloop==0.19.0
     # via uvicorn
 watchfiles==0.21.0


Test patch:
diff --git a/tests/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index 9acebf0..0efa001 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -1,6 +1,5 @@
 import json
 
-import httpx_sse
 import pytest
 from fastapi.testclient import TestClient
 
@@ -33,12 +32,6 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer):
     with open(document_path, "w") as file:
         file.write("!\n")
 
-    # Reset starlette_sse AppStatus for each run
-    # See https://github.com/sysid/sse-starlette/issues/59
-    from sse_starlette.sse import AppStatus
-
-    AppStatus.should_exit_event = None
-
     with TestClient(app(config=config, ignore_unavailable_components=False)) as client:
         authenticate(client)
 
@@ -104,15 +97,12 @@ def test_e2e(tmp_local_root, multiple_answer_chunks, stream_answer):
 
         prompt = "?"
         if stream_answer:
-            chunks = []
-            with httpx_sse.connect_sse(
-                client,
+            with client.stream(
                 "POST",
                 f"/chats/{chat['id']}/answer",
                 json={"prompt": prompt, "stream": True},
-            ) as event_source:
-                for sse in event_source.iter_sse():
-                    chunks.append(json.loads(sse.data))
+            ) as response:
+                chunks = [json.loads(chunk) for chunk in response.iter_lines()]
             message = chunks[0]
             assert all(chunk["sources"] is None for chunk in chunks[1:])
             message["content"] = "".join(chunk["content"] for chunk in chunks)