diff --git a/surfsense_backend/alembic/versions/104_add_chunks_document_id_index.py b/surfsense_backend/alembic/versions/104_add_chunks_document_id_index.py new file mode 100644 index 000000000..4e6dc6a15 --- /dev/null +++ b/surfsense_backend/alembic/versions/104_add_chunks_document_id_index.py @@ -0,0 +1,41 @@ +"""104_add_chunks_document_id_index + +Revision ID: 104 +Revises: 103 +Create Date: 2026-03-09 + +Adds a B-tree index on chunks.document_id to speed up chunk lookups +during hybrid search (both retrievers fetch chunks by document_id +after RRF ranking selects the top documents). +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + +revision: str = "104" +down_revision: str | None = "103" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.execute( + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_indexes + WHERE tablename = 'chunks' AND indexname = 'ix_chunks_document_id' + ) THEN + CREATE INDEX ix_chunks_document_id ON chunks(document_id); + END IF; + END$$; + """ + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS ix_chunks_document_id") diff --git a/surfsense_backend/app/agents/new_chat/tools/search_surfsense_docs.py b/surfsense_backend/app/agents/new_chat/tools/search_surfsense_docs.py index ec86c3ffa..b8b1527c7 100644 --- a/surfsense_backend/app/agents/new_chat/tools/search_surfsense_docs.py +++ b/surfsense_backend/app/agents/new_chat/tools/search_surfsense_docs.py @@ -8,6 +8,7 @@ The documentation is indexed at deployment time from MDX files and stored in dedicated tables (surfsense_docs_documents, surfsense_docs_chunks). """ +import asyncio import json from langchain_core.tools import tool @@ -100,7 +101,7 @@ async def search_surfsense_docs_async( Formatted string with relevant documentation content """ # Get embedding for the query - query_embedding = embed_text(query) + query_embedding = await asyncio.to_thread(embed_text, query) # Vector similarity search on chunks, joining with documents stmt = ( diff --git a/surfsense_backend/app/agents/new_chat/tools/shared_memory.py b/surfsense_backend/app/agents/new_chat/tools/shared_memory.py index ba69f1ce8..c826d808f 100644 --- a/surfsense_backend/app/agents/new_chat/tools/shared_memory.py +++ b/surfsense_backend/app/agents/new_chat/tools/shared_memory.py @@ -1,5 +1,6 @@ """Shared (team) memory backend for search-space-scoped AI context.""" +import asyncio import logging from typing import Any from uuid import UUID @@ -64,7 +65,7 @@ async def save_shared_memory( count = await get_shared_memory_count(db_session, search_space_id) if count >= MAX_MEMORIES_PER_SEARCH_SPACE: await delete_oldest_shared_memory(db_session, search_space_id) - embedding = embed_text(content) + embedding = await asyncio.to_thread(embed_text, content) row = SharedMemory( search_space_id=search_space_id, created_by_id=_to_uuid(created_by_id), @@ -108,7 +109,7 @@ async def recall_shared_memory( if category and category in valid_categories: stmt = stmt.where(SharedMemory.category == MemoryCategory(category)) if query: - query_embedding = embed_text(query) + query_embedding = await asyncio.to_thread(embed_text, query) stmt = stmt.order_by( SharedMemory.embedding.op("<=>")(query_embedding) ).limit(top_k) diff --git a/surfsense_backend/app/agents/new_chat/tools/user_memory.py b/surfsense_backend/app/agents/new_chat/tools/user_memory.py index 8aa516454..81e849856 100644 --- a/surfsense_backend/app/agents/new_chat/tools/user_memory.py +++ b/surfsense_backend/app/agents/new_chat/tools/user_memory.py @@ -9,6 +9,7 @@ Features: - recall_memory: Retrieve relevant memories using semantic search """ +import asyncio import logging from typing import Any from uuid import UUID @@ -177,8 +178,7 @@ def create_save_memory_tool( # Delete oldest memory to make room await delete_oldest_memory(db_session, user_id, search_space_id) - # Generate embedding for the memory - embedding = embed_text(content) + embedding = await asyncio.to_thread(embed_text, content) # Create new memory using ORM # The pgvector Vector column type handles embedding conversion automatically @@ -267,8 +267,7 @@ def create_recall_memory_tool( uuid_user_id = _to_uuid(user_id) if query: - # Semantic search using embeddings - query_embedding = embed_text(query) + query_embedding = await asyncio.to_thread(embed_text, query) # Build query with vector similarity stmt = ( diff --git a/surfsense_backend/app/config/__init__.py b/surfsense_backend/app/config/__init__.py index 68c65a818..a03ef5f8a 100644 --- a/surfsense_backend/app/config/__init__.py +++ b/surfsense_backend/app/config/__init__.py @@ -327,6 +327,7 @@ class Config: EMBEDDING_MODEL, **embedding_kwargs, ) + is_local_embedding_model = "://" not in (EMBEDDING_MODEL or "") chunker_instance = RecursiveChunker( chunk_size=getattr(embedding_model_instance, "max_seq_length", 512) ) diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index dc355dd94..04d1328a6 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -960,7 +960,7 @@ class Chunk(BaseModel, TimestampMixin): embedding = Column(Vector(config.embedding_model_instance.dimension)) document_id = Column( - Integer, ForeignKey("documents.id", ondelete="CASCADE"), nullable=False + Integer, ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, index=True ) document = relationship("Document", back_populates="chunks") diff --git a/surfsense_backend/app/indexing_pipeline/document_embedder.py b/surfsense_backend/app/indexing_pipeline/document_embedder.py index adec24434..f545d9097 100644 --- a/surfsense_backend/app/indexing_pipeline/document_embedder.py +++ b/surfsense_backend/app/indexing_pipeline/document_embedder.py @@ -1,3 +1,3 @@ -from app.utils.document_converters import embed_text +from app.utils.document_converters import embed_text, embed_texts -__all__ = ["embed_text"] +__all__ = ["embed_text", "embed_texts"] diff --git a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py index 9460f900c..0fadfc42f 100644 --- a/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py +++ b/surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.db import Chunk, Document, DocumentStatus from app.indexing_pipeline.connector_document import ConnectorDocument from app.indexing_pipeline.document_chunker import chunk_text -from app.indexing_pipeline.document_embedder import embed_text +from app.indexing_pipeline.document_embedder import embed_texts from app.indexing_pipeline.document_hashing import ( compute_content_hash, compute_unique_identifier_hash, @@ -195,25 +195,23 @@ class IndexingPipelineService: else: content = connector_doc.source_markdown - t_step = time.perf_counter() - embedding = embed_text(content) - perf.debug( - "[indexing] embed_text (summary) doc=%d in %.3fs", - document.id, - time.perf_counter() - t_step, - ) - await self.session.execute( delete(Chunk).where(Chunk.document_id == document.id) ) t_step = time.perf_counter() + chunk_texts = chunk_text( + connector_doc.source_markdown, + use_code_chunker=connector_doc.should_use_code_chunker, + ) + + texts_to_embed = [content, *chunk_texts] + embeddings = embed_texts(texts_to_embed) + summary_embedding, *chunk_embeddings = embeddings + chunks = [ - Chunk(content=text, embedding=embed_text(text)) - for text in chunk_text( - connector_doc.source_markdown, - use_code_chunker=connector_doc.should_use_code_chunker, - ) + Chunk(content=text, embedding=emb) + for text, emb in zip(chunk_texts, chunk_embeddings) ] perf.info( "[indexing] chunk+embed doc=%d chunks=%d in %.3fs", @@ -223,7 +221,7 @@ class IndexingPipelineService: ) document.content = content - document.embedding = embedding + document.embedding = summary_embedding attach_chunks_to_document(document, chunks) document.updated_at = datetime.now(UTC) document.status = DocumentStatus.ready() diff --git a/surfsense_backend/app/services/connector_service.py b/surfsense_backend/app/services/connector_service.py index 9db0cc9e1..870e175d3 100644 --- a/surfsense_backend/app/services/connector_service.py +++ b/surfsense_backend/app/services/connector_service.py @@ -264,7 +264,9 @@ class ConnectorService: # Reuse caller-provided embedding or compute once for both retrievers. if query_embedding is None: t_embed = time.perf_counter() - query_embedding = config.embedding_model_instance.embed(query_text) + query_embedding = await asyncio.to_thread( + config.embedding_model_instance.embed, query_text + ) perf.info( "[connector_svc] _combined_rrf embedding in %.3fs type=%s", time.perf_counter() - t_embed, diff --git a/surfsense_backend/app/utils/document_converters.py b/surfsense_backend/app/utils/document_converters.py index 8049b0de5..6a59990f5 100644 --- a/surfsense_backend/app/utils/document_converters.py +++ b/surfsense_backend/app/utils/document_converters.py @@ -55,6 +55,23 @@ def embed_text(text: str) -> np.ndarray: return config.embedding_model_instance.embed(truncate_for_embedding(text)) +def embed_texts(texts: list[str]) -> list[np.ndarray]: + """Batch-embed multiple texts in a single call. + + Each text is truncated to fit the model's context window before embedding. + For API-based models (``://`` in the model string) this uses + ``embed_batch`` to collapse many network round-trips into one. + For local models (SentenceTransformers) it falls back to sequential + ``embed`` calls to avoid padding overhead. + """ + if not texts: + return [] + truncated = [truncate_for_embedding(t) for t in texts] + if config.is_local_embedding_model: + return [config.embedding_model_instance.embed(t) for t in truncated] + return config.embedding_model_instance.embed_batch(truncated) + + def get_model_context_window(model_name: str) -> int: """Get the total context window size for a model (input + output tokens).""" try: @@ -209,12 +226,11 @@ async def create_document_chunks(content: str) -> list[Chunk]: Returns: List of Chunk objects with embeddings """ + chunk_texts = [c.text for c in config.chunker_instance.chunk(content)] + chunk_embeddings = embed_texts(chunk_texts) return [ - Chunk( - content=chunk.text, - embedding=embed_text(chunk.text), - ) - for chunk in config.chunker_instance.chunk(content) + Chunk(content=text, embedding=emb) + for text, emb in zip(chunk_texts, chunk_embeddings) ] diff --git a/surfsense_backend/tests/integration/conftest.py b/surfsense_backend/tests/integration/conftest.py index 8b92a5aa8..4e43ea302 100644 --- a/surfsense_backend/tests/integration/conftest.py +++ b/surfsense_backend/tests/integration/conftest.py @@ -129,10 +129,12 @@ def patched_summarize_raises(monkeypatch) -> AsyncMock: @pytest.fixture -def patched_embed_text(monkeypatch) -> MagicMock: - mock = MagicMock(return_value=[0.1] * _EMBEDDING_DIM) +def patched_embed_texts(monkeypatch) -> MagicMock: + mock = MagicMock( + side_effect=lambda texts: [[0.1] * _EMBEDDING_DIM for _ in texts] + ) monkeypatch.setattr( - "app.indexing_pipeline.indexing_pipeline_service.embed_text", + "app.indexing_pipeline.indexing_pipeline_service.embed_texts", mock, ) return mock diff --git a/surfsense_backend/tests/integration/document_upload/conftest.py b/surfsense_backend/tests/integration/document_upload/conftest.py index 41639fc2f..45cfef7ac 100644 --- a/surfsense_backend/tests/integration/document_upload/conftest.py +++ b/surfsense_backend/tests/integration/document_upload/conftest.py @@ -265,8 +265,8 @@ def _mock_external_apis(monkeypatch): AsyncMock(return_value="Mocked summary."), ) monkeypatch.setattr( - "app.indexing_pipeline.indexing_pipeline_service.embed_text", - MagicMock(return_value=[0.1] * _EMBEDDING_DIM), + "app.indexing_pipeline.indexing_pipeline_service.embed_texts", + MagicMock(side_effect=lambda texts: [[0.1] * _EMBEDDING_DIM for _ in texts]), ) monkeypatch.setattr( "app.indexing_pipeline.indexing_pipeline_service.chunk_text", diff --git a/surfsense_backend/tests/integration/indexing_pipeline/adapters/test_file_upload_adapter.py b/surfsense_backend/tests/integration/indexing_pipeline/adapters/test_file_upload_adapter.py index fa0fe5787..9fc802aa6 100644 --- a/surfsense_backend/tests/integration/indexing_pipeline/adapters/test_file_upload_adapter.py +++ b/surfsense_backend/tests/integration/indexing_pipeline/adapters/test_file_upload_adapter.py @@ -8,7 +8,7 @@ pytestmark = pytest.mark.integration @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_sets_status_ready(db_session, db_search_space, db_user, mocker): """Document status is READY after successful indexing.""" @@ -31,7 +31,7 @@ async def test_sets_status_ready(db_session, db_search_space, db_user, mocker): @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_content_is_summary(db_session, db_search_space, db_user, mocker): """Document content is set to the LLM-generated summary.""" @@ -55,7 +55,7 @@ async def test_content_is_summary(db_session, db_search_space, db_user, mocker): @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_chunks_written_to_db(db_session, db_search_space, db_user, mocker): """Chunks derived from the source markdown are persisted in the DB.""" @@ -84,7 +84,7 @@ async def test_chunks_written_to_db(db_session, db_search_space, db_user, mocker @pytest.mark.usefixtures( - "patched_summarize_raises", "patched_embed_text", "patched_chunk_text" + "patched_summarize_raises", "patched_embed_texts", "patched_chunk_text" ) async def test_raises_on_indexing_failure(db_session, db_search_space, db_user, mocker): """RuntimeError is raised when the indexing step fails so the caller can fire a failure notification.""" @@ -107,7 +107,7 @@ async def test_raises_on_indexing_failure(db_session, db_search_space, db_user, @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_reindex_updates_content(db_session, db_search_space, db_user, mocker): """Document content is updated to the new summary after reindexing.""" @@ -136,7 +136,7 @@ async def test_reindex_updates_content(db_session, db_search_space, db_user, moc @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_reindex_updates_content_hash( db_session, db_search_space, db_user, mocker @@ -168,7 +168,7 @@ async def test_reindex_updates_content_hash( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_reindex_sets_status_ready(db_session, db_search_space, db_user, mocker): """Document status is READY after successful reindexing.""" @@ -196,7 +196,7 @@ async def test_reindex_sets_status_ready(db_session, db_search_space, db_user, m assert DocumentStatus.is_state(document.status, DocumentStatus.READY) -@pytest.mark.usefixtures("patched_summarize", "patched_embed_text") +@pytest.mark.usefixtures("patched_summarize", "patched_embed_texts") async def test_reindex_replaces_chunks(db_session, db_search_space, db_user, mocker): """Reindexing replaces old chunks with new content rather than appending.""" mocker.patch( @@ -235,7 +235,7 @@ async def test_reindex_replaces_chunks(db_session, db_search_space, db_user, moc @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_reindex_clears_reindexing_flag( db_session, db_search_space, db_user, mocker @@ -266,7 +266,7 @@ async def test_reindex_clears_reindexing_flag( assert document.content_needs_reindexing is False -@pytest.mark.usefixtures("patched_embed_text", "patched_chunk_text") +@pytest.mark.usefixtures("patched_embed_texts", "patched_chunk_text") async def test_reindex_raises_on_failure(db_session, db_search_space, db_user, mocker): """RuntimeError is raised when reindexing fails so the caller can handle it.""" mocker.patch( diff --git a/surfsense_backend/tests/integration/indexing_pipeline/test_index_document.py b/surfsense_backend/tests/integration/indexing_pipeline/test_index_document.py index 2e8ee4d92..a82148f96 100644 --- a/surfsense_backend/tests/integration/indexing_pipeline/test_index_document.py +++ b/surfsense_backend/tests/integration/indexing_pipeline/test_index_document.py @@ -11,7 +11,7 @@ pytestmark = pytest.mark.integration @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_sets_status_ready( db_session, @@ -38,7 +38,7 @@ async def test_sets_status_ready( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_content_is_summary_when_should_summarize_true( db_session, @@ -65,7 +65,7 @@ async def test_content_is_summary_when_should_summarize_true( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_content_is_source_markdown_when_should_summarize_false( db_session, @@ -95,7 +95,7 @@ async def test_content_is_source_markdown_when_should_summarize_false( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_chunks_written_to_db( db_session, @@ -123,7 +123,7 @@ async def test_chunks_written_to_db( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_embedding_written_to_db( db_session, @@ -151,7 +151,7 @@ async def test_embedding_written_to_db( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_updated_at_advances_after_indexing( db_session, @@ -183,7 +183,7 @@ async def test_updated_at_advances_after_indexing( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_no_llm_falls_back_to_source_markdown( db_session, @@ -214,7 +214,7 @@ async def test_no_llm_falls_back_to_source_markdown( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_fallback_summary_used_when_llm_unavailable( db_session, @@ -245,7 +245,7 @@ async def test_fallback_summary_used_when_llm_unavailable( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_reindex_replaces_old_chunks( db_session, @@ -282,7 +282,7 @@ async def test_reindex_replaces_old_chunks( @pytest.mark.usefixtures( - "patched_summarize_raises", "patched_embed_text", "patched_chunk_text" + "patched_summarize_raises", "patched_embed_texts", "patched_chunk_text" ) async def test_llm_error_sets_status_failed( db_session, @@ -309,7 +309,7 @@ async def test_llm_error_sets_status_failed( @pytest.mark.usefixtures( - "patched_summarize_raises", "patched_embed_text", "patched_chunk_text" + "patched_summarize_raises", "patched_embed_texts", "patched_chunk_text" ) async def test_llm_error_leaves_no_partial_data( db_session, diff --git a/surfsense_backend/tests/integration/indexing_pipeline/test_prepare_for_indexing.py b/surfsense_backend/tests/integration/indexing_pipeline/test_prepare_for_indexing.py index 837b02c9f..776180b9a 100644 --- a/surfsense_backend/tests/integration/indexing_pipeline/test_prepare_for_indexing.py +++ b/surfsense_backend/tests/integration/indexing_pipeline/test_prepare_for_indexing.py @@ -33,7 +33,7 @@ async def test_new_document_is_persisted_with_pending_status( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_unchanged_ready_document_is_skipped( db_session, @@ -56,7 +56,7 @@ async def test_unchanged_ready_document_is_skipped( @pytest.mark.usefixtures( - "patched_summarize", "patched_embed_text", "patched_chunk_text" + "patched_summarize", "patched_embed_texts", "patched_chunk_text" ) async def test_title_only_change_updates_title_in_db( db_session, @@ -339,7 +339,7 @@ async def test_same_content_from_different_source_is_skipped( @pytest.mark.usefixtures( - "patched_summarize_raises", "patched_embed_text", "patched_chunk_text" + "patched_summarize_raises", "patched_embed_texts", "patched_chunk_text" ) async def test_failed_document_with_unchanged_content_is_requeued( db_session, diff --git a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/composio-drive-config.tsx b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/composio-drive-config.tsx index 239125565..d8967a83e 100644 --- a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/composio-drive-config.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/composio-drive-config.tsx @@ -1,6 +1,8 @@ "use client"; import { + ChevronDown, + ChevronRight, File, FileSpreadsheet, FileText, @@ -12,7 +14,6 @@ import { import type { FC } from "react"; import { useEffect, useState } from "react"; import { ComposioDriveFolderTree } from "@/components/connectors/composio-drive-folder-tree"; -import { Button } from "@/components/ui/button"; import { Label } from "@/components/ui/label"; import { Select, @@ -108,9 +109,11 @@ export const ComposioDriveConfig: FC = ({ const [selectedFolders, setSelectedFolders] = useState(existingFolders); const [selectedFiles, setSelectedFiles] = useState(existingFiles); - const [showFolderSelector, setShowFolderSelector] = useState(false); const [indexingOptions, setIndexingOptions] = useState(existingIndexingOptions); + const [isEditMode] = useState(() => existingFolders.length > 0 || existingFiles.length > 0); + const [isFolderTreeOpen, setIsFolderTreeOpen] = useState(!isEditMode); + // Update selected folders and files when connector config changes useEffect(() => { const folders = (connector.config?.selected_folders as SelectedFolder[] | undefined) || []; @@ -241,8 +244,21 @@ export const ComposioDriveConfig: FC = ({ )} - {showFolderSelector ? ( -
+ {isEditMode ? ( +
+ + {isFolderTreeOpen && ( = ({ selectedFiles={selectedFiles} onSelectFiles={handleSelectFiles} /> - -
- ) : ( - - )} + )} +
+ ) : ( + + )} {/* Indexing Options */} diff --git a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/google-drive-config.tsx b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/google-drive-config.tsx index 383f6ce0e..8480ab53c 100644 --- a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/google-drive-config.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/google-drive-config.tsx @@ -1,6 +1,8 @@ "use client"; import { + ChevronDown, + ChevronRight, File, FileSpreadsheet, FileText, @@ -12,7 +14,6 @@ import { import type { FC } from "react"; import { useEffect, useState } from "react"; import { GoogleDriveFolderTree } from "@/components/connectors/google-drive-folder-tree"; -import { Button } from "@/components/ui/button"; import { Label } from "@/components/ui/label"; import { Select, @@ -97,9 +98,11 @@ export const GoogleDriveConfig: FC = ({ connector, onConfi const [selectedFolders, setSelectedFolders] = useState(existingFolders); const [selectedFiles, setSelectedFiles] = useState(existingFiles); - const [showFolderSelector, setShowFolderSelector] = useState(false); const [indexingOptions, setIndexingOptions] = useState(existingIndexingOptions); + const [isEditMode] = useState(() => existingFolders.length > 0 || existingFiles.length > 0); + const [isFolderTreeOpen, setIsFolderTreeOpen] = useState(!isEditMode); + // Update selected folders and files when connector config changes useEffect(() => { const folders = (connector.config?.selected_folders as SelectedFolder[] | undefined) || []; @@ -225,8 +228,21 @@ export const GoogleDriveConfig: FC = ({ connector, onConfi )} - {showFolderSelector ? ( -
+ {isEditMode ? ( +
+ + {isFolderTreeOpen && ( = ({ connector, onConfi selectedFiles={selectedFiles} onSelectFiles={handleSelectFiles} /> - -
- ) : ( - - )} + )} +
+ ) : ( + + )} {/* Indexing Options */}