apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-07-10 22:32:16 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat: added ai file sorting

Browse source
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-04-14 01:43:30 -07:00
parent fa0b47dfca
commit 4bee367d4a
51 changed files with 1703 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

3
surfsense_backend/app/agents/new_chat/middleware/dedup_tool_calls.py
View file

@ -93,7 +93,8 @@ class DedupHITLToolCallsMiddleware(AgentMiddleware): # type: ignore[type-arg]
@staticmethod
def _dedup(
state: AgentState, dedup_keys: dict[str, str] # type: ignore[type-arg]
state: AgentState,
dedup_keys: dict[str, str], # type: ignore[type-arg]
) -> dict[str, Any] | None:
messages = state.get("messages")
if not messages:

2
surfsense_backend/app/agents/new_chat/middleware/filesystem.py
View file

@ -593,7 +593,7 @@ class SurfSenseFilesystemMiddleware(FilesystemMiddleware):
runtime: ToolRuntime[None, FilesystemState],
timeout: int | None,
) -> str:
sandbox, is_new = await get_or_create_sandbox(self._thread_id)
sandbox, _is_new = await get_or_create_sandbox(self._thread_id)
# NOTE: sync_files_to_sandbox is intentionally disabled.
# The virtual FS contains XML-wrapped KB documents whose paths
# would double-nest under SANDBOX_DOCUMENTS_ROOT (e.g.

201
surfsense_backend/app/agents/new_chat/middleware/knowledge_search.py
View file

@ -58,6 +58,14 @@ class KBSearchPlan(BaseModel):
default=None,
description="Optional ISO end date or datetime for KB search filtering.",
)
is_recency_query: bool = Field(
default=False,
description=(
"True when the user's intent is primarily about recency or temporal "
"ordering (e.g. 'latest', 'newest', 'most recent', 'last uploaded') "
"rather than topical relevance."
),
)
def _extract_text_from_message(message: BaseMessage) -> str:
@ -245,7 +253,7 @@ def _build_kb_planner_prompt(
return (
"You optimize internal knowledge-base search inputs for document retrieval.\n"
"Return JSON only with this exact shape:\n"
'{"optimized_query":"string","start_date":"ISO string or null","end_date":"ISO string or null"}\n\n'
'{"optimized_query":"string","start_date":"ISO string or null","end_date":"ISO string or null","is_recency_query":bool}\n\n'
"Rules:\n"
"- Preserve the user's intent.\n"
"- Rewrite the query to improve retrieval using concrete entities, acronyms, projects, tools, people, and document-specific terms when helpful.\n"
@ -253,6 +261,11 @@ def _build_kb_planner_prompt(
"- Only use date filters when the latest user request or recent dialogue clearly implies a time range.\n"
"- If you use date filters, prefer returning both bounds.\n"
"- If no date filter is useful, return null for both dates.\n"
'- Set "is_recency_query" to true ONLY when the user\'s primary intent is about '
"recency or temporal ordering rather than topical relevance. Examples: "
'"latest file", "newest upload", "most recent document", "what did I save last", '
'"show me files from today", "last thing I added". '
"When true, results will be sorted by date instead of relevance.\n"
"- Do not include markdown, prose, or explanations.\n\n"
f"Today's UTC date: {today}\n\n"
f"Recent conversation:\n{recent_conversation or '(none)'}\n\n"
@ -506,6 +519,135 @@ def _resolve_search_types(
return list(expanded) if expanded else None
_RECENCY_MAX_CHUNKS_PER_DOC = 5
async def browse_recent_documents(
*,
search_space_id: int,
document_type: list[str] | None = None,
top_k: int = 10,
start_date: datetime | None = None,
end_date: datetime | None = None,
) -> list[dict[str, Any]]:
"""Return documents ordered by recency (newest first), no relevance ranking.
Used when the user's intent is temporal ("latest file", "most recent upload")
and hybrid search would produce poor results because the query has no
meaningful topical signal.
"""
from sqlalchemy import func, select
from app.db import DocumentType
async with shielded_async_session() as session:
base_conditions = [
Document.search_space_id == search_space_id,
func.coalesce(Document.status["state"].astext, "ready") != "deleting",
]
if document_type is not None:
import contextlib
doc_type_enums = []
for dt in document_type:
if isinstance(dt, str):
with contextlib.suppress(KeyError):
doc_type_enums.append(DocumentType[dt])
else:
doc_type_enums.append(dt)
if doc_type_enums:
if len(doc_type_enums) == 1:
base_conditions.append(Document.document_type == doc_type_enums[0])
else:
base_conditions.append(Document.document_type.in_(doc_type_enums))
if start_date is not None:
base_conditions.append(Document.updated_at >= start_date)
if end_date is not None:
base_conditions.append(Document.updated_at <= end_date)
doc_query = (
select(Document)
.where(*base_conditions)
.order_by(Document.updated_at.desc())
.limit(top_k)
)
result = await session.execute(doc_query)
documents = result.scalars().unique().all()
if not documents:
return []
doc_ids = [d.id for d in documents]
numbered = (
select(
Chunk.id.label("chunk_id"),
Chunk.document_id,
Chunk.content,
func.row_number()
.over(partition_by=Chunk.document_id, order_by=Chunk.id)
.label("rn"),
)
.where(Chunk.document_id.in_(doc_ids))
.subquery("numbered")
)
chunk_query = (
select(numbered.c.chunk_id, numbered.c.document_id, numbered.c.content)
.where(numbered.c.rn <= _RECENCY_MAX_CHUNKS_PER_DOC)
.order_by(numbered.c.document_id, numbered.c.chunk_id)
)
chunk_result = await session.execute(chunk_query)
fetched_chunks = chunk_result.all()
doc_chunks: dict[int, list[dict[str, Any]]] = {d.id: [] for d in documents}
for row in fetched_chunks:
if row.document_id in doc_chunks:
doc_chunks[row.document_id].append(
{"chunk_id": row.chunk_id, "content": row.content}
)
results: list[dict[str, Any]] = []
for doc in documents:
chunks_list = doc_chunks.get(doc.id, [])
metadata = doc.document_metadata or {}
results.append(
{
"document_id": doc.id,
"content": "\n\n".join(
c["content"] for c in chunks_list if c.get("content")
),
"score": 0.0,
"chunks": chunks_list,
"matched_chunk_ids": [],
"document": {
"id": doc.id,
"title": doc.title,
"document_type": (
doc.document_type.value
if getattr(doc, "document_type", None)
else None
),
"metadata": metadata,
},
"source": (
doc.document_type.value
if getattr(doc, "document_type", None)
else None
),
}
)
logger.info(
"browse_recent_documents: %d docs returned for space=%d",
len(results),
search_space_id,
)
return results
async def search_knowledge_base(
*,
query: str,
@ -704,10 +846,13 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg]
*,
messages: Sequence[BaseMessage],
user_text: str,
) -> tuple[str, datetime | None, datetime | None]:
"""Rewrite the KB query and infer optional date filters with the LLM."""
) -> tuple[str, datetime | None, datetime | None, bool]:
"""Rewrite the KB query and infer optional date filters with the LLM.
Returns (optimized_query, start_date, end_date, is_recency_query).
"""
if self.llm is None:
return user_text, None, None
return user_text, None, None, False
recent_conversation = _render_recent_conversation(
messages,
@ -734,15 +879,18 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg]
plan.start_date,
plan.end_date,
)
is_recency = plan.is_recency_query
_perf_log.info(
"[kb_fs_middleware] planner in %.3fs query=%r optimized=%r start=%s end=%s",
"[kb_fs_middleware] planner in %.3fs query=%r optimized=%r "
"start=%s end=%s recency=%s",
loop.time() - t0,
user_text[:80],
optimized_query[:120],
start_date.isoformat() if start_date else None,
end_date.isoformat() if end_date else None,
is_recency,
)
return optimized_query, start_date, end_date
return optimized_query, start_date, end_date, is_recency
except (json.JSONDecodeError, ValidationError, ValueError) as exc:
logger.warning(
"KB planner returned invalid output, using raw query: %s", exc
@ -750,7 +898,7 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg]
except Exception as exc: # pragma: no cover - defensive fallback
logger.warning("KB planner failed, using raw query: %s", exc)
return user_text, None, None
return user_text, None, None, False
def before_agent( # type: ignore[override]
self,
@ -789,7 +937,12 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg]
t0 = _perf_log and asyncio.get_event_loop().time()
existing_files = state.get("files")
planned_query, start_date, end_date = await self._plan_search_inputs(
(
planned_query,
start_date,
end_date,
is_recency,
) = await self._plan_search_inputs(
messages=messages,
user_text=user_text,
)
@ -805,16 +958,28 @@ class KnowledgeBaseSearchMiddleware(AgentMiddleware): # type: ignore[type-arg]
# messages within the same agent instance.
self.mentioned_document_ids = []
# --- 2. Run KB hybrid search ---
search_results = await search_knowledge_base(
query=planned_query,
search_space_id=self.search_space_id,
available_connectors=self.available_connectors,
available_document_types=self.available_document_types,
top_k=self.top_k,
start_date=start_date,
end_date=end_date,
)
# --- 2. Run KB search (recency browse or hybrid) ---
if is_recency:
doc_types = _resolve_search_types(
self.available_connectors, self.available_document_types
)
search_results = await browse_recent_documents(
search_space_id=self.search_space_id,
document_type=doc_types,
top_k=self.top_k,
start_date=start_date,
end_date=end_date,
)
else:
search_results = await search_knowledge_base(
query=planned_query,
search_space_id=self.search_space_id,
available_connectors=self.available_connectors,
available_document_types=self.available_document_types,
top_k=self.top_k,
start_date=start_date,
end_date=end_date,
)
# --- 3. Merge: mentioned first, then search (dedup by doc id) ---
seen_doc_ids: set[int] = set()

5
surfsense_backend/app/agents/new_chat/sandbox.py
View file

@ -138,7 +138,9 @@ def _find_or_create(thread_id: str) -> tuple[_TimeoutAwareSandbox, bool]:
try:
client.delete(sandbox)
except Exception:
logger.debug("Could not delete broken sandbox %s", sandbox.id, exc_info=True)
logger.debug(
"Could not delete broken sandbox %s", sandbox.id, exc_info=True
)
sandbox = client.create(_sandbox_create_params(labels))
is_new = True
logger.info("Created replacement sandbox: %s", sandbox.id)
@ -203,6 +205,7 @@ async def get_or_create_sandbox(
def _schedule_sandbox_delete(sandbox: _TimeoutAwareSandbox) -> None:
"""Best-effort background deletion of an evicted sandbox."""
def _delete() -> None:
try:
client = _get_client()

2
surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py
View file

@ -2,10 +2,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.confluence_history import ConfluenceHistoryConnector
from app.services.confluence import ConfluenceToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py
View file

@ -2,10 +2,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.confluence_history import ConfluenceHistoryConnector
from app.services.confluence import ConfluenceToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py
View file

@ -2,10 +2,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.confluence_history import ConfluenceHistoryConnector
from app.services.confluence import ConfluenceToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/dropbox/create_file.py
View file

@ -5,10 +5,10 @@ from pathlib import Path
from typing import Any, Literal
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.dropbox.client import DropboxClient
from app.db import SearchSourceConnector, SearchSourceConnectorType

2
surfsense_backend/app/agents/new_chat/tools/dropbox/trash_file.py
View file

@ -2,11 +2,11 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy import String, and_, cast, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.dropbox.client import DropboxClient
from app.db import (
Document,

2
surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py
View file

@ -6,9 +6,9 @@ from email.mime.text import MIMEText
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.services.gmail import GmailToolMetadataService
logger = logging.getLogger(__name__)

2
surfsense_backend/app/agents/new_chat/tools/gmail/send_email.py
View file

@ -6,9 +6,9 @@ from email.mime.text import MIMEText
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.services.gmail import GmailToolMetadataService
logger = logging.getLogger(__name__)

2
surfsense_backend/app/agents/new_chat/tools/gmail/trash_email.py
View file

@ -4,9 +4,9 @@ from datetime import datetime
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.services.gmail import GmailToolMetadataService
logger = logging.getLogger(__name__)

2
surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py
View file

@ -6,9 +6,9 @@ from email.mime.text import MIMEText
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.services.gmail import GmailToolMetadataService
logger = logging.getLogger(__name__)

4
surfsense_backend/app/agents/new_chat/tools/google_calendar/update_event.py
View file

@ -150,7 +150,9 @@ def create_update_calendar_event_tool(
final_new_end_datetime = result.params.get(
"new_end_datetime", new_end_datetime
)
final_new_description = result.params.get("new_description", new_description)
final_new_description = result.params.get(
"new_description", new_description
)
final_new_location = result.params.get("new_location", new_location)
final_new_attendees = result.params.get("new_attendees", new_attendees)

4
surfsense_backend/app/agents/new_chat/tools/hitl.py
View file

@ -58,7 +58,9 @@ def _parse_decision(approval: Any) -> tuple[str, dict[str, Any]]:
raise ValueError("No approval decision received")
decision = decisions[0]
decision_type: str = decision.get("type") or decision.get("decision_type") or "approve"
decision_type: str = (
decision.get("type") or decision.get("decision_type") or "approve"
)
edited_params: dict[str, Any] = {}
edited_action = decision.get("edited_action")

2
surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py
View file

@ -3,10 +3,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.jira_history import JiraHistoryConnector
from app.services.jira import JiraToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py
View file

@ -3,10 +3,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.jira_history import JiraHistoryConnector
from app.services.jira import JiraToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py
View file

@ -3,10 +3,10 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import flag_modified
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.jira_history import JiraHistoryConnector
from app.services.jira import JiraToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/linear/create_issue.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.linear_connector import LinearAPIError, LinearConnector
from app.services.linear import LinearToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/linear/delete_issue.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.linear_connector import LinearAPIError, LinearConnector
from app.services.linear import LinearToolMetadataService

10
surfsense_backend/app/agents/new_chat/tools/linear/update_issue.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.linear_connector import LinearAPIError, LinearConnector
from app.services.linear import LinearKBSyncService, LinearToolMetadataService
@ -157,9 +157,13 @@ def create_update_linear_issue_tool(
final_issue_id = result.params.get("issue_id", issue_id)
final_document_id = result.params.get("document_id", document_id)
final_new_title = result.params.get("new_title", new_title)
final_new_description = result.params.get("new_description", new_description)
final_new_description = result.params.get(
"new_description", new_description
)
final_new_state_id = result.params.get("new_state_id", new_state_id)
final_new_assignee_id = result.params.get("new_assignee_id", new_assignee_id)
final_new_assignee_id = result.params.get(
"new_assignee_id", new_assignee_id
)
final_new_priority = result.params.get("new_priority", new_priority)
final_new_label_ids: list[str] | None = result.params.get(
"new_label_ids", new_label_ids

2
surfsense_backend/app/agents/new_chat/tools/notion/create_page.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector
from app.services.notion import NotionToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/notion/delete_page.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector
from app.services.notion.tool_metadata_service import NotionToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/notion/update_page.py
View file

@ -2,9 +2,9 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.notion_history import NotionAPIError, NotionHistoryConnector
from app.services.notion import NotionToolMetadataService

2
surfsense_backend/app/agents/new_chat/tools/onedrive/create_file.py
View file

@ -5,10 +5,10 @@ from pathlib import Path
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.onedrive.client import OneDriveClient
from app.db import SearchSourceConnector, SearchSourceConnectorType

2
surfsense_backend/app/agents/new_chat/tools/onedrive/trash_file.py
View file

@ -2,11 +2,11 @@ import logging
from typing import Any
from langchain_core.tools import tool
from app.agents.new_chat.tools.hitl import request_approval
from sqlalchemy import String, and_, cast, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from app.agents.new_chat.tools.hitl import request_approval
from app.connectors.onedrive.client import OneDriveClient
from app.db import (
Document,

4
surfsense_backend/app/db.py
View file

@ -1321,6 +1321,10 @@ class SearchSpace(BaseModel, TimestampMixin):
Integer, nullable=True, default=0
) # For vision/screenshot analysis, defaults to Auto mode
ai_file_sort_enabled = Column(
Boolean, nullable=False, default=False, server_default="false"
)
user_id = Column(
UUID(as_uuid=True), ForeignKey("user.id", ondelete="CASCADE"), nullable=False
)

25
surfsense_backend/app/indexing_pipeline/indexing_pipeline_service.py
View file

@ -422,6 +422,8 @@ class IndexingPipelineService:
)
log_index_success(ctx, chunk_count=len(chunks))
await self._enqueue_ai_sort_if_enabled(document)
except RETRYABLE_LLM_ERRORS as e:
log_retryable_llm_error(ctx, e)
await rollback_and_persist_failure(
@ -457,6 +459,29 @@ class IndexingPipelineService:
return document
async def _enqueue_ai_sort_if_enabled(self, document: Document) -> None:
"""Fire-and-forget: enqueue incremental AI sort if the search space has it enabled."""
try:
from app.db import SearchSpace
result = await self.session.execute(
select(SearchSpace.ai_file_sort_enabled).where(
SearchSpace.id == document.search_space_id
)
)
enabled = result.scalar()
if not enabled:
return
from app.tasks.celery_tasks.document_tasks import ai_sort_document_task
user_id = str(document.created_by_id) if document.created_by_id else ""
ai_sort_document_task.delay(document.search_space_id, user_id, document.id)
except Exception:
logging.getLogger(__name__).warning(
"Failed to enqueue AI sort for document %s", document.id, exc_info=True
)
async def index_batch_parallel(
self,
connector_docs: list[ConnectorDocument],

4
surfsense_backend/app/routes/export_routes.py
View file

@ -20,7 +20,9 @@ router = APIRouter()
@router.get("/search-spaces/{search_space_id}/export")
async def export_knowledge_base(
search_space_id: int,
folder_id: int | None = Query(None, description="Export only this folder's subtree"),
folder_id: int | None = Query(
None, description="Export only this folder's subtree"
),
session: AsyncSession = Depends(get_async_session),
user: User = Depends(current_active_user),
):

37
surfsense_backend/app/routes/search_spaces_routes.py
View file

@ -216,6 +216,7 @@ async def read_search_spaces(
user_id=space.user_id,
citations_enabled=space.citations_enabled,
qna_custom_instructions=space.qna_custom_instructions,
ai_file_sort_enabled=space.ai_file_sort_enabled,
member_count=member_count,
is_owner=is_owner,
)
@ -384,6 +385,42 @@ async def edit_team_memory(
return db_search_space
@router.post("/searchspaces/{search_space_id}/ai-sort")
async def trigger_ai_sort(
search_space_id: int,
session: AsyncSession = Depends(get_async_session),
user: User = Depends(current_active_user),
):
"""Trigger a full AI file sort for all documents in the search space."""
try:
await check_permission(
session,
user,
search_space_id,
Permission.SETTINGS_UPDATE.value,
"You don't have permission to trigger AI sort on this search space",
)
result = await session.execute(
select(SearchSpace).filter(SearchSpace.id == search_space_id)
)
db_search_space = result.scalars().first()
if not db_search_space:
raise HTTPException(status_code=404, detail="Search space not found")
from app.tasks.celery_tasks.document_tasks import ai_sort_search_space_task
ai_sort_search_space_task.delay(search_space_id, str(user.id))
return {"message": "AI sort started"}
except HTTPException:
raise
except Exception as e:
logger.error(f"Failed to trigger AI sort: {e!s}", exc_info=True)
raise HTTPException(
status_code=500, detail=f"Failed to trigger AI sort: {e!s}"
) from e
@router.delete("/searchspaces/{search_space_id}", response_model=dict)
async def delete_search_space(
search_space_id: int,

2
surfsense_backend/app/schemas/search_space.py
View file

@ -22,6 +22,7 @@ class SearchSpaceUpdate(BaseModel):
citations_enabled: bool | None = None
qna_custom_instructions: str | None = None
shared_memory_md: str | None = None
ai_file_sort_enabled: bool | None = None
class SearchSpaceRead(SearchSpaceBase, IDModel, TimestampModel):
@ -31,6 +32,7 @@ class SearchSpaceRead(SearchSpaceBase, IDModel, TimestampModel):
citations_enabled: bool
qna_custom_instructions: str | None = None
shared_memory_md: str | None = None
ai_file_sort_enabled: bool = False
model_config = ConfigDict(from_attributes=True)

329
surfsense_backend/app/services/ai_file_sort_service.py Normal file
View file

@ -0,0 +1,329 @@
"""AI File Sort Service: builds connector-type/date/category/subcategory folder paths."""
from __future__ import annotations
import json
import logging
import re
from datetime import UTC, datetime
from langchain_core.messages import HumanMessage
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy.orm import selectinload
from app.db import (
Chunk,
Document,
DocumentType,
SearchSourceConnector,
SearchSourceConnectorType,
)
from app.services.folder_service import ensure_folder_hierarchy_with_depth_validation
logger = logging.getLogger(__name__)
_DOCTYPE_TO_CONNECTOR_LABEL: dict[str, str] = {
DocumentType.EXTENSION: "Browser Extension",
DocumentType.CRAWLED_URL: "Web Crawl",
DocumentType.FILE: "File Upload",
DocumentType.SLACK_CONNECTOR: "Slack",
DocumentType.TEAMS_CONNECTOR: "Teams",
DocumentType.ONEDRIVE_FILE: "OneDrive",
DocumentType.NOTION_CONNECTOR: "Notion",
DocumentType.YOUTUBE_VIDEO: "YouTube",
DocumentType.GITHUB_CONNECTOR: "GitHub",
DocumentType.LINEAR_CONNECTOR: "Linear",
DocumentType.DISCORD_CONNECTOR: "Discord",
DocumentType.JIRA_CONNECTOR: "Jira",
DocumentType.CONFLUENCE_CONNECTOR: "Confluence",
DocumentType.CLICKUP_CONNECTOR: "ClickUp",
DocumentType.GOOGLE_CALENDAR_CONNECTOR: "Google Calendar",
DocumentType.GOOGLE_GMAIL_CONNECTOR: "Gmail",
DocumentType.GOOGLE_DRIVE_FILE: "Google Drive",
DocumentType.AIRTABLE_CONNECTOR: "Airtable",
DocumentType.LUMA_CONNECTOR: "Luma",
DocumentType.ELASTICSEARCH_CONNECTOR: "Elasticsearch",
DocumentType.BOOKSTACK_CONNECTOR: "BookStack",
DocumentType.CIRCLEBACK: "Circleback",
DocumentType.OBSIDIAN_CONNECTOR: "Obsidian",
DocumentType.NOTE: "Notes",
DocumentType.DROPBOX_FILE: "Dropbox",
DocumentType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR: "Google Drive (Composio)",
DocumentType.COMPOSIO_GMAIL_CONNECTOR: "Gmail (Composio)",
DocumentType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR: "Google Calendar (Composio)",
DocumentType.LOCAL_FOLDER_FILE: "Local Folder",
}
_CONNECTOR_TYPE_LABEL: dict[str, str] = {
SearchSourceConnectorType.SERPER_API: "Serper Search",
SearchSourceConnectorType.TAVILY_API: "Tavily Search",
SearchSourceConnectorType.SEARXNG_API: "SearXNG Search",
SearchSourceConnectorType.LINKUP_API: "Linkup Search",
SearchSourceConnectorType.BAIDU_SEARCH_API: "Baidu Search",
SearchSourceConnectorType.SLACK_CONNECTOR: "Slack",
SearchSourceConnectorType.TEAMS_CONNECTOR: "Teams",
SearchSourceConnectorType.ONEDRIVE_CONNECTOR: "OneDrive",
SearchSourceConnectorType.NOTION_CONNECTOR: "Notion",
SearchSourceConnectorType.GITHUB_CONNECTOR: "GitHub",
SearchSourceConnectorType.LINEAR_CONNECTOR: "Linear",
SearchSourceConnectorType.DISCORD_CONNECTOR: "Discord",
SearchSourceConnectorType.JIRA_CONNECTOR: "Jira",
SearchSourceConnectorType.CONFLUENCE_CONNECTOR: "Confluence",
SearchSourceConnectorType.CLICKUP_CONNECTOR: "ClickUp",
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR: "Google Calendar",
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR: "Gmail",
SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR: "Google Drive",
SearchSourceConnectorType.AIRTABLE_CONNECTOR: "Airtable",
SearchSourceConnectorType.LUMA_CONNECTOR: "Luma",
SearchSourceConnectorType.ELASTICSEARCH_CONNECTOR: "Elasticsearch",
SearchSourceConnectorType.WEBCRAWLER_CONNECTOR: "Web Crawl",
SearchSourceConnectorType.BOOKSTACK_CONNECTOR: "BookStack",
SearchSourceConnectorType.CIRCLEBACK_CONNECTOR: "Circleback",
SearchSourceConnectorType.OBSIDIAN_CONNECTOR: "Obsidian",
SearchSourceConnectorType.MCP_CONNECTOR: "MCP",
SearchSourceConnectorType.DROPBOX_CONNECTOR: "Dropbox",
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR: "Google Drive (Composio)",
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR: "Gmail (Composio)",
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR: "Google Calendar (Composio)",
}
_MAX_CONTENT_CHARS = 4000
_MAX_CHUNKS_FOR_CONTEXT = 5
_CATEGORY_PROMPT = (
"Based on the document information below, classify it into a broad category "
"and a more specific subcategory.\n\n"
"Rules:\n"
"- category: 1-2 word broad theme (e.g. Science, Finance, Engineering, Communication, Media)\n"
"- subcategory: 1-2 word specific topic within the category "
"(e.g. Physics, Tax Reports, Backend, Team Updates)\n"
"- Use nouns only. Do not include generic terms like 'General' or 'Miscellaneous'.\n\n"
"Title: {title}\n\n"
"Content: {summary}\n\n"
'Respond with ONLY a JSON object: {{"category": "...", "subcategory": "..."}}'
)
_SAFE_NAME_RE = re.compile(r"[^a-zA-Z0-9 _\-()]")
_FALLBACK_CATEGORY = "Uncategorized"
_FALLBACK_SUBCATEGORY = "General"
def resolve_root_folder_label(
document: Document, connector: SearchSourceConnector | None
) -> str:
if connector is not None:
return _CONNECTOR_TYPE_LABEL.get(
connector.connector_type, str(connector.connector_type)
)
return _DOCTYPE_TO_CONNECTOR_LABEL.get(
document.document_type, str(document.document_type)
)
def resolve_date_folder(document: Document) -> str:
ts = document.updated_at or document.created_at
if ts is None:
ts = datetime.now(UTC)
return ts.strftime("%Y-%m-%d")
def sanitize_category_folder_name(
value: str | None, fallback: str = _FALLBACK_CATEGORY
) -> str:
if not value or not value.strip():
return fallback
cleaned = _SAFE_NAME_RE.sub("", value.strip())
cleaned = " ".join(cleaned.split())
if not cleaned:
return fallback
return cleaned[:50]
async def _resolve_document_text(
session: AsyncSession,
document: Document,
) -> str:
"""Build the best available text representation for taxonomy generation.
Prefers ``document.content``; falls back to joining the first few chunks
when content is empty or too short to be useful.
"""
text = (document.content or "").strip()
if len(text) >= 100:
return text[:_MAX_CONTENT_CHARS]
stmt = (
select(Chunk.content)
.where(Chunk.document_id == document.id)
.order_by(Chunk.id)
.limit(_MAX_CHUNKS_FOR_CONTEXT)
)
result = await session.execute(stmt)
chunk_texts = [row[0] for row in result.all() if row[0]]
if chunk_texts:
combined = "\n\n".join(chunk_texts)
return combined[:_MAX_CONTENT_CHARS]
return text[:_MAX_CONTENT_CHARS]
def _get_cached_taxonomy(document: Document) -> tuple[str, str] | None:
"""Return (category, subcategory) from document metadata cache, or None."""
meta = document.document_metadata
if not isinstance(meta, dict):
return None
cat = meta.get("ai_sort_category")
subcat = meta.get("ai_sort_subcategory")
if cat and subcat and isinstance(cat, str) and isinstance(subcat, str):
return cat, subcat
return None
def _set_cached_taxonomy(document: Document, category: str, subcategory: str) -> None:
"""Persist the AI taxonomy on document metadata for deterministic re-sorts."""
meta = dict(document.document_metadata or {})
meta["ai_sort_category"] = category
meta["ai_sort_subcategory"] = subcategory
document.document_metadata = meta
async def generate_ai_taxonomy(
title: str,
summary_or_content: str,
llm,
) -> tuple[str, str]:
"""Return (category, subcategory) using a single structured LLM call."""
text = (summary_or_content or "").strip()
if not text:
return _FALLBACK_CATEGORY, _FALLBACK_SUBCATEGORY
if len(text) > _MAX_CONTENT_CHARS:
text = text[:_MAX_CONTENT_CHARS]
prompt = _CATEGORY_PROMPT.format(title=title or "Untitled", summary=text)
try:
result = await llm.ainvoke([HumanMessage(content=prompt)])
raw = result.content.strip()
if raw.startswith("```"):
raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
parsed = json.loads(raw)
category = sanitize_category_folder_name(
parsed.get("category"), _FALLBACK_CATEGORY
)
subcategory = sanitize_category_folder_name(
parsed.get("subcategory"), _FALLBACK_SUBCATEGORY
)
return category, subcategory
except Exception:
logger.warning("AI taxonomy generation failed, using fallback", exc_info=True)
return _FALLBACK_CATEGORY, _FALLBACK_SUBCATEGORY
def _build_path_segments(
root_label: str,
date_label: str,
category: str,
subcategory: str,
) -> list[dict]:
return [
{"name": root_label, "metadata": {"ai_sort": True, "ai_sort_level": 1}},
{"name": date_label, "metadata": {"ai_sort": True, "ai_sort_level": 2}},
{"name": category, "metadata": {"ai_sort": True, "ai_sort_level": 3}},
{"name": subcategory, "metadata": {"ai_sort": True, "ai_sort_level": 4}},
]
async def _resolve_taxonomy(
session: AsyncSession,
document: Document,
llm,
) -> tuple[str, str]:
"""Return (category, subcategory), reusing cached values when available."""
cached = _get_cached_taxonomy(document)
if cached is not None:
return cached
content_text = await _resolve_document_text(session, document)
category, subcategory = await generate_ai_taxonomy(
document.title, content_text, llm
)
_set_cached_taxonomy(document, category, subcategory)
return category, subcategory
async def ai_sort_document(
session: AsyncSession,
document: Document,
llm,
) -> Document:
"""Sort a single document into the 4-level AI folder hierarchy."""
connector: SearchSourceConnector | None = None
if document.connector_id is not None:
connector = await session.get(SearchSourceConnector, document.connector_id)
root_label = resolve_root_folder_label(document, connector)
date_label = resolve_date_folder(document)
category, subcategory = await _resolve_taxonomy(session, document, llm)
segments = _build_path_segments(root_label, date_label, category, subcategory)
leaf_folder = await ensure_folder_hierarchy_with_depth_validation(
session,
document.search_space_id,
segments,
)
document.folder_id = leaf_folder.id
await session.flush()
return document
async def ai_sort_all_documents(
session: AsyncSession,
search_space_id: int,
llm,
) -> tuple[int, int]:
"""Sort all documents in a search space. Returns (sorted_count, failed_count)."""
stmt = (
select(Document)
.where(Document.search_space_id == search_space_id)
.options(selectinload(Document.connector))
)
result = await session.execute(stmt)
documents = list(result.scalars().all())
sorted_count = 0
failed_count = 0
for doc in documents:
try:
connector = doc.connector
root_label = resolve_root_folder_label(doc, connector)
date_label = resolve_date_folder(doc)
category, subcategory = await _resolve_taxonomy(session, doc, llm)
segments = _build_path_segments(
root_label, date_label, category, subcategory
)
leaf_folder = await ensure_folder_hierarchy_with_depth_validation(
session,
search_space_id,
segments,
)
doc.folder_id = leaf_folder.id
sorted_count += 1
except Exception:
logger.error("Failed to AI-sort document %s", doc.id, exc_info=True)
failed_count += 1
await session.commit()
logger.info(
"AI sort complete for search_space=%d: sorted=%d, failed=%d",
search_space_id,
sorted_count,
failed_count,
)
return sorted_count, failed_count

52
surfsense_backend/app/services/folder_service.py
View file

@ -142,6 +142,58 @@ async def generate_folder_position(
return generate_key_between(last_position, None)
async def ensure_folder_hierarchy_with_depth_validation(
session: AsyncSession,
search_space_id: int,
path_segments: list[dict],
) -> Folder:
"""Create or return a nested folder chain, validating depth at each step.
Each item in ``path_segments`` is a dict with:
- ``name`` (str): folder display name
- ``metadata`` (dict | None): optional ``folder_metadata`` JSONB payload
Returns the deepest (leaf) Folder in the chain.
"""
parent_id: int | None = None
current_folder: Folder | None = None
for segment in path_segments:
name = segment["name"]
metadata = segment.get("metadata")
stmt = select(Folder).where(
Folder.search_space_id == search_space_id,
Folder.name == name,
Folder.parent_id == parent_id
if parent_id is not None
else Folder.parent_id.is_(None),
)
result = await session.execute(stmt)
folder = result.scalar_one_or_none()
if folder is None:
await validate_folder_depth(session, parent_id, subtree_depth=0)
position = await generate_folder_position(
session, search_space_id, parent_id
)
folder = Folder(
name=name,
search_space_id=search_space_id,
parent_id=parent_id,
position=position,
folder_metadata=metadata,
)
session.add(folder)
await session.flush()
current_folder = folder
parent_id = folder.id
assert current_folder is not None, "path_segments must not be empty"
return current_folder
async def get_folder_subtree_ids(session: AsyncSession, folder_id: int) -> list[int]:
"""Return all folder IDs in the subtree rooted at folder_id (inclusive)."""
result = await session.execute(

119
surfsense_backend/app/tasks/celery_tasks/document_tasks.py
View file

@ -4,6 +4,7 @@ import asyncio
import contextlib
import logging
import os
import time
from uuid import UUID
from app.celery_app import celery_app
@ -1551,3 +1552,121 @@ async def _index_uploaded_folder_files_async(
heartbeat_task.cancel()
if notification_id is not None:
_stop_heartbeat(notification_id)
# ===== AI File Sort tasks =====
AI_SORT_LOCK_TTL_SECONDS = 600 # 10 minutes
_ai_sort_redis = None
def _get_ai_sort_redis():
import redis
global _ai_sort_redis
if _ai_sort_redis is None:
_ai_sort_redis = redis.from_url(config.REDIS_APP_URL, decode_responses=True)
return _ai_sort_redis
def _ai_sort_lock_key(search_space_id: int) -> str:
return f"ai_sort:search_space:{search_space_id}:lock"
@celery_app.task(name="ai_sort_search_space", bind=True, max_retries=1)
def ai_sort_search_space_task(self, search_space_id: int, user_id: str):
"""Full AI sort for all documents in a search space."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(_ai_sort_search_space_async(search_space_id, user_id))
finally:
loop.close()
async def _ai_sort_search_space_async(search_space_id: int, user_id: str):
r = _get_ai_sort_redis()
lock_key = _ai_sort_lock_key(search_space_id)
if not r.set(lock_key, "running", nx=True, ex=AI_SORT_LOCK_TTL_SECONDS):
logger.info(
"AI sort already running for search_space=%d, skipping",
search_space_id,
)
return
t_start = time.perf_counter()
try:
from app.services.ai_file_sort_service import ai_sort_all_documents
from app.services.llm_service import get_document_summary_llm
async with get_celery_session_maker()() as session:
llm = await get_document_summary_llm(
session, search_space_id, disable_streaming=True
)
if llm is None:
logger.warning(
"No LLM configured for search_space=%d, skipping AI sort",
search_space_id,
)
return
sorted_count, failed_count = await ai_sort_all_documents(
session, search_space_id, llm
)
elapsed = time.perf_counter() - t_start
logger.info(
"AI sort search_space=%d done in %.1fs: sorted=%d failed=%d",
search_space_id,
elapsed,
sorted_count,
failed_count,
)
finally:
r.delete(lock_key)
@celery_app.task(
name="ai_sort_document", bind=True, max_retries=2, default_retry_delay=10
)
def ai_sort_document_task(self, search_space_id: int, user_id: str, document_id: int):
"""Incremental AI sort for a single document after indexing."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(
_ai_sort_document_async(search_space_id, user_id, document_id)
)
finally:
loop.close()
async def _ai_sort_document_async(search_space_id: int, user_id: str, document_id: int):
from app.db import Document
from app.services.ai_file_sort_service import ai_sort_document
from app.services.llm_service import get_document_summary_llm
async with get_celery_session_maker()() as session:
document = await session.get(Document, document_id)
if document is None:
logger.warning("Document %d not found, skipping AI sort", document_id)
return
llm = await get_document_summary_llm(
session, search_space_id, disable_streaming=True
)
if llm is None:
logger.warning(
"No LLM for search_space=%d, skipping AI sort of doc=%d",
search_space_id,
document_id,
)
return
await ai_sort_document(session, document, llm)
await session.commit()
logger.info(
"AI sorted document=%d into search_space=%d",
document_id,
search_space_id,
)

9
surfsense_backend/app/tasks/chat/stream_new_chat.py
View file

@ -61,6 +61,7 @@ from app.services.new_streaming_service import VercelStreamingService
from app.utils.content_utils import bootstrap_history_from_db
from app.utils.perf import get_perf_logger, log_system_snapshot, trim_native_heap
_background_tasks: set[asyncio.Task] = set()
_perf_log = get_perf_logger()
@ -1552,7 +1553,7 @@ async def stream_new_chat(
# Shared threads write to team memory; private threads write to user memory.
if not stream_result.agent_called_update_memory:
if visibility == ChatVisibility.SEARCH_SPACE:
asyncio.create_task(
task = asyncio.create_task(
extract_and_save_team_memory(
user_message=user_query,
search_space_id=search_space_id,
@ -1560,14 +1561,18 @@ async def stream_new_chat(
author_display_name=current_user_display_name,
)
)
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
elif user_id:
asyncio.create_task(
task = asyncio.create_task(
extract_and_save_memory(
user_message=user_query,
user_id=user_id,
llm=llm,
)
)
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
# Finish the step and message
yield streaming_service.format_finish_step()

3
surfsense_backend/app/tasks/connector_indexers/google_drive_indexer.py
View file

@ -961,6 +961,7 @@ async def index_google_drive_files(
vision_llm = None
if connector_enable_vision_llm:
from app.services.llm_service import get_vision_llm
vision_llm = await get_vision_llm(session, search_space_id)
drive_client = GoogleDriveClient(
session, connector_id, credentials=pre_built_credentials
@ -1168,6 +1169,7 @@ async def index_google_drive_single_file(
vision_llm = None
if connector_enable_vision_llm:
from app.services.llm_service import get_vision_llm
vision_llm = await get_vision_llm(session, search_space_id)
drive_client = GoogleDriveClient(
session, connector_id, credentials=pre_built_credentials
@ -1306,6 +1308,7 @@ async def index_google_drive_selected_files(
vision_llm = None
if connector_enable_vision_llm:
from app.services.llm_service import get_vision_llm
vision_llm = await get_vision_llm(session, search_space_id)
drive_client = GoogleDriveClient(
session, connector_id, credentials=pre_built_credentials

4
surfsense_backend/app/tasks/connector_indexers/local_folder_indexer.py
View file

@ -1360,7 +1360,9 @@ async def index_uploaded_files(
try:
content, content_hash = await _compute_file_content_hash(
temp_path, filename, search_space_id,
temp_path,
filename,
search_space_id,
vision_llm=vision_llm_instance,
)
except Exception as e:

1
surfsense_backend/app/tasks/connector_indexers/onedrive_indexer.py
View file

@ -656,6 +656,7 @@ async def index_onedrive_files(
vision_llm = None
if connector_enable_vision_llm:
from app.services.llm_service import get_vision_llm
vision_llm = await get_vision_llm(session, search_space_id)
onedrive_client = OneDriveClient(session, connector_id)