api-shared 0.0.1__tar.gz

api_shared-0.0.1/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.3
+Name: api-shared
+Version: 0.0.1
+Summary: Shared dependencies for both api and the worker services.
+Requires-Dist: pyyaml>=6.0.1
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: logfire[redis,httpx,system-metrics]>=3.12.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: opentelemetry-distro[otlp]>=0.52b0
+Requires-Dist: opentelemetry-instrumentation-logging>=0.52b0
+Requires-Dist: opentelemetry-instrumentation-redis>=0.52b0
+Requires-Dist: pydantic-settings>=2.8.1
+Requires-Dist: taskiq-aio-pika>=0.4.1
+Requires-Dist: taskiq-redis>=1.0.4
+Requires-Dist: taskiq[opentelemetry]>=0.12.1
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+

api_shared-0.0.1/pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "api-shared"
+version = "0.0.1"
+description = "Shared dependencies for both api and the worker services."
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "pyyaml>=6.0.1",
+    "httpx>=0.28.1",
+    "logfire[redis,httpx,system-metrics]>=3.12.0",
+    "loguru>=0.7.3",
+    "opentelemetry-distro[otlp]>=0.52b0",
+    "opentelemetry-instrumentation-logging>=0.52b0",
+    "opentelemetry-instrumentation-redis>=0.52b0",
+    "pydantic-settings>=2.8.1",
+    "taskiq-aio-pika>=0.4.1",
+    "taskiq-redis>=1.0.4",
+    "taskiq[opentelemetry]>=0.12.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.9,<0.14"]
+build-backend = "uv_build"

api_shared-0.0.1/src/api_shared/broker.py

