diff --git a/surfsense_backend/app/agents/new_chat/chat_deepagent.py b/surfsense_backend/app/agents/new_chat/chat_deepagent.py index a00cb4b7b..c69ba1063 100644 --- a/surfsense_backend/app/agents/new_chat/chat_deepagent.py +++ b/surfsense_backend/app/agents/new_chat/chat_deepagent.py @@ -20,10 +20,10 @@ from langgraph.types import Checkpointer from sqlalchemy.ext.asyncio import AsyncSession from app.agents.new_chat.context import SurfSenseContextSchema +from app.agents.new_chat.llm_config import AgentConfig from app.agents.new_chat.middleware.dedup_tool_calls import ( DedupHITLToolCallsMiddleware, ) -from app.agents.new_chat.llm_config import AgentConfig from app.agents.new_chat.system_prompt import ( build_configurable_system_prompt, build_surfsense_system_prompt, diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py index b7decbe0c..50ef24a1b 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/create_page.py @@ -120,6 +120,7 @@ def create_create_confluence_page_tool( return {"status": "error", "message": "A space must be selected."} from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType actual_connector_id = final_connector_id diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py index 285b1058d..ba1dae653 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/delete_page.py @@ -124,6 +124,7 @@ def create_delete_confluence_page_tool( final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType if not final_connector_id: diff --git a/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py b/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py index e43aef59d..36dd091b4 100644 --- a/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py +++ b/surfsense_backend/app/agents/new_chat/tools/confluence/update_page.py @@ -136,6 +136,7 @@ def create_update_confluence_page_tool( final_document_id = final_params.get("document_id", document_id) from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType if not final_connector_id: diff --git a/surfsense_backend/app/agents/new_chat/tools/google_calendar/__init__.py b/surfsense_backend/app/agents/new_chat/tools/google_calendar/__init__.py index 4fea648f0..d1ce4e795 100644 --- a/surfsense_backend/app/agents/new_chat/tools/google_calendar/__init__.py +++ b/surfsense_backend/app/agents/new_chat/tools/google_calendar/__init__.py @@ -1,15 +1,15 @@ from app.agents.new_chat.tools.google_calendar.create_event import ( create_create_calendar_event_tool, ) -from app.agents.new_chat.tools.google_calendar.update_event import ( - create_update_calendar_event_tool, -) from app.agents.new_chat.tools.google_calendar.delete_event import ( create_delete_calendar_event_tool, ) +from app.agents.new_chat.tools.google_calendar.update_event import ( + create_update_calendar_event_tool, +) __all__ = [ "create_create_calendar_event_tool", - "create_update_calendar_event_tool", "create_delete_calendar_event_tool", + "create_update_calendar_event_tool", ] diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py index 273cb163b..2fe49b3df 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/create_issue.py @@ -128,6 +128,7 @@ def create_create_jira_issue_tool( return {"status": "error", "message": "A project must be selected."} from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType actual_connector_id = final_connector_id diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py index 0f656fbb0..2f8c370ad 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/delete_issue.py @@ -121,6 +121,7 @@ def create_delete_jira_issue_tool( final_delete_from_kb = final_params.get("delete_from_kb", delete_from_kb) from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType if not final_connector_id: diff --git a/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py b/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py index f25eb826f..3150eee1a 100644 --- a/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py +++ b/surfsense_backend/app/agents/new_chat/tools/jira/update_issue.py @@ -131,6 +131,7 @@ def create_update_jira_issue_tool( final_document_id = final_params.get("document_id", document_id) from sqlalchemy.future import select + from app.db import SearchSourceConnector, SearchSourceConnectorType if not final_connector_id: diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 9374fd9d0..e738e06c6 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/surfsense_backend/app/agents/new_chat/tools/registry.py @@ -45,6 +45,11 @@ from langchain_core.tools import BaseTool from app.db import ChatVisibility +from .confluence import ( + create_create_confluence_page_tool, + create_delete_confluence_page_tool, + create_update_confluence_page_tool, +) from .display_image import create_display_image_tool from .generate_image import create_generate_image_tool from .gmail import ( @@ -62,6 +67,11 @@ from .google_drive import ( create_create_google_drive_file_tool, create_delete_google_drive_file_tool, ) +from .jira import ( + create_create_jira_issue_tool, + create_delete_jira_issue_tool, + create_update_jira_issue_tool, +) from .knowledge_base import create_search_knowledge_base_tool from .linear import ( create_create_linear_issue_tool, @@ -70,16 +80,6 @@ from .linear import ( ) from .link_preview import create_link_preview_tool from .mcp_tool import load_mcp_tools -from .jira import ( - create_create_jira_issue_tool, - create_delete_jira_issue_tool, - create_update_jira_issue_tool, -) -from .confluence import ( - create_create_confluence_page_tool, - create_delete_confluence_page_tool, - create_update_confluence_page_tool, -) from .notion import ( create_create_notion_page_tool, create_delete_notion_page_tool, diff --git a/surfsense_backend/app/connectors/notion_history.py b/surfsense_backend/app/connectors/notion_history.py index 3553f3cc0..55fef7240 100644 --- a/surfsense_backend/app/connectors/notion_history.py +++ b/surfsense_backend/app/connectors/notion_history.py @@ -1,13 +1,12 @@ import asyncio import contextlib import logging -import re from collections.abc import Awaitable, Callable from typing import Any, TypeVar from notion_client import AsyncClient -from notion_markdown import to_notion from notion_client.errors import APIResponseError +from notion_markdown import to_notion from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select diff --git a/surfsense_backend/app/routes/linear_add_connector_route.py b/surfsense_backend/app/routes/linear_add_connector_route.py index 310e3736a..9345ae495 100644 --- a/surfsense_backend/app/routes/linear_add_connector_route.py +++ b/surfsense_backend/app/routes/linear_add_connector_route.py @@ -12,10 +12,9 @@ import httpx from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse from pydantic import ValidationError +from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession - -from sqlalchemy import select from sqlalchemy.orm.attributes import flag_modified from app.config import config diff --git a/surfsense_backend/app/routes/notion_add_connector_route.py b/surfsense_backend/app/routes/notion_add_connector_route.py index d862a5855..16e80ebcb 100644 --- a/surfsense_backend/app/routes/notion_add_connector_route.py +++ b/surfsense_backend/app/routes/notion_add_connector_route.py @@ -12,10 +12,9 @@ import httpx from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse from pydantic import ValidationError +from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession - -from sqlalchemy import select from sqlalchemy.orm.attributes import flag_modified from app.config import config diff --git a/surfsense_backend/app/routes/search_source_connectors_routes.py b/surfsense_backend/app/routes/search_source_connectors_routes.py index 769186867..c12621f60 100644 --- a/surfsense_backend/app/routes/search_source_connectors_routes.py +++ b/surfsense_backend/app/routes/search_source_connectors_routes.py @@ -72,6 +72,7 @@ from app.tasks.connector_indexers import ( index_slack_messages, ) from app.users import current_active_user +from app.utils.connector_naming import ensure_unique_connector_name from app.utils.indexing_locks import ( acquire_connector_indexing_lock, release_connector_indexing_lock, @@ -81,7 +82,6 @@ from app.utils.periodic_scheduler import ( delete_periodic_schedule, update_periodic_schedule, ) -from app.utils.connector_naming import ensure_unique_connector_name from app.utils.rbac import check_permission # Set up logging diff --git a/surfsense_backend/app/utils/google_credentials.py b/surfsense_backend/app/utils/google_credentials.py index b490768bc..202ea7fa3 100644 --- a/surfsense_backend/app/utils/google_credentials.py +++ b/surfsense_backend/app/utils/google_credentials.py @@ -1,7 +1,7 @@ """Shared Google OAuth credential utilities for native and Composio connectors.""" import logging -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from google.oauth2.credentials import Credentials @@ -31,11 +31,11 @@ def build_composio_credentials(connected_account_id: str) -> Credentials: def composio_refresh_handler(request, scopes): fresh_token = service.get_access_token(connected_account_id) - expiry = datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(minutes=55) + expiry = datetime.now(UTC).replace(tzinfo=None) + timedelta(minutes=55) return fresh_token, expiry return Credentials( token=access_token, - expiry=datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(minutes=55), + expiry=datetime.now(UTC).replace(tzinfo=None) + timedelta(minutes=55), refresh_handler=composio_refresh_handler, ) diff --git a/surfsense_backend/tests/integration/google_unification/conftest.py b/surfsense_backend/tests/integration/google_unification/conftest.py index 4483cf7dd..de68c7acb 100644 --- a/surfsense_backend/tests/integration/google_unification/conftest.py +++ b/surfsense_backend/tests/integration/google_unification/conftest.py @@ -67,7 +67,7 @@ def make_chunk(*, content: str, document_id: int) -> Chunk: @pytest_asyncio.fixture async def seed_google_docs( - db_session: AsyncSession, db_user: "User", db_search_space: "SearchSpace" + db_session: AsyncSession, db_user: User, db_search_space: SearchSpace ): """Insert a native Drive doc, a legacy Composio Drive doc, and a FILE doc. @@ -269,7 +269,7 @@ def mock_task_logger(): async def seed_connector( async_engine, *, - connector_type: "SearchSourceConnectorType", + connector_type: SearchSourceConnectorType, config: dict, name_prefix: str = "test", ): diff --git a/surfsense_backend/tests/integration/google_unification/test_drive_indexer_credentials.py b/surfsense_backend/tests/integration/google_unification/test_drive_indexer_credentials.py index f17fc69f5..5bb0b6137 100644 --- a/surfsense_backend/tests/integration/google_unification/test_drive_indexer_credentials.py +++ b/surfsense_backend/tests/integration/google_unification/test_drive_indexer_credentials.py @@ -7,7 +7,7 @@ mocked at their system boundaries. from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest import pytest_asyncio diff --git a/surfsense_backend/tests/integration/google_unification/test_hybrid_search_type_filtering.py b/surfsense_backend/tests/integration/google_unification/test_hybrid_search_type_filtering.py index 5a29b8333..402d3ee0b 100644 --- a/surfsense_backend/tests/integration/google_unification/test_hybrid_search_type_filtering.py +++ b/surfsense_backend/tests/integration/google_unification/test_hybrid_search_type_filtering.py @@ -8,7 +8,6 @@ which is the foundation of the Google unification changes. import pytest -from app.config import config as app_config from app.retriever.chunks_hybrid_search import ChucksHybridSearchRetriever from .conftest import DUMMY_EMBEDDING diff --git a/surfsense_backend/tests/unit/google_unification/test_composio_credentials.py b/surfsense_backend/tests/unit/google_unification/test_composio_credentials.py index 988e7854f..589b5c8b6 100644 --- a/surfsense_backend/tests/unit/google_unification/test_composio_credentials.py +++ b/surfsense_backend/tests/unit/google_unification/test_composio_credentials.py @@ -5,11 +5,10 @@ returned ``google.oauth2.credentials.Credentials`` object is correctly configured with a token and a working refresh handler. """ -from datetime import datetime, timezone +from datetime import UTC, datetime from unittest.mock import MagicMock, patch import pytest - from google.oauth2.credentials import Credentials pytestmark = pytest.mark.unit @@ -29,7 +28,7 @@ def test_returns_credentials_with_token_and_expiry(MockComposioService): assert isinstance(creds, Credentials) assert creds.token == "fake-access-token" assert creds.expiry is not None - assert creds.expiry > datetime.now(timezone.utc).replace(tzinfo=None) + assert creds.expiry > datetime.now(UTC).replace(tzinfo=None) @patch("app.services.composio_service.ComposioService") @@ -53,5 +52,5 @@ def test_refresh_handler_fetches_fresh_token(MockComposioService): new_token, new_expiry = refresh_handler(request=None, scopes=None) assert new_token == "refreshed-token" - assert new_expiry > datetime.now(timezone.utc).replace(tzinfo=None) + assert new_expiry > datetime.now(UTC).replace(tzinfo=None) assert mock_service.get_access_token.call_count == 2 diff --git a/surfsense_backend/tests/unit/google_unification/test_connector_credential_acceptance.py b/surfsense_backend/tests/unit/google_unification/test_connector_credential_acceptance.py index bb7bce8ed..da7548cc2 100644 --- a/surfsense_backend/tests/unit/google_unification/test_connector_credential_acceptance.py +++ b/surfsense_backend/tests/unit/google_unification/test_connector_credential_acceptance.py @@ -10,7 +10,7 @@ allows Composio credentials through without raising ValueError or persisting to from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -21,7 +21,7 @@ pytestmark = pytest.mark.unit def _utcnow_naive() -> datetime: """Return current UTC time as a naive datetime (matches google-auth convention).""" - return datetime.now(timezone.utc).replace(tzinfo=None) + return datetime.now(UTC).replace(tzinfo=None) def _composio_credentials(*, expired: bool = False) -> Credentials: