digitalkin 0.2.11__tar.gz → 0.2.13__tar.gz

{digitalkin-0.2.11 → digitalkin-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.11
+Version: 0.2.13
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,29 +452,36 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.8
+Requires-Dist: digitalkin-proto>=0.1.10
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: openai>=1.76.2
 Requires-Dist: pydantic>=2.11.4
 Provides-Extra: dev
-Requires-Dist: pytest>=8.3.4; extra == "dev"
-Requires-Dist: pytest-asyncio>=0.26.0; extra == "dev"
-Requires-Dist: pytest-cov>=6.1.0; extra == "dev"
-Requires-Dist: typos>=1.31.2; extra == "dev"
-Requires-Dist: ruff>=0.11.7; extra == "dev"
+Requires-Dist: typos>=1.32.0; extra == "dev"
+Requires-Dist: ruff>=0.11.9; extra == "dev"
 Requires-Dist: mypy>=1.15.0; extra == "dev"
 Requires-Dist: pyright>=1.1.400; extra == "dev"
 Requires-Dist: pre-commit>=4.2.0; extra == "dev"
 Requires-Dist: bump2version>=1.0.1; extra == "dev"
 Requires-Dist: build>=1.2.2; extra == "dev"
 Requires-Dist: twine>=6.1.0; extra == "dev"
-Requires-Dist: cryptography>=44.0.2; extra == "dev"
-Requires-Dist: grpcio-testing>=1.71.0; extra == "dev"
-Requires-Dist: freezegun>=1.5.1; extra == "dev"
+Requires-Dist: cryptography>=44.0.3; extra == "dev"
+Requires-Dist: taskiq[reload]>=0.11.17; extra == "dev"
 Provides-Extra: examples
-Requires-Dist: openai>=1.66.3; extra == "examples"
+Requires-Dist: openai>=1.75.0; extra == "examples"
+Provides-Extra: tests
+Requires-Dist: freezegun>=1.5.1; extra == "tests"
+Requires-Dist: hdrhistogram>=0.10.3; extra == "tests"
+Requires-Dist: grpcio-testing>=1.71.0; extra == "tests"
+Requires-Dist: psutil>=7.0.0; extra == "tests"
+Requires-Dist: pytest>=8.3.4; extra == "tests"
+Requires-Dist: pytest-asyncio>=0.26.0; extra == "tests"
+Requires-Dist: pytest-cov>=6.1.0; extra == "tests"
+Provides-Extra: taskiq
+Requires-Dist: rstream>=0.20.9; extra == "taskiq"
+Requires-Dist: taskiq-aio-pika>=0.4.2; extra == "taskiq"
+Requires-Dist: taskiq-redis>=1.0.8; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK
@@ -509,6 +516,16 @@ To install the DigitalKin SDK, simply run:
 pip install digitalkin
 ```
 
+**Optional Taskiq Integration**: Asynchronous task execution powered by Taskiq, backed by RabbitMQ and Redis
+To enable the Rabbitmq streaming capabilities, run:
+
+```sh
+sudo rabbitmq-plugins enable rabbitmq_stream
+
+# Core + Taskiq integration (RabbitMQ broker)
+pip install digitalkin[taskiq]
+```
+
 ## 🛠️ Usage
 
 ### Basic Import
@@ -519,6 +536,20 @@ Start by importing the necessary modules:
 import digitalkin
 ```
 
+## Features
+
+### Taskiq with RabbitMQ
+
+TaskIQ intergration allows the module to scale for heavy CPU tasks by having the request's stateless module in a new instance.
+
+- **Decoupled Scalability**: RabbitMQ brokers messages, letting producers and consumers scale independently.
+- **Reliability**: Durable queues, acknowledgements, and dead-lettering ensure tasks aren’t lost.
+- **Concurrency Control**: Taskiq’s worker pool manages parallel execution without custom schedulers.
+- **Flexibility**: Built-in retries, exponential backoff, and Redis result-backend for resilient workflows.
+- **Ecosystem**: Battle-tested `aio-pika` AMQP client plus Taskiq’s decorator-based API.
+
+By combining Taskiq’s async API with RabbitMQ’s guarantees, you get a robust, production-ready queue with minimal boilerplate.
+
 ## 👷‍♂️ Development
 
 ### Prerequisites

{digitalkin-0.2.11 → digitalkin-0.2.13}/README.md

@@ -30,6 +30,16 @@ To install the DigitalKin SDK, simply run:
 pip install digitalkin
 ```
 
+**Optional Taskiq Integration**: Asynchronous task execution powered by Taskiq, backed by RabbitMQ and Redis
+To enable the Rabbitmq streaming capabilities, run:
+
+```sh
+sudo rabbitmq-plugins enable rabbitmq_stream
+
+# Core + Taskiq integration (RabbitMQ broker)
+pip install digitalkin[taskiq]
+```
+
 ## 🛠️ Usage
 
 ### Basic Import
@@ -40,6 +50,20 @@ Start by importing the necessary modules:
 import digitalkin
 ```
 
+## Features
+
+### Taskiq with RabbitMQ
+
+TaskIQ intergration allows the module to scale for heavy CPU tasks by having the request's stateless module in a new instance.
+
+- **Decoupled Scalability**: RabbitMQ brokers messages, letting producers and consumers scale independently.
+- **Reliability**: Durable queues, acknowledgements, and dead-lettering ensure tasks aren’t lost.
+- **Concurrency Control**: Taskiq’s worker pool manages parallel execution without custom schedulers.
+- **Flexibility**: Built-in retries, exponential backoff, and Redis result-backend for resilient workflows.
+- **Ecosystem**: Battle-tested `aio-pika` AMQP client plus Taskiq’s decorator-based API.
+
+By combining Taskiq’s async API with RabbitMQ’s guarantees, you get a robust, production-ready queue with minimal boilerplate.
+
 ## 👷‍♂️ Development
 
 ### Prerequisites

digitalkin-0.2.13/examples/modules/cpu_intensive_module.py

@@ -0,0 +1,271 @@
+"""Simple module calling an LLM."""
+
+import logging
+from collections.abc import Callable
+from typing import Any, ClassVar, Literal
+
+from pydantic import BaseModel, Field
+
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerConfig, ServerMode
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_models import ServicesStrategy
+from digitalkin.services.setup.setup_strategy import SetupData
+
+# Configure logging with clear formatting
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+class MessageInputPayload(BaseModel):
+    """Message trigger model for the CPU Archetype module."""
+
+    payload_type: Literal["message"] = "message"
+    user_prompt: str = Field(
+        ...,
+        title="User Prompt",
+        description="The prompt provided by the user for processing.",
+    )
+
+
+class InputFile(BaseModel):
+    """File model for the CPU Archetype module."""
+
+    name: str = Field(
+        ...,
+        title="File Name",
+        description="The name of the file to be processed.",
+    )
+    content: bytes = Field(
+        ...,
+        title="File Content",
+        description="The content of the file to be processed.",
+    )
+
+    file_type: str = Field(
+        ...,
+        title="File Type",
+        description="The type of the file to be processed.",
+    )
+
+
+class FileInputPayload(BaseModel):
+    """File input model for the CPU Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[InputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files to be processed.",
+    )
+
+
+class CPUInput(BaseModel):
+    """Input model defining what data the module expects."""
+
+    payload: MessageInputPayload | FileInputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or list of file input.",
+    )
+
+
+class MessageOutputPayload(BaseModel):
+    """Message output model for the CPU Archetype module."""
+
+    payload_type: Literal["message"] = "message"
+    user_response: str = Field(
+        ...,
+        title="User Response",
+        description="The response generated by the assistant based on the user prompt.",
+    )
+
+
+class OutputFile(BaseModel):
+    """File model for the CPU Archetype module."""
+
+    name: str = Field(
+        ...,
+        title="File Name",
+        description="The name of the file to be processed.",
+    )
+    url: str | None = Field(
+        ...,
+        title="File URL",
+        description="The URL of the file to be processed.",
+    )
+
+    message: str | None = Field(
+        None,
+        title="Message",
+        description="Optional message associated with the file.",
+    )
+
+
+class FileOutputPayload(BaseModel):
+    """File output model for the CPU Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[OutputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files generated by the assistant.",
+    )
+
+
+class CPUOutput(BaseModel):
+    """Output model defining what data the module produces."""
+
+    payload: MessageOutputPayload | FileOutputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or file response.",
+    )
+
+
+class CPUSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    model_name: str = Field(
+        ...,
+        title="Model Name",
+        description="The name of the CPU model to use for processing.",
+    )
+    developer_prompt: str = Field(
+        ...,
+        title="Developer Prompt",
+        description="The developer prompt new versions of system prompt, it defines the behavior of the assistant.",
+    )
+    temperature: float = Field(
+        0.7,
+        title="Temperature",
+        description="Controls the randomness of the model's output. Higher values make output more random.",
+    )
+    max_tokens: int = Field(
+        100,
+        title="Max Tokens",
+        description="The maximum number of tokens to generate in the response.",
+    )
+
+
+class CPUToolSecret(BaseModel):
+    """Secret model defining module configuration parameters."""
+
+
+server_config = ServerConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    max_workers=10,
+    credentials=None,
+)
+
+
+client_config = ClientConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    credentials=None,
+)
+
+
+class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret]):
+    """A CPU endpoint tool module module."""
+
+    name = "CPUIntensiveModule"
+    description = "A module that interacts with CPU API to process text"
+
+    # Define the schema formats for the module
+    input_format = CPUInput
+    output_format = CPUOutput
+    setup_format = CPUSetup
+    secret_format = CPUToolSecret
+
+    # Define module metadata for discovery
+    metadata: ClassVar[dict[str, Any]] = {
+        "name": "CPUIntensiveModule",
+        "description": "Transforms input text using a streaming LLM response.",
+        "version": "1.0.0",
+        "tags": ["text", "transformation", "encryption", "streaming"],
+    }
+    # Define services_config_params with default values
+    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {
+        "storage": {
+            "config": {"chat_history": None},
+            "client_config": client_config,
+        },
+        "filesystem": {
+            "config": {},
+            "client_config": client_config,
+        },
+        "cost": {
+            "config": {},
+            "client_config": client_config,
+        },
+    }
+
+    async def initialize(self, setup_data: SetupData) -> None:
+        """Initialize the module capabilities.
+
+        This method is called when the module is loaded by the server.
+        Use it to set up module-specific resources or configurations.
+        """
+
+    async def run(
+        self,
+        input_data: CPUInput,
+        setup_data: CPUSetup,
+        callback: Callable,
+    ) -> None:
+        """Run the module.
+
+        Args:
+            input_data: Input data for the module
+            setup_data: Setup data for the module
+            callback: Callback function to report progress
+
+        Raises:
+            ValueError: If the payload type is unknown
+        """
+        # Validate the input data
+        input_model = self.input_format.model_validate(input_data)
+        self.setup_format.model_validate(setup_data)
+        logger.debug("Running with input data: %s", input_model)
+
+        if not hasattr(input_model, "payload"):
+            error_msg = "Input data is missing 'payload' field"
+            raise ValueError(error_msg)
+
+        if not hasattr(input_model.payload, "payload_type"):
+            error_msg = "Input payload is missing 'type' field"
+            raise ValueError(error_msg)
+
+        total = 0
+        input = MessageInputPayload.model_validate(input_model.payload).user_prompt
+
+        for i in range(int(input)):
+            total += i * i
+            if i % 100 == 0 or i == int(input) - 1:
+                message_output_payload = MessageOutputPayload(
+                    payload_type="message",
+                    user_response=f"result iteration {i}: {total}",
+                )
+                output_model = self.output_format.model_validate({"payload": message_output_payload})
+                await callback(output_data=output_model)
+        logger.info("Job %s completed", self.job_id)
+
+    async def cleanup(self) -> None:
+        """Clean up any resources when the module is stopped.
+
+        This method is called when the module is being shut down.
+        Use it to close connections, free resources, etc.
+        """
+        logger.info("Cleaning up module %s", self.metadata["name"])
+        # Release any resources here if needed.