digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/grpc_servers/module_server.py

@@ -1,37 +1,32 @@
 """Module gRPC server implementation for DigitalKin."""
 
-from pathlib import Path
-import uuid
+from typing import TYPE_CHECKING
 
-import grpc
+from agentic_mesh_protocol.module.v1 import (
+    module_service_pb2,
+    module_service_pb2_grpc,
+)
 
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.module_servicer import ModuleServicer
-from digitalkin.grpc_servers.utils.exceptions import ServerError
-from digitalkin.grpc_servers.utils.models import (
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import (
     ClientConfig,
     ModuleServerConfig,
     SecurityMode,
 )
 from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.registry import GrpcRegistry
 
-from digitalkin_proto.digitalkin.module.v2 import (
-    module_service_pb2,
-    module_service_pb2_grpc,
-)
-from digitalkin_proto.digitalkin.module_registry.v2 import (
-    metadata_pb2,
-    module_registry_service_pb2_grpc,
-    registration_pb2,
-)
-from digitalkin.logger import logger
+if TYPE_CHECKING:
+    from digitalkin.services.registry import RegistryStrategy
 
 
 class ModuleServer(BaseServer):
     """gRPC server for a DigitalKin module.
 
     This server exposes the module's functionality through the ModuleService gRPC interface.
-    It can optionally register itself with a ModuleRegistry server.
+    It can optionally register itself with a Registry server.
 
     Attributes:
         module: The module instance being served.
@@ -50,13 +45,20 @@ class ModuleServer(BaseServer):
 
         Args:
             module_class: The module instance to be served.
-            config: Server configuration including registry address if auto-registration is desired.
+            server_config: Server configuration including registry address if auto-registration is desired.
+            client_config: Client configuration used by services.
         """
         super().__init__(server_config)
         self.module_class = module_class
         self.server_config = server_config
         self.client_config = client_config
         self.module_servicer: ModuleServicer | None = None
+        self.registry: RegistryStrategy | None = None
+
+        self._registry_client_config: ClientConfig | None = None
+        if self.server_config.registry_address:
+            self._registry_client_config = self._build_registry_client_config()
+        self._prepare_registry_config()
 
     def _register_servicers(self) -> None:
         """Register the module servicer with the gRPC server.
@@ -77,207 +79,130 @@ class ModuleServer(BaseServer):
         )
         logger.debug("Registered Module servicer")
 
+    def _prepare_registry_config(self) -> None:
+        """Prepare registry client config on module_class before server starts.
+
+        This ensures ServicesConfig created by JobManager will have registry config,
+        allowing spawned module instances to inherit the registry configuration.
+        """
+        if not self._registry_client_config:
+            return
+
+        self.module_class.services_config_params["registry"] = {"client_config": self._registry_client_config}
+
+    def _build_registry_client_config(self) -> ClientConfig:
+        """Build ClientConfig for registry from server_config.registry_address.
+
+        Returns:
+            ClientConfig configured for registry connection.
+        """
+        host, port = self.server_config.registry_address.rsplit(":", 1)
+        return ClientConfig(
+            host=host,
+            port=int(port),
+            mode=self.server_config.mode,
+            security=self.client_config.security if self.client_config else SecurityMode.INSECURE,
+            credentials=self.client_config.credentials if self.client_config else None,
+        )
+
+    def _init_registry(self) -> None:
+        """Initialize server-level registry client for registration.
+
+        Note: services_config_params["registry"] is already set in _prepare_registry_config()
+        which runs in __init__(). This method only creates the server-level client instance.
+        """
+        if not self._registry_client_config:
+            return
+
+        self.registry = GrpcRegistry("", "", "", self._registry_client_config)
+
     def start(self) -> None:
         """Start the module server and register with the registry if configured."""
-        logger.info("Starting module server",extra={"server_config": self.server_config})
+        logger.info("Starting module server", extra={"server_config": self.server_config})
         super().start()
 
-        logger.debug("Starting module server",extra={"server_config": self.server_config})
-        # If a registry address is provided, register the module
-        if self.server_config.registry_address:
-            try:
-                self._register_with_registry()
-            except Exception:
-                logger.exception("Failed to register with registry")
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.debug(
-                "Setup post init started",extra={"client_config": self.client_config}
-            )
+            logger.debug("Setup post init started", extra={"client_config": self.client_config})
             self.module_servicer.setup.__post_init__(self.client_config)
 
     async def start_async(self) -> None:
         """Start the module server and register with the registry if configured."""
-        logger.info("Starting module server",extra={"server_config": self.server_config})
+        logger.info("Starting module server", extra={"server_config": self.server_config})
         await super().start_async()
-        # If a registry address is provided, register the module
-        if self.server_config.registry_address:
-            try:
-                self._register_with_registry()
-            except Exception:
-                logger.exception("Failed to register with registry")
+
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.info(
-                "Setup post init started",extra={"client_config": self.client_config}
-            )
-            await self.module_servicer.job_manager._start()
+            logger.info("Setup post init started", extra={"client_config": self.client_config})
+            await self.module_servicer.job_manager.start()
             self.module_servicer.setup.__post_init__(self.client_config)
 
     def stop(self, grace: float | None = None) -> None:
-        """Stop the module server and deregister from the registry if needed."""
-        # If registered with a registry, deregister
-        if self.server_config.registry_address:
-            try:
-                self._deregister_from_registry()
-            except ServerError:
-                logger.exception("Failed to deregister from registry")
+        """Stop the module server.
 
+        Modules become inactive when they stop sending heartbeats
+        """
         super().stop(grace)
 
     def _register_with_registry(self) -> None:
-        """Register this module with the registry server.
-
-        Raises:
-            ServerError: If communication with the registry server fails.
-        """
-        logger.debug(
-            "Registering module with registry at %s",
-            self.server_config.registry_address,
-        )
-
-        # Create appropriate channel based on security mode
-        channel = self._create_registry_channel()
-
-        with channel:
-            # Create a stub (client)
-            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
-
-            # Determine module type
-            module_type = self._determine_module_type()
-
-            metadata = metadata_pb2.Metadata(
-                name=self.module_class.metadata["name"],
-                tags=[
-                    metadata_pb2.Tag(tag=tag)
-                    for tag in self.module_class.metadata["tags"]
-                ],
-                description=self.module_class.metadata["description"],
+        """Register this module with the registry server."""
+        if not self.registry:
+            logger.debug("No registry configured, skipping registration")
+            return
+
+        module_id = self.module_class.get_module_id()
+        version = self.module_class.metadata.get("version", "0.0.0")
+
+        if not module_id or module_id == "unknown":
+            logger.warning(
+                "Module has no valid module_id, skipping registration",
+                extra={"module_class": self.module_class.__name__},
             )
-
-            self.module_class.metadata["module_id"] = (
-                f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
-            )
-            # Create registration request
-            request = registration_pb2.RegisterRequest(
-                module_id=self.module_class.metadata["module_id"],
-                version=self.module_class.metadata["version"],
-                module_type=module_type,
-                address=self.server_config.address,
-                metadata=metadata,
-            )
-
-            try:
-                # Call the register method
-                logger.debug(
-                    "Request sent to registry for module: %s:%s",
-                    self.module_class.metadata["name"],
-                    self.module_class.metadata["module_id"],
-                )
-                response = stub.RegisterModule(request)
-
-                if response.success:
-                    logger.debug("Module registered successfully")
-                else:
-                    logger.error("Module registration failed")
-            except grpc.RpcError:
-                logger.exception("RPC error during registration:")
-                raise ServerError
-
-    def _deregister_from_registry(self) -> None:
-        """Deregister this module from the registry server.
-
-        Raises:
-            ServerError: If communication with the registry server fails.
-        """
-        logger.debug(
-            "Deregistering module from registry at %s",
-            self.server_config.registry_address,
+            return
+
+        logger.info(
+            "Attempting to register module with registry",
+            extra={
+                "module_id": module_id,
+                "address": self.server_config.host,
+                "port": self.server_config.port,
+                "version": version,
+                "registry_address": self.server_config.registry_address,
+            },
         )
 
-        # Create appropriate channel based on security mode
-        channel = self._create_registry_channel()
-
-        with channel:
-            # Create a stub (client)
-            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
-
-            # Create deregistration request
-            request = registration_pb2.DeregisterRequest(
-                module_id=self.module_class.metadata["module_id"],
-            )
-        try:
-            # Call the deregister method
-            response = stub.DeregisterModule(request)
-
-            if response.success:
-                logger.debug("Module deregistered successfull")
-            else:
-                logger.error("Module deregistration failed")
-        except grpc.RpcError:
-            logger.exception("RPC error during deregistration")
-            raise ServerError
-
-    def _create_registry_channel(self) -> grpc.Channel:
-        """Create an appropriate channel to the registry server.
-
-        Returns:
-            A gRPC channel for communication with the registry.
+        result = self.registry.register(
+            module_id=module_id,
+            address=self.server_config.host,
+            port=self.server_config.port,
+            version=version,
+        )
 
-        Raises:
-            ValueError: If credentials are required but not provided.
-        """
-        if (
-            self.client_config is not None
-            and self.client_config.security == SecurityMode.SECURE
-            and self.client_config.credentials
-        ):
-            # Secure channel
-            # Secure channel
-            root_certificates = Path(
-                self.client_config.credentials.root_cert_path
-            ).read_bytes()
-
-            # mTLS channel
-            private_key = None
-            certificate_chain = None
-            if (
-                self.client_config.credentials.client_cert_path is not None
-                and self.client_config.credentials.client_key_path is not None
-            ):
-                private_key = Path(
-                    self.client_config.credentials.client_key_path
-                ).read_bytes()
-                certificate_chain = Path(
-                    self.client_config.credentials.client_cert_path
-                ).read_bytes()
-
-            # Create channel credentials
-            channel_credentials = grpc.ssl_channel_credentials(
-                root_certificates=root_certificates,
-                certificate_chain=certificate_chain,
-                private_key=private_key,
+        if result:
+            logger.info(
+                "Module registered successfully",
+                extra={
+                    "module_id": result.module_id,
+                    "address": self.server_config.host,
+                    "port": self.server_config.port,
+                    "registry_address": self.server_config.registry_address,
+                },
             )
-            return grpc.secure_channel(
-                self.server_config.registry_address, channel_credentials
+        else:
+            logger.warning(
+                "Module registration returned None (module may not exist in registry)",
+                extra={
+                    "module_id": module_id,
+                    "registry_address": self.server_config.registry_address,
+                },
             )
-        # Insecure channel
-        return grpc.insecure_channel(self.server_config.registry_address)
-
-    def _determine_module_type(self) -> str:
-        """Determine the module type based on its class.
-
-        Returns:
-            A string representing the module type.
-        """
-        module_type = "UNKNOWN"
-        class_name = self.module_class.__name__
-
-        if class_name == "ToolModule":
-            module_type = "TOOL"
-        elif class_name == "TriggerModule":
-            module_type = "TRIGGER"
-        elif class_name == "ArchetypeModule":
-            module_type = "KIN"
-
-        return module_type

digitalkin/grpc_servers/module_servicer.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 import grpc
-from digitalkin_proto.digitalkin.module.v2 import (
+from agentic_mesh_protocol.module.v1 import (
     information_pb2,
     lifecycle_pb2,
     module_service_pb2_grpc,
@@ -13,12 +13,13 @@ from digitalkin_proto.digitalkin.module.v2 import (
 )
 from google.protobuf import json_format, struct_pb2
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
+from digitalkin.models.core.job_manager_models import JobManagerMode
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import JobManagerMode
+from digitalkin.services.registry import GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode
 from digitalkin.services.setup.default_setup import DefaultSetup
 from digitalkin.services.setup.grpc_setup import GrpcSetup
@@ -40,6 +41,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
     args: Namespace
     setup: SetupStrategy
     job_manager: BaseJobManager
+    _registry_cache: RegistryStrategy | None = None
 
     def _add_parser_args(self, parser: ArgumentParser) -> None:
         super()._add_parser_args(parser)
@@ -76,12 +78,32 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         self.job_manager = job_manager_class(module_class, self.args.services_mode)
 
         logger.debug(
-            "ModuleServicer initialized with job manager: %s | %s",
+            "ModuleServicer initialized with job manager: %s",
             self.args.job_manager_mode,
-            self.job_manager,
+            extra={"job_manager": self.job_manager},
         )
         self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
+    def _get_registry(self) -> RegistryStrategy | None:
+        """Get a cached registry instance if configured.
+
+        Returns:
+            Cached GrpcRegistry instance if registry config exists, None otherwise.
+        """
+        if self._registry_cache is not None:
+            return self._registry_cache
+
+        registry_config = self.module_class.services_config_params.get("registry")
+        if not registry_config:
+            return None
+
+        client_config = registry_config.get("client_config")
+        if not client_config:
+            return None
+
+        self._registry_cache = GrpcRegistry("", "", "", client_config)
+        return self._registry_cache
+
     async def ConfigSetupModule(  # noqa: N802
         self,
         request: lifecycle_pb2.ConfigSetupModuleRequest,
@@ -108,11 +130,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                 "mission_id": request.mission_id,
             },
         )
-        # Process the module input
-        # TODO: Secret should be used here as well
         setup_version = request.setup_version
         config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
-        setup_version_data = self.module_class.create_setup_model(
+        setup_version_data = await self.module_class.create_setup_model(
             json_format.MessageToDict(request.setup_version.content),
             config_fields=True,
         )
@@ -139,8 +159,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
 
         updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
-        logger.info("Setup updated")
-        logger.debug(f"Updated setup data: {updated_setup_data=}")
+        logger.info("Setup updated", extra={"job_id": job_id})
+        logger.debug("Updated setup data", extra={"job_id": job_id, "setup_data": updated_setup_data})
         setup_version.content = json_format.ParseDict(
             updated_setup_data,
             struct_pb2.Struct(),
@@ -172,7 +192,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         # Process the module input
         # TODO: Check failure of input data format
-        input_data = self.module_class.create_input_model(dict(request.input.items()))
+        input_data = self.module_class.create_input_model(json_format.MessageToDict(request.input))
+
         setup_data_class = self.setup.get_setup(
             setup_dict={
                 "setup_id": request.setup_id,
@@ -184,7 +205,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             msg = "No setup data returned."
             raise ServicerError(msg)
 
-        setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
+        setup_data = await self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # create a task to run the module in background
         job_id = await self.job_manager.create_module_instance_job(
@@ -201,27 +222,37 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             yield lifecycle_pb2.StartModuleResponse(success=False)
             return
 
-        async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
-            async for message in stream:
-                if message.get("error", None) is not None:
-                    context.set_code(message["error"]["code"])
-                    context.set_details(message["error"]["error_message"])
-                    yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
-                    break
-
-                if message.get("exception", None) is not None:
-                    logger.error("Error in output_data")
-                    context.set_code(message["short_description"])
-                    context.set_details(message["exception"])
-                    yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
-                    break
-
-                if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
-                    yield lifecycle_pb2.StartModuleResponse(success=True, job_id=job_id)
-                    break
-
-                proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
-                yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+        try:
+            async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
+                async for message in stream:
+                    if message.get("error", None) is not None:
+                        logger.error("Error in output_data", extra={"message": message})
+                        context.set_code(message["error"]["code"])
+                        context.set_details(message["error"]["error_message"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    if message.get("exception", None) is not None:
+                        logger.error("Exception in output_data", extra={"message": message})
+                        context.set_code(message["short_description"])
+                        context.set_details(message["exception"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    logger.info("Yielding message from job %s: %s", job_id, message)
+                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+
+                    if message.get("root", {}).get("protocol") == "end_of_stream":
+                        logger.info(
+                            "End of stream signal received",
+                            extra={"job_id": job_id, "mission_id": request.mission_id},
+                        )
+                        break
+        finally:
+            await self.job_manager.wait_for_completion(job_id)
+            await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
+
         logger.info("Job %s finished", job_id)
 
     async def StopModule(  # noqa: N802
@@ -238,17 +269,19 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         Returns:
             A response indicating success or failure.
         """
-        logger.debug("StopModule called for module: '%s'", self.module_class.__name__)
+        logger.debug(
+            "StopModule called",
+            extra={"module_class": self.module_class.__name__, "job_id": request.job_id},
+        )
 
         response: bool = await self.job_manager.stop_module(request.job_id)
         if not response:
-            message = f"Job {request.job_id} not found"
-            logger.warning(message)
+            logger.warning("Job not found for stop request", extra={"job_id": request.job_id})
             context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
+            context.set_details(f"Job {request.job_id} not found")
             return lifecycle_pb2.StopModuleResponse(success=False)
 
-        logger.debug("Job %s stopped successfully", request.job_id)
+        logger.debug("Job stopped successfully", extra={"job_id": request.job_id})
         return lifecycle_pb2.StopModuleResponse(success=True)
 
     async def GetModuleStatus(  # noqa: N802
@@ -339,7 +372,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get input schema if available
         try:
             # Convert schema to proto format
-            input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
+            input_schema_proto = await self.module_class.get_input_format(
+                llm_format=request.llm_format,
+            )
             input_format_struct = json_format.Parse(
                 text=input_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -375,7 +410,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get output schema if available
         try:
             # Convert schema to proto format
-            output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
+            output_schema_proto = await self.module_class.get_output_format(
+                llm_format=request.llm_format,
+            )
             output_format_struct = json_format.Parse(
                 text=output_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -411,7 +448,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
+            setup_schema_proto = await self.module_class.get_setup_format(llm_format=request.llm_format)
             setup_format_struct = json_format.Parse(
                 text=setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -428,7 +465,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             setup_schema=setup_format_struct,
         )
 
-    def GetModuleSecret(  # noqa: N802
+    async def GetModuleSecret(  # noqa: N802
         self,
         request: information_pb2.GetModuleSecretRequest,
         context: grpc.ServicerContext,
@@ -447,7 +484,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get secret schema if available
         try:
             # Convert schema to proto format
-            secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_schema_proto = await self.module_class.get_secret_format(llm_format=request.llm_format)
             secret_format_struct = json_format.Parse(
                 text=secret_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -483,7 +520,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_schema_proto = await self.module_class.get_config_setup_format(llm_format=request.llm_format)
             config_setup_format_struct = json_format.Parse(
                 text=config_setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member

digitalkin/grpc_servers/utils/__init__.py

@@ -0,0 +1 @@
+"""gRPC servers utilities package."""

digitalkin/grpc_servers/utils/exceptions.py

@@ -27,11 +27,3 @@ class ServerStateError(ServerError):
 
 class ReflectionError(ServerError):
     """Error related to gRPC reflection service."""
-
-
-class HealthCheckError(ServerError):
-    """Error related to gRPC health check service."""
-
-
-class OptionalFeatureNotImplementedError(NotImplementedError):
-    """Raised when an optional feature is not implemented, but was requested."""