digitalkin 0.3.1.dev2__tar.gz → 0.3.2.dev2__tar.gz

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.1.dev2
+Version: 0.3.2.dev2
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,17 +452,17 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto==0.2.0.dev5
+Requires-Dist: agentic-mesh-protocol==0.2.1.dev0
 Requires-Dist: grpcio-health-checking>=1.76.0
 Requires-Dist: grpcio-reflection>=1.76.0
 Requires-Dist: grpcio-status>=1.76.0
 Requires-Dist: pydantic>=2.12.5
-Requires-Dist: surrealdb>=1.0.6
+Requires-Dist: surrealdb>=1.0.7
 Provides-Extra: taskiq
 Requires-Dist: rstream>=0.40.0; extra == "taskiq"
 Requires-Dist: taskiq-aio-pika>=0.5.0; extra == "taskiq"
-Requires-Dist: taskiq-redis>=1.1.2; extra == "taskiq"
-Requires-Dist: taskiq[reload]>=0.12.0; extra == "taskiq"
+Requires-Dist: taskiq-redis>=1.2.0; extra == "taskiq"
+Requires-Dist: taskiq[reload]>=0.12.1; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/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
 
 

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/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
 
 

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/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

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/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-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/modules/dynamic_setup_module.py

@@ -209,8 +209,6 @@ class DynamicModuleOutput(DataModel[MessageOutputTrigger]):
 class DynamicModuleSecret(BaseModel):
     """Secret model (empty for this example)."""
 
-    pass
-
 
 # =============================================================================
 # Module Implementation
@@ -285,13 +283,7 @@ class DynamicSetupModule(
 
 async def demonstrate_dynamic_schema() -> None:
     """Demonstrate the dynamic schema functionality."""
-    print("=" * 60)
-    print("Dynamic Schema Demonstration")
-    print("=" * 60)
-
     # 1. Show schema WITHOUT force (dynamic fields not resolved)
-    print("\n1. Schema without force=True (fetchers NOT called):")
-    print("-" * 40)
 
     model_no_force = await DynamicAgentSetup.get_clean_model(
         config_fields=True,
@@ -302,13 +294,10 @@ async def demonstrate_dynamic_schema() -> None:
 
     # Check if enum is present
     model_name_schema = schema_no_force.get("properties", {}).get("model_name", {})
-    print(f"model_name has enum: {'enum' in model_name_schema}")
     if "enum" in model_name_schema:
-        print(f"  enum values: {model_name_schema['enum']}")
+        pass
 
     # 2. Show schema WITH force (dynamic fields resolved)
-    print("\n2. Schema with force=True (fetchers called):")
-    print("-" * 40)
 
     model_with_force = await DynamicAgentSetup.get_clean_model(
         config_fields=True,
@@ -319,43 +308,30 @@ async def demonstrate_dynamic_schema() -> None:
 
     # Check enum values after force
     model_name_schema = schema_with_force.get("properties", {}).get("model_name", {})
-    print(f"model_name has enum: {'enum' in model_name_schema}")
     if "enum" in model_name_schema:
-        print(f"  enum values: {model_name_schema['enum']}")
+        pass
 
     language_schema = schema_with_force.get("properties", {}).get("language", {})
-    print(f"language has enum: {'enum' in language_schema}")
     if "enum" in language_schema:
-        print(f"  enum values: {language_schema['enum']}")
+        pass
 
     # 3. Show that static json_schema_extra is preserved
-    print("\n3. Static json_schema_extra preserved:")
-    print("-" * 40)
-    print(f"model_name ui:widget: {model_name_schema.get('ui:widget', 'NOT FOUND')}")
 
     # 4. Show field filtering
-    print("\n4. Field filtering demonstration:")
-    print("-" * 40)
 
     # Config fields only (hidden excluded)
-    config_model = await DynamicAgentSetup.get_clean_model(
+    await DynamicAgentSetup.get_clean_model(
         config_fields=True,
         hidden_fields=False,
         force=False,
     )
-    print(f"Config fields (hidden=False): {list(config_model.model_fields.keys())}")
 
     # All fields including hidden
-    all_model = await DynamicAgentSetup.get_clean_model(
+    await DynamicAgentSetup.get_clean_model(
         config_fields=True,
         hidden_fields=True,
         force=False,
     )
-    print(f"All fields (hidden=True): {list(all_model.model_fields.keys())}")
-
-    print("\n" + "=" * 60)
-    print("Demonstration complete!")
-    print("=" * 60)
 
 
 if __name__ == "__main__":

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/services/filesystem_module.py

@@ -3,16 +3,16 @@
 import asyncio
 import datetime
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from pydantic import BaseModel, Field
 
 from digitalkin.logger import logger
 from digitalkin.models.module import ModuleStatus
 from digitalkin.modules.archetype_module import ArchetypeModule
+from digitalkin.services.filesystem.filesystem_strategy import FileFilter, UploadFileData
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
-from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord, FileFilter, UploadFileData
 
 
 class ExampleInput(BaseModel):
@@ -62,7 +62,10 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"cost": {"config": {}}, "storage": {"config": {}}} # Filesystem has no config but it's enabled
+    services_config_params = {
+        "cost": {"config": {}},
+        "storage": {"config": {}},
+    }  # Filesystem has no config but it's enabled
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -128,7 +131,6 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
             callback(record.model_dump())
         # Call the callback with the output data
 
-
         # Wait a bit to simulate processing time
         await asyncio.sleep(1)
 
@@ -173,7 +175,7 @@ async def test_module() -> None:
 
     # Check the storage
     if module.status == ModuleStatus.STOPPED:
-        files, nb_results = module.filesystem.get_files(
+        files, _nb_results = module.filesystem.get_files(
             filters=FileFilter(name="example_output.txt", context="test-mission-123"),
         )
         for file in files:

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/examples/services/storage_module.py

@@ -64,7 +64,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"storage": {"config": {"example": ExampleOutput}},"cost": {"config":{}}}
+    services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "cost": {"config": {}}}
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -184,7 +184,9 @@ async def test_module() -> None:
 def test_storage_directly() -> None:
     """Test the storage service directly."""
     # Initialize storage service
-    storage = ServicesConfig().storage(mission_id="test-mission",setup_version_id="test-setup-123", config={"example": ExampleStorage})
+    storage = ServicesConfig().storage(
+        mission_id="test-mission", setup_version_id="test-setup-123", config={"example": ExampleStorage}
+    )
 
     # Create a test record
     storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
 
-    version = "0.3.1.dev2"
+    version = "0.3.2.dev2"
     classifiers = [
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
@@ -28,20 +28,20 @@
     ]
 
     dependencies = [
-        "digitalkin-proto==0.2.0.dev5",
+        "agentic-mesh-protocol==0.2.1.dev0",
         "grpcio-health-checking>=1.76.0",
         "grpcio-reflection>=1.76.0",
         "grpcio-status>=1.76.0",
         "pydantic>=2.12.5",
-        "surrealdb>=1.0.6",
+        "surrealdb>=1.0.7",
     ]
 
     [project.optional-dependencies]
         taskiq = [
             "rstream>=0.40.0",
             "taskiq-aio-pika>=0.5.0",
-            "taskiq-redis>=1.1.2",
-            "taskiq[reload]>=0.12.0",
+            "taskiq-redis>=1.2.0",
+            "taskiq[reload]>=0.12.1",
         ]
 
     [project.urls]
@@ -56,7 +56,7 @@
 [dependency-groups]
     dev = [
         "typos>=1.40.0",
-        "ruff>=0.14.7",
+        "ruff>=0.14.8",
         "mypy>=1.19.0",
         "pyright>=1.1.407",
         "pre-commit>=4.5.0",
@@ -70,7 +70,7 @@
         "hdrhistogram>=0.10.3",
         "grpcio-testing>=1.76.0",
         "psutil>=7.1.3",
-        "pytest>=9.0.1",
+        "pytest>=9.0.2",
         "pytest-asyncio>=1.3.0",
         "pytest-cov>=7.0.0",
         "pytest-html==4.1.1",
@@ -97,11 +97,10 @@
         "mkdocs-material[imaging]>=9.7.0",
         "mkdocs-minify-plugin>=0.8.0",
         "mkdocs-section-index>=0.3.10",
-        "mkdocstrings-python>=2.0.0",
+        "mkdocstrings-python>=2.0.1",
         "mkdocs-open-in-new-tab>=1.0.8",
         "tomli>=2.3.0",
     ]
-
 [tool.setuptools]
     package-dir = { "" = "src" }
 

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/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.dev2"

{digitalkin-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/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-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/digitalkin/core/job_manager/single_job_manager.py

@@ -15,8 +15,8 @@ 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.module import ModuleCodeModel
+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-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/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-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/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-0.3.1.dev2 → digitalkin-0.3.2.dev2}/src/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.