digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

base_server/server_async_insecure.py

@@ -9,8 +9,9 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -30,7 +31,7 @@ class AsyncGreeterImpl(Greeter):
 
     async def SayHello(self, request, context):  # noqa: N802
         """Asynchronous implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -40,7 +41,7 @@ class AsyncGreeterImpl(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -97,7 +98,7 @@ async def main_async() -> int:
         # as the KeyboardInterrupt usually breaks out of asyncio.run()
         logger.info("Server stopping due to keyboard interrupt...")
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
     finally:
         # Clean up resources if server was started
@@ -116,7 +117,7 @@ def main():
         logger.info("Server stopped by keyboard interrupt")
         return 0  # Clean exit
     except Exception as e:
-        logger.exception(f"Fatal error: {e}")
+        logger.exception("Fatal error: %s", e)
         return 1
 
 

base_server/server_async_secure.py

@@ -9,13 +9,14 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import (
     SecurityMode,
     ServerConfig,
     ServerCredentials,
     ServerMode,
 )
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -35,7 +36,7 @@ class AsyncGreeterImpl(Greeter):
 
     async def SayHello(self, request, context):  # noqa: N802
         """Asynchronous implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -45,7 +46,7 @@ class AsyncGreeterImpl(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -115,7 +116,7 @@ async def main_async() -> int:
         # as the KeyboardInterrupt usually breaks out of asyncio.run()
         logger.info("Server stopping due to keyboard interrupt...")
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
     finally:
         # Clean up resources if server was started
@@ -134,7 +135,7 @@ def main():
         logger.info("Server stopped by keyboard interrupt")
         return 0  # Clean exit
     except Exception as e:
-        logger.exception(f"Fatal error: {e}")
+        logger.exception("Fatal error: %s", e)
         return 1
 
 

base_server/server_sync_insecure.py

@@ -8,8 +8,9 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -29,7 +30,7 @@ class SyncGreeterServicer(Greeter):
 
     def SayHello(self, request, context):  # noqa: N802
         """Implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -39,7 +40,7 @@ class SyncGreeterServicer(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -92,7 +93,7 @@ def main() -> int:
             server.stop()
 
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
 
     return 0

base_server/server_sync_secure.py

@@ -8,13 +8,14 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import (
     SecurityMode,
     ServerConfig,
     ServerCredentials,
     ServerMode,
 )
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -34,7 +35,7 @@ class SyncGreeterServicer(Greeter):
 
     def SayHello(self, request, context):  # noqa: N802
         """Implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -44,7 +45,7 @@ class SyncGreeterServicer(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -111,7 +112,7 @@ def main() -> int:
             server.stop()
 
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
 
     return 0

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.1.dev2"
+    __version__ = "0.3.2.dev14"

digitalkin/core/job_manager/base_job_manager.py

@@ -8,8 +8,8 @@ from typing import Any, Generic
 from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
+from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode

digitalkin/core/job_manager/single_job_manager.py

@@ -15,7 +15,7 @@ from digitalkin.core.task_manager.local_task_manager import LocalTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.base_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
@@ -86,7 +86,10 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 message=f"Module {job_id} did not respond within 30 seconds",
             )
         finally:
-            logger.info(f"{job_id=}: {session.queue.empty()}")
+            logger.debug(
+                "Config setup response retrieved",
+                extra={"job_id": job_id, "queue_empty": session.queue.empty()},
+            )
 
     async def create_config_setup_instance_job(
         self,
@@ -126,7 +129,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         except Exception:
             # Remove the module from the manager in case of an error.
             del self.tasks_sessions[job_id]
-            logger.exception("Failed to start module %s: %s", job_id)
+            logger.exception("Failed to start module", extra={"job_id": job_id})
             raise
         else:
             return job_id
@@ -140,7 +143,8 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             job_id: The unique identifier of the job.
             output_data: The output data produced by the job.
         """
-        await self.tasks_sessions[job_id].queue.put(output_data.model_dump())
+        session = self.tasks_sessions[job_id]
+        await session.queue.put(output_data.model_dump())
 
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
@@ -259,6 +263,18 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
         return job_id
 
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
+        """Clean a task's session.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+
+        Returns:
+            bool: True if the task was successfully cleaned, False otherwise.
+        """
+        return await self._task_manager.clean_session(task_id, mission_id)
+
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module job.
 
@@ -271,20 +287,23 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         Raises:
             Exception: If an error occurs while stopping the module.
         """
-        logger.info(f"STOP required for {job_id=}")
+        logger.info("Stop module requested", extra={"job_id": job_id})
 
         async with self._lock:
             session = self.tasks_sessions.get(job_id)
 
             if not session:
-                logger.warning(f"session with id: {job_id} not found")
+                logger.warning("Session not found", extra={"job_id": job_id})
                 return False
             try:
                 await session.module.stop()
                 await self.cancel_task(job_id, session.mission_id)
-                logger.debug(f"session {job_id} ({session.module.name}) stopped successfully")
-            except Exception as e:
-                logger.error(f"Error while stopping module {job_id}: {e}")
+                logger.debug(
+                    "Module stopped successfully",
+                    extra={"job_id": job_id, "mission_id": session.mission_id},
+                )
+            except Exception:
+                logger.exception("Error stopping module", extra={"job_id": job_id})
                 raise
             else:
                 return True

digitalkin/core/job_manager/taskiq_broker.py

@@ -21,8 +21,9 @@ from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.core.task_manager.task_executor import TaskExecutor
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
-from digitalkin.models.core.job_manager_models import StreamCodeModel
-from digitalkin.models.module.module_types import OutputModelT
+from digitalkin.models.module.module import ModuleCodeModel
+from digitalkin.models.module.module_types import DataModel, OutputModelT
+from digitalkin.models.module.utility import EndOfStreamOutput
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
@@ -141,7 +142,7 @@ async def cleanup_global_resources() -> None:
         logger.warning("Failed to shutdown Taskiq broker: %s", e)
 
 
-async def send_message_to_stream(job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+async def send_message_to_stream(job_id: str, output_data: OutputModelT | ModuleCodeModel) -> None:  # type: ignore[type-var]
     """Callback define to add a message frame to the Rstream.
 
     Args:
@@ -186,7 +187,7 @@ async def run_start_module(
     module_class.discover()
 
     job_id = context.message.task_id
-    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
+    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)  # type: ignore[type-var]
     module = ModuleFactory.create_module_instance(module_class, job_id, mission_id, setup_id, setup_version_id)
 
     channel = None
@@ -201,7 +202,7 @@ async def run_start_module(
         # Create a proper done callback that handles errors
         async def send_end_of_stream(_: Any) -> None:  # noqa: ANN401
             try:
-                await callback(StreamCodeModel(code="__END_OF_STREAM__"))
+                await callback(DataModel(root=EndOfStreamOutput()))
             except Exception as e:
                 logger.error("Error sending end of stream: %s", e, exc_info=True)
 
@@ -272,7 +273,7 @@ async def run_config_module(
     logger.debug("Services config: %s | Module config: %s", services_config, module_class.services_config)
 
     job_id = context.message.task_id
-    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
+    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)  # type: ignore[type-var]
     module = ModuleFactory.create_module_instance(module_class, job_id, mission_id, setup_id, setup_version_id)
 
     # Override environment variables temporarily to use manager's SurrealDB

digitalkin/core/job_manager/taskiq_job_manager.py

@@ -24,7 +24,7 @@ from digitalkin.core.job_manager.taskiq_broker import STREAM, STREAM_RETENTION,
 from digitalkin.core.task_manager.remote_task_manager import RemoteTaskManager
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
 

digitalkin/core/task_manager/surrealdb_repository.py

@@ -4,7 +4,7 @@ import asyncio
 import datetime
 import os
 from collections.abc import AsyncGenerator
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, cast
 from uuid import UUID
 
 from surrealdb import AsyncHttpSurrealConnection, AsyncSurreal, AsyncWsSurrealConnection, RecordID
@@ -91,7 +91,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Connecting to SurrealDB at %s", self.url)
         self.db = AsyncSurreal(self.url)  # type: ignore
         await self.db.signin({"username": self.username, "password": self.password})
-        await self.db.use(self.namespace, self.database)
+        await self.db.use(self.namespace, self.database)  # type: ignore[arg-type]
         logger.debug("Successfully connected to SurrealDB")
 
     async def close(self) -> None:
@@ -112,7 +112,7 @@ class SurrealDBConnection(Generic[TSurreal]):
             # Process results and track failures
             failed_queries = []
             for live_id, result in zip(live_query_ids, results):
-                if isinstance(result, (ConnectionError, TimeoutError, Exception)):
+                if isinstance(result, ConnectionError | TimeoutError | Exception):
                     failed_queries.append((live_id, str(result)))
                 else:
                     self._live_queries.discard(live_id)
@@ -146,7 +146,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Creating record in %s with data: %s", table_name, data)
         result = await self.db.create(table_name, data)
         logger.debug("create result: %s", result)
-        return result
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def merge(
         self,
@@ -170,7 +170,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Updating record in %s with data: %s", record_id, data)
         result = await self.db.merge(record_id, data)
         logger.debug("update result: %s", result)
-        return result
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def update(
         self,
@@ -194,7 +194,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Updating record in %s with data: %s", record_id, data)
         result = await self.db.update(record_id, data)
         logger.debug("update result: %s", result)
-        return result
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def execute_query(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
         """Execute a custom SurrealQL query.
@@ -209,7 +209,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("execute_query: %s with params: %s", query, params)
         result = await self.db.query(query, params or {})
         logger.debug("execute_query result: %s", result)
-        return [result] if isinstance(result, dict) else result
+        return cast("list[dict[str, Any]]", [result] if isinstance(result, dict) else result)
 
     async def select_by_task_id(self, table: str, value: str) -> dict[str, Any]:
         """Fetch a record from a table by a unique field.