nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/standalone/auth.py

@@ -81,9 +81,7 @@ class AuthHeaderAuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(
-        self, request: HTTPConnection
-    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -94,9 +92,7 @@ class AuthHeaderAuthenticationBackend(NucliaCloudAuthenticationBackend):
         user = request.headers[self.settings.auth_policy_user_header]
         nuclia_user: BaseUser = NucliaUser(username=user)
 
-        auth_creds = AuthCredentials(
-            get_mapped_roles(settings=self.settings, data={"user": user})
-        )
+        auth_creds = AuthCredentials(get_mapped_roles(settings=self.settings, data={"user": user}))
 
         return auth_creds, nuclia_user
 
@@ -113,9 +109,7 @@ class OAuth2AuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(
-        self, request: HTTPConnection
-    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -133,9 +127,7 @@ class OAuth2AuthenticationBackend(NucliaCloudAuthenticationBackend):
         try:
             token_data = orjson.loads(base64.b64decode(token_split[1] + "==="))
         except Exception:
-            logger.warning(
-                f"Could not parse jwt bearer token value: {token}", exc_info=True
-            )
+            logger.warning(f"Could not parse jwt bearer token value: {token}", exc_info=True)
             return None
 
         if "sub" not in token_data:
@@ -168,9 +160,7 @@ class BasicAuthAuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(
-        self, request: HTTPConnection
-    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -186,9 +176,7 @@ class BasicAuthAuthenticationBackend(NucliaCloudAuthenticationBackend):
         user = token.split(":")[0]
 
         nuclia_user: BaseUser = NucliaUser(username=user)
-        auth_creds = AuthCredentials(
-            get_mapped_roles(settings=self.settings, data={"user": user})
-        )
+        auth_creds = AuthCredentials(get_mapped_roles(settings=self.settings, data={"user": user}))
 
         return auth_creds, nuclia_user
 
@@ -201,9 +189,7 @@ class UpstreamNaiveAuthenticationBackend(NucliaCloudAuthenticationBackend):
             user_header=settings.auth_policy_user_header,
         )
 
-    async def authenticate(
-        self, request: HTTPConnection
-    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp

nucliadb/standalone/config.py

@@ -37,19 +37,16 @@ def config_standalone_driver(nucliadb_args: Settings):
     )
 
     # update global settings with arg values
-    for fieldname in DriverSettings.__fields__.keys():
+    for fieldname in DriverSettings.model_fields.keys():
         setattr(ingest_settings, fieldname, getattr(nucliadb_args, fieldname))
-    for fieldname in StorageSettings.__fields__.keys():
+    for fieldname in StorageSettings.model_fields.keys():
         setattr(storage_settings, fieldname, getattr(nucliadb_args, fieldname))
 
     if ingest_settings.driver == DriverConfig.NOT_SET:
         # no driver specified, for standalone, we force defaulting to local here
-        ingest_settings.driver = DriverConfig.LOCAL
+        ingest_settings.driver = DriverConfig.PG
 
-    if (
-        ingest_settings.driver == DriverConfig.LOCAL
-        and ingest_settings.driver_local_url is None
-    ):
+    if ingest_settings.driver == DriverConfig.LOCAL and ingest_settings.driver_local_url is None:
         # also provide default path for local driver when none provided
         ingest_settings.driver_local_url = "./data/main"
 
@@ -57,11 +54,11 @@ def config_standalone_driver(nucliadb_args: Settings):
         # no driver specified, for standalone, we try to automate some settings here
         storage_settings.file_backend = FileBackendConfig.LOCAL
 
-    if (
-        storage_settings.file_backend == FileBackendConfig.LOCAL
-        and storage_settings.local_files is None
-    ):
-        storage_settings.local_files = "./data/blob"
+    if storage_settings.file_backend == FileBackendConfig.LOCAL:
+        if storage_settings.local_files is None:
+            storage_settings.local_files = "./data/blob"
+        if storage_settings.local_indexing_bucket is None:
+            storage_settings.local_indexing_bucket = "indexer"
 
     if ingest_settings.driver_local_url is not None and not os.path.isdir(
         ingest_settings.driver_local_url

nucliadb/standalone/introspect.py

@@ -34,7 +34,7 @@ from pydantic import BaseModel
 
 from nucliadb.common.cluster import manager as cluster_manager
 from nucliadb.standalone.settings import Settings
-from nucliadb_telemetry.settings import LogSettings
+from nucliadb_telemetry.settings import LogOutputType, LogSettings
 
 MB = 1024 * 1024
 CHUNK_SIZE = 2 * MB
@@ -70,7 +70,7 @@ class NodeInfo(BaseModel):
     id: str
     address: str
     shard_count: int
-    primary_id: Optional[str]
+    primary_id: Optional[str] = None
 
 
 class ClusterInfo(BaseModel):
@@ -86,7 +86,7 @@ async def stream_tar(app: FastAPI) -> AsyncGenerator[bytes, None]:
             await add_cluster_info(temp_dir, tar)
             settings: Settings = app.settings.copy()  # type: ignore
             await add_settings(temp_dir, tar, settings)
-            if settings.log_output_type == "file":
+            if settings.log_output_type == LogOutputType.FILE:
                 await add_logs(tar)
 
         async for chunk in stream_out_tar(tar_file):
@@ -164,7 +164,7 @@ def _add_cluster_info_to_tar(temp_dir: str, tar: tarfile.TarFile):
     )
     cluster_info_file = os.path.join(temp_dir, "cluster_info.txt")
     with open(cluster_info_file, "w") as f:
-        f.write(cluster_info.json(indent=4))
+        f.write(cluster_info.model_dump_json(indent=4))
     tar.add(cluster_info_file, arcname="cluster_info.txt")
 
 
@@ -177,7 +177,7 @@ def _add_settings_to_tar(temp_dir: str, tar: tarfile.TarFile, settings: Settings
     remove_sensitive_settings(settings)
     settings_file = os.path.join(temp_dir, "settings.json")
     with open(settings_file, "w") as f:
-        f.write(settings.json(indent=4))
+        f.write(settings.model_dump_json(indent=4))
     tar.add(settings_file, arcname="settings.json")
 
 

nucliadb/standalone/lifecycle.py

@@ -18,48 +18,49 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
 
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
+from nucliadb.common.context.fastapi import inject_app_context
 from nucliadb.ingest.app import initialize_grpc as initialize_ingest_grpc
 from nucliadb.ingest.app import initialize_pull_workers
 from nucliadb.ingest.settings import settings as ingest_settings
-from nucliadb.reader.lifecycle import finalize as finalize_reader
-from nucliadb.reader.lifecycle import initialize as initialize_reader
-from nucliadb.search.lifecycle import finalize as finalize_search
-from nucliadb.search.lifecycle import initialize as initialize_search
-from nucliadb.train.lifecycle import finalize as finalize_train
-from nucliadb.train.lifecycle import initialize as initialize_train
-from nucliadb.writer.lifecycle import finalize as finalize_writer
-from nucliadb.writer.lifecycle import initialize as initialize_writer
+from nucliadb.reader.lifecycle import lifespan as reader_lifespan
+from nucliadb.search.lifecycle import lifespan as search_lifespan
+from nucliadb.train.lifecycle import lifespan as train_lifespan
+from nucliadb.writer.lifecycle import lifespan as writer_lifespan
 from nucliadb_utils.utilities import finalize_utilities
 
 SYNC_FINALIZERS = []
 
 
-async def initialize():
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     if ingest_settings.disable_pull_worker:
         finalizers = await initialize_ingest_grpc()
     else:
         finalizers = await initialize_pull_workers()
     SYNC_FINALIZERS.extend(finalizers)
-    await initialize_writer()
-    await initialize_reader()
-    await initialize_search()
-    await initialize_train()
-    await setup_cluster()
 
+    async with (
+        writer_lifespan(app),
+        reader_lifespan(app),
+        search_lifespan(app),
+        train_lifespan(app),
+        inject_app_context(app),
+    ):
+        await setup_cluster()
+
+        yield
 
-async def finalize():
-    for finalizer in SYNC_FINALIZERS:
-        if asyncio.iscoroutinefunction(finalizer):
-            await finalizer()
-        else:
-            finalizer()
-    SYNC_FINALIZERS.clear()
+        for finalizer in SYNC_FINALIZERS:
+            if asyncio.iscoroutinefunction(finalizer):
+                await finalizer()
+            else:
+                finalizer()
+        SYNC_FINALIZERS.clear()
 
-    await finalize_writer()
-    await finalize_reader()
-    await finalize_search()
-    await finalize_train()
     await finalize_utilities()
     await teardown_cluster()

nucliadb/standalone/migrations.py

@@ -0,0 +1,58 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+import asyncio
+import sys
+
+from nucliadb.common import locking
+from nucliadb.common.cluster.standalone.utils import is_worker_node
+from nucliadb.migrator.command import run as run_migrator
+
+
+def run_migrations():
+    """
+    Run migrations for the standalone mode.
+    """
+    loop = asyncio.new_event_loop()
+    loop.run_until_complete(safe_run_migrations())
+    loop.close()
+
+
+async def safe_run_migrations():
+    """
+    Run migrations for the standalone mode, only if the node is a worker node.
+    The worker node will keep blocked until the migrations are run -- it relies
+    on the migrator's internal distributed lock.
+    """
+    if not is_worker_node():
+        return
+
+    sys.stdout.write(
+        """-------------------------------------------------
+|   Running Migrations for NucliaDB Standalone
+-------------------------------------------------
+"""
+    )
+    while True:
+        try:
+            await run_migrator(forever=False)
+            break
+        except locking.ResourceLocked:
+            sys.stdout.write("Another worker is already running migrations. Waiting...\n")
+            continue

nucliadb/standalone/purge.py

@@ -20,21 +20,22 @@
 # Standalone purge command
 import asyncio
 
-import pydantic_argparse
+import argdantic
 
 from nucliadb.standalone.config import config_nucliadb
 from nucliadb.standalone.settings import Settings
 
+parser: argdantic.ArgParser = argdantic.ArgParser()
+
+
+@parser.command(singleton=True, name="NucliaDB", help="NucliaDB Starting script")
+def setting(settings: Settings) -> Settings:
+    return settings
+
 
 def purge():
     from nucliadb.purge import main
 
-    parser = pydantic_argparse.ArgumentParser(
-        model=Settings,
-        prog="NucliaDB",
-        description="NucliaDB Starting script",
-    )
-    nucliadb_args = parser.parse_typed_args()
-
+    nucliadb_args = parser()
     config_nucliadb(nucliadb_args)
     asyncio.run(main())

nucliadb/standalone/run.py

@@ -23,7 +23,7 @@ import os
 import sys
 from typing import Optional
 
-import pydantic_argparse
+import argdantic
 import uvicorn  # type: ignore
 from fastapi import FastAPI
 
@@ -31,6 +31,7 @@ from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.ingest.settings import settings as ingest_settings
 from nucliadb.standalone import versions
 from nucliadb.standalone.config import config_nucliadb
+from nucliadb.standalone.migrations import run_migrations
 from nucliadb.standalone.settings import Settings
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.fastapi import instrument_app
@@ -41,14 +42,17 @@ from nucliadb_utils.settings import nuclia_settings, storage_settings
 logger = logging.getLogger(__name__)
 
 
+parser: argdantic.ArgParser = argdantic.ArgParser()
+
+
+@parser.command(singleton=True, name="NucliaDB", help="NucliaDB Starting script")
+def setting_parser(settings: Settings) -> Settings:
+    return settings
+
+
 def setup() -> Settings:
     errors.setup_error_handling(versions.get_installed_version("nucliadb"))
-    parser = pydantic_argparse.ArgumentParser(
-        model=Settings,
-        prog="NucliaDB",
-        description="NucliaDB Starting script",
-    )
-    nucliadb_args = parser.parse_typed_args()
+    nucliadb_args = parser()
 
     log_settings = LogSettings(
         # change default settings for standalone
@@ -80,6 +84,7 @@ def get_server(settings: Settings) -> tuple[FastAPI, uvicorn.Server]:
 
 def run():
     settings = setup()
+    run_migrations()
     app, server = get_server(settings)
     instrument_app(app, excluded_urls=["/"], metrics=True)
 
@@ -92,29 +97,31 @@ def run():
     settings_to_output = {
         "API": f"http://{settings.http_host}:{settings.http_port}/api",
         "Admin UI": f"http://{settings.http_host}:{settings.http_port}/admin",
-        "Key-value backend": ingest_settings.driver,
-        "Blog storage backend": storage_settings.file_backend,
-        "Cluster discovery mode": cluster_settings.cluster_discovery_mode,
+        "Key-value backend": ingest_settings.driver.value,
+        "Blob storage backend": storage_settings.file_backend.value,
+        "Cluster discovery mode": cluster_settings.cluster_discovery_mode.value,
         "Node replicas": cluster_settings.node_replicas,
         "Index data path": os.path.realpath(cluster_settings.data_path),
         "Node port": cluster_settings.standalone_node_port,
-        "Auth policy": settings.auth_policy,
-        "Log output type": settings.log_output_type,
-        "Node role": cluster_settings.standalone_node_role,
+        "Auth policy": settings.auth_policy.value,
+        "Node role": cluster_settings.standalone_node_role.value,
+    }
+    log_settings = {
+        "Log output type": settings.log_output_type.value,
+        "Log format type": settings.log_format_type.value,
+        "Log level": settings.log_level.value,
     }
     if settings.log_output_type == LogOutputType.FILE:
         log_folder = os.path.realpath(os.path.dirname(LogSettings().access_log))
-        settings_to_output["Log folder path"] = log_folder
+        log_settings["Log folder path"] = log_folder
+    settings_to_output.update(log_settings)
 
     if nuclia_settings.nuclia_service_account:
         settings_to_output["NUA API key"] = "Configured ✔"
         settings_to_output["NUA API zone"] = nuclia_settings.nuclia_zone
 
     settings_to_output_fmted = "\n".join(
-        [
-            f"||      - {k}:{' ' * (27 - len(k))}{v}"
-            for k, v in settings_to_output.items()
-        ]
+        [f"||      - {k}:{' ' * (27 - len(k))}{v}" for k, v in settings_to_output.items()]
     )
 
     installed_version = versions.installed_nucliadb()

nucliadb/standalone/settings.py

@@ -27,42 +27,40 @@ from nucliadb.ingest.settings import DriverSettings
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_telemetry.settings import LogFormatType, LogLevel, LogOutputType
 from nucliadb_utils.settings import StorageSettings
+from nucliadb_utils.storages.settings import Settings as ExtendedStorageSettings
 
 
-class StandaloneDiscoveryMode(str, Enum):
+class StandaloneDiscoveryMode(Enum):
     DEFAULT = "default"
     MANUAL = "manual"
     KUBERNETES = "kubernetes"
     SINGLE_NODE = "single_node"
 
 
-class AuthPolicy(str, Enum):
+class AuthPolicy(Enum):
     UPSTREAM_NAIVE = "upstream_naive"
     UPSTREAM_AUTH_HEADER = "upstream_auth_header"
     UPSTREAM_OAUTH2 = "upstream_oauth2"
     UPSTREAM_BASICAUTH = "upstream_basicauth"
 
 
-class Settings(DriverSettings, StorageSettings):
+class Settings(DriverSettings, StorageSettings, ExtendedStorageSettings):
     # be consistent here with DATA_PATH env var
-    data_path: str = pydantic.Field(
-        "./data/node", description="Path to node index files"
-    )
+    data_path: str = pydantic.Field("./data/node", description="Path to node index files")
 
     # all settings here are mapped in to other env var settings used
     # in the app. These are helper settings to make things easier to
     # use with standalone app vs cluster app.
     nua_api_key: Optional[str] = pydantic.Field(
-        description="Nuclia Understanding API Key. Read how to generate a NUA Key here: https://docs.nuclia.dev/docs/docs/using/understanding/intro#get-a-nua-key"  # noqa
+        default=None,
+        description="Nuclia Understanding API Key. Read how to generate a NUA Key here: https://docs.nuclia.dev/docs/rag/advanced/understanding/intro#get-a-nua-key",  # noqa
     )
-    zone: Optional[str] = pydantic.Field(description="Nuclia Understanding API Zone ID")
+    zone: Optional[str] = pydantic.Field(default=None, description="Nuclia Understanding API Zone ID")
     http_host: str = pydantic.Field(default="0.0.0.0", description="HTTP Port")
     http_port: int = pydantic.Field(default=8080, description="HTTP Port")
     ingest_grpc_port: int = pydantic.Field(default=8030, description="Ingest GRPC Port")
     train_grpc_port: int = pydantic.Field(default=8031, description="Train GRPC Port")
-    standalone_node_port: int = pydantic.Field(
-        default=10009, description="Node GRPC Port"
-    )
+    standalone_node_port: int = pydantic.Field(default=10009, description="Node GRPC Port")
 
     auth_policy: AuthPolicy = pydantic.Field(
         default=AuthPolicy.UPSTREAM_NAIVE,
@@ -90,9 +88,7 @@ class Settings(DriverSettings, StorageSettings):
         description="Default role to assign to user that is authenticated \
                     upstream. Not used with `upstream_naive` auth policy.",
     )
-    auth_policy_role_mapping: Optional[
-        dict[str, dict[str, list[NucliaDBRoles]]]
-    ] = pydantic.Field(
+    auth_policy_role_mapping: Optional[dict[str, dict[str, list[NucliaDBRoles]]]] = pydantic.Field(
         default=None,
         description="""
 Role mapping for `upstream_auth_header`, `upstream_oauth2` and `upstream_basicauth` auth policies.

nucliadb/standalone/versions.py

@@ -18,10 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import enum
+import importlib.metadata
 import logging
 from typing import Optional
 
-import pkg_resources
 from cachetools import TTLCache
 
 from nucliadb.common.http_clients.pypi import PyPi
@@ -65,14 +65,24 @@ def is_newer_release(installed: str, latest: str) -> bool:
     >>> is_newer_release("1.2.3", "1.2.3.post1")
     False
     """
-    parsed_installed = pkg_resources.parse_version(_release(installed))
-    parsed_latest = pkg_resources.parse_version(_release(latest))
+
+    def parse_version(version: str) -> tuple[int, int, int]:
+        parts = version.split(".")
+        if len(parts) > 3:
+            raise ValueError(f"Invalid version string: {version}")
+        major = int(parts[0]) if len(parts) >= 1 else 0
+        minor = int(parts[1]) if len(parts) >= 2 else 0
+        patch = int(parts[2]) if len(parts) == 3 else 0
+        return (major, minor, patch)
+
+    parsed_installed = parse_version(_release(installed))
+    parsed_latest = parse_version(_release(latest))
     return parsed_latest > parsed_installed
 
 
 def _release(version: str) -> str:
     """
-    Strips the .postX part of the version so that wecan compare major.minor.patch only.
+    Strips the .postX part of the version so that we can compare major.minor.patch only.
 
     >>> _release("1.2.3")
     '1.2.3'
@@ -83,7 +93,7 @@ def _release(version: str) -> str:
 
 
 def get_installed_version(package_name: str) -> str:
-    return pkg_resources.get_distribution(package_name).version
+    return importlib.metadata.distribution(package_name).version
 
 
 async def get_latest_version(package: str) -> Optional[str]:

nucliadb/tasks/consumer.py

@@ -60,7 +60,9 @@ class NatsTaskConsumer:
     async def initialize(self, context: ApplicationContext):
         self.context = context
         await create_nats_stream_if_not_exists(
-            self.context, self.stream.name, subjects=[self.stream.subject]  # type: ignore
+            self.context,
+            self.stream.name,  # type: ignore
+            subjects=[self.stream.subject],  # type: ignore
         )
         await self._setup_nats_subscription()
         self.initialized = True
@@ -128,11 +130,9 @@ class NatsTaskConsumer:
             f"Message received: subject:{subject}, seqid: {seqid}, reply: {reply}",
             extra={"consumer_name": self.name},
         )
-        async with MessageProgressUpdater(
-            msg, nats_consumer_settings.nats_ack_wait * 0.66
-        ):
+        async with MessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66):
             try:
-                task_msg = self.msg_type.parse_raw(msg.data)
+                task_msg = self.msg_type.model_validate_json(msg.data)
             except pydantic.ValidationError as e:
                 errors.capture_exception(e)
                 logger.error(
@@ -144,9 +144,7 @@ class NatsTaskConsumer:
                 await msg.ack()
                 return
 
-            logger.info(
-                f"Starting task consumption", extra={"consumer_name": self.name}
-            )
+            logger.info(f"Starting task consumption", extra={"consumer_name": self.name})
             try:
                 await self.callback(self.context, task_msg)  # type: ignore
             except asyncio.CancelledError:
@@ -200,9 +198,7 @@ def create_consumer(
     return consumer
 
 
-async def start_consumer(
-    task_name: str, context: ApplicationContext
-) -> NatsTaskConsumer:
+async def start_consumer(task_name: str, context: ApplicationContext) -> NatsTaskConsumer:
     """
     Returns an initialized consumer for the given task name, ready to consume messages from the task stream.
     """
@@ -214,7 +210,7 @@ async def start_consumer(
         name=f"{task_name}_consumer",
         stream=task.stream,
         callback=task.callback,  # type: ignore
-        msg_type=task.msg_type,  # type: ignore
+        msg_type=task.msg_type,
         max_concurrent_messages=task.max_concurrent_messages,
     )
     await consumer.initialize(context)

nucliadb/tasks/producer.py

@@ -44,11 +44,13 @@ class NatsTaskProducer:
     async def initialize(self, context: ApplicationContext):
         self.context = context
         await create_nats_stream_if_not_exists(
-            self.context, self.stream.name, subjects=[self.stream.subject]  # type: ignore
+            self.context,
+            self.stream.name,  # type: ignore
+            subjects=[self.stream.subject],  # type: ignore
         )
         self.initialized = True
 
-    async def __call__(self, msg: MsgType) -> int:  # type: ignore
+    async def __call__(self, msg: MsgType) -> int:
         """
         Publish message to the producer's nats stream.
         Returns the sequence number of the published message.
@@ -57,7 +59,8 @@ class NatsTaskProducer:
             raise RuntimeError("NatsTaskProducer not initialized")
         try:
             pub_ack = await self.context.nats_manager.js.publish(  # type: ignore
-                self.stream.subject, msg.json().encode("utf-8")  # type: ignore
+                self.stream.subject,  # type: ignore
+                msg.model_dump_json().encode("utf-8"),  # type: ignore
             )
             logger.info(
                 "Message sent to Nats",
@@ -93,8 +96,6 @@ async def get_producer(task_name: str, context: ApplicationContext) -> NatsTaskP
         task = get_registered_task(task_name)
     except KeyError:
         raise ValueError(f"Task {task_name} not registered")
-    producer = create_producer(
-        name=f"{task_name}_producer", stream=task.stream, msg_type=task.msg_type
-    )
+    producer = create_producer(name=f"{task_name}_producer", stream=task.stream, msg_type=task.msg_type)
     await producer.initialize(context)
     return producer