apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-06-24 21:38:09 +02:00
Code Issues Projects Releases Packages Wiki Activity

chore: ran linting

Browse source
This commit is contained in:
Anish Sarkar 2026-03-22 00:43:53 +05:30
parent 44f59bfb61
commit 2bc6a0c3bc
24 changed files with 166 additions and 160 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -144,7 +144,7 @@ def create_create_gmail_draft_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_GMAIL_TYPES = [
_gmail_types = [
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
]
@ -155,7 +155,7 @@ def create_create_gmail_draft_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()
@ -170,7 +170,7 @@ def create_create_gmail_draft_tool(
select(SearchSourceConnector).filter(
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()

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

@ -145,7 +145,7 @@ def create_send_gmail_email_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_GMAIL_TYPES = [
_gmail_types = [
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
]
@ -156,7 +156,7 @@ def create_send_gmail_email_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()
@ -171,7 +171,7 @@ def create_send_gmail_email_tool(
select(SearchSourceConnector).filter(
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()

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

@ -162,7 +162,7 @@ def create_trash_gmail_email_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_GMAIL_TYPES = [
_gmail_types = [
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
]
@ -172,7 +172,7 @@ def create_trash_gmail_email_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()

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

@ -192,7 +192,7 @@ def create_update_gmail_draft_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_GMAIL_TYPES = [
_gmail_types = [
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
]
@ -202,7 +202,7 @@ def create_update_gmail_draft_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
SearchSourceConnector.connector_type.in_(_gmail_types),
)
)
connector = result.scalars().first()
@ -419,7 +419,9 @@ async def _find_draft_id_by_message(gmail_service: Any, message_id: str) -> str
response = await asyncio.get_event_loop().run_in_executor(
None,
lambda: gmail_service.users().drafts().list(**kwargs).execute(),
lambda kwargs=kwargs: (
gmail_service.users().drafts().list(**kwargs).execute()
),
)
for draft in response.get("drafts", []):

6
surfsense_backend/app/agents/new_chat/tools/google_calendar/create_event.py
View file

@ -155,7 +155,7 @@ def create_create_calendar_event_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_CALENDAR_TYPES = [
_calendar_types = [
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
]
@ -166,7 +166,7 @@ def create_create_calendar_event_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_CALENDAR_TYPES),
SearchSourceConnector.connector_type.in_(_calendar_types),
)
)
connector = result.scalars().first()
@ -181,7 +181,7 @@ def create_create_calendar_event_tool(
select(SearchSourceConnector).filter(
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_CALENDAR_TYPES),
SearchSourceConnector.connector_type.in_(_calendar_types),
)
)
connector = result.scalars().first()

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

@ -161,7 +161,7 @@ def create_delete_calendar_event_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_CALENDAR_TYPES = [
_calendar_types = [
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
]
@ -171,7 +171,7 @@ def create_delete_calendar_event_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_CALENDAR_TYPES),
SearchSourceConnector.connector_type.in_(_calendar_types),
)
)
connector = result.scalars().first()

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

@ -178,7 +178,7 @@ def create_update_calendar_event_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_CALENDAR_TYPES = [
_calendar_types = [
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
]
@ -188,7 +188,7 @@ def create_update_calendar_event_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_CALENDAR_TYPES),
SearchSourceConnector.connector_type.in_(_calendar_types),
)
)
connector = result.scalars().first()

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

@ -166,7 +166,7 @@ def create_create_google_drive_file_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_DRIVE_TYPES = [
_drive_types = [
SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
]
@ -177,7 +177,7 @@ def create_create_google_drive_file_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_DRIVE_TYPES),
SearchSourceConnector.connector_type.in_(_drive_types),
)
)
connector = result.scalars().first()
@ -192,7 +192,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.in_(_DRIVE_TYPES),
SearchSourceConnector.connector_type.in_(_drive_types),
)
)
connector = result.scalars().first()

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

