ergon-framework-python 0.1.0__py3-none-any.whl

ergon/connector/sqs/models.py

@@ -0,0 +1,36 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class SQSClient(BaseModel):
+    region_name: str = Field(description="AWS region name (e.g. us-east-1)")
+    aws_access_key_id: Optional[str] = Field(default=None, description="AWS access key ID. None = use env/IAM role")
+    aws_secret_access_key: Optional[str] = Field(
+        default=None, description="AWS secret access key. None = use env/IAM role"
+    )
+    aws_session_token: Optional[str] = Field(default=None, description="AWS session token for temporary credentials")
+    endpoint_url: Optional[str] = Field(default=None, description="Custom endpoint URL (e.g. LocalStack)")
+    queue_url: Optional[str] = Field(default=None, description="Default SQS queue URL, can be overridden per-call")
+
+
+class SQSConsumerConfig(BaseModel):
+    queue_url: Optional[str] = Field(default=None, description="Override default queue URL for consuming")
+    wait_time_seconds: int = Field(default=20, description="Long polling wait time in seconds (0-20)")
+    visibility_timeout: Optional[int] = Field(
+        default=None, description="Override queue's default visibility timeout in seconds"
+    )
+    attribute_names: List[str] = Field(default=["All"], description="System attribute names to include in response")
+    message_attribute_names: List[str] = Field(
+        default=["All"], description="Message attribute names to include in response"
+    )
+
+
+class SQSProducerConfig(BaseModel):
+    queue_url: Optional[str] = Field(default=None, description="Override default queue URL for producing")
+    delay_seconds: int = Field(default=0, description="Delay in seconds before message becomes visible (0-900)")
+    message_group_id: Optional[str] = Field(default=None, description="Message group ID for FIFO queues")
+    message_deduplication_id: Optional[str] = Field(default=None, description="Deduplication ID for FIFO queues")
+    message_attributes: Optional[Dict[str, Dict]] = Field(
+        default=None, description="Custom message attributes to attach to sent messages"
+    )

ergon/connector/sqs/service.py

