diff --git a/surfsense_backend/app/services/llm_service.py b/surfsense_backend/app/services/llm_service.py index 24bc4138b..e562d0d35 100644 --- a/surfsense_backend/app/services/llm_service.py +++ b/surfsense_backend/app/services/llm_service.py @@ -162,7 +162,7 @@ async def validate_llm_config( async def get_search_space_llm_instance( - session: AsyncSession, search_space_id: int, role: str + session: AsyncSession, search_space_id: int, role: str, disable_streaming: bool = False ) -> ChatLiteLLM | ChatLiteLLMRouter | None: """ Get a ChatLiteLLM instance for a specific search space and role. @@ -218,7 +218,7 @@ async def get_search_space_llm_instance( logger.debug( f"Using Auto mode (LLM Router) for search space {search_space_id}, role {role}" ) - return ChatLiteLLMRouter() + return ChatLiteLLMRouter(disable_streaming=disable_streaming) except Exception as e: logger.error(f"Failed to create ChatLiteLLMRouter: {e}") return None @@ -284,6 +284,9 @@ async def get_search_space_llm_instance( if global_config.get("litellm_params"): litellm_kwargs.update(global_config["litellm_params"]) + if disable_streaming: + litellm_kwargs["disable_streaming"] = True + return ChatLiteLLM(**litellm_kwargs) # Get the LLM configuration from database (NewLLMConfig) @@ -357,6 +360,9 @@ async def get_search_space_llm_instance( if llm_config.litellm_params: litellm_kwargs.update(llm_config.litellm_params) + if disable_streaming: + litellm_kwargs["disable_streaming"] = True + return ChatLiteLLM(**litellm_kwargs) except Exception as e: @@ -374,20 +380,20 @@ async def get_agent_llm( async def get_document_summary_llm( - session: AsyncSession, search_space_id: int + session: AsyncSession, search_space_id: int, disable_streaming: bool = False ) -> ChatLiteLLM | ChatLiteLLMRouter | None: """Get the search space's document summary LLM instance.""" return await get_search_space_llm_instance( - session, search_space_id, LLMRole.DOCUMENT_SUMMARY + session, search_space_id, LLMRole.DOCUMENT_SUMMARY, disable_streaming=disable_streaming ) # Backward-compatible alias (LLM preferences are now per-search-space, not per-user) async def get_user_long_context_llm( - session: AsyncSession, user_id: str, search_space_id: int + session: AsyncSession, user_id: str, search_space_id: int, disable_streaming: bool = False ) -> ChatLiteLLM | ChatLiteLLMRouter | None: """ Deprecated: Use get_document_summary_llm instead. The user_id parameter is ignored as LLM preferences are now per-search-space. """ - return await get_document_summary_llm(session, search_space_id) + return await get_document_summary_llm(session, search_space_id, disable_streaming=disable_streaming) diff --git a/surfsense_backend/app/services/notion/kb_sync_service.py b/surfsense_backend/app/services/notion/kb_sync_service.py index 3cacdcfe5..ea6328d46 100644 --- a/surfsense_backend/app/services/notion/kb_sync_service.py +++ b/surfsense_backend/app/services/notion/kb_sync_service.py @@ -104,7 +104,7 @@ class NotionKBSyncService: logger.debug("Generating summary and embeddings") user_llm = await get_user_long_context_llm( - self.db_session, user_id, search_space_id + self.db_session, user_id, search_space_id, disable_streaming=True # disable streaming to avoid leaking into the chat ) if user_llm: diff --git a/surfsense_web/components/tool-ui/update-notion-page.tsx b/surfsense_web/components/tool-ui/update-notion-page.tsx index 0242057ff..de8fdd359 100644 --- a/surfsense_web/components/tool-ui/update-notion-page.tsx +++ b/surfsense_web/components/tool-ui/update-notion-page.tsx @@ -195,7 +195,7 @@ function ApprovalCard({
Notion Account
- {account.workspace_icon} {account.workspace_name} + {account.workspace_name}
)}