diff --git a/chat_app.py b/chat_app.py index 8846cf9..5bd2393 100644 --- a/chat_app.py +++ b/chat_app.py @@ -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) diff --git a/pyproject.toml b/pyproject.toml index fdf5507..5063f08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/llm.py b/src/llm.py index b8605ec..d24b122 100644 --- a/src/llm.py +++ b/src/llm.py @@ -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. diff --git a/tests/test_llm.py b/tests/test_llm.py index 958982a..5d4440d 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -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", []) diff --git a/uv.lock b/uv.lock index 4d8b456..7c79a34 100644 --- a/uv.lock +++ b/uv.lock @@ -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"