@@ -0,0 +1,219 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+import boto3
+
+from .models import SQSClient
+
+logger = logging.getLogger(__name__)
+
+
+class SQSService:
+    def __init__(self, client: SQSClient) -> None:
+        self.client = client
+
+        kwargs: Dict[str, Any] = {"region_name": client.region_name}
+        if client.aws_access_key_id:
+            kwargs["aws_access_key_id"] = client.aws_access_key_id
+        if client.aws_secret_access_key:
+            kwargs["aws_secret_access_key"] = client.aws_secret_access_key
+        if client.aws_session_token:
+            kwargs["aws_session_token"] = client.aws_session_token
+        if client.endpoint_url:
+            kwargs["endpoint_url"] = client.endpoint_url
+
+        self._sqs = boto3.client("sqs", **kwargs)
+
+    # ---------- Receive ----------
+
+    def receive_messages(
+        self,
+        queue_url: Optional[str] = None,
+        max_number_of_messages: int = 1,
+        wait_time_seconds: int = 20,
+        visibility_timeout: Optional[int] = None,
+        attribute_names: Optional[List[str]] = None,
+        message_attribute_names: Optional[List[str]] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Receive up to `max_number_of_messages` from the queue.
+
+        SQS caps MaxNumberOfMessages at 10 per API call, so for larger
+        batch sizes this loops until the requested count is reached or
+        the queue returns no more messages.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        collected: List[Dict[str, Any]] = []
+        remaining = max_number_of_messages
+
+        first_call = True
+        while remaining > 0:
+            fetch_count = min(remaining, 10)
+
+            params: Dict[str, Any] = {
+                "QueueUrl": url,
+                "MaxNumberOfMessages": fetch_count,
+                "WaitTimeSeconds": wait_time_seconds if first_call else 0,
+                "AttributeNames": attribute_names or ["All"],
+                "MessageAttributeNames": message_attribute_names or ["All"],
+            }
+            if visibility_timeout is not None:
+                params["VisibilityTimeout"] = visibility_timeout
+
+            response = self._sqs.receive_message(**params)
+            messages = response.get("Messages", [])
+            first_call = False
+
+            if not messages:
+                break
+
+            collected.extend(messages)
+            remaining -= len(messages)
+
+            if len(messages) < fetch_count:
+                break
+
+        return collected
+
+    # ---------- Send ----------
+
+    def send_message(
+        self,
+        message_body: Any,
+        queue_url: Optional[str] = None,
+        delay_seconds: int = 0,
+        message_attributes: Optional[Dict[str, Dict]] = None,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        body = json.dumps(message_body) if not isinstance(message_body, str) else message_body
+
+        params: Dict[str, Any] = {
+            "QueueUrl": url,
+            "MessageBody": body,
+            "DelaySeconds": delay_seconds,
+        }
+        if message_attributes:
+            params["MessageAttributes"] = message_attributes
+        if message_group_id:
+            params["MessageGroupId"] = message_group_id
+        if message_deduplication_id:
+            params["MessageDeduplicationId"] = message_deduplication_id
+
+        return self._sqs.send_message(**params)
+
+    def send_message_batch(
+        self,
+        entries: List[Dict[str, Any]],
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Send messages in batches of up to 10 (SQS limit).
+        Each entry must have at minimum 'Id' and 'MessageBody'.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        successful = []
+        failed = []
+
+        for i in range(0, len(entries), 10):
+            chunk = entries[i : i + 10]
+            response = self._sqs.send_message_batch(QueueUrl=url, Entries=chunk)
+            successful.extend(response.get("Successful", []))
+            failed.extend(response.get("Failed", []))
+
+        return {"Successful": successful, "Failed": failed}
+
+    # ---------- Delete ----------
+
+    def delete_message(self, receipt_handle: str, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        self._sqs.delete_message(QueueUrl=url, ReceiptHandle=receipt_handle)
+
+    def delete_message_batch(
+        self,
+        entries: List[Dict[str, str]],
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Delete messages in batches of up to 10.
+        Each entry must have 'Id' and 'ReceiptHandle'.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        successful = []
+        failed = []
+
+        for i in range(0, len(entries), 10):
+            chunk = entries[i : i + 10]
+            response = self._sqs.delete_message_batch(QueueUrl=url, Entries=chunk)
+            successful.extend(response.get("Successful", []))
+            failed.extend(response.get("Failed", []))
+
+        return {"Successful": successful, "Failed": failed}
+
+    # ---------- Queue Management ----------
+
+    def list_queues(self, prefix: Optional[str] = None) -> List[str]:
+        params: Dict[str, Any] = {}
+        if prefix:
+            params["QueueNamePrefix"] = prefix
+
+        response = self._sqs.list_queues(**params)
+        return response.get("QueueUrls", [])
+
+    def get_queue_attributes(
+        self,
+        attribute_names: Optional[List[str]] = None,
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, str]:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        response = self._sqs.get_queue_attributes(
+            QueueUrl=url,
+            AttributeNames=attribute_names or ["All"],
+        )
+        return response.get("Attributes", {})
+
+    def create_queue(
+        self,
+        queue_name: str,
+        attributes: Optional[Dict[str, str]] = None,
+    ) -> str:
+        params: Dict[str, Any] = {"QueueName": queue_name}
+        if attributes:
+            params["Attributes"] = attributes
+
+        response = self._sqs.create_queue(**params)
+        return response["QueueUrl"]
+
+    def delete_queue(self, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        self._sqs.delete_queue(QueueUrl=url)
+
+    def purge_queue(self, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        self._sqs.purge_queue(QueueUrl=url)

ergon/connector/transaction.py

@@ -0,0 +1,14 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Transaction(BaseModel):
+    """
+    Generic transaction model.
+    """
+
+    id: str
+    payload: Any
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    model_config = ConfigDict(frozen=True)

ergon/service/__init__.py

@@ -0,0 +1,5 @@
+from .service import ServiceConfig
+
+__all__ = [
+    "ServiceConfig",
+]

ergon/service/service.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class ServiceConfig(BaseModel):
+    """
+    Generic service configuration.
+
+    service: class implementing a service
+    args: positional arguments passed to service's __init__
+    kwargs: keyword arguments passed to service's __init__
+    """
+
+    service: type
+    args: tuple[Any, ...] = ()
+    kwargs: dict[str, Any] = {}

ergon/task/__init__.py

@@ -0,0 +1,13 @@
+from . import exceptions, helpers, mixins, policies, utils
+from .base import TaskConfig
+from .manager import manager
+
+__all__ = [
+    "manager",
+    "TaskConfig",
+    "mixins",
+    "policies",
+    "exceptions",
+    "helpers",
+    "utils",
+]

ergon/task/base.py

@@ -0,0 +1,222 @@
+import socket
+from abc import ABC, ABCMeta, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+from .. import telemetry
+from ..connector import Connector, ConnectorConfig
+from ..service import ServiceConfig
+
+
+class TaskConfig(BaseModel):
+    """
+    Declarative configuration for a task.
+    Specifies:
+      - which Task class to run
+      - input/output connectors
+      - logger configuration
+      - worker scaling
+    """
+
+    name: str
+    task: object
+    max_workers: int = 1
+    connectors: Dict[str, ConnectorConfig] = Field(default_factory=dict)
+    services: Dict[str, ServiceConfig] = Field(default_factory=dict)
+    policies: List[Any] = Field(default_factory=list)
+    logging: Optional[telemetry.logging.LoggingConfig] = None
+    metrics: Optional[telemetry.metrics.MetricsConfig] = None
+    tracing: Optional[telemetry.tracing.TracingConfig] = None
+
+    @model_validator(mode="after")
+    def validate(self) -> "TaskConfig":
+        if not self.connectors and not self.services:
+            raise ValueError("Task must define at least one connector or service.")
+
+        if self.max_workers < 1:
+            raise ValueError("max_workers must be at least 1.")
+
+        return self
+
+
+class TaskMeta(ABCMeta):
+    """
+    Metaclass that injects connectors, services and execution policies into task instances.
+    Consumed by BaseTask and BaseAsyncTask.
+    """
+
+    def __call__(
+        cls,
+        connectors: Dict[str, Connector],
+        services: Dict[str, Any],
+        policies: List[Any],
+        worker_id: Optional[int] = None,
+        task_config: Optional[TaskConfig] = None,
+        *args,
+        **kwargs,
+    ):
+        # ---------------------------------------------------------
+        # Create instance normally
+        # ---------------------------------------------------------
+        self = super().__call__(
+            connectors=connectors,
+            services=services,
+            policies=policies,
+            worker_id=worker_id,
+            task_config=task_config,
+            *args,
+            **kwargs,
+        )
+
+        # ---------------------------------------------------------
+        # Store DI containers
+        # ---------------------------------------------------------
+        self.connectors = connectors
+        self.services = services
+        self.policies = policies
+        self.worker_id = worker_id
+        self.task_config = task_config
+
+        # ---------------------------------------------------------
+        # Expose connectors as attributes
+        # Example: pipefy_connector, rabbitmq_connector
+        # ---------------------------------------------------------
+        for name, conn in connectors.items():
+            setattr(self, f"{name}_connector", conn)
+
+        # ---------------------------------------------------------
+        # Expose services as attributes
+        # Example: openai_service, s3_service
+        # ---------------------------------------------------------
+        for name, service in services.items():
+            setattr(self, f"{name}_service", service)
+
+        # ---------------------------------------------------------
+        # Expose policies as attributes
+        # Example: policy_name
+        # ---------------------------------------------------------
+        for policy in policies:
+            name = policy.name or policy.__class__.__name__
+            setattr(self, f"{name}_policy", policy)
+
+        # ---------------------------------------------------------
+        # Task identity
+        # ---------------------------------------------------------
+        self.name = getattr(self, "name", cls.__name__)
+
+        return self
+
+
+# ---------- BASE TASKS ----------
+
+
+class BaseTask(ABC, metaclass=TaskMeta):
+    """
+    Base class for all tasks in the Ergon Task Framework.
+
+    Behavior (produce/consume) is added via mixins:
+    - ProduceMixin
+    - ConsumeMixin
+    """
+
+    name: str = "base"
+    connectors: Dict[str, Connector]
+    services: Dict[str, Any]
+    policies: List[Any]
+    worker_id: Optional[int]
+    task_config: Optional["TaskConfig"]
+
+    def __init__(
+        self,
+        connectors: Dict[str, Connector],
+        services: Dict[str, Any],
+        policies: List[Any],
+        worker_id: Optional[int] = None,
+        task_config: Optional["TaskConfig"] = None,
+        *args,
+        **kwargs,
+    ):
+        # Don't pass args/kwargs to object.__init__()
+        super().__init__()
+
+    @abstractmethod
+    def execute(self) -> Any:
+        """Main entry point for running the task."""
+        raise NotImplementedError
+
+    def exit(self):
+        """
+        Gracefully exit the task.
+
+        Tasks may override this to:
+        - close resources
+        - flush buffers
+        - notify services
+
+        IMPORTANT: By default, this does NOT kill the interpreter.
+        """
+        return
+
+
+class BaseAsyncTask(ABC, metaclass=TaskMeta):
+    """
+    Asynchronous base class for all async tasks in the Ergon Task Framework.
+
+    Behavior (async produce/consume) will be added via async mixins.
+    """
+
+    name: str = "base_async"
+    connectors: Dict[str, Any]
+    services: Dict[str, Any]
+    policies: List[Any]
+    worker_id: Optional[int]
+    task_config: Optional["TaskConfig"]
+
+    def __init__(
+        self,
+        connectors: Dict[str, Any],
+        services: Dict[str, Any],
+        policies: List[Any],
+        worker_id: Optional[int] = None,
+        task_config: Optional["TaskConfig"] = None,
+        *args,
+        **kwargs,
+    ):
+        # Don't pass args/kwargs to object.__init__()
+        super().__init__()
+
+    @abstractmethod
+    async def execute(self) -> Any:
+        """Main async entry point for running the task."""
+        raise NotImplementedError
+
+    async def exit(self):
+        """Optional async cleanup hook."""
+        return
+
+
+class TaskExecMetadata(BaseModel):
+    task_name: str
+    execution_id: str
+    execution_start_time: str
+    pid: int
+    worker_id: Optional[int]
+    host_name: str = Field(default_factory=lambda: TaskExecMetadata._get_host_name())
+    host_ip: str = Field(default_factory=lambda: TaskExecMetadata._get_host_ip())
+
+    # ---------- STATIC HELPERS ----------
+    @staticmethod
+    def _get_host_ip() -> str:
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        try:
+            s.connect(("8.8.8.8", 80))
+            return s.getsockname()[0]
+        except Exception:
+            return "127.0.0.1"
+        finally:
+            s.close()
+
+    @staticmethod
+    def _get_host_name() -> str:
+        return socket.gethostname()