@@ -0,0 +1,225 @@
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field, ValidationError
+from taskiq import (
+    AsyncBroker,
+    AsyncResultBackend,
+    InMemoryBroker,
+    SmartRetryMiddleware,
+)
+from taskiq.instrumentation import TaskiqInstrumentor
+from taskiq.schedule_sources import LabelScheduleSource
+from taskiq.scheduler.scheduler import TaskiqScheduler
+from taskiq_aio_pika import AioPikaBroker
+from taskiq_redis import RedisAsyncResultBackend
+
+from api_shared.core.settings import Environment, OLTPLogMethod, settings
+
+if settings.TASKIQ_DASHBOARD_URL:
+    from api_shared.middlewares.dashboard import DashboardMiddleware
+
+if settings.OLTP_LOG_METHOD != OLTPLogMethod.NONE:
+    TaskiqInstrumentor().instrument()
+
+
+class BrokerConfigSchema(BaseModel):
+    """Pydantic model for broker configuration validation.
+
+    Attributes:
+        queue: Queue name for this broker's tasks.
+        routing_key: Routing key pattern for message routing.
+        exchange: Exchange name for this broker.
+        description: Optional description of the broker's purpose.
+    """
+
+    queue: str = Field(..., description="Queue name for this broker's tasks")
+    routing_key: str = Field(
+        default="#", description="Routing key pattern for message routing"
+    )
+    exchange: str = Field(..., description="Exchange name for this broker")
+    description: str | None = Field(
+        default=None, description="Description of the broker's purpose"
+    )
+
+
+class BrokersConfigSchema(BaseModel):
+    """Pydantic model for the complete brokers YAML file.
+
+    Attributes:
+        brokers: Dictionary mapping broker names to their configurations.
+    """
+
+    brokers: dict[str, BrokerConfigSchema] = Field(
+        ..., description="Broker configurations"
+    )
+
+
+class BrokerManager:
+    """Manages taskiq broker instances and provides access to them."""
+
+    def __init__(self):
+        """Initialize the broker manager and create all configured brokers."""
+        self._brokers: dict[str, AsyncBroker] = {}
+        self._scheduler: TaskiqScheduler | None = None
+        self._broker_configs: dict[str, BrokerConfigSchema] = {}
+        self._load_broker_configs()
+        self._initialize_brokers()
+
+    def _load_broker_configs(self) -> None:
+        """Load and validate broker configurations from YAML file.
+
+        Raises:
+            FileNotFoundError: If the broker config file doesn't exist.
+            ValidationError: If the YAML structure is invalid.
+        """
+        if settings.TASKIQ_BROKERS_CONFIG_FILE:
+            config_path = Path(settings.TASKIQ_BROKERS_CONFIG_FILE)
+        else:
+            config_path = (
+                Path(__file__).parent.parent.parent / "configs" / "brokers.yml"
+            )
+
+        if not config_path.exists():
+            raise FileNotFoundError(
+                f"Broker configuration file not found: {config_path}. "
+                f"Create the file or set TASKIQ_BROKERS_CONFIG_FILE environment variable."
+            )
+
+        with config_path.open("r") as f:
+            raw_config = yaml.safe_load(f)
+
+        try:
+            validated_config = BrokersConfigSchema.model_validate(raw_config)
+        except ValidationError as e:
+            raise ValueError(
+                f"Invalid broker configuration in {config_path}:\n{e}"
+            ) from e
+
+        self._broker_configs = validated_config.brokers
+
+    def _create_broker(
+        self, broker_name: str, broker_config: BrokerConfigSchema
+    ) -> AsyncBroker:
+        """Create a configured taskiq broker instance.
+
+        Args:
+            broker_name: Name identifier for the broker.
+            broker_config: Broker configuration object.
+
+        Returns:
+            Configured AsyncBroker instance.
+        """
+        result_backend: AsyncResultBackend[Any] = RedisAsyncResultBackend(
+            redis_url=str(settings.REDIS_URL.with_path(f"/{settings.REDIS_TASK_DB}")),
+        )
+
+        middlewares = [
+            SmartRetryMiddleware(
+                default_retry_count=5,
+                default_delay=10,
+                use_jitter=True,
+                use_delay_exponent=True,
+                max_delay_exponent=120,
+            ),
+        ]
+
+        if settings.TASKIQ_DASHBOARD_URL:
+            middlewares.append(
+                DashboardMiddleware(
+                    url=settings.TASKIQ_DASHBOARD_URL,
+                    api_token=settings.TASKIQ_DASHBOARD_API_TOKEN,
+                    broker_name=broker_name,
+                )
+            )
+
+        return (
+            AioPikaBroker(
+                str(settings.RABBITMQ_URL),
+                queue_name=broker_config.queue,
+                routing_key=broker_config.routing_key,
+                exchange_name=broker_config.exchange,
+            )
+            .with_result_backend(result_backend)
+            .with_middlewares(*middlewares)
+        )
+
+    def _initialize_brokers(self) -> None:
+        """Initialize all configured brokers."""
+        # NOTE: For testing, create in-memory brokers for all configured brokers
+        if settings.ENVIRONMENT == Environment.TEST:
+            self._brokers = {
+                broker_name: InMemoryBroker() for broker_name in self._broker_configs
+            }
+        else:
+            for broker_name, broker_config in self._broker_configs.items():
+                self._brokers[broker_name] = self._create_broker(
+                    broker_name, broker_config
+                )
+
+        if settings.ENVIRONMENT != Environment.TEST:
+            workers_broker = self._brokers.get("general")
+            if workers_broker:
+                self._scheduler = TaskiqScheduler(
+                    broker=workers_broker,
+                    sources=[LabelScheduleSource(workers_broker)],
+                )
+
+    def get_broker(self, name: str) -> AsyncBroker:
+        """Get a broker by name.
+
+        Args:
+            name: Name of the broker to retrieve.
+
+        Returns:
+            AsyncBroker instance.
+
+        Raises:
+            RuntimeError: If the broker is not enabled.
+        """
+        broker = self._brokers.get(name)
+        if broker is None:
+            raise RuntimeError(
+                f"Broker '{name}' is not enabled. Enable it in the broker configuration file."
+            )
+        return broker
+
+    def get_all_brokers(self) -> dict[str, AsyncBroker]:
+        """Get all configured brokers.
+
+        Returns:
+            Dictionary mapping broker names to broker instances.
+        """
+        return self._brokers.copy()
+
+    async def startup_all(self) -> None:
+        """Start up all configured brokers.
+
+        Only starts brokers that are not in worker process mode.
+        """
+        for broker_instance in self._brokers.values():
+            if not broker_instance.is_worker_process:
+                await broker_instance.startup()
+
+    async def shutdown_all(self) -> None:
+        """Shut down all configured brokers.
+
+        Only shuts down brokers that are not in worker process mode.
+        """
+        for broker_instance in self._brokers.values():
+            if not broker_instance.is_worker_process:
+                await broker_instance.shutdown()
+
+    @property
+    def scheduler(self) -> TaskiqScheduler | None:
+        """Get the scheduler instance if available.
+
+        Returns:
+            TaskiqScheduler instance or None if not configured.
+        """
+        return self._scheduler
+
+
+# Create singleton instance
+broker_manager = BrokerManager()

api_shared-0.0.1/src/api_shared/core/settings.py

@@ -0,0 +1,209 @@
+import ipaddress
+from enum import StrEnum
+from typing import Annotated
+
+from pydantic import (
+    AfterValidator,
+    AnyHttpUrl,
+    Field,
+    PlainValidator,
+    TypeAdapter,
+)
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from yarl import URL
+
+AnyHttpUrlAdapter = TypeAdapter(AnyHttpUrl)
+
+CustomHttpUrlStr = Annotated[
+    str,
+    PlainValidator(AnyHttpUrlAdapter.validate_strings),
+    AfterValidator(lambda x: str(x).rstrip("/")),
+]
+
+
+class Environment(StrEnum):
+    """Possible environments."""
+
+    DEV = "dev"
+    TEST = "test"
+    PROD = "prod"
+
+
+class OLTPLogMethod(StrEnum):
+    NONE = "none"
+    MANUAL = "manual"
+    LOGFIRE = "logfire"
+
+
+class RunMode(StrEnum):
+    NONE = "none"
+    API = "api"
+    WORKER = "worker"
+
+
+def _should_use_http_scheme(host: str) -> bool:
+    """Check if the host should use HTTP scheme instead of HTTPS.
+
+    Uses HTTP for:
+    - IP addresses (IPv4 or IPv6)
+    - Simple hostnames without dots (like docker hostnames: redis, postgres, etc.)
+
+    Uses HTTPS for:
+    - Domain names with dots (like redis.example.com)
+
+    Args:
+        host: The host string to check.
+
+    Returns:
+        True if should use HTTP scheme, False if should use HTTPS.
+    """
+    try:
+        ipaddress.ip_address(host)
+        return True
+    except ValueError:
+        pass
+
+    return "." not in host
+
+
+class SharedBaseSettings(BaseSettings):
+    """Base settings class with common configuration shared across all services."""
+
+    ENVIRONMENT: Environment = Field(
+        default=Environment.DEV,
+        description="Application environment (dev, test, prod).",
+    )
+    OLTP_LOG_METHOD: OLTPLogMethod = Field(
+        default=OLTPLogMethod.NONE,
+        description="OpenTelemetry logging method (none, manual, logfire).",
+    )
+    OTLP_ENDPOINT: CustomHttpUrlStr | None = Field(
+        default=None,
+        description="OpenTelemetry GRPC endpoint for OTLP exporter.",
+    )
+    OLTP_STD_LOGGING_ENABLED: bool = Field(
+        default=False,
+        description="Enable standard logging integration with OpenTelemetry.",
+    )
+    RABBITMQ_HOST: str = Field(
+        default="localhost",
+        description="RabbitMQ server hostname or IP address.",
+    )
+    RABBITMQ_PORT: int = Field(
+        default=5672,
+        description="RabbitMQ server port.",
+    )
+    RABBITMQ_USERNAME: str = Field(
+        description="RabbitMQ authentication username.",
+    )
+    RABBITMQ_PASSWORD: str = Field(
+        description="RabbitMQ authentication password.",
+    )
+    RABBITMQ_VHOST: str = Field(
+        default="/",
+        description="RabbitMQ virtual host.",
+    )
+    REDIS_PORT: int = Field(
+        default=6379,
+        description="Redis server port.",
+    )
+    REDIS_HOST: str = Field(
+        default="localhost",
+        description="Redis server hostname or IP address.",
+    )
+    REDIS_USER: str | None = Field(
+        default=None,
+        description="Redis authentication username.",
+    )
+    REDIS_PASS: str = Field(
+        description="Redis authentication password.",
+    )
+    REDIS_BASE: str | None = Field(
+        default=None,
+        description="Redis database base path.",
+    )
+    REDIS_TASK_DB: int = Field(
+        default=1,
+        ge=1,
+        le=16,
+        description="Redis database number for taskiq result backend. Must be between 1-16.",
+    )
+    RUN_MODE: RunMode = Field(
+        default=RunMode.NONE,
+        description="Application run mode api or worker).",
+    )
+    TASKIQ_DASHBOARD_HOST: str | None = Field(
+        default=None,
+        description="Taskiq dashboard server hostname or IP address.",
+    )
+    TASKIQ_DASHBOARD_PORT: int = Field(
+        default=8001,
+        description="Taskiq dashboard server port.",
+    )
+    TASKIQ_DASHBOARD_API_TOKEN: str = Field(
+        default="supersecret",
+        description="API token for Taskiq dashboard authentication.",
+    )
+    TASKIQ_BROKERS_CONFIG_FILE: str | None = Field(
+        default=None,
+        description="Path to YAML file containing broker configurations.",
+    )
+
+    @property
+    def TASKIQ_DASHBOARD_URL(self) -> str | None:
+        """Assemble Taskiq Dashboard URL from settings.
+
+        Returns:
+            Taskiq Dashboard URL.
+        """
+        if self.TASKIQ_DASHBOARD_HOST is None:
+            return None
+
+        return f"http://{self.TASKIQ_DASHBOARD_HOST}:{self.TASKIQ_DASHBOARD_PORT}"
+
+    @property
+    def REDIS_URL(self) -> URL:
+        """Assemble REDIS URL from settings.
+
+        Returns:
+            Redis URL.
+        """
+        path = f"/{self.REDIS_BASE}" if self.REDIS_BASE is not None else ""
+        scheme = "redis" if _should_use_http_scheme(self.REDIS_HOST) else "rediss"
+
+        return URL.build(
+            scheme=scheme,
+            host=self.REDIS_HOST,
+            port=self.REDIS_PORT,
+            user=self.REDIS_USER,
+            password=self.REDIS_PASS,
+            path=path,
+        )
+
+    @property
+    def RABBITMQ_URL(self) -> URL:
+        """Assemble RabbitMQ URL from settings.
+
+        Returns:
+            RabbitMQ URL.
+        """
+        scheme = "amqp" if _should_use_http_scheme(self.RABBITMQ_HOST) else "amqps"
+
+        return URL.build(
+            scheme=scheme,
+            host=self.RABBITMQ_HOST,
+            port=self.RABBITMQ_PORT,
+            user=self.RABBITMQ_USERNAME,
+            password=self.RABBITMQ_PASSWORD,
+            path=self.RABBITMQ_VHOST,
+        )
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        # env_prefix="API_TEMPLATE_SHARED_",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
+settings = SharedBaseSettings()  # type: ignore[call-arg]

