fix: resolve mypy type error

Author: italojohnny

src/backend/base/langflow/graph/graph/base.py

@@ -44,9 +44,8 @@
 from langflow.utils.async_helpers import run_until_complete
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Generator, Iterable, AsyncGenerator
-
-    from sqlmodel.ext.asyncio.session import AsyncSession
+    from collections.abc import Callable, Generator, Iterable
+    from contextlib import AbstractAsyncContextManager
 
     from langflow.api.v1.schemas import InputValueRequest
     from langflow.custom.custom_component.component import Component
@@ -653,7 +652,7 @@ def set_run_id(self, run_id: uuid.UUID | None = None) -> None:
 
         self._run_id = str(run_id)
 
-    async def initialize_run(self, session_scope: AsyncGenerator[AsyncSession, None] | None) -> None:
+    async def initialize_run(self, session_scope: Callable[[], AbstractAsyncContextManager[Any]] | None = None) -> None:
         if not self._run_id:
             self.set_run_id()
 

src/backend/base/langflow/services/tracing/service.py

@@ -14,10 +14,10 @@
 from langflow.services.deps import get_variable_service
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
+    from collections.abc import Callable
+    from contextlib import AbstractAsyncContextManager
 
     from langchain.callbacks.base import BaseCallbackHandler
-    from sqlmodel.ext.asyncio.session import AsyncSession
 
     from langflow.custom.custom_component.component import Component
     from langflow.graph.vertex.base import Vertex
@@ -222,7 +222,7 @@ async def get_varaibles_from_db(self, session_scope, user_id, variable_names) ->
 
     async def start_tracers(
         self,
-        session_scope: AsyncGenerator[AsyncSession, None],
+        session_scope: Callable[[], AbstractAsyncContextManager[Any]],
         run_id: UUID,
         run_name: str,
         user_id: str | None,