digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/services/filesystem/grpc_filesystem.py

@@ -0,0 +1,317 @@
+"""gRPC filesystem implementation."""
+
+from typing import Any, Literal
+
+from agentic_mesh_protocol.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
+from google.protobuf import struct_pb2
+from google.protobuf.json_format import MessageToDict
+
+from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+from digitalkin.grpc_servers.utils.grpc_error_handler import GrpcErrorHandlerMixin
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig
+from digitalkin.services.filesystem.filesystem_strategy import (
+    FileFilter,
+    FilesystemRecord,
+    FilesystemServiceError,
+    FilesystemStrategy,
+    UploadFileData,
+)
+
+
+class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
+    """Default state filesystem strategy."""
+
+    @staticmethod
+    def _file_type_to_enum(file_type: str) -> filesystem_pb2.FileType:
+        """Convert a file type string to a FileType enum.
+
+        Args:
+            file_type: The file type string to convert
+
+        Returns:
+            filesystem_pb2.FileType: The converted file type enum
+        """
+        if not file_type.upper().startswith("FILE_TYPE_"):
+            file_type = f"FILE_TYPE_{file_type.upper()}"
+        try:
+            return getattr(filesystem_pb2.FileType, file_type.upper())
+        except AttributeError:
+            return filesystem_pb2.FileType.FILE_TYPE_UNSPECIFIED
+
+    @staticmethod
+    def _file_status_to_enum(file_status: str) -> filesystem_pb2.FileStatus:
+        """Convert a file status string to a FileStatus enum.
+
+        Args:
+            file_status: The file status string to convert
+
+        Returns:
+            filesystem_pb2.FileStatus: The converted file status enum
+        """
+        if not file_status.upper().startswith("FILE_STATUS_"):
+            file_status = f"FILE_STATUS_{file_status.upper()}"
+        try:
+            return getattr(filesystem_pb2.FileStatus, file_status.upper())
+        except AttributeError:
+            return filesystem_pb2.FileStatus.FILE_STATUS_UNSPECIFIED
+
+    @staticmethod
+    def _file_proto_to_data(file: filesystem_pb2.File) -> FilesystemRecord:
+        """Convert a File proto message to FilesystemRecord.
+
+        Args:
+            file: The File proto message to convert
+
+        Returns:
+            FilesystemRecord: The converted data
+        """
+        return FilesystemRecord(
+            id=file.file_id,
+            context=file.context,
+            name=file.name,
+            file_type=filesystem_pb2.FileType.Name(file.file_type),
+            content_type=file.content_type,
+            size_bytes=file.size_bytes,
+            checksum=file.checksum,
+            metadata=MessageToDict(file.metadata),
+            storage_uri=file.storage_uri,
+            file_url=file.file_url,
+            status=filesystem_pb2.FileStatus.Name(file.status),
+            content=file.content,
+        )
+
+    def _filter_to_proto(self, filters: FileFilter) -> filesystem_pb2.FileFilter:
+        """Convert a FileFilter to a FileFilter proto message.
+
+        Args:
+            filters: The FileFilter to convert
+
+        Returns:
+            filesystem_pb2.FileFilter: The converted FileFilter proto message
+        """
+        return filesystem_pb2.FileFilter(
+            **filters.model_dump(exclude={"file_types", "status"}),
+            file_types=[self._file_type_to_enum(file_type) for file_type in filters.file_types]
+            if filters.file_types
+            else None,
+            status=self._file_status_to_enum(filters.status) if filters.status else None,
+        )
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        client_config: ClientConfig,
+        config: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize the gRPC filesystem strategy.
+
+        Args:
+            mission_id: The ID of the mission this strategy is associated with
+            setup_id: The ID of the setup
+            setup_version_id: The ID of the setup version this strategy is associated with
+            client_config: Configuration for the gRPC client connection
+            config: Configuration for the filesystem strategy
+        """
+        super().__init__(mission_id, setup_id, setup_version_id, config)
+        self.service_name = "FilesystemService"
+        channel = self._init_channel(client_config)
+        self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
+        logger.debug("Channel client 'Filesystem' initialized successfully")
+
+    def upload_files(
+        self,
+        files: list[UploadFileData],
+    ) -> tuple[list[FilesystemRecord], int, int]:
+        """Upload multiple files to the filesystem.
+
+        Args:
+            files: List of tuples containing (content, name, file_type, content_type, metadata, replace_if_exists)
+
+        Returns:
+            tuple[list[FilesystemRecord], int, int]: List of uploaded files, total uploaded count, total failed count
+        """
+        logger.debug("Uploading %d files", len(files))
+        with self.handle_grpc_errors("UploadFiles", FilesystemServiceError):
+            upload_files: list[filesystem_pb2.UploadFileData] = []
+            for file in files:
+                metadata_struct: struct_pb2.Struct | None = None
+                if file.metadata:
+                    metadata_struct = struct_pb2.Struct()
+                    metadata_struct.update(file.metadata)
+                upload_files.append(
+                    filesystem_pb2.UploadFileData(
+                        context=self.mission_id,
+                        name=file.name,
+                        file_type=self._file_type_to_enum(file.file_type),
+                        content_type=file.content_type or "application/octet-stream",
+                        content=file.content,
+                        metadata=metadata_struct,
+                        status=filesystem_pb2.FileStatus.FILE_STATUS_UPLOADING,
+                        replace_if_exists=file.replace_if_exists,
+                    )
+                )
+            request = filesystem_pb2.UploadFilesRequest(files=upload_files)
+            response: filesystem_pb2.UploadFilesResponse = self.exec_grpc_query("UploadFiles", request)
+            results = [self._file_proto_to_data(result.file) for result in response.results if result.HasField("file")]
+            logger.debug("Uploaded files: %s", results)
+            return results, response.total_uploaded, response.total_failed
+
+    def get_file(
+        self,
+        file_id: str,
+        context: Literal["mission", "setup"] = "mission",
+        *,
+        include_content: bool = False,
+    ) -> FilesystemRecord:
+        """Get a file from the filesystem.
+
+        Args:
+            file_id: The ID of the file to be retrieved
+            context: The context of the files (mission or setup)
+            include_content: Whether to include file content in response
+
+        Returns:
+            FilesystemRecord: Metadata about the retrieved file
+
+        Raises:
+            FilesystemServiceError: If there is an error retrieving the file
+        """
+        match context:
+            case "setup":
+                context_id = self.setup_id
+            case "mission":
+                context_id = self.mission_id
+        with self.handle_grpc_errors("GetFile", FilesystemServiceError):
+            request = filesystem_pb2.GetFileRequest(
+                context=context_id,
+                file_id=file_id,
+                include_content=include_content,
+            )
+
+            response: filesystem_pb2.GetFileResponse = self.exec_grpc_query("GetFile", request)
+
+            return self._file_proto_to_data(response.file)
+
+    def update_file(
+        self,
+        file_id: str,
+        content: bytes | None = None,
+        file_type: Literal[
+            "UNSPECIFIED",
+            "DOCUMENT",
+            "IMAGE",
+            "VIDEO",
+            "AUDIO",
+            "ARCHIVE",
+            "CODE",
+            "OTHER",
+        ]
+        | None = None,
+        content_type: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        new_name: str | None = None,
+        status: str | None = None,
+    ) -> FilesystemRecord:
+        """Update a file in the filesystem.
+
+        Args:
+            file_id: The id of the file to be updated
+            content: Optional new content of the file
+            file_type: Optional new type of data
+            content_type: Optional new MIME type
+            metadata: Optional new metadata (will merge with existing)
+            new_name: Optional new name for the file
+            status: Optional new status for the file
+
+        Returns:
+            FilesystemRecord: Metadata about the updated file
+
+        Raises:
+            FilesystemServiceError: If there is an error during update
+        """
+        with self.handle_grpc_errors("UpdateFile", FilesystemServiceError):
+            request = filesystem_pb2.UpdateFileRequest(
+                context=self.mission_id,
+                file_id=file_id,
+                content=content,
+                file_type=self._file_type_to_enum(file_type) if file_type else None,
+                content_type=content_type,
+                new_name=new_name,
+                status=self._file_status_to_enum(status) if status else None,
+            )
+
+            if metadata:
+                request.metadata.update(metadata)
+
+            response: filesystem_pb2.UpdateFileResponse = self.exec_grpc_query("UpdateFile", request)
+            return self._file_proto_to_data(response.result.file)
+
+    def delete_files(
+        self,
+        filters: FileFilter,
+        *,
+        permanent: bool = False,
+        force: bool = False,
+    ) -> tuple[dict[str, bool], int, int]:
+        """Delete multiple files from the filesystem.
+
+        Args:
+            filters: Filter criteria for the files
+            permanent: Whether to permanently delete the files
+            force: Whether to force delete even if files are in use
+
+        Returns:
+            tuple[dict[str, bool], int, int]: Results per file, total deleted count, total failed count
+        """
+        with self.handle_grpc_errors("DeleteFiles", FilesystemServiceError):
+            request = filesystem_pb2.DeleteFilesRequest(
+                context=self.mission_id,
+                filters=self._filter_to_proto(filters),
+                permanent=permanent,
+                force=force,
+            )
+
+            response: filesystem_pb2.DeleteFilesResponse = self.exec_grpc_query("DeleteFiles", request)
+            return dict(response.results), response.total_deleted, response.total_failed
+
+    def get_files(
+        self,
+        filters: FileFilter,
+        *,
+        list_size: int = 100,
+        offset: int = 0,
+        order: str | None = None,
+        include_content: bool = False,
+    ) -> tuple[list[FilesystemRecord], int]:
+        """Get multiple files from the filesystem.
+
+        Args:
+            filters: Filter criteria for the files
+            list_size: Number of files to return per page
+            offset: Offset to start from
+            order: Field to order results by
+            include_content: Whether to include file content in response
+
+        Returns:
+            tuple[list[FilesystemRecord], int]: List of files and total count
+        """
+        match filters.context:
+            case "setup":
+                context_id = self.setup_id
+            case "mission":
+                context_id = self.mission_id
+        with self.handle_grpc_errors("GetFiles", FilesystemServiceError):
+            request = filesystem_pb2.GetFilesRequest(
+                context=context_id,
+                filters=self._filter_to_proto(filters),
+                include_content=include_content,
+                list_size=list_size,
+                offset=offset,
+                order=order,
+            )
+            response: filesystem_pb2.GetFilesResponse = self.exec_grpc_query("GetFiles", request)
+
+            return [self._file_proto_to_data(file) for file in response.files], response.total_count

digitalkin/services/identity/__init__.py

@@ -0,0 +1,6 @@
+"""This module is responsible for handling the identity service."""
+
+from digitalkin.services.identity.default_identity import DefaultIdentity
+from digitalkin.services.identity.identity_strategy import IdentityStrategy
+
+__all__ = ["DefaultIdentity", "IdentityStrategy"]

digitalkin/services/identity/default_identity.py

@@ -0,0 +1,15 @@
+"""Default identity."""
+
+from digitalkin.services.identity.identity_strategy import IdentityStrategy
+
+
+class DefaultIdentity(IdentityStrategy):
+    """DefaultIdentity is the default identity strategy."""
+
+    async def get_identity(self) -> str:  # noqa: PLR6301
+        """Get the identity.
+
+        Returns:
+            str: The identity
+        """
+        return "default_identity"

digitalkin/services/identity/identity_strategy.py

@@ -0,0 +1,14 @@
+"""This module contains the abstract base class for identity strategies."""
+
+from abc import ABC, abstractmethod
+
+from digitalkin.services.base_strategy import BaseStrategy
+
+
+class IdentityStrategy(BaseStrategy, ABC):
+    """IdentityStrategy is the abstract base class for all identity strategies."""
+
+    @abstractmethod
+    async def get_identity(self) -> str:
+        """Get the identity."""
+        raise NotImplementedError

digitalkin/services/registry/__init__.py

@@ -0,0 +1,27 @@
+"""This module is responsible for handling the registry service."""
+
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    ModuleStatusInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.default_registry import DefaultRegistry
+from digitalkin.services.registry.exceptions import (
+    RegistryModuleNotFoundError,
+    RegistryServiceError,
+)
+from digitalkin.services.registry.grpc_registry import GrpcRegistry
+from digitalkin.services.registry.registry_strategy import RegistryStrategy
+
+__all__ = [
+    "DefaultRegistry",
+    "GrpcRegistry",
+    "ModuleInfo",
+    "ModuleStatusInfo",
+    "RegistryModuleNotFoundError",
+    "RegistryModuleStatus",
+    "RegistryModuleType",
+    "RegistryServiceError",
+    "RegistryStrategy",
+]

digitalkin/services/registry/default_registry.py

@@ -0,0 +1,141 @@
+"""Default registry implementation."""
+
+from typing import ClassVar
+
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    ModuleStatusInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.exceptions import RegistryModuleNotFoundError
+from digitalkin.services.registry.registry_strategy import RegistryStrategy
+
+
+class DefaultRegistry(RegistryStrategy):
+    """Default registry strategy using in-memory storage."""
+
+    _modules: ClassVar[dict[str, ModuleInfo]] = {}
+
+    def discover_by_id(self, module_id: str) -> ModuleInfo:
+        """Get module info by ID.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleInfo with module details.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+        return self._modules[module_id]
+
+    def search(
+        self,
+        name: str | None = None,
+        module_type: str | None = None,
+        organization_id: str | None = None,  # noqa: ARG002
+    ) -> list[ModuleInfo]:
+        """Search for modules by criteria.
+
+        Args:
+            name: Filter by name (partial match).
+            module_type: Filter by type (archetype, tool).
+            organization_id: Filter by organization (not used in local storage).
+
+        Returns:
+            List of matching modules.
+        """
+        results = list(self._modules.values())
+
+        if name:
+            results = [m for m in results if name in m.name]
+
+        if module_type:
+            results = [m for m in results if m.module_type == module_type]
+
+        return results
+
+    def get_status(self, module_id: str) -> ModuleStatusInfo:
+        """Get module status.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleStatusInfo with current status.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+
+        module = self._modules[module_id]
+        return ModuleStatusInfo(
+            module_id=module_id,
+            status=module.status or RegistryModuleStatus.UNSPECIFIED,
+        )
+
+    def register(
+        self,
+        module_id: str,
+        address: str,
+        port: int,
+        version: str,
+    ) -> ModuleInfo | None:
+        """Register a module with the registry.
+
+        Note: Updates existing module or creates new one in local storage.
+
+        Args:
+            module_id: Unique module identifier.
+            address: Network address.
+            port: Network port.
+            version: Module version.
+
+        Returns:
+            ModuleInfo if successful, None otherwise.
+        """
+        existing = self._modules.get(module_id)
+        self._modules[module_id] = ModuleInfo(
+            module_id=module_id,
+            module_type=existing.module_type if existing else RegistryModuleType.UNSPECIFIED,
+            address=address,
+            port=port,
+            version=version,
+            name=existing.name if existing else module_id,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return self._modules[module_id]
+
+    def heartbeat(self, module_id: str) -> RegistryModuleStatus:
+        """Send heartbeat to keep module active.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            Current module status after heartbeat.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+
+        module = self._modules[module_id]
+        # Update status to ACTIVE on heartbeat
+        self._modules[module_id] = ModuleInfo(
+            module_id=module.module_id,
+            module_type=module.module_type,
+            address=module.address,
+            port=module.port,
+            version=module.version,
+            name=module.name,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return RegistryModuleStatus.ACTIVE

digitalkin/services/registry/exceptions.py

@@ -0,0 +1,47 @@
+"""Registry-specific exceptions.
+
+This module contains custom exceptions for registry service operations.
+"""
+
+
+class RegistryServiceError(Exception):
+    """Base exception for registry service errors."""
+
+
+class RegistryModuleNotFoundError(RegistryServiceError):
+    """Raised when a module is not found in the registry."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that was not found.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' not found in registry")
+
+
+class ModuleAlreadyExistsError(RegistryServiceError):
+    """Raised when attempting to register an already-registered module."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that already exists.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' already registered")
+
+
+class InvalidStatusError(RegistryServiceError):
+    """Raised when an invalid status is provided."""
+
+    def __init__(self, status: int) -> None:
+        """Initialize the exception.
+
+        Args:
+            status: The invalid status value.
+        """
+        self.status = status
+        super().__init__(f"Invalid module status: {status}")