api_shared-0.0.1/src/api_shared/middlewares/dashboard.py

@@ -0,0 +1,169 @@
+# NOTE: From https://github.com/danfimov/taskiq-dashboard/blob/main/taskiq_dashboard/interface/middleware.py
+import asyncio
+from datetime import datetime, timezone
+from logging import getLogger
+from typing import Any
+from urllib.parse import urljoin
+
+import httpx
+from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.compat import model_dump
+from taskiq.message import TaskiqMessage
+from taskiq.result import TaskiqResult
+
+logger = getLogger("taskiq_dashboard.admin_middleware")
+
+
+class DashboardMiddleware(TaskiqMiddleware):
+    """A Taskiq middleware that reports task lifecycle events to an external admin dashboard API.
+
+    This middleware sends HTTP POST requests to a configured endpoint when tasks
+    are queued, started, or completed. It can be used for task monitoring, auditing,
+    or visualization in external systems.
+
+    Attributes:
+        url (str): Base URL of the admin API.
+        api_token (str): Token used for authenticating with the API.
+        timeout (float): Timeout (in seconds) for API requests.
+        broker_name (str): Name of the broker instance to include in the payload. Defaults to 'default_broker'.
+        _pending (set[asyncio.Task]): Set of currently running background request tasks.
+        _client (httpx.AsyncClient | None): HTTP client session used for sending requests.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        api_token: str,
+        timeout: float = 5.0,
+        broker_name: str = "default_broker",
+    ) -> None:
+        super().__init__()
+        self.url = url
+        self.timeout = timeout
+        self.api_token = api_token
+        self.broker_name = broker_name
+        self._pending: set[asyncio.Task[Any]] = set()
+        self._client: httpx.AsyncClient | None = None
+
+    @staticmethod
+    def _now_iso() -> str:
+        return datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
+
+    def _get_client(self) -> httpx.AsyncClient:
+        """Create and cache session."""
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=self.timeout)
+        return self._client
+
+    async def startup(self) -> None:
+        """Startup method to initialize httpx.AsyncClient."""
+        self._client = self._get_client()
+
+    async def shutdown(self) -> None:
+        """Shutdown method to run all pending requests and close the session."""
+        if self._pending:
+            await asyncio.gather(*self._pending, return_exceptions=True)
+        if self._client is not None:
+            await self._client.aclose()
+
+    async def _spawn_request(
+        self,
+        endpoint: str,
+        payload: dict[str, Any],
+    ) -> None:
+        """Fire and forget helper.
+
+        Start an async POST to the admin API, keep the resulting Task in _pending
+        so it can be awaited/cleaned during graceful shutdown.
+        """
+
+        async def _send() -> None:
+            client = self._get_client()
+            try:
+                resp = await client.post(
+                    urljoin(self.url, endpoint),
+                    headers={"access-token": self.api_token},
+                    json=payload,
+                )
+                resp.raise_for_status()
+                if not resp.is_success:
+                    logger.error("POST %s - %s", endpoint, resp.status_code)
+            except httpx.HTTPStatusError:
+                logger.exception("POST %s failed with HTTP error", endpoint)
+            except httpx.RequestError:
+                logger.exception("POST %s failed with request error", endpoint)
+
+        task = asyncio.create_task(_send())
+        self._pending.add(task)
+        task.add_done_callback(self._pending.discard)
+
+    async def post_send(self, message: TaskiqMessage) -> None:
+        """
+        This hook is executed right after the task is sent.
+
+        This is a client-side hook. It executes right
+        after the message is kicked in broker.
+
+        :param message: kicked message.
+        """
+        dict_message: dict[str, Any] = model_dump(message)
+        await self._spawn_request(
+            f"api/tasks/{message.task_id}/queued",
+            {
+                "args": dict_message["args"],
+                "kwargs": dict_message["kwargs"],
+                "labels": dict_message["labels"],
+                "queuedAt": self._now_iso(),
+                "taskName": message.task_name,
+                "worker": self.broker_name,
+            },
+        )
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+        """
+        This hook is called before executing task.
+
+        This is a worker-side hook, which means it
+        executes in the worker process.
+
+        :param message: incoming parsed taskiq message.
+        :return: modified message.
+        """
+        dict_message: dict[str, Any] = model_dump(message)
+        await self._spawn_request(
+            f"api/tasks/{message.task_id}/started",
+            {
+                "args": dict_message["args"],
+                "kwargs": dict_message["kwargs"],
+                "labels": dict_message["labels"],
+                "startedAt": self._now_iso(),
+                "taskName": message.task_name,
+                "worker": self.broker_name,
+            },
+        )
+        return message
+
+    async def post_execute(
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult[Any],
+    ) -> None:
+        """
+        This hook executes after task is complete.
+
+        This is a worker-side hook. It's called
+        in worker process.
+
+        :param message: incoming message.
+        :param result: result of execution for current task.
+        """
+        dict_result: dict[str, Any] = model_dump(result)
+        await self._spawn_request(
+            f"api/tasks/{message.task_id}/executed",
+            {
+                "finishedAt": self._now_iso(),
+                "executionTime": result.execution_time,
+                "error": None if result.error is None else repr(result.error),
+                "returnValue": {"return_value": dict_result["return_value"]},
+            },
+        )

api_shared-0.0.1/src/api_shared/tasks/general/__init__.py

@@ -0,0 +1,31 @@
+"""Worker tasks that use the 'general' broker."""
+
+from api_shared.broker import broker_manager
+
+# This will raise RuntimeError if general broker is not enabled
+workers_broker = broker_manager.get_broker("general")
+
+from api_shared.tasks.general.complex_task import (
+    LongRunningProcessResult,
+    long_running_process,
+)
+from api_shared.tasks.general.dummy import add_one, add_one_with_retry
+from api_shared.tasks.general.failing_task import failing_process
+from api_shared.tasks.general.pydantic_parse_task import (
+    NestedModel,
+    PydanticParseInput,
+    PydanticParseResult,
+    pydantic_parse_check,
+)
+
+__all__ = [
+    "LongRunningProcessResult",
+    "NestedModel",
+    "PydanticParseInput",
+    "PydanticParseResult",
+    "add_one",
+    "add_one_with_retry",
+    "failing_process",
+    "long_running_process",
+    "pydantic_parse_check",
+]

api_shared-0.0.1/src/api_shared/tasks/general/complex_task.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+
+from api_shared.broker import broker_manager
+
+broker = broker_manager.get_broker("general")
+
+
+class LongRunningProcessResult(BaseModel):
+    start_time: float
+    end_time: float
+    elapsed: float
+    status: str
+
+
+@broker.task(task_name="long_running_process")
+async def long_running_process(duration: int = 5) -> LongRunningProcessResult:
+    """
+    Simulates a long-running process by sleeping.
+    """
+    raise NotImplementedError("This task is implemented in the worker package")

api_shared-0.0.1/src/api_shared/tasks/general/dummy.py

@@ -0,0 +1,46 @@
+import random
+from datetime import datetime
+
+from loguru import logger
+
+from api_shared.broker import broker_manager
+
+broker = broker_manager.get_broker("general")
+
+
+@broker.task
+async def add_one(value: int) -> int:
+    return value + 1
+
+
+@broker.task(retry_on_error=True, max_retries=5, delay=15)
+async def add_one_with_retry(value: int) -> int:
+    # Randomly fail 50% of the time
+    if random.random() < 0.5:  # noqa: PLR2004
+        raise RuntimeError("Random failure in add_one_with_retry")
+
+    return value + 1
+
+
+@broker.task(
+    schedule=[
+        {
+            "cron": "*/2 * * * *",  # type: str, either cron or time should be specified. Runs every 2 minutes.
+            "cron_offset": None,  # type: str | timedelta | None, can be omitted.
+            "time": None,  # type: datetime | None, either cron or time should be specified.
+            "args": [1],  # type List[Any] | None, can be omitted.
+            "kwargs": {},  # type: Dict[str, Any] | None, can be omitted.
+            "labels": {},  # type: Dict[str, Any] | None, can be omitted.
+        }
+    ]
+)
+async def add_one_scheduled(value: int) -> int:
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    logger.info(f"Current time: {current_time}")
+
+    return value + 1
+
+
+@broker.task
+async def parse_int(val: str) -> int:
+    return int(val)

api_shared-0.0.1/src/api_shared/tasks/general/failing_task.py

@@ -0,0 +1,13 @@
+from api_shared.broker import broker_manager
+
+broker = broker_manager.get_broker("general")
+
+
+@broker.task(task_name="failing_process")
+async def failing_process(
+    error_message: str = "This is a deliberate error",
+) -> None:
+    """
+    A task that intentionally fails to demonstrate error handling.
+    """
+    raise NotImplementedError("This task is implemented in the worker package")

api_shared-0.0.1/src/api_shared/tasks/general/pydantic_parse_task.py

@@ -0,0 +1,33 @@
+from pydantic import BaseModel
+
+from api_shared.broker import broker_manager
+
+broker = broker_manager.get_broker("general")
+
+
+class NestedModel(BaseModel):
+    name: str
+    value: int
+    tags: list[str]
+
+
+class PydanticParseInput(BaseModel):
+    text: str
+    count: int
+    nested: NestedModel
+
+
+class PydanticParseResult(BaseModel):
+    received_text: str
+    received_count: int
+    received_nested: NestedModel
+    doubled_count: int
+    status: str
+
+
+@broker.task(task_name="pydantic_parse_check")
+async def pydantic_parse_check(data: PydanticParseInput) -> PydanticParseResult:
+    """
+    Tests taskiq's ability to parse and serialize Pydantic BaseModels.
+    """
+    raise NotImplementedError("This task is implemented in the worker package")

api_shared-0.0.1/src/api_shared/tasks/general/redis.py

@@ -0,0 +1,21 @@
+# from typing import Annotated
+
+# from app.api.redis.deps import get_redis_pool
+# from loguru import logger
+# from redis.asyncio import ConnectionPool, Redis
+# from taskiq import TaskiqDepends
+
+# from api_shared.broker import broker_manager
+
+# broker = broker_manager.get_broker("general")
+
+
+# @broker.task
+# async def my_redis_task(
+#     key: str,
+#     val: str,
+#     pool: Annotated[ConnectionPool, TaskiqDepends(get_redis_pool)],
+# ):
+#     async with Redis(connection_pool=pool) as redis:
+#         await redis.set(name=key, value=val)
+#         logger.debug(f"Set key {key} with value {val} in Redis using connection pool")

api_shared-0.0.1/src/api_shared/tasks/ml/__init__.py

@@ -0,0 +1,20 @@
+"""ML tasks that use the 'ml' broker."""
+
+from api_shared.broker import broker_manager
+
+# This will raise RuntimeError if ML broker is not enabled
+ml_broker = broker_manager.get_broker("ml")
+
+from api_shared.tasks.ml.ml_tasks import (
+    MLInferenceResult,
+    MLTrainingResult,
+    ml_inference_task,
+    train_model_task,
+)
+
+__all__ = [
+    "MLInferenceResult",
+    "MLTrainingResult",
+    "ml_inference_task",
+    "train_model_task",
+]

api_shared-0.0.1/src/api_shared/tasks/ml/ml_tasks.py

@@ -0,0 +1,60 @@
+from pydantic import BaseModel
+
+from api_shared.broker import broker_manager
+
+ml_broker = broker_manager.get_broker("ml")
+
+
+class MLInferenceResult(BaseModel):
+    model_id: str
+    predictions: list[float]
+    confidence: float
+    status: str
+
+
+@ml_broker.task(task_name="ml_inference")
+async def ml_inference_task(model_id: str, input_data: dict) -> MLInferenceResult:
+    """
+    Perform ML inference on input data.
+
+    This task is implemented in the ML worker package.
+    It will be processed by ML workers listening to the 'taskiq_ml' queue.
+
+    Args:
+        model_id: Identifier for the ML model to use.
+        input_data: Input data for inference.
+
+    Returns:
+        Inference results with predictions and confidence scores.
+    """
+    raise NotImplementedError("This task is implemented in the ML worker package")
+
+
+class MLTrainingResult(BaseModel):
+    dataset_id: str
+    model_id: str
+    training_metrics: dict[str, float | int]
+    status: str
+
+
+@ml_broker.task(task_name="train_model")
+async def train_model_task(
+    dataset_id: str,
+    model_config: dict,
+    hyperparameters: dict,
+) -> MLTrainingResult:
+    """
+    Train an ML model with given configuration.
+
+    This task is implemented in the ML worker package.
+    It will be processed by ML workers listening to the 'taskiq_ml' queue.
+
+    Args:
+        dataset_id: Identifier for the training dataset.
+        model_config: Model architecture configuration.
+        hyperparameters: Training hyperparameters.
+
+    Returns:
+        Training results and model metadata.
+    """
+    raise NotImplementedError("This task is implemented in the ML worker package")