Merge branch 'main' into main

.github/workflows/ci.yaml

@@ -8,16 +8,19 @@ on:
     branches: [main, dev]
 
 jobs:
+  pytest:
+    uses: ./.github/workflows/pytest.yaml
+    secrets: inherit
   mypy:
     uses: ./.github/workflows/mypy.yaml
     secrets: inherit
-  tests:
-    uses: ./.github/workflows/tests.yaml
+  e2e-tests:
+    uses: ./.github/workflows/e2e-tests.yaml
     secrets: inherit
   ci:
     runs-on: ubuntu-latest
     name: Run CI
-    needs: [mypy, tests]
+    needs: [mypy, pytest, e2e-tests]
     steps:
       - name: 'Done'
         run: echo "Done"

.github/workflows/tests.yaml → .github/workflows/e2e-tests.yaml

@@ -1,4 +1,4 @@
-name: Tests
+name: E2ETests
 
 on: [workflow_call]
 

.github/workflows/pytest.yaml

@@ -0,0 +1,23 @@
+name: Pytest
+
+on: [workflow_call]
+
+jobs:
+  mypy:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./backend
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+          cache: 'pip'
+      - name: Install Poetry
+        run: pip install poetry
+      - name: Install dependencies
+        run: poetry install --with tests --with mypy --with custom-data
+      - name: Run Pytest
+        run: poetry run pytest --cov=chainlit/

.gitignore

@@ -55,3 +55,6 @@ dist-ssr
 *.njsproj
 *.sln
 *.sw?
+
+.aider*
+.coverage

backend/chainlit/data/__init__.py

@@ -27,6 +27,7 @@
     ThreadFilter,
 )
 from chainlit.user import PersistedUser, User
+from httpx import HTTPStatusError, RequestError
 from literalai import Attachment
 from literalai import Score as LiteralScore
 from literalai import Step as LiteralStep
@@ -294,6 +295,14 @@ async def upsert_feedback(
             )
             return created.id or ""
 
+    async def safely_send_steps(self, steps):
+        try:
+            await self.client.api.send_steps(steps)
+        except HTTPStatusError as e:
+            logger.error(f"HTTP Request: error sending steps: {e.response.status_code}")
+        except RequestError as e:
+            logger.error(f"HTTP Request: error for {e.request.url!r}.")
+
     @queue_until_user_message()
     async def create_element(self, element: "Element"):
         metadata = {
@@ -322,7 +331,7 @@ async def create_element(self, element: "Element"):
             )
             object_key = uploaded["object_key"]
 
-        await self.client.api.send_steps(
+        await self.safely_send_steps(
             [
                 {
                     "id": element.for_id,
@@ -384,7 +393,7 @@ async def create_step(self, step_dict: "StepDict"):
         if step_dict.get("isError"):
             step["error"] = step_dict.get("output")
 
-        await self.client.api.send_steps([step])
+        await self.safely_send_steps([step])
 
     @queue_until_user_message()
     async def update_step(self, step_dict: "StepDict"):

backend/chainlit/element.py

@@ -235,14 +235,10 @@ def __post_init__(self) -> None:
         if not isinstance(self.figure, Figure):
             raise TypeError("figure must be a matplotlib.figure.Figure")
 
-        options = {
-            "dpi": 200,
-            "bbox_inches": "tight",
-            "backend": "Agg",
-            "format": "png",
-        }
         image = BytesIO()
-        self.figure.savefig(image, **options)
+        self.figure.savefig(
+            image, dpi=200, bbox_inches="tight", backend="Agg", format="png"
+        )
         self.content = image.getvalue()
 
         super().__post_init__()

backend/pyproject.toml

@@ -48,22 +48,21 @@ lazify = "^0.4.0"
 packaging = "^23.1"
 python-multipart = "^0.0.9"
 pyjwt = "^2.8.0"
-# Fix https://stackoverflow.com/questions/77582651/why-doesnt-numpy-support-pep517-builds
-numpy = [
-    { version = "^1.26", python = ">=3.9" },
-    { version = "^1.24.4", python = "<3.9" },
-]
+numpy = "^1.26"
 
 [tool.poetry.group.tests]
 optional = true
 
 [tool.poetry.group.tests.dependencies]
+pytest = "^8.3.2"
+pytest-asyncio = "^0.23.8"
+pytest-cov = "^5.0.0"
 openai = "^1.11.1"
 langchain = "^0.1.5"
 llama-index = "^0.10.45"
-tenacity = "8.3.0"          # pin until fixed https://github.com/langchain-ai/langchain/issues/22972
-transformers = "^4.30.1"
-matplotlib = "3.7.1"
+tenacity = "^8.4.1"
+transformers = "^4.38"
+matplotlib = "^3.7.1"
 farm-haystack = "^1.18.0"
 plotly = "^5.18.0"
 slack_bolt = "^1.18.1"
@@ -84,33 +83,27 @@ types-aiofiles = "^23.1.0.5"
 mypy-boto3-dynamodb = "^1.34.113"
 
 [tool.mypy]
-python_version = "3.8"
+python_version = "3.9"
 
 
 [[tool.mypy.overrides]]
 module = [
-    "anthropic",
     "boto3.dynamodb.types",
-    "huggingface_hub.inference_api",
-    "fastapi_socketio",
     "botbuilder.*",
     "filetype",
     "haystack.*",
     "langflow",
     "lazify",
-    "matplotlib.*",                       # remove when 3.8.0 is out, it should export types
-    "plotly.*",
+    "plotly",
     "nest_asyncio",
     "socketio.*",
     "uptrace",
     "syncer",
-    "vertexai.language_models",
-    "vertexai.preview.generative_models",
-    "google.cloud",
     "azure.storage.filedatalake",
 ]
 ignore_missing_imports = true
 
+
 [tool.poetry.group.custom-data]
 optional = true
 
@@ -125,3 +118,8 @@ azure-storage-file-datalake = "^12.14.0"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"

backend/tests/__init__.py

@@ -0,0 +1 @@
+

backend/tests/test_context.py

@@ -0,0 +1,80 @@
+from unittest.mock import Mock
+
+import pytest
+from chainlit.context import (
+    ChainlitContext,
+    ChainlitContextException,
+    get_context,
+    init_http_context,
+    init_ws_context,
+)
+from chainlit.emitter import BaseChainlitEmitter, ChainlitEmitter
+from chainlit.session import HTTPSession, WebsocketSession
+
+
+@pytest.fixture
+def mock_websocket_session():
+    return Mock(spec=WebsocketSession)
+
+
+@pytest.fixture
+def mock_http_session():
+    return Mock(spec=HTTPSession)
+
+
+@pytest.fixture
+def mock_emitter():
+    return Mock(spec=BaseChainlitEmitter)
+
+
+async def test_chainlit_context_init_with_websocket(
+    mock_websocket_session, mock_emitter
+):
+    context = ChainlitContext(mock_websocket_session, mock_emitter)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_websocket_session
+    assert context.active_steps == []
+
+
+async def test_chainlit_context_init_with_http(mock_http_session):
+    context = ChainlitContext(mock_http_session)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_http_session
+    assert context.active_steps == []
+
+
+async def test_init_ws_context(mock_websocket_session):
+    context = init_ws_context(mock_websocket_session)
+    assert isinstance(context, ChainlitContext)
+    assert context.session == mock_websocket_session
+    assert isinstance(context.emitter, ChainlitEmitter)
+
+
+async def test_init_http_context():
+    context = init_http_context()
+    assert isinstance(context, ChainlitContext)
+    assert isinstance(context.session, HTTPSession)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+
+
+async def test_get_context():
+    with pytest.raises(ChainlitContextException):
+        get_context()
+
+    init_http_context()  # Initialize a context
+    context = get_context()
+    assert isinstance(context, ChainlitContext)
+
+
+async def test_current_step_and_run():
+    context = init_http_context()
+    assert context.current_step is None
+    assert context.current_run is None
+
+    # Mock a step
+    mock_step = Mock()
+    mock_step.name = "on_chat_start"
+    context.active_steps.append(mock_step)
+
+    assert context.current_step == mock_step
+    assert context.current_run == mock_step