Merge pull request road-core#333 from TamiTakamiya/TamiTakamiya/event_stream_format_support

Event Stream Format support

Author: tisnik

examples/rcsconfig.yaml

@@ -73,6 +73,7 @@ ols_config:
   default_provider: my_bam
   default_model: ibm/granite-3-8b-instruct
   expire_llm_is_ready_persistent_state: -1
+  enable_event_stream_format: true
   # query_filters:
   #   - name: foo_filter
   #     pattern: '\b(?:foo)\b'

ols/app/endpoints/streaming_ols.py

@@ -129,13 +129,19 @@ async def invalid_response_generator() -> AsyncGenerator[str, None]:
     yield prompts.INVALID_QUERY_RESP
 
 
+def format_stream_data(d: dict) -> str:
+    """Format outbound data in the Event Stream Format if required."""
+    data = json.dumps(d)
+    return f"data: {data}\n\n" if config.ols_config.enable_event_stream_format else data
+
+
 def stream_start_event(conversation_id: str) -> str:
     """Yield the start of the data stream.
 
     Args:
         conversation_id: The conversation ID (UUID).
     """
-    return json.dumps(
+    return format_stream_data(
         {
             "event": "start",
             "data": {
@@ -157,7 +163,7 @@ def stream_end_event(
         token_counter: Token counter for the whole stream.
     """
     if media_type == constants.MEDIA_TYPE_JSON:
-        return json.dumps(
+        return format_stream_data(
             {
                 "event": "end",
                 "data": {
@@ -203,7 +209,7 @@ def prompt_too_long_error(error: PromptTooLongError, media_type: str) -> str:
     logger.error("Prompt is too long: %s", error)
     if media_type == MEDIA_TYPE_TEXT:
         return f"Prompt is too long: {error}"
-    return json.dumps(
+    return format_stream_data(
         {
             "event": "error",
             "data": {
@@ -230,7 +236,7 @@ def generic_llm_error(error: Exception, media_type: str) -> str:
 
     if media_type == MEDIA_TYPE_TEXT:
         return f"{response}: {cause}"
-    return json.dumps(
+    return format_stream_data(
         {
             "event": "error",
             "data": {
@@ -254,7 +260,12 @@ def build_yield_item(item: str, idx: int, media_type: str) -> str:
     """
     if media_type == MEDIA_TYPE_TEXT:
         return item
-    return json.dumps({"event": "token", "data": {"id": idx, "token": item}})
+    return format_stream_data(
+        {
+            "event": "token",
+            "data": {"id": idx, "token": item},
+        }
+    )
 
 
 def store_data(

ols/app/models/config.py

@@ -870,6 +870,8 @@ class OLSConfig(BaseModel):
     extra_ca: list[FilePath] = []
     certificate_directory: Optional[str] = None
 
+    enable_event_stream_format: bool = False
+
     def __init__(
         self, data: Optional[dict] = None, ignore_missing_certs: bool = False
     ) -> None:
@@ -920,6 +922,7 @@ def __init__(
         self.tls_security_profile = TLSSecurityProfile(
             data.get("tlsSecurityProfile", None)
         )
+        self.enable_event_stream_format = data.get("enable_event_stream_format", False)
 
     def __eq__(self, other: object) -> bool:
         """Compare two objects for equality."""

tests/unit/app/endpoints/test_streaming_ols.py

@@ -8,6 +8,7 @@
 from ols.app.endpoints.streaming_ols import (
     build_referenced_docs,
     build_yield_item,
+    format_stream_data,
     generic_llm_error,
     invalid_response_generator,
     prompt_too_long_error,
@@ -149,3 +150,26 @@ def test_build_referenced_docs():
         {"doc_title": "title_1", "doc_url": "url_1"},
         {"doc_title": "title_2", "doc_url": "url_2"},
     ]
+
+
+@pytest.mark.usefixtures("_load_config")
+def test_format_stream_data():
+    """Test format_stream_data function."""
+    stream_data = {
+        "event": "token",
+        "data": {"id": "ABC-123456", "token": "***TOKEN***"},
+    }
+    stringified_data = json.dumps(stream_data)
+    data_in_event_stream_data_format = f"data: {stringified_data}\n\n"
+
+    saved_value = config.ols_config.enable_event_stream_format
+    try:
+        config.ols_config.enable_event_stream_format = False
+        output = format_stream_data((stream_data))
+        assert output == stringified_data
+
+        config.ols_config.enable_event_stream_format = True
+        output = format_stream_data((stream_data))
+        assert output == data_in_event_stream_data_format
+    finally:
+        config.ols_config.enable_event_stream_format = saved_value