Add S3StorageClient for elements persistence

Author: GFJHogue

bin/chat-chainlit.py

@@ -10,9 +10,9 @@
 from agent.graph import AgentGraph
 from agent.profiles import ProfileName, get_chat_profiles
 from agent.profiles.base import OutputState
-from util.chainlit_helpers import (is_feature_enabled, message_rate_limited,
-                                   save_openai_metrics, static_messages,
-                                   update_search_results)
+from util.chainlit_helpers import (PrefixedS3StorageClient, is_feature_enabled,
+                                   message_rate_limited, save_openai_metrics,
+                                   static_messages, update_search_results)
 from util.config_yml import Config, TriggerEvent
 from util.logging import logging
 
@@ -22,12 +22,26 @@
 profiles: list[ProfileName] = config.profiles if config else [ProfileName.React_to_Me]
 llm_graph = AgentGraph(profiles)
 
-if os.getenv("POSTGRES_CHAINLIT_DB"):
-    CHAINLIT_DB_URI = f"postgresql+psycopg://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}@postgres:5432/{os.getenv('POSTGRES_CHAINLIT_DB')}?sslmode=disable"
+POSTGRES_CHAINLIT_DB = os.getenv("POSTGRES_CHAINLIT_DB")
+POSTGRES_USER = os.getenv("POSTGRES_USER")
+POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD")
+S3_BUCKET = os.getenv("S3_BUCKET")
+S3_CHAINLIT_PREFIX = os.getenv("S3_CHAINLIT_PREFIX")
+
+if POSTGRES_CHAINLIT_DB and POSTGRES_USER and POSTGRES_PASSWORD:
+    CHAINLIT_DB_URI = f"postgresql+psycopg://{POSTGRES_USER}:{POSTGRES_PASSWORD}@postgres:5432/{POSTGRES_CHAINLIT_DB}?sslmode=disable"
+
+    if S3_BUCKET and S3_CHAINLIT_PREFIX:
+        storage_client = PrefixedS3StorageClient(S3_BUCKET, S3_CHAINLIT_PREFIX)
+    else:
+        storage_client = None
 
     @cl.data_layer
     def get_data_layer() -> BaseDataLayer:
-        return SQLAlchemyDataLayer(conninfo=CHAINLIT_DB_URI)
+        return SQLAlchemyDataLayer(
+            conninfo=CHAINLIT_DB_URI,
+            storage_provider=storage_client,
+        )
 
 else:
     logging.warning("POSTGRES_CHAINLIT_DB undefined; Chainlit persistence disabled.")
@@ -100,15 +114,16 @@ async def main(message: cl.Message) -> None:
         thread_id=thread_id,
         enable_postprocess=enable_postprocess,
     )
+    assistant_message: cl.Message | None = chainlit_cb.final_stream
 
     if (
         enable_postprocess
-        and chainlit_cb.final_stream
+        and assistant_message
         and len(result["additional_content"]["search_results"]) > 0
     ):
         await update_search_results(
             result["additional_content"]["search_results"],
-            chainlit_cb.final_stream,
+            assistant_message,
         )
 
     await static_messages(config, after_messages=message_count)

docker-compose.yml

@@ -8,6 +8,8 @@ services:
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD}
       - POSTGRES_CHAINLIT_DB=${POSTGRES_CHAINLIT_DB}
       - POSTGRES_LANGGRAPH_DB=${POSTGRES_LANGGRAPH_DB}
+      - S3_BUCKET=${S3_BUCKET}
+      - S3_CHAINLIT_PREFIX=${S3_CHAINLIT_PREFIX}
       - LOG_LEVEL=${LOG_LEVEL}
       - UVICORN_LOG_LEVEL=${LOG_LEVEL}
       - CHAT_ENV=${CHAT_ENV}

src/util/chainlit_helpers.py

@@ -1,9 +1,11 @@
 import os
 from datetime import datetime
+from pathlib import PurePosixPath
 from typing import Any, Iterable
 
 import chainlit as cl
 from chainlit.data import get_data_layer
+from chainlit.data.storage_clients.s3 import S3StorageClient
 from langchain_community.callbacks import OpenAICallbackHandler
 
 from util.config_yml import Config, TriggerEvent
@@ -12,6 +14,30 @@
 guest_user_metadata: dict[str, Any] = {}
 
 
+class PrefixedS3StorageClient(S3StorageClient):
+    def __init__(self, bucket: str, prefix: str, **kwargs: Any) -> None:
+        super().__init__(bucket, **kwargs)
+        self._prefix = PurePosixPath(prefix)
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: bytes | str,
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> dict[str, Any]:
+        object_key = str(self._prefix / object_key)
+        return await super().upload_file(object_key, data, mime, overwrite)
+
+    async def delete_file(self, object_key: str) -> bool:
+        object_key = str(self._prefix / object_key)
+        return await super().delete_file(object_key)
+
+    async def get_read_url(self, object_key: str) -> str:
+        object_key = str(self._prefix / object_key)
+        return await super().get_read_url(object_key)
+
+
 def get_user_id() -> str | None:
     user: cl.User | None = cl.user_session.get("user")
     return user.identifier if user else None
@@ -136,7 +162,7 @@ async def update_search_results(
         name="SearchResults",
         props={"results": search_results},
     )
-    message.elements = [search_results_element]  # type: ignore
+    message.elements.append(search_results_element)  # type: ignore[arg-type]
     await message.update()