From 9005e992c0e3234855e6e14fb980008f9aa5152c Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 18 Jun 2026 23:32:36 +0530 Subject: [PATCH 1/3] feat: add unit tests for LLM bundle streaming functionality - Introduced a new test file to validate the LLM bundle construction for streaming flows in chat. - Implemented tests to ensure that both DB-backed and global models enable streaming correctly. - Utilized mocking to isolate dependencies and verify the expected behavior of the LLM constructor. --- .../tasks/chat/streaming/test_llm_bundle.py | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 surfsense_backend/tests/unit/tasks/chat/streaming/test_llm_bundle.py diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_llm_bundle.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_llm_bundle.py new file mode 100644 index 000000000..cecf8be5d --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_llm_bundle.py @@ -0,0 +1,154 @@ +"""Contracts for chat LLM construction in streaming flows. + +``stream_new_chat`` / ``stream_resume_chat`` depend on LangChain receiving +token chunks from ``ChatLiteLLM``. ``langchain-litellm`` defaults +``streaming`` to ``False``, so the shared bundle loader must opt in +explicitly for both DB-backed and global model paths. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +import app.tasks.chat.streaming.flows.shared.llm_bundle as llm_bundle + +pytestmark = pytest.mark.unit + + +class _CapturedChatLiteLLM: + calls: list[dict[str, Any]] = [] + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self.__class__.calls.append(kwargs) + + +@pytest.fixture(autouse=True) +def _patch_common_bundle_dependencies(monkeypatch: pytest.MonkeyPatch): + """Keep these tests focused on the LLM constructor contract.""" + + _CapturedChatLiteLLM.calls = [] + + async def _fake_search_space(_session: Any, _search_space_id: int) -> SimpleNamespace: + return SimpleNamespace(id=42, user_id="user-1") + + monkeypatch.setattr(llm_bundle, "_load_search_space", _fake_search_space) + monkeypatch.setattr(llm_bundle, "SanitizedChatLiteLLM", _CapturedChatLiteLLM) + monkeypatch.setattr(llm_bundle, "register_model_usage_metadata", lambda **_kw: None) + monkeypatch.setattr( + llm_bundle, + "has_capability", + lambda _model, capability: capability in {"chat", "vision"}, + ) + + return None + + +async def test_load_llm_bundle_enables_streaming_for_db_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connection = SimpleNamespace( + provider="openai", + api_key="sk-test", + base_url=None, + extra={"litellm_params": {"temperature": 0.1}}, + ) + model = SimpleNamespace( + id=7, + model_id="gpt-4o-mini", + display_name="GPT 4o Mini", + connection=connection, + ) + + async def _fake_db_model(_session: Any, *, model_id: int, search_space: Any) -> Any: + assert model_id == 7 + assert search_space.id == 42 + return model + + monkeypatch.setattr(llm_bundle, "_load_db_model", _fake_db_model) + monkeypatch.setattr( + llm_bundle, + "to_litellm", + lambda _conn, _model_id: ( + "openai/gpt-4o-mini", + {"api_key": "sk-test", "temperature": 0.1}, + ), + ) + + llm, agent_config, error = await llm_bundle.load_llm_bundle( + object(), + config_id=7, + search_space_id=42, + ) + + assert error is None + assert llm is not None + assert agent_config is not None + assert _CapturedChatLiteLLM.calls == [ + { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "temperature": 0.1, + "streaming": True, + } + ] + + +async def test_load_llm_bundle_enables_streaming_for_global_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + global_model = { + "id": -11, + "connection_id": -101, + "model_id": "claude-sonnet-4-5", + "display_name": "Claude Sonnet", + "billing_tier": "premium", + } + global_connection = { + "id": -101, + "provider": "anthropic", + "api_key": "sk-ant-test", + "base_url": None, + "extra": {"litellm_params": {"temperature": 0.2}}, + } + monkeypatch.setattr( + llm_bundle.config, + "GLOBAL_MODELS", + [global_model], + raising=False, + ) + monkeypatch.setattr( + llm_bundle.config, + "GLOBAL_CONNECTIONS", + [global_connection], + raising=False, + ) + monkeypatch.setattr( + llm_bundle, + "to_litellm", + lambda _conn, _model_id: ( + "anthropic/claude-sonnet-4-5", + {"api_key": "sk-ant-test", "temperature": 0.2}, + ), + ) + + llm, agent_config, error = await llm_bundle.load_llm_bundle( + object(), + config_id=-11, + search_space_id=42, + ) + + assert error is None + assert llm is not None + assert agent_config is not None + assert _CapturedChatLiteLLM.calls == [ + { + "model": "anthropic/claude-sonnet-4-5", + "api_key": "sk-ant-test", + "temperature": 0.2, + "streaming": True, + } + ] From 6e970be220e00465775e086215ec49760e01cfad Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 18 Jun 2026 23:39:55 +0530 Subject: [PATCH 2/3] feat: enable streaming in LLM bundle construction - Updated the LLM bundle construction to include a streaming option for both DB-backed and global models. - Modified the `litellm_kwargs` to set the streaming parameter to True, enhancing the functionality for chat streaming flows. --- .../app/tasks/chat/streaming/flows/shared/llm_bundle.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/shared/llm_bundle.py b/surfsense_backend/app/tasks/chat/streaming/flows/shared/llm_bundle.py index 9304f1698..6f905e8f4 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/shared/llm_bundle.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/shared/llm_bundle.py @@ -130,7 +130,9 @@ async def load_llm_bundle( billing_tier="free", ) return ( - SanitizedChatLiteLLM(model=model_string, **litellm_kwargs), + SanitizedChatLiteLLM( + model=model_string, **{**litellm_kwargs, "streaming": True} + ), agent_config, None, ) @@ -174,7 +176,9 @@ async def load_llm_bundle( billing_tier=str(global_model.get("billing_tier", "free")).lower(), ) return ( - SanitizedChatLiteLLM(model=model_string, **litellm_kwargs), + SanitizedChatLiteLLM( + model=model_string, **{**litellm_kwargs, "streaming": True} + ), agent_config, None, ) From bb664a1f325fd2828cd0626968442d8029b8863a Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 18 Jun 2026 23:56:58 +0530 Subject: [PATCH 3/3] fix: enable smooth rendering in MarkdownText component - Updated the MarkdownTextPrimitive component to enable smooth rendering by default. - Adjusted the props to streamline the rendering process for improved user experience. --- surfsense_web/components/assistant-ui/chat-viewport.tsx | 4 ++-- surfsense_web/components/assistant-ui/markdown-text.tsx | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/surfsense_web/components/assistant-ui/chat-viewport.tsx b/surfsense_web/components/assistant-ui/chat-viewport.tsx index dedada7a5..83308b642 100644 --- a/surfsense_web/components/assistant-ui/chat-viewport.tsx +++ b/surfsense_web/components/assistant-ui/chat-viewport.tsx @@ -27,8 +27,8 @@ export interface ChatViewportProps { export const ChatViewport: FC = ({ children, footer }) => ( { return (