@ -162,7 +162,7 @@ def create_delete_google_drive_file_tool(
from app.db import SearchSourceConnector, SearchSourceConnectorType
_DRIVE_TYPES = [
_drive_types = [
SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
]
@ -172,7 +172,7 @@ def create_delete_google_drive_file_tool(
SearchSourceConnector.id == final_connector_id,
SearchSourceConnector.search_space_id == search_space_id,
SearchSourceConnector.user_id == user_id,
SearchSourceConnector.connector_type.in_(_DRIVE_TYPES),
SearchSourceConnector.connector_type.in_(_drive_types),
)
)
connector = result.scalars().first()

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

@ -9,6 +9,7 @@ This module provides:
"""
import asyncio
import contextlib
import json
import re
import time
@ -89,10 +90,8 @@ async def _browse_recent_documents(
doc_type_enums = []
for dt in type_list:
if isinstance(dt, str):
try:
with contextlib.suppress(KeyError):
doc_type_enums.append(DocumentType[dt])
except KeyError:
pass
else:
doc_type_enums.append(dt)
if not doc_type_enums:

11
surfsense_backend/app/connectors/google_calendar_connector.py
View file

@ -59,11 +59,12 @@ class GoogleCalendarConnector:
"""
has_standard_refresh = bool(self._credentials.refresh_token)
if has_standard_refresh:
if not all([self._credentials.client_id, self._credentials.client_secret]):
raise ValueError(
"Google OAuth credentials (client_id, client_secret) must be set"
)
if has_standard_refresh and not all(
[self._credentials.client_id, self._credentials.client_secret]
):
raise ValueError(
"Google OAuth credentials (client_id, client_secret) must be set"
)
if self._credentials and not self._credentials.expired:
return self._credentials

1
surfsense_backend/app/connectors/google_drive/client.py
View file

@ -3,6 +3,7 @@
import io
from typing import Any
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from googleapiclient.http import MediaIoBaseUpload

11
surfsense_backend/app/connectors/google_gmail_connector.py
View file

@ -88,11 +88,12 @@ class GoogleGmailConnector:
"""
has_standard_refresh = bool(self._credentials.refresh_token)
if has_standard_refresh:
if not all([self._credentials.client_id, self._credentials.client_secret]):
raise ValueError(
"Google OAuth credentials (client_id, client_secret) must be set"
)
if has_standard_refresh and not all(
[self._credentials.client_id, self._credentials.client_secret]
):
raise ValueError(
"Google OAuth credentials (client_id, client_secret) must be set"
)
if self._credentials and not self._credentials.expired:
return self._credentials

5
surfsense_backend/app/retriever/chunks_hybrid_search.py
View file

@ -1,4 +1,5 @@
import asyncio
import contextlib
import time
from datetime import datetime
@ -225,10 +226,8 @@ class ChucksHybridSearchRetriever:
doc_type_enums = []
for dt in type_list:
if isinstance(dt, str):
try:
with contextlib.suppress(KeyError):
doc_type_enums.append(DocumentType[dt])
except KeyError:
pass
else:
doc_type_enums.append(dt)
if not doc_type_enums:

5
surfsense_backend/app/retriever/documents_hybrid_search.py
View file

@ -1,3 +1,4 @@
import contextlib
import time
from datetime import datetime
@ -205,10 +206,8 @@ class DocumentHybridSearchRetriever:
doc_type_enums = []
for dt in type_list:
if isinstance(dt, str):
try:
with contextlib.suppress(KeyError):
doc_type_enums.append(DocumentType[dt])
except KeyError:
pass
else:
doc_type_enums.append(dt)
if not doc_type_enums:

8
surfsense_backend/app/services/gmail/tool_metadata_service.py
View file

@ -225,7 +225,9 @@ class GmailToolMetadataService:
service = build("gmail", "v1", credentials=creds)
profile = await asyncio.get_event_loop().run_in_executor(
None,
lambda: service.users().getProfile(userId="me").execute(),
lambda service=service: (
service.users().getProfile(userId="me").execute()
),
)
acc_dict["email"] = profile.get("emailAddress", "")
except Exception:
@ -334,7 +336,9 @@ class GmailToolMetadataService:
kwargs["pageToken"] = page_token
response = await asyncio.get_event_loop().run_in_executor(
None,
lambda: service.users().drafts().list(**kwargs).execute(),
lambda kwargs=kwargs: (
service.users().drafts().list(**kwargs).execute()
),
)
for draft in response.get("drafts", []):
if draft.get("message", {}).get("id") == message_id:

18
surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py
View file

@ -16,15 +16,6 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.connectors.google_calendar_connector import GoogleCalendarConnector
from app.db import Document, DocumentStatus, DocumentType, SearchSourceConnectorType
from app.utils.google_credentials import (
COMPOSIO_GOOGLE_CONNECTOR_TYPES,
build_composio_credentials,
)
ACCEPTED_CALENDAR_CONNECTOR_TYPES = {
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
}
from app.services.llm_service import get_user_long_context_llm
from app.services.task_logging_service import TaskLoggingService
from app.utils.document_converters import (
@ -34,6 +25,10 @@ from app.utils.document_converters import (
generate_document_summary,
generate_unique_identifier_hash,
)
from app.utils.google_credentials import (
COMPOSIO_GOOGLE_CONNECTOR_TYPES,
build_composio_credentials,
)
from .base import (
check_document_by_unique_identifier,
@ -46,6 +41,11 @@ from .base import (
update_connector_last_indexed,
)
ACCEPTED_CALENDAR_CONNECTOR_TYPES = {
SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
}
# Type hint for heartbeat callback
HeartbeatCallbackType = Callable[[int], Awaitable[None]]

18
surfsense_backend/app/tasks/connector_indexers/google_gmail_indexer.py
View file

@ -21,15 +21,6 @@ from app.db import (
DocumentType,
SearchSourceConnectorType,
)
from app.utils.google_credentials import (
COMPOSIO_GOOGLE_CONNECTOR_TYPES,
build_composio_credentials,
)
ACCEPTED_GMAIL_CONNECTOR_TYPES = {
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
}
from app.services.llm_service import get_user_long_context_llm
from app.services.task_logging_service import TaskLoggingService
from app.utils.document_converters import (
@ -39,6 +30,10 @@ from app.utils.document_converters import (
generate_document_summary,
generate_unique_identifier_hash,
)
from app.utils.google_credentials import (
COMPOSIO_GOOGLE_CONNECTOR_TYPES,
build_composio_credentials,
)
from .base import (
calculate_date_range,
@ -51,6 +46,11 @@ from .base import (
update_connector_last_indexed,
)
ACCEPTED_GMAIL_CONNECTOR_TYPES = {
SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
}
# Type hint for heartbeat callback
HeartbeatCallbackType = Callable[[int], Awaitable[None]]

4
surfsense_backend/tests/integration/google_unification/test_search_includes_legacy_docs.py
View file

@ -23,7 +23,7 @@ async def test_search_google_drive_includes_legacy_composio_docs(
async with patched_session_factory() as session:
service = ConnectorService(session, search_space_id=space_id)
result_object, raw_docs = await service.search_google_drive(
_, raw_docs = await service.search_google_drive(
user_query="quarterly budget",
search_space_id=space_id,
top_k=10,
@ -55,7 +55,7 @@ async def test_search_files_does_not_include_google_types(
async with patched_session_factory() as session:
service = ConnectorService(session, search_space_id=space_id)
result_object, raw_docs = await service.search_files(
_, raw_docs = await service.search_files(
user_query="quarterly budget",
search_space_id=space_id,
top_k=10,

8
surfsense_backend/tests/unit/google_unification/test_composio_credentials.py
View file

@ -15,11 +15,11 @@ pytestmark = pytest.mark.unit
@patch("app.services.composio_service.ComposioService")
def test_returns_credentials_with_token_and_expiry(MockComposioService):
def test_returns_credentials_with_token_and_expiry(mock_composio_service):
"""build_composio_credentials returns a Credentials object with the Composio access token."""
mock_service = MagicMock()
mock_service.get_access_token.return_value = "fake-access-token"
MockComposioService.return_value = mock_service
mock_composio_service.return_value = mock_service
from app.utils.google_credentials import build_composio_credentials
@ -32,14 +32,14 @@ def test_returns_credentials_with_token_and_expiry(MockComposioService):
@patch("app.services.composio_service.ComposioService")
def test_refresh_handler_fetches_fresh_token(MockComposioService):
def test_refresh_handler_fetches_fresh_token(mock_composio_service):
"""The refresh_handler on the returned Credentials fetches a new token from Composio."""
mock_service = MagicMock()
mock_service.get_access_token.side_effect = [
"initial-token",
"refreshed-token",
]
MockComposioService.return_value = mock_service
mock_composio_service.return_value = mock_service
from app.utils.google_credentials import build_composio_credentials

2
surfsense_backend/tests/unit/google_unification/test_connector_credential_acceptance.py
View file

@ -208,7 +208,7 @@ async def test_drive_client_uses_prebuilt_composio_credentials(mock_build):
credentials=creds,
)
files, next_token, error = await client.list_files()
files, _next_token, error = await client.list_files()
assert error is None
assert files == []