Merge pull request #77 from wagtail/feature/add-litellm-chat-streamin…

…g-test

Feature/add litellm chat streaming test

pyproject.toml

@@ -27,7 +27,7 @@ requires-python = ">=3.11"
 dependencies = [
     "Django>=4.2",
     "Wagtail>=5.2",
-    "litellm>=1.41.15",
+    "litellm>=1.43.2",
     "openai>=1.28.1",
     "aiohttp>=3.9.0b0; python_version >= '3.12'",
 ]

tests/test_ai_utils/test_backends/test_litellm.py

@@ -156,6 +156,23 @@ def test_litellm_chat(make_chat_backend):
     assert response.choices[0] == response_text
 
 
+@if_litellm_installed
+def test_litellm_streaming_chat(make_chat_backend):
+    backend = make_chat_backend()
+    input_text = "Little trotty wagtail, he waddled in the mud."
+    response_text = "And left his little footmarks, trample where he would."
+    messages: List[ChatMessage] = [
+        {"content": message, "role": "user"} for message in input_text
+    ]
+    response = backend.chat(messages=messages, stream=True, mock_response=response_text)
+
+    full_text = ""
+    for chunk in response:
+        full_text += chunk["content"]
+
+    assert full_text == response_text
+
+
 @if_litellm_installed
 async def test_litellm_async_chat(make_chat_backend):
     backend = make_chat_backend()