Feat/prompt optimization refactor

packages/nvidia_nat_core/pyproject.toml

@@ -71,6 +71,7 @@ dependencies = [
   "nest-asyncio2~=1.7",
   "networkx~=3.4",
   "nvidia-nat-atif == {version}",
+  "nvidia-nat-prompt-optimization == {version}",
   "numpy~=2.3",
   "openinference-semantic-conventions>=0.1.14,<1.0.0",
   "optuna~=4.4",
@@ -120,6 +121,7 @@ test = [
 [tool.uv.sources]
 nvidia-nat-atif = { path = "../nvidia_nat_atif", editable = true }
 nvidia-nat-eval = { path = "../nvidia_nat_eval", editable = true }
+nvidia-nat-prompt-optimization = { path = "../nvidia_nat_prompt_optimization", editable = true }
 nvidia-nat-test = { path = "../nvidia_nat_test", editable = true }
 
 [project.entry-points.'nat.components']

packages/nvidia_nat_core/src/nat/data_models/optimizer.py

@@ -18,6 +18,7 @@
 from pathlib import Path
 from typing import Literal
 
+from nvidia_prompt_optimization import PromptMutationInput
 from pydantic import BaseModel
 from pydantic import Field
 
@@ -40,12 +41,9 @@ class OptimizerMetric(BaseModel):
     weight: float = Field(description="Weight of the metric in the optimization process.", default=1.0)
 
 
-class PromptOptimizerInputSchema(BaseModel):
-    """Input schema for prompt optimizer mutator/recombiner helper functions."""
-
-    original_prompt: str
-    objective: str
-    oracle_feedback: str | None = None
+# Compatibility alias for existing NAT imports while the standalone prompt optimizer
+# package becomes the source of truth for prompt callback payloads.
+PromptOptimizerInputSchema = PromptMutationInput
 
 
 class SamplerType(StrEnum):

packages/nvidia_nat_core/tests/nat/data_models/test_optimizer.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nvidia_prompt_optimization import PromptMutationInput
+
+from nat.data_models.optimizer import PromptOptimizerInputSchema
+
+
+def test_prompt_optimizer_input_schema_is_standalone_contract():
+    assert PromptOptimizerInputSchema is PromptMutationInput

packages/nvidia_nat_langchain/pyproject.toml

@@ -58,6 +58,7 @@ dependencies = [
   "nvidia-nat-core == {version}",
   "nvidia-nat-eval == {version}",
   "nvidia-nat-opentelemetry == {version}",
+  "nvidia-nat-prompt-optimization == {version}",
   "langchain>=1.2.3,<2.0.0",
   "langchain-aws>=1.1.0,<2.0.0",
   "langchain-classic>=1.0.1,<2.0.0",
@@ -88,6 +89,7 @@ config-settings = { editable_mode = "compat" }
 nvidia-nat-core = { path = "../nvidia_nat_core", editable = true }
 nvidia-nat-eval = { path = "../nvidia_nat_eval", editable = true }
 nvidia-nat-opentelemetry = { path = "../nvidia_nat_opentelemetry", editable = true }
+nvidia-nat-prompt-optimization = { path = "../nvidia_nat_prompt_optimization", editable = true }
 nvidia-nat-test = { path = "../nvidia_nat_test", editable = true }
 
 [project.entry-points.'nat.components']

packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/prompt_optimizer/register.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from nvidia_prompt_optimization import PromptMutationInput
+from nvidia_prompt_optimization import PromptRecombinationInput
 from pydantic import Field
 
 from nat.builder.builder import Builder
@@ -21,7 +23,8 @@
 from nat.cli.register_workflow import register_function
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.function import FunctionBaseConfig
-from nat.data_models.optimizer import PromptOptimizerInputSchema
+
+PromptOptimizerInputSchema = PromptMutationInput
 
 
 class PromptOptimizerConfig(FunctionBaseConfig, name="prompt_init"):
@@ -66,7 +69,7 @@ async def prompt_optimizer_function(config: PromptOptimizerConfig, builder: Buil
         input_variables=["original_prompt", "objective", "oracle_feedback_section"],
         validate_template=True)
 
-    async def _inner(input_message: PromptOptimizerInputSchema) -> str:
+    async def _inner(input_message: PromptMutationInput) -> str:
         """
         Optimize the prompt using the provided LLM.
         """
@@ -130,8 +133,8 @@ async def prompt_recombiner_function(config: PromptRecombinerConfig, builder: Bu
 
     base_prompt: str = (await template.ainvoke(input={"system_objective": config.system_objective})).to_string()
 
-    class RecombineSchema(PromptOptimizerInputSchema):
-        parent_b: str | None = None
+    class RecombineSchema(PromptRecombinationInput):
+        pass
 
     async def _inner(input_message: RecombineSchema) -> str:
         parent_a = input_message.original_prompt

packages/nvidia_nat_langchain/tests/agent/prompt_optimizer/test_prompt_optimizer_register.py

@@ -13,13 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from nvidia_prompt_optimization import PromptMutationInput
+
 from nat.plugins.langchain.agent.prompt_optimizer.prompt import oracle_feedback_template
 from nat.plugins.langchain.agent.prompt_optimizer.register import PromptOptimizerInputSchema
 
 
 class TestPromptOptimizerInputSchema:
     """Tests for PromptOptimizerInputSchema."""
 
+    def test_schema_is_standalone_contract(self):
+        assert PromptOptimizerInputSchema is PromptMutationInput
+
     def test_oracle_feedback_is_optional(self):
         """Oracle feedback defaults to None."""
         schema = PromptOptimizerInputSchema(

packages/nvidia_nat_prompt_optimization/README.md

@@ -0,0 +1,3 @@
+# NVIDIA Prompt Optimization
+
+Standalone prompt optimization contracts for NVIDIA NeMo Agent Toolkit.

packages/nvidia_nat_prompt_optimization/pyproject.toml

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=64", "setuptools-scm>=8", "setuptools_dynamic_dependencies>=1.0.0"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["nvidia_prompt_optimization*"]
+
+[tool.setuptools_scm]
+git_describe_command = "git describe --long --first-parent"
+root = "../.."
+
+[project]
+name = "nvidia-nat-prompt-optimization"
+dynamic = ["version", "dependencies", "optional-dependencies"]
+requires-python = ">=3.11,<3.14"
+description = "Standalone prompt optimization contracts for NVIDIA NeMo Agent Toolkit"
+readme = "README.md"
+keywords = ["ai", "agents", "prompt", "optimization"]
+license = { text = "Apache-2.0" }
+authors = [{ name = "NVIDIA Corporation" }]
+maintainers = [{ name = "NVIDIA Corporation" }]
+classifiers = [
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+]
+
+[project.urls]
+documentation = "https://docs.nvidia.com/nemo/agent-toolkit/latest/"
+source = "https://github.com/NVIDIA/NeMo-Agent-Toolkit"
+
+[tool.setuptools_dynamic_dependencies]
+dependencies = [
+  "pydantic~=2.11",
+]
+
+[tool.setuptools_dynamic_dependencies.optional-dependencies]
+test = [
+  "nvidia-nat-test == {version}",
+]
+
+[tool.uv]
+build-constraint-dependencies = ["setuptools>=64", "setuptools-scm>=8", "setuptools_dynamic_dependencies>=1.0.0"]
+managed = true
+config-settings = { editable_mode = "compat" }
+
+[tool.uv.sources]
+nvidia-nat-test = { path = "../nvidia_nat_test", editable = true }

packages/nvidia_nat_prompt_optimization/src/nvidia_prompt_optimization/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone prompt optimization contracts."""
+
+from .contracts import PromptEvaluationItem
+from .contracts import PromptEvaluationResult
+from .contracts import PromptEvaluator
+from .contracts import PromptMutationInput
+from .contracts import PromptMutator
+from .contracts import PromptOptimizationHistoryEntry
+from .contracts import PromptOptimizationResult
+from .contracts import PromptRecombinationInput
+from .contracts import PromptRecombiner
+from .contracts import PromptSet
+
+__all__ = [
+    "PromptEvaluationItem",
+    "PromptEvaluationResult",
+    "PromptEvaluator",
+    "PromptMutationInput",
+    "PromptMutator",
+    "PromptOptimizationHistoryEntry",
+    "PromptOptimizationResult",
+    "PromptRecombinationInput",
+    "PromptRecombiner",
+    "PromptSet",
+]