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/4] 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/4] 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/4] 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 ( Date: Fri, 19 Jun 2026 01:41:21 +0530 Subject: [PATCH 4/4] refactor: enhance chat UI components for mobile responsiveness - Updated the layout of the ComposerAction and ChatHeader components to improve mobile compatibility. - Added a new prop to ImageModelSelector for mobile-specific rendering. - Adjusted ModelSelector to conditionally render elements based on mobile view, enhancing user experience on smaller screens. --- .../components/assistant-ui/thread.tsx | 6 +++--- .../components/new-chat/chat-header.tsx | 4 ++-- .../new-chat/image-model-selector.tsx | 20 ++++++++++++++----- .../components/new-chat/model-selector.tsx | 12 +++++++---- 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 9f74895d1..c8da125f4 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -1577,7 +1577,7 @@ const ComposerAction: FC = ({ Select a model )} -
+
= ({ variant="default" size="icon" className={cn( - "aui-composer-send size-9 rounded-full", + "aui-composer-send size-9 shrink-0 rounded-full", isSendDisabled && "cursor-not-allowed opacity-50" )} aria-label="Send message" @@ -1617,7 +1617,7 @@ const ComposerAction: FC = ({ type="button" variant="default" size="icon" - className="aui-composer-cancel size-9 rounded-full" + className="aui-composer-cancel size-9 shrink-0 rounded-full" aria-label="Stop generating" > diff --git a/surfsense_web/components/new-chat/chat-header.tsx b/surfsense_web/components/new-chat/chat-header.tsx index 9882530d4..99d56eb02 100644 --- a/surfsense_web/components/new-chat/chat-header.tsx +++ b/surfsense_web/components/new-chat/chat-header.tsx @@ -11,13 +11,13 @@ interface ChatHeaderProps { export function ChatHeader({ searchSpaceId, className, onChatModelSelected }: ChatHeaderProps) { return ( -
+
- +
); } diff --git a/surfsense_web/components/new-chat/image-model-selector.tsx b/surfsense_web/components/new-chat/image-model-selector.tsx index e90a46c09..5cd898afc 100644 --- a/surfsense_web/components/new-chat/image-model-selector.tsx +++ b/surfsense_web/components/new-chat/image-model-selector.tsx @@ -33,6 +33,7 @@ import { providerDisplay } from "../settings/model-connections/provider-metadata interface ImageModelSelectorProps { searchSpaceId: number; className?: string; + mobileIconOnly?: boolean; } type ImageModel = ModelRead & { @@ -95,7 +96,11 @@ function groupedModels(models: ImageModel[]) { }, {}); } -export function ImageModelSelector({ searchSpaceId, className }: ImageModelSelectorProps) { +export function ImageModelSelector({ + searchSpaceId, + className, + mobileIconOnly = false, +}: ImageModelSelectorProps) { const router = useRouter(); const isMobile = useIsMobile(); const [open, setOpen] = useState(false); @@ -126,6 +131,7 @@ export function ImageModelSelector({ searchSpaceId, className }: ImageModelSelec const groups = useMemo(() => groupedModels(visibleImageModels), [visibleImageModels]); const loading = globalLoading || connectionsLoading; const hasSearchQuery = search.trim().length > 0; + const showIconOnlyTrigger = isMobile && mobileIconOnly; function handleOpenChange(nextOpen: boolean) { if (!nextOpen) setSearch(""); @@ -252,12 +258,14 @@ export function ImageModelSelector({ searchSpaceId, className }: ImageModelSelec type="button" variant="ghost" size="sm" + aria-label="Select image model" className={cn( "h-8 min-w-0 gap-2 rounded-md px-3 text-muted-foreground transition-colors", "select-none", "hover:bg-foreground/10 hover:text-foreground", "data-[state=open]:bg-foreground/10 data-[state=open]:text-foreground", - className + className, + showIconOnlyTrigger && "h-9 w-auto shrink-0 justify-center gap-1 px-2" )} > {selected ? ( @@ -265,9 +273,11 @@ export function ImageModelSelector({ searchSpaceId, className }: ImageModelSelec ) : ( )} - - {selected ? modelName(selected) : "Auto"} - + {showIconOnlyTrigger ? null : ( + + {selected ? modelName(selected) : "Auto"} + + )} ); diff --git a/surfsense_web/components/new-chat/model-selector.tsx b/surfsense_web/components/new-chat/model-selector.tsx index 22d86aa92..c10bfd862 100644 --- a/surfsense_web/components/new-chat/model-selector.tsx +++ b/surfsense_web/components/new-chat/model-selector.tsx @@ -131,6 +131,7 @@ export function ModelSelector({ const groups = useMemo(() => groupedModels(visibleChatModels), [visibleChatModels]); const loading = globalLoading || connectionsLoading; const hasSearchQuery = search.trim().length > 0; + const showIconOnlyTrigger = isMobile; function handleOpenChange(nextOpen: boolean) { if (!nextOpen) setSearch(""); @@ -276,15 +277,18 @@ export function ModelSelector({ "select-none", "hover:bg-foreground/10 hover:text-foreground", "data-[state=open]:bg-foreground/10 data-[state=open]:text-foreground", - className + className, + showIconOnlyTrigger && "h-9 w-auto shrink-0 justify-center gap-1 px-2" )} > {selected ? getProviderIcon(selected.provider, { className: "size-4 shrink-0" }) : getProviderIcon(AUTO_PROVIDER_ICON_KEY, { className: "size-4 shrink-0" })} - - {selected ? modelName(selected) : "Auto"} - + {showIconOnlyTrigger ? null : ( + + {selected ? modelName(selected) : "Auto"} + + )} );