fix(chat): use native async LLM call to prevent session reset on message submit
The previous implementation used asyncio.to_thread(llm.complete) to avoid blocking the event loop, but Chainlit's contextvars context is not reliably propagated across thread boundaries, causing the session to reset and clear the message history on each user submission. Changes: - Add llm.acomplete() using litellm.acompletion() (native coroutine) - Replace asyncio.to_thread() in on_message with await llm.acomplete() - Store the welcome message in state.messages so it is replayed on reconnect - Persist state to cl.user_session immediately after appending the user's message (before the LLM call) so reconnect detection has the latest history - Add pytest-asyncio dev dependency and asyncio_mode = "auto" config - Add 6 async tests for acomplete() in tests/test_llm.py https://claude.ai/code/session_01SUWzMzFvSfWiHXA2p6rPg9
This commit is contained in:
Claude
+10
-9
@@ -18,7 +18,6 @@ Environment variables (see .env.example):
|
||||
DATA_DIR Directory for completed registration JSON (default: data/)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
@@ -101,6 +100,8 @@ async def on_chat_start() -> None:
|
||||
# Brand new session
|
||||
session_id = str(uuid.uuid4())
|
||||
state = ConversationState(conversation_id=session_id)
|
||||
# Store the welcome in history so it's replayed if the session reconnects.
|
||||
state.messages.append(ChatMessage(role="assistant", content=_WELCOME_DE))
|
||||
cl.user_session.set("state", state.to_dict())
|
||||
logger.info("Chat session started: %s", session_id)
|
||||
await cl.Message(content=_WELCOME_DE).send()
|
||||
@@ -114,20 +115,20 @@ async def on_message(message: cl.Message) -> None:
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
state.last_activity = now
|
||||
|
||||
# Append parent's message to history
|
||||
# Append parent's message to history and persist immediately so that any
|
||||
# WebSocket reconnect during the LLM call can replay the full conversation.
|
||||
state.messages.append(ChatMessage(role="user", content=message.content))
|
||||
cl.user_session.set("state", state.to_dict())
|
||||
|
||||
# --- Build system prompt ---
|
||||
system = build_system_prompt(_kb, state)
|
||||
|
||||
# --- Call LLM in a thread so the async event loop (and WebSocket) stay alive ---
|
||||
# litellm.completion is synchronous; running it directly in an async handler
|
||||
# blocks the event loop for the full response duration and causes WebSocket
|
||||
# timeouts that trigger on_chat_start again (clearing the screen).
|
||||
# --- Call LLM natively async to keep the event loop free ---
|
||||
# litellm.acompletion() is a true coroutine — it does not block the event
|
||||
# loop and does not require asyncio.to_thread (which can drop Chainlit's
|
||||
# contextvars, causing the session to reset mid-conversation).
|
||||
try:
|
||||
full_content = await asyncio.to_thread(
|
||||
llm.complete, _config.ai_model, system, state.messages
|
||||
)
|
||||
full_content = await llm.acomplete(_config.ai_model, system, state.messages)
|
||||
except Exception:
|
||||
logger.exception("LLM call failed for session %s", state.conversation_id)
|
||||
error_text = fallback_message(state.language)
|
||||
|
||||
@@ -26,5 +26,9 @@ packages = ["src"]
|
||||
[tool.uv]
|
||||
dev-dependencies = [
|
||||
"pytest>=8.0.0",
|
||||
"pytest-asyncio>=1.3.0",
|
||||
"pytest-mock>=3.14.0",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
|
||||
+21
@@ -5,6 +5,27 @@ from collections.abc import Generator
|
||||
import litellm
|
||||
|
||||
|
||||
async def acomplete(model: str, system: str, messages: list) -> str:
|
||||
"""Call any LLM asynchronously and return the response text.
|
||||
|
||||
This is the async equivalent of ``complete()`` — use this from async
|
||||
handlers (e.g. Chainlit's ``@cl.on_message``) to avoid blocking the
|
||||
event loop and losing framework context variables.
|
||||
|
||||
Args:
|
||||
model: litellm model string, e.g. "anthropic/claude-opus-4-6".
|
||||
system: System prompt text.
|
||||
messages: List of objects with .role and .content attributes.
|
||||
|
||||
Returns:
|
||||
The model's reply as a plain string.
|
||||
"""
|
||||
api_messages = [{"role": "system", "content": system}]
|
||||
api_messages += [{"role": m.role, "content": m.content} for m in messages]
|
||||
response = await litellm.acompletion(model=model, messages=api_messages, max_tokens=2048)
|
||||
return response.choices[0].message.content
|
||||
|
||||
|
||||
def complete(model: str, system: str, messages: list) -> str:
|
||||
"""Call any LLM and return the response text.
|
||||
|
||||
|
||||
@@ -6,6 +6,17 @@ from src import llm
|
||||
from src.models.conversation import ChatMessage
|
||||
|
||||
|
||||
def _make_async_mock(mocker, content: str):
|
||||
"""Return an awaitable mock that resolves to a response with the given content."""
|
||||
mock_response = mocker.MagicMock()
|
||||
mock_response.choices[0].message.content = content
|
||||
|
||||
async def _coro(*args, **kwargs):
|
||||
return mock_response
|
||||
|
||||
return mocker.patch("litellm.acompletion", side_effect=_coro), mock_response
|
||||
|
||||
|
||||
class TestLlmComplete:
|
||||
def test_returns_model_reply(self, mocker):
|
||||
mock_response = mocker.MagicMock()
|
||||
@@ -122,3 +133,63 @@ class TestLlmStreamComplete:
|
||||
|
||||
messages = mock_completion.call_args.kwargs["messages"]
|
||||
assert messages[1] == {"role": "user", "content": "Hallo"}
|
||||
|
||||
|
||||
class TestLlmAComplete:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_model_reply(self, mocker):
|
||||
mock_acompletion, _ = _make_async_mock(mocker, "Hallo! Wie heisst dein Kind?")
|
||||
|
||||
result = await llm.acomplete("anthropic/claude-opus-4-6", "system prompt", [])
|
||||
|
||||
assert result == "Hallo! Wie heisst dein Kind?"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passes_model_to_litellm(self, mocker):
|
||||
mock_acompletion, _ = _make_async_mock(mocker, "ok")
|
||||
|
||||
await llm.acomplete("openai/gpt-4o", "system", [])
|
||||
|
||||
call_kwargs = mock_acompletion.call_args.kwargs
|
||||
assert call_kwargs["model"] == "openai/gpt-4o"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_system_prompt_prepended_as_system_message(self, mocker):
|
||||
mock_acompletion, _ = _make_async_mock(mocker, "ok")
|
||||
|
||||
await llm.acomplete("anthropic/claude-opus-4-6", "You are helpful.", [])
|
||||
|
||||
messages = mock_acompletion.call_args.kwargs["messages"]
|
||||
assert messages[0] == {"role": "system", "content": "You are helpful."}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_messages_appended_after_system(self, mocker):
|
||||
mock_acompletion, _ = _make_async_mock(mocker, "ok")
|
||||
|
||||
chat = [
|
||||
ChatMessage(role="user", content="Hallo"),
|
||||
ChatMessage(role="assistant", content="Guten Tag"),
|
||||
]
|
||||
await llm.acomplete("anthropic/claude-opus-4-6", "system", chat)
|
||||
|
||||
messages = mock_acompletion.call_args.kwargs["messages"]
|
||||
assert messages[1] == {"role": "user", "content": "Hallo"}
|
||||
assert messages[2] == {"role": "assistant", "content": "Guten Tag"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_tokens_passed(self, mocker):
|
||||
mock_acompletion, _ = _make_async_mock(mocker, "ok")
|
||||
|
||||
await llm.acomplete("anthropic/claude-opus-4-6", "system", [])
|
||||
|
||||
assert mock_acompletion.call_args.kwargs["max_tokens"] == 2048
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exception_propagates(self, mocker):
|
||||
async def _raise(*args, **kwargs):
|
||||
raise RuntimeError("API error")
|
||||
|
||||
mocker.patch("litellm.acompletion", side_effect=_raise)
|
||||
|
||||
with pytest.raises(RuntimeError, match="API error"):
|
||||
await llm.acomplete("anthropic/claude-opus-4-6", "system", [])
|
||||
|
||||
@@ -1082,6 +1082,7 @@ dependencies = [
|
||||
[package.dev-dependencies]
|
||||
dev = [
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-mock" },
|
||||
]
|
||||
|
||||
@@ -1096,6 +1097,7 @@ requires-dist = [
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "pytest", specifier = ">=8.0.0" },
|
||||
{ name = "pytest-asyncio", specifier = ">=1.3.0" },
|
||||
{ name = "pytest-mock", specifier = ">=3.14.0" },
|
||||
]
|
||||
|
||||
@@ -2139,6 +2141,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-asyncio"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/90/2c/8af215c0f776415f3590cac4f9086ccefd6fd463befeae41cd4d3f193e5a/pytest_asyncio-1.3.0.tar.gz", hash = "sha256:d7f52f36d231b80ee124cd216ffb19369aa168fc10095013c6b014a34d3ee9e5", size = 50087, upload-time = "2025-11-10T16:07:47.256Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-mock"
|
||||
version = "3.15.1"
|
||||
|
||||
Reference in New Issue
Block a user