apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-01 03:46:25 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

refactor: unify all 3 google Composio and non-Composio connector types and pipelines keeping same credential adapters

Browse source
This commit is contained in:
Anish Sarkar 2026-03-19 05:08:21 +05:30
parent 6c37b563c0
commit 83152e8e7e
24 changed files with 633 additions and 3596 deletions
Download patch file Download diff file Expand all files Collapse all files

24
surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py
View file

@ -157,14 +157,18 @@ def create_create_google_drive_file_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_DRIVE_TYPES = [
SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
]
if final_connector_id is not None:
result = await db_session.execute(
select(SearchSourceConnector).filter(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnector.connector_type.in_(_DRIVE_TYPES),
)
)
connector = result.scalars().first()
@ -179,8 +183,7 @@ def create_create_google_drive_file_tool(
select(SearchSourceConnector).filter(
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnector.connector_type.in_(_DRIVE_TYPES),
)
)
connector = result.scalars().first()
@ -194,8 +197,19 @@ def create_create_google_drive_file_tool(
logger.info(
f"Creating Google Drive file: name='{final_name}', type='{final_file_type}', connector={actual_connector_id}"
)
pre_built_creds = None
if connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR:
from app.utils.google_credentials import build_composio_credentials
cca_id = connector.config.get("composio_connected_account_id")
if cca_id:
pre_built_creds = build_composio_credentials(cca_id)
client = GoogleDriveClient(
session=db_session, connector_id=actual_connector_id
session=db_session,
connector_id=actual_connector_id,
credentials=pre_built_creds,
)
try:
created = await client.create_file(

23
surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py
View file

@ -151,13 +151,17 @@ def create_delete_google_drive_file_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_DRIVE_TYPES = [
SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
]
result = await db_session.execute(
select(SearchSourceConnector).filter(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type
== SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnector.connector_type.in_(_DRIVE_TYPES),
)
)
connector = result.scalars().first()
@ -170,7 +174,20 @@ def create_delete_google_drive_file_tool(
logger.info(
f"Deleting Google Drive file: file_id='{final_file_id}', connector={final_connector_id}"
)
client = GoogleDriveClient(session=db_session, connector_id=connector.id)
pre_built_creds = None
if connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR:
from app.utils.google_credentials import build_composio_credentials
cca_id = connector.config.get("composio_connected_account_id")
if cca_id:
pre_built_creds = build_composio_credentials(cca_id)
client = GoogleDriveClient(
session=db_session,
connector_id=connector.id,
credentials=pre_built_creds,
)
try:
await client.trash_file(file_id=final_file_id)
except HttpError as http_err:

48
surfsense_backend/app/agents/new_chat/tools/knowledge_base.py
View file

@ -60,7 +60,7 @@ def _is_degenerate_query(query: str) -> bool:
async def _browse_recent_documents(
search_space_id: int,
document_type: str | None,
document_type: str | list[str] | None,
top_k: int,
start_date: datetime | None,
end_date: datetime | None,
@ -83,14 +83,22 @@ async def _browse_recent_documents(
base_conditions = [Document.search_space_id == search_space_id]
if document_type is not None:
if isinstance(document_type, str):
try:
doc_type_enum = DocumentType[document_type]
base_conditions.append(Document.document_type == doc_type_enum)
except KeyError:
return []
type_list = document_type if isinstance(document_type, list) else [document_type]
doc_type_enums = []
for dt in type_list:
if isinstance(dt, str):
try:
doc_type_enums.append(DocumentType[dt])
except KeyError:
pass
else:
doc_type_enums.append(dt)
if not doc_type_enums:
return []
if len(doc_type_enums) == 1:
base_conditions.append(Document.document_type == doc_type_enums[0])
else:
base_conditions.append(Document.document_type == document_type)
base_conditions.append(Document.document_type.in_(doc_type_enums))
if start_date is not None:
base_conditions.append(Document.updated_at >= start_date)
@ -195,10 +203,6 @@ _ALL_CONNECTORS: list[str] = [
"CRAWLED_URL",
"CIRCLEBACK",
"OBSIDIAN_CONNECTOR",
# Composio connectors
"COMPOSIO_GOOGLE_DRIVE_CONNECTOR",
"COMPOSIO_GMAIL_CONNECTOR",
"COMPOSIO_GOOGLE_CALENDAR_CONNECTOR",
]
# Human-readable descriptions for each connector type
@ -228,10 +232,6 @@ CONNECTOR_DESCRIPTIONS: dict[str, str] = {
"BOOKSTACK_CONNECTOR": "BookStack pages (personal documentation)",
"CIRCLEBACK": "Circleback meeting notes, transcripts, and action items",
"OBSIDIAN_CONNECTOR": "Obsidian vault notes and markdown files (personal notes)",
# Composio connectors
"COMPOSIO_GOOGLE_DRIVE_CONNECTOR": "Google Drive files via Composio (personal cloud storage)",
"COMPOSIO_GMAIL_CONNECTOR": "Gmail emails via Composio (personal emails)",
"COMPOSIO_GOOGLE_CALENDAR_CONNECTOR": "Google Calendar events via Composio (personal calendar)",
}
@ -654,6 +654,20 @@ async def search_knowledge_base_async(
)
browse_connectors = connectors if connectors else [None] # type: ignore[list-item]
# Expand native Google types to include legacy Composio equivalents
# so old documents remain searchable until re-indexed.
_LEGACY_ALIASES: dict[str, str] = {
"GOOGLE_DRIVE_FILE": "COMPOSIO_GOOGLE_DRIVE_CONNECTOR",
"GOOGLE_GMAIL_CONNECTOR": "COMPOSIO_GMAIL_CONNECTOR",
"GOOGLE_CALENDAR_CONNECTOR": "COMPOSIO_GOOGLE_CALENDAR_CONNECTOR",
}
expanded_browse = []
for c in browse_connectors:
if c is not None and c in _LEGACY_ALIASES:
expanded_browse.append([c, _LEGACY_ALIASES[c]])
else:
expanded_browse.append(c)
browse_results = await asyncio.gather(
*[
_browse_recent_documents(
@ -663,7 +677,7 @@ async def search_knowledge_base_async(
start_date=resolved_start_date,
end_date=resolved_end_date,
)
for c in browse_connectors
for c in expanded_browse
]
)
for docs in browse_results: