nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aiq/__init__.py +1 -1
- nat/{front_ends/mcp → agent/auto_memory_wrapper}/__init__.py +1 -1
- nat/agent/auto_memory_wrapper/agent.py +278 -0
- nat/agent/auto_memory_wrapper/register.py +227 -0
- nat/agent/auto_memory_wrapper/state.py +30 -0
- nat/agent/base.py +1 -1
- nat/agent/dual_node.py +1 -1
- nat/agent/prompt_optimizer/prompt.py +1 -1
- nat/agent/prompt_optimizer/register.py +1 -1
- nat/agent/react_agent/agent.py +16 -9
- nat/agent/react_agent/output_parser.py +2 -2
- nat/agent/react_agent/prompt.py +3 -2
- nat/agent/react_agent/register.py +2 -2
- nat/agent/react_agent/register_per_user_agent.py +104 -0
- nat/agent/reasoning_agent/reasoning_agent.py +1 -1
- nat/agent/register.py +3 -1
- nat/agent/responses_api_agent/__init__.py +1 -1
- nat/agent/responses_api_agent/register.py +1 -1
- nat/agent/rewoo_agent/agent.py +9 -4
- nat/agent/rewoo_agent/prompt.py +1 -1
- nat/agent/rewoo_agent/register.py +1 -1
- nat/agent/tool_calling_agent/agent.py +5 -4
- nat/agent/tool_calling_agent/register.py +1 -1
- nat/authentication/__init__.py +1 -1
- nat/authentication/api_key/__init__.py +1 -1
- nat/authentication/api_key/api_key_auth_provider.py +1 -1
- nat/authentication/api_key/api_key_auth_provider_config.py +22 -7
- nat/authentication/api_key/register.py +1 -1
- nat/authentication/credential_validator/__init__.py +1 -1
- nat/authentication/credential_validator/bearer_token_validator.py +1 -1
- nat/authentication/exceptions/__init__.py +1 -1
- nat/authentication/exceptions/api_key_exceptions.py +1 -1
- nat/authentication/http_basic_auth/http_basic_auth_provider.py +1 -1
- nat/authentication/http_basic_auth/register.py +1 -1
- nat/authentication/interfaces.py +1 -1
- nat/authentication/oauth2/__init__.py +1 -1
- nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +1 -1
- nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +1 -1
- nat/authentication/oauth2/oauth2_resource_server_config.py +1 -1
- nat/authentication/oauth2/register.py +1 -1
- nat/authentication/register.py +1 -1
- nat/builder/builder.py +511 -1
- nat/builder/child_builder.py +385 -0
- nat/builder/component_utils.py +28 -4
- nat/builder/context.py +17 -1
- nat/builder/embedder.py +1 -1
- nat/builder/eval_builder.py +19 -7
- nat/builder/evaluator.py +1 -1
- nat/builder/framework_enum.py +2 -1
- nat/builder/front_end.py +1 -1
- nat/builder/function.py +40 -3
- nat/builder/function_base.py +1 -1
- nat/builder/function_info.py +1 -1
- nat/builder/intermediate_step_manager.py +1 -1
- nat/builder/llm.py +1 -1
- nat/builder/per_user_workflow_builder.py +843 -0
- nat/builder/retriever.py +1 -1
- nat/builder/sync_builder.py +571 -0
- nat/builder/user_interaction_manager.py +1 -1
- nat/builder/workflow.py +1 -1
- nat/builder/workflow_builder.py +536 -424
- nat/cli/__init__.py +1 -1
- nat/cli/cli_utils/config_override.py +1 -1
- nat/cli/cli_utils/validation.py +32 -1
- nat/cli/commands/configure/channel/add.py +1 -1
- nat/cli/commands/configure/channel/channel.py +1 -1
- nat/cli/commands/configure/channel/remove.py +1 -1
- nat/cli/commands/configure/channel/update.py +1 -1
- nat/cli/commands/configure/configure.py +1 -1
- nat/cli/commands/evaluate.py +87 -13
- nat/cli/commands/finetune.py +132 -0
- nat/cli/commands/info/__init__.py +1 -1
- nat/cli/commands/info/info.py +1 -1
- nat/cli/commands/info/list_channels.py +1 -1
- nat/cli/commands/info/list_components.py +1 -1
- nat/cli/commands/object_store/__init__.py +1 -1
- nat/cli/commands/object_store/object_store.py +1 -1
- nat/cli/commands/optimize.py +1 -1
- nat/cli/commands/{mcp → red_teaming}/__init__.py +1 -1
- nat/cli/commands/red_teaming/red_teaming.py +138 -0
- nat/cli/commands/red_teaming/red_teaming_utils.py +73 -0
- nat/cli/commands/registry/__init__.py +1 -1
- nat/cli/commands/registry/publish.py +1 -1
- nat/cli/commands/registry/pull.py +1 -1
- nat/cli/commands/registry/registry.py +1 -1
- nat/cli/commands/registry/remove.py +1 -1
- nat/cli/commands/registry/search.py +1 -1
- nat/cli/commands/sizing/__init__.py +1 -1
- nat/cli/commands/sizing/calc.py +1 -1
- nat/cli/commands/sizing/sizing.py +1 -1
- nat/cli/commands/start.py +1 -1
- nat/cli/commands/uninstall.py +1 -1
- nat/cli/commands/validate.py +1 -1
- nat/cli/commands/workflow/__init__.py +1 -1
- nat/cli/commands/workflow/workflow.py +1 -1
- nat/cli/commands/workflow/workflow_commands.py +3 -2
- nat/cli/entrypoint.py +15 -37
- nat/cli/main.py +2 -2
- nat/cli/plugin_loader.py +69 -0
- nat/cli/register_workflow.py +183 -5
- nat/cli/type_registry.py +169 -3
- nat/control_flow/register.py +1 -1
- nat/control_flow/router_agent/agent.py +1 -1
- nat/control_flow/router_agent/prompt.py +1 -1
- nat/control_flow/router_agent/register.py +1 -1
- nat/control_flow/sequential_executor.py +28 -7
- nat/data_models/__init__.py +1 -1
- nat/data_models/agent.py +1 -1
- nat/data_models/api_server.py +38 -3
- nat/data_models/authentication.py +1 -1
- nat/data_models/common.py +1 -1
- nat/data_models/component.py +7 -1
- nat/data_models/component_ref.py +34 -1
- nat/data_models/config.py +62 -1
- nat/data_models/dataset_handler.py +15 -2
- nat/data_models/discovery_metadata.py +1 -1
- nat/data_models/embedder.py +1 -1
- nat/data_models/evaluate.py +6 -1
- nat/data_models/evaluator.py +1 -1
- nat/data_models/finetuning.py +260 -0
- nat/data_models/front_end.py +1 -1
- nat/data_models/function.py +1 -1
- nat/data_models/function_dependencies.py +1 -1
- nat/data_models/gated_field_mixin.py +1 -1
- nat/data_models/interactive.py +1 -1
- nat/data_models/intermediate_step.py +29 -2
- nat/data_models/invocation_node.py +1 -1
- nat/data_models/llm.py +1 -1
- nat/data_models/logging.py +1 -1
- nat/data_models/memory.py +1 -1
- nat/data_models/middleware.py +3 -1
- nat/data_models/object_store.py +1 -1
- nat/data_models/openai_mcp.py +1 -1
- nat/data_models/optimizable.py +1 -1
- nat/data_models/optimizer.py +1 -1
- nat/data_models/profiler.py +1 -1
- nat/data_models/registry_handler.py +1 -1
- nat/data_models/retriever.py +1 -1
- nat/data_models/retry_mixin.py +1 -1
- nat/data_models/runtime_enum.py +1 -1
- nat/data_models/span.py +1 -1
- nat/data_models/step_adaptor.py +1 -1
- nat/data_models/streaming.py +1 -1
- nat/data_models/swe_bench_model.py +1 -1
- nat/data_models/telemetry_exporter.py +1 -1
- nat/data_models/thinking_mixin.py +1 -1
- nat/data_models/ttc_strategy.py +1 -1
- nat/embedder/azure_openai_embedder.py +1 -1
- nat/embedder/nim_embedder.py +1 -1
- nat/embedder/openai_embedder.py +1 -1
- nat/embedder/register.py +1 -1
- nat/eval/__init__.py +1 -1
- nat/eval/config.py +8 -1
- nat/eval/dataset_handler/dataset_downloader.py +1 -1
- nat/eval/dataset_handler/dataset_filter.py +1 -1
- nat/eval/dataset_handler/dataset_handler.py +4 -2
- nat/eval/evaluate.py +217 -80
- nat/eval/evaluator/__init__.py +1 -1
- nat/eval/evaluator/base_evaluator.py +2 -2
- nat/eval/evaluator/evaluator_model.py +3 -2
- nat/eval/intermediate_step_adapter.py +1 -1
- nat/eval/llm_validator.py +336 -0
- nat/eval/rag_evaluator/evaluate.py +17 -10
- nat/eval/rag_evaluator/register.py +1 -1
- nat/eval/red_teaming_evaluator/__init__.py +14 -0
- nat/eval/red_teaming_evaluator/data_models.py +66 -0
- nat/eval/red_teaming_evaluator/evaluate.py +327 -0
- nat/eval/red_teaming_evaluator/filter_conditions.py +75 -0
- nat/eval/red_teaming_evaluator/register.py +55 -0
- nat/eval/register.py +2 -1
- nat/eval/remote_workflow.py +1 -1
- nat/eval/runners/__init__.py +1 -1
- nat/eval/runners/config.py +1 -1
- nat/eval/runners/multi_eval_runner.py +1 -1
- nat/eval/runners/red_teaming_runner/__init__.py +24 -0
- nat/eval/runners/red_teaming_runner/config.py +282 -0
- nat/eval/runners/red_teaming_runner/report_utils.py +707 -0
- nat/eval/runners/red_teaming_runner/runner.py +867 -0
- nat/eval/runtime_evaluator/__init__.py +1 -1
- nat/eval/runtime_evaluator/evaluate.py +1 -1
- nat/eval/runtime_evaluator/register.py +1 -1
- nat/eval/runtime_event_subscriber.py +1 -1
- nat/eval/swe_bench_evaluator/evaluate.py +1 -1
- nat/eval/swe_bench_evaluator/register.py +1 -1
- nat/eval/trajectory_evaluator/evaluate.py +2 -2
- nat/eval/trajectory_evaluator/register.py +1 -1
- nat/eval/tunable_rag_evaluator/evaluate.py +5 -5
- nat/eval/tunable_rag_evaluator/register.py +1 -1
- nat/eval/usage_stats.py +1 -1
- nat/eval/utils/eval_trace_ctx.py +1 -1
- nat/eval/utils/output_uploader.py +1 -1
- nat/eval/utils/tqdm_position_registry.py +1 -1
- nat/eval/utils/weave_eval.py +1 -1
- nat/experimental/decorators/experimental_warning_decorator.py +1 -1
- nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +1 -1
- nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +1 -1
- nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +1 -1
- nat/experimental/test_time_compute/functions/execute_score_select_function.py +1 -1
- nat/experimental/test_time_compute/functions/multi_llm_judge_function.py +88 -0
- nat/experimental/test_time_compute/functions/plan_select_execute_function.py +1 -1
- nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +1 -1
- nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +1 -1
- nat/experimental/test_time_compute/models/editor_config.py +1 -1
- nat/experimental/test_time_compute/models/scoring_config.py +1 -1
- nat/experimental/test_time_compute/models/search_config.py +20 -2
- nat/experimental/test_time_compute/models/selection_config.py +33 -2
- nat/experimental/test_time_compute/models/stage_enums.py +1 -1
- nat/experimental/test_time_compute/models/strategy_base.py +1 -1
- nat/experimental/test_time_compute/models/tool_use_config.py +1 -1
- nat/experimental/test_time_compute/models/ttc_item.py +1 -1
- nat/experimental/test_time_compute/register.py +4 -1
- nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +1 -1
- nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +1 -1
- nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +1 -1
- nat/experimental/test_time_compute/search/multi_llm_generation.py +115 -0
- nat/experimental/test_time_compute/search/multi_llm_planner.py +1 -1
- nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +1 -1
- nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +1 -1
- nat/experimental/test_time_compute/selection/best_of_n_selector.py +1 -1
- nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +1 -1
- nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +1 -1
- nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +1 -1
- nat/experimental/test_time_compute/selection/llm_judge_selection.py +127 -0
- nat/experimental/test_time_compute/selection/threshold_selector.py +1 -1
- nat/finetuning/__init__.py +24 -0
- nat/finetuning/finetuning_runtime.py +143 -0
- nat/finetuning/interfaces/__init__.py +24 -0
- nat/finetuning/interfaces/finetuning_runner.py +261 -0
- nat/finetuning/interfaces/trainer_adapter.py +103 -0
- nat/finetuning/interfaces/trajectory_builder.py +115 -0
- nat/finetuning/utils/__init__.py +15 -0
- nat/finetuning/utils/parsers/__init__.py +15 -0
- nat/finetuning/utils/parsers/adk_parser.py +141 -0
- nat/finetuning/utils/parsers/base_parser.py +238 -0
- nat/finetuning/utils/parsers/common.py +91 -0
- nat/finetuning/utils/parsers/langchain_parser.py +267 -0
- nat/finetuning/utils/parsers/llama_index_parser.py +218 -0
- nat/front_ends/__init__.py +1 -1
- nat/front_ends/console/__init__.py +1 -1
- nat/front_ends/console/authentication_flow_handler.py +1 -1
- nat/front_ends/console/console_front_end_config.py +4 -1
- nat/front_ends/console/console_front_end_plugin.py +5 -4
- nat/front_ends/console/register.py +1 -1
- nat/front_ends/cron/__init__.py +1 -1
- nat/front_ends/fastapi/__init__.py +1 -1
- nat/front_ends/fastapi/async_job.py +128 -0
- nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +1 -1
- nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +13 -9
- nat/front_ends/fastapi/dask_client_mixin.py +1 -1
- nat/front_ends/fastapi/fastapi_front_end_config.py +1 -1
- nat/front_ends/fastapi/fastapi_front_end_controller.py +1 -1
- nat/front_ends/fastapi/fastapi_front_end_plugin.py +25 -30
- nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +195 -60
- nat/front_ends/fastapi/html_snippets/__init__.py +1 -1
- nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +1 -1
- nat/front_ends/fastapi/intermediate_steps_subscriber.py +12 -1
- nat/front_ends/fastapi/job_store.py +23 -11
- nat/front_ends/fastapi/main.py +1 -1
- nat/front_ends/fastapi/message_handler.py +27 -4
- nat/front_ends/fastapi/message_validator.py +54 -2
- nat/front_ends/fastapi/register.py +1 -1
- nat/front_ends/fastapi/response_helpers.py +16 -15
- nat/front_ends/fastapi/step_adaptor.py +1 -1
- nat/front_ends/fastapi/utils.py +1 -1
- nat/front_ends/register.py +1 -2
- nat/front_ends/simple_base/__init__.py +1 -1
- nat/front_ends/simple_base/simple_front_end_plugin_base.py +6 -4
- nat/llm/aws_bedrock_llm.py +1 -1
- nat/llm/azure_openai_llm.py +10 -1
- nat/llm/dynamo_llm.py +363 -0
- nat/llm/huggingface_llm.py +177 -0
- nat/llm/litellm_llm.py +1 -1
- nat/llm/nim_llm.py +1 -1
- nat/llm/openai_llm.py +1 -1
- nat/llm/register.py +3 -1
- nat/llm/utils/__init__.py +1 -1
- nat/llm/utils/env_config_value.py +1 -1
- nat/llm/utils/error.py +1 -1
- nat/llm/utils/thinking.py +1 -1
- nat/memory/__init__.py +1 -1
- nat/memory/interfaces.py +1 -1
- nat/memory/models.py +1 -1
- nat/meta/pypi.md +1 -1
- nat/middleware/__init__.py +5 -5
- nat/middleware/cache/__init__.py +14 -0
- nat/middleware/{cache_middleware.py → cache/cache_middleware.py} +39 -42
- nat/middleware/cache/cache_middleware_config.py +44 -0
- nat/middleware/cache/register.py +33 -0
- nat/middleware/defense/__init__.py +14 -0
- nat/middleware/defense/defense_middleware.py +362 -0
- nat/middleware/defense/defense_middleware_content_guard.py +455 -0
- nat/middleware/defense/defense_middleware_data_models.py +91 -0
- nat/middleware/defense/defense_middleware_output_verifier.py +440 -0
- nat/middleware/defense/defense_middleware_pii.py +356 -0
- nat/middleware/defense/register.py +82 -0
- nat/middleware/dynamic/__init__.py +14 -0
- nat/middleware/dynamic/dynamic_function_middleware.py +962 -0
- nat/middleware/dynamic/dynamic_middleware_config.py +132 -0
- nat/middleware/dynamic/register.py +34 -0
- nat/middleware/function_middleware.py +236 -52
- nat/middleware/logging/__init__.py +14 -0
- nat/middleware/logging/logging_middleware.py +67 -0
- nat/middleware/logging/logging_middleware_config.py +28 -0
- nat/middleware/logging/register.py +33 -0
- nat/middleware/middleware.py +142 -28
- nat/middleware/red_teaming/__init__.py +14 -0
- nat/middleware/red_teaming/red_teaming_middleware.py +344 -0
- nat/middleware/red_teaming/red_teaming_middleware_config.py +112 -0
- nat/middleware/red_teaming/register.py +47 -0
- nat/middleware/register.py +7 -20
- nat/middleware/utils/__init__.py +14 -0
- nat/middleware/utils/workflow_inventory.py +155 -0
- nat/object_store/__init__.py +1 -1
- nat/object_store/in_memory_object_store.py +1 -1
- nat/object_store/interfaces.py +1 -1
- nat/object_store/models.py +1 -1
- nat/object_store/register.py +1 -1
- nat/observability/__init__.py +1 -1
- nat/observability/exporter/__init__.py +1 -1
- nat/observability/exporter/base_exporter.py +1 -1
- nat/observability/exporter/exporter.py +1 -1
- nat/observability/exporter/file_exporter.py +1 -1
- nat/observability/exporter/processing_exporter.py +1 -1
- nat/observability/exporter/raw_exporter.py +1 -1
- nat/observability/exporter/span_exporter.py +7 -1
- nat/observability/exporter_manager.py +1 -1
- nat/observability/mixin/__init__.py +1 -1
- nat/observability/mixin/batch_config_mixin.py +1 -1
- nat/observability/mixin/collector_config_mixin.py +1 -1
- nat/observability/mixin/file_mixin.py +1 -1
- nat/observability/mixin/file_mode.py +1 -1
- nat/observability/mixin/redaction_config_mixin.py +1 -1
- nat/observability/mixin/resource_conflict_mixin.py +1 -1
- nat/observability/mixin/serialize_mixin.py +1 -1
- nat/observability/mixin/tagging_config_mixin.py +1 -1
- nat/observability/mixin/type_introspection_mixin.py +1 -1
- nat/observability/processor/__init__.py +1 -1
- nat/observability/processor/batching_processor.py +1 -1
- nat/observability/processor/callback_processor.py +1 -1
- nat/observability/processor/falsy_batch_filter_processor.py +1 -1
- nat/observability/processor/intermediate_step_serializer.py +1 -1
- nat/observability/processor/processor.py +1 -1
- nat/observability/processor/processor_factory.py +1 -1
- nat/observability/processor/redaction/__init__.py +1 -1
- nat/observability/processor/redaction/contextual_redaction_processor.py +1 -1
- nat/observability/processor/redaction/contextual_span_redaction_processor.py +1 -1
- nat/observability/processor/redaction/redaction_processor.py +1 -1
- nat/observability/processor/redaction/span_header_redaction_processor.py +1 -1
- nat/observability/processor/span_tagging_processor.py +1 -1
- nat/observability/register.py +1 -1
- nat/observability/utils/__init__.py +1 -1
- nat/observability/utils/dict_utils.py +1 -1
- nat/observability/utils/time_utils.py +1 -1
- nat/profiler/calc/__init__.py +1 -1
- nat/profiler/calc/calc_runner.py +3 -3
- nat/profiler/calc/calculations.py +1 -1
- nat/profiler/calc/data_models.py +1 -1
- nat/profiler/calc/plot.py +30 -3
- nat/profiler/callbacks/agno_callback_handler.py +1 -1
- nat/profiler/callbacks/base_callback_class.py +1 -1
- nat/profiler/callbacks/langchain_callback_handler.py +33 -3
- nat/profiler/callbacks/llama_index_callback_handler.py +13 -10
- nat/profiler/callbacks/semantic_kernel_callback_handler.py +1 -1
- nat/profiler/callbacks/token_usage_base_model.py +1 -1
- nat/profiler/data_frame_row.py +1 -1
- nat/profiler/data_models.py +1 -1
- nat/profiler/decorators/framework_wrapper.py +16 -1
- nat/profiler/decorators/function_tracking.py +1 -1
- nat/profiler/forecasting/config.py +1 -1
- nat/profiler/forecasting/model_trainer.py +1 -1
- nat/profiler/forecasting/models/__init__.py +1 -1
- nat/profiler/forecasting/models/forecasting_base_model.py +1 -1
- nat/profiler/forecasting/models/linear_model.py +1 -1
- nat/profiler/forecasting/models/random_forest_regressor.py +1 -1
- nat/profiler/inference_metrics_model.py +1 -1
- nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +1 -1
- nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +1 -1
- nat/profiler/inference_optimization/data_models.py +1 -1
- nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +1 -1
- nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +1 -1
- nat/profiler/inference_optimization/llm_metrics.py +1 -1
- nat/profiler/inference_optimization/prompt_caching.py +1 -1
- nat/profiler/inference_optimization/token_uniqueness.py +1 -1
- nat/profiler/inference_optimization/workflow_runtimes.py +1 -1
- nat/profiler/intermediate_property_adapter.py +1 -1
- nat/profiler/parameter_optimization/optimizable_utils.py +1 -1
- nat/profiler/parameter_optimization/optimizer_runtime.py +1 -1
- nat/profiler/parameter_optimization/parameter_optimizer.py +1 -1
- nat/profiler/parameter_optimization/parameter_selection.py +1 -1
- nat/profiler/parameter_optimization/pareto_visualizer.py +1 -1
- nat/profiler/parameter_optimization/prompt_optimizer.py +1 -1
- nat/profiler/parameter_optimization/update_helpers.py +1 -1
- nat/profiler/profile_runner.py +1 -1
- nat/profiler/utils.py +1 -1
- nat/registry_handlers/local/local_handler.py +1 -1
- nat/registry_handlers/local/register_local.py +1 -1
- nat/registry_handlers/metadata_factory.py +1 -1
- nat/registry_handlers/package_utils.py +1 -1
- nat/registry_handlers/pypi/pypi_handler.py +1 -1
- nat/registry_handlers/pypi/register_pypi.py +1 -1
- nat/registry_handlers/register.py +1 -1
- nat/registry_handlers/registry_handler_base.py +1 -1
- nat/registry_handlers/rest/register_rest.py +1 -1
- nat/registry_handlers/rest/rest_handler.py +1 -1
- nat/registry_handlers/schemas/headers.py +1 -1
- nat/registry_handlers/schemas/package.py +1 -1
- nat/registry_handlers/schemas/publish.py +1 -1
- nat/registry_handlers/schemas/pull.py +1 -1
- nat/registry_handlers/schemas/remove.py +1 -1
- nat/registry_handlers/schemas/search.py +1 -1
- nat/registry_handlers/schemas/status.py +1 -1
- nat/retriever/interface.py +1 -1
- nat/retriever/milvus/__init__.py +1 -1
- nat/retriever/milvus/register.py +1 -1
- nat/retriever/milvus/retriever.py +1 -1
- nat/retriever/models.py +1 -1
- nat/retriever/nemo_retriever/__init__.py +1 -1
- nat/retriever/nemo_retriever/register.py +1 -1
- nat/retriever/nemo_retriever/retriever.py +5 -5
- nat/retriever/register.py +1 -1
- nat/runtime/__init__.py +1 -1
- nat/runtime/loader.py +10 -3
- nat/runtime/metrics.py +180 -0
- nat/runtime/runner.py +1 -5
- nat/runtime/session.py +451 -32
- nat/runtime/user_metadata.py +1 -1
- nat/settings/global_settings.py +1 -1
- nat/tool/chat_completion.py +1 -1
- nat/tool/code_execution/README.md +1 -1
- nat/tool/code_execution/code_sandbox.py +1 -1
- nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +1 -1
- nat/tool/code_execution/local_sandbox/__init__.py +1 -1
- nat/tool/code_execution/local_sandbox/local_sandbox_server.py +1 -1
- nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +1 -1
- nat/tool/code_execution/register.py +1 -1
- nat/tool/code_execution/utils.py +1 -1
- nat/tool/datetime_tools.py +1 -1
- nat/tool/document_search.py +1 -1
- nat/tool/github_tools.py +1 -1
- nat/tool/memory_tools/add_memory_tool.py +1 -1
- nat/tool/memory_tools/delete_memory_tool.py +1 -1
- nat/tool/memory_tools/get_memory_tool.py +1 -1
- nat/tool/nvidia_rag.py +2 -2
- nat/tool/register.py +1 -1
- nat/tool/retriever.py +1 -1
- nat/tool/server_tools.py +1 -1
- nat/utils/__init__.py +8 -5
- nat/utils/callable_utils.py +1 -1
- nat/utils/data_models/schema_validator.py +1 -1
- nat/utils/debugging_utils.py +1 -1
- nat/utils/decorators.py +1 -1
- nat/utils/dump_distro_mapping.py +1 -1
- nat/utils/exception_handlers/automatic_retries.py +3 -3
- nat/utils/exception_handlers/schemas.py +1 -1
- nat/utils/io/model_processing.py +1 -1
- nat/utils/io/supress_logs.py +33 -0
- nat/utils/io/yaml_tools.py +1 -1
- nat/utils/log_levels.py +1 -1
- nat/utils/log_utils.py +13 -1
- nat/utils/metadata_utils.py +1 -1
- nat/utils/optional_imports.py +1 -1
- nat/utils/producer_consumer_queue.py +1 -1
- nat/utils/reactive/base/observable_base.py +1 -1
- nat/utils/reactive/base/observer_base.py +1 -1
- nat/utils/reactive/base/subject_base.py +1 -1
- nat/utils/reactive/observable.py +1 -1
- nat/utils/reactive/observer.py +1 -1
- nat/utils/reactive/subject.py +1 -1
- nat/utils/reactive/subscription.py +1 -1
- nat/utils/responses_api.py +1 -1
- nat/utils/settings/global_settings.py +1 -1
- nat/utils/string_utils.py +1 -1
- nat/utils/type_converter.py +18 -5
- nat/utils/type_utils.py +1 -1
- nat/utils/url_utils.py +1 -1
- {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/METADATA +39 -14
- nvidia_nat-1.4.0a20260113.dist-info/RECORD +547 -0
- nvidia_nat-1.4.0a20260113.dist-info/entry_points.txt +38 -0
- nat/cli/commands/mcp/mcp.py +0 -986
- nat/front_ends/mcp/introspection_token_verifier.py +0 -73
- nat/front_ends/mcp/mcp_front_end_config.py +0 -109
- nat/front_ends/mcp/mcp_front_end_plugin.py +0 -155
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +0 -388
- nat/front_ends/mcp/memory_profiler.py +0 -320
- nat/front_ends/mcp/register.py +0 -27
- nat/front_ends/mcp/tool_converter.py +0 -321
- nvidia_nat-1.4.0a20251120.dist-info/RECORD +0 -488
- nvidia_nat-1.4.0a20251120.dist-info/entry_points.txt +0 -23
- {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/WHEEL +0 -0
- {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE.md +0 -0
- {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
"""Configuration for dynamic middleware."""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
from pydantic import BaseModel
|
|
20
|
+
from pydantic import Field
|
|
21
|
+
from pydantic import model_validator
|
|
22
|
+
|
|
23
|
+
from nat.data_models.component import ComponentGroup
|
|
24
|
+
from nat.data_models.component_ref import AuthenticationRef
|
|
25
|
+
from nat.data_models.component_ref import EmbedderRef
|
|
26
|
+
from nat.data_models.component_ref import LLMRef
|
|
27
|
+
from nat.data_models.component_ref import MemoryRef
|
|
28
|
+
from nat.data_models.component_ref import ObjectStoreRef
|
|
29
|
+
from nat.data_models.component_ref import RetrieverRef
|
|
30
|
+
from nat.data_models.middleware import FunctionMiddlewareBaseConfig
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class AllowedComponentFunctions(BaseModel):
|
|
34
|
+
"""Component functions allowed for auto-registration.
|
|
35
|
+
|
|
36
|
+
Default allowlists are provided for each component type. User-provided
|
|
37
|
+
values are automatically merged with defaults.
|
|
38
|
+
Set to None or omit to use only defaults.
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
llms: set[str] | None = Field(
|
|
42
|
+
default=None, description="Additional LLM functions that should be allowed to register with middleware.")
|
|
43
|
+
embedders: set[str] | None = Field(
|
|
44
|
+
default=None, description="Additional Embedder functions that should be allowed to register with middleware.")
|
|
45
|
+
retrievers: set[str] | None = Field(
|
|
46
|
+
default=None, description="Additional Retriever functions that should be allowed to register with middleware.")
|
|
47
|
+
memory: set[str] | None = Field(
|
|
48
|
+
default=None, description="Additional Memory functions that should be allowed to register with middleware.")
|
|
49
|
+
object_stores: set[str] | None = Field(
|
|
50
|
+
default=None,
|
|
51
|
+
description="Additional Object Store functions that should be allowed to register with middleware.")
|
|
52
|
+
authentication: set[str] | None = Field(
|
|
53
|
+
default=None,
|
|
54
|
+
description="Additional Authentication functions that should be allowed to register with middleware.")
|
|
55
|
+
|
|
56
|
+
@model_validator(mode='after')
|
|
57
|
+
def merge_with_defaults(self):
|
|
58
|
+
"""Merge user-provided values with defaults from COMPONENT_FUNCTION_ALLOWLISTS."""
|
|
59
|
+
from nat.middleware.utils.workflow_inventory import COMPONENT_FUNCTION_ALLOWLISTS
|
|
60
|
+
|
|
61
|
+
def merge(component_group: ComponentGroup, user_set: set[str] | None) -> set[str]:
|
|
62
|
+
defaults = COMPONENT_FUNCTION_ALLOWLISTS[component_group]
|
|
63
|
+
if user_set is None:
|
|
64
|
+
return defaults.copy()
|
|
65
|
+
return defaults | user_set
|
|
66
|
+
|
|
67
|
+
self.llms = merge(ComponentGroup.LLMS, self.llms)
|
|
68
|
+
self.embedders = merge(ComponentGroup.EMBEDDERS, self.embedders)
|
|
69
|
+
self.retrievers = merge(ComponentGroup.RETRIEVERS, self.retrievers)
|
|
70
|
+
self.memory = merge(ComponentGroup.MEMORY, self.memory)
|
|
71
|
+
self.object_stores = merge(ComponentGroup.OBJECT_STORES, self.object_stores)
|
|
72
|
+
self.authentication = merge(ComponentGroup.AUTHENTICATION, self.authentication)
|
|
73
|
+
|
|
74
|
+
return self
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
class DynamicMiddlewareConfig(FunctionMiddlewareBaseConfig, name="dynamic_middleware"):
|
|
78
|
+
"""Configuration for dynamic middleware.
|
|
79
|
+
|
|
80
|
+
Controls which components and functions to intercept, and which policies to apply.
|
|
81
|
+
Supports explicit component references and auto-discovery flags.
|
|
82
|
+
"""
|
|
83
|
+
|
|
84
|
+
# === Component References ===
|
|
85
|
+
|
|
86
|
+
llms: list[LLMRef] | None = Field(default=None, description="LLMs to intercept")
|
|
87
|
+
|
|
88
|
+
embedders: list[EmbedderRef] | None = Field(default=None, description="Embedders component functions to intercept")
|
|
89
|
+
|
|
90
|
+
retrievers: list[RetrieverRef] | None = Field(default=None,
|
|
91
|
+
description="Retrievers component functions to intercept")
|
|
92
|
+
|
|
93
|
+
memory: list[MemoryRef] | None = Field(default=None, description="Memory component functions to intercept")
|
|
94
|
+
|
|
95
|
+
object_stores: list[ObjectStoreRef] | None = Field(default=None,
|
|
96
|
+
description="Object stores component functions to intercept")
|
|
97
|
+
|
|
98
|
+
auth_providers: list[AuthenticationRef] | None = Field(
|
|
99
|
+
default=None, description="Authentication providers component functions to intercept")
|
|
100
|
+
|
|
101
|
+
# === Component and Function Auto-Discovery Flags ===
|
|
102
|
+
|
|
103
|
+
register_llms: bool | None = Field(default=False,
|
|
104
|
+
description="Auto-discover and register all LLMs component functions")
|
|
105
|
+
|
|
106
|
+
register_embedders: bool | None = Field(default=False,
|
|
107
|
+
description="Auto-discover and register all embedders component functions")
|
|
108
|
+
|
|
109
|
+
register_retrievers: bool | None = Field(
|
|
110
|
+
default=False, description="Auto-discover and register all retrievers component functions")
|
|
111
|
+
|
|
112
|
+
register_memory: bool | None = Field(
|
|
113
|
+
default=False, description="Auto-discover and register all memory providers component functions")
|
|
114
|
+
|
|
115
|
+
register_object_stores: bool | None = Field(
|
|
116
|
+
default=False, description="Auto-discover and register all object stores component functions")
|
|
117
|
+
|
|
118
|
+
register_auth_providers: bool | None = Field(
|
|
119
|
+
default=False, description="Auto-discover and register all authentication providers component functions")
|
|
120
|
+
|
|
121
|
+
register_workflow_functions: bool | None = Field(default=False,
|
|
122
|
+
description="Auto-discover and register all workflow functions")
|
|
123
|
+
|
|
124
|
+
# === Enable/Disable ===
|
|
125
|
+
|
|
126
|
+
enabled: bool = Field(default=True, description="Whether this middleware is active")
|
|
127
|
+
|
|
128
|
+
# === Component Function Allowlists ===
|
|
129
|
+
|
|
130
|
+
allowed_component_functions: AllowedComponentFunctions | None = Field(
|
|
131
|
+
default=None,
|
|
132
|
+
description="Functions allowed for auto-registration. Omit to use defaults, provide to extend them")
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
"""Registration for dynamic middleware."""
|
|
16
|
+
|
|
17
|
+
from nat.builder.builder import Builder
|
|
18
|
+
from nat.cli.register_workflow import register_middleware
|
|
19
|
+
from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware
|
|
20
|
+
from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@register_middleware(config_type=DynamicMiddlewareConfig)
|
|
24
|
+
async def dynamic_middleware(config: DynamicMiddlewareConfig, builder: Builder):
|
|
25
|
+
"""Build a dynamic middleware from configuration.
|
|
26
|
+
|
|
27
|
+
Args:
|
|
28
|
+
config: The dynamic middleware configuration
|
|
29
|
+
builder: The workflow builder
|
|
30
|
+
|
|
31
|
+
Yields:
|
|
32
|
+
A configured dynamic middleware instance
|
|
33
|
+
"""
|
|
34
|
+
yield DynamicFunctionMiddleware(config=config, builder=builder)
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
2
|
# SPDX-License-Identifier: Apache-2.0
|
|
3
3
|
#
|
|
4
4
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
@@ -37,81 +37,250 @@ from typing import Any
|
|
|
37
37
|
from nat.middleware.middleware import CallNext
|
|
38
38
|
from nat.middleware.middleware import CallNextStream
|
|
39
39
|
from nat.middleware.middleware import FunctionMiddlewareContext
|
|
40
|
+
from nat.middleware.middleware import InvocationContext
|
|
40
41
|
from nat.middleware.middleware import Middleware
|
|
41
42
|
|
|
42
43
|
|
|
43
44
|
class FunctionMiddleware(Middleware):
|
|
44
|
-
"""
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
45
|
+
"""Base class for function middleware with pre/post-invoke hooks.
|
|
46
|
+
|
|
47
|
+
Middleware intercepts function calls and can:
|
|
48
|
+
- Transform inputs before execution (pre_invoke)
|
|
49
|
+
- Transform outputs after execution (post_invoke)
|
|
50
|
+
- Override function_middleware_invoke for full control
|
|
51
|
+
|
|
52
|
+
Lifecycle:
|
|
53
|
+
- Framework checks ``enabled`` property before calling any methods
|
|
54
|
+
- If disabled, middleware is skipped entirely (no methods called)
|
|
55
|
+
- Users do NOT need to check ``enabled`` in their implementations
|
|
56
|
+
|
|
57
|
+
Inherited abstract members that must be implemented:
|
|
58
|
+
- enabled: Property that returns whether middleware should run
|
|
59
|
+
- pre_invoke: Transform inputs before function execution
|
|
60
|
+
- post_invoke: Transform outputs after function execution
|
|
61
|
+
|
|
62
|
+
Context Flow:
|
|
63
|
+
- FunctionMiddlewareContext (frozen): Static function metadata only
|
|
64
|
+
- InvocationContext: Unified context for both pre and post invoke phases
|
|
65
|
+
- Pre-invoke: output is None, modify modified_args/modified_kwargs
|
|
66
|
+
- Post-invoke: output has the result, modify output to transform
|
|
67
|
+
|
|
68
|
+
Example::
|
|
69
|
+
|
|
70
|
+
class LoggingMiddleware(FunctionMiddleware):
|
|
71
|
+
def __init__(self, config: LoggingConfig):
|
|
72
|
+
super().__init__()
|
|
73
|
+
self._config = config
|
|
74
|
+
|
|
75
|
+
@property
|
|
76
|
+
def enabled(self) -> bool:
|
|
77
|
+
return self._config.enabled
|
|
78
|
+
|
|
79
|
+
async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
80
|
+
logger.info(f"Calling {context.function_context.name} with {context.modified_args}")
|
|
81
|
+
logger.info(f"Original args: {context.original_args}")
|
|
82
|
+
return None # Pass through unchanged
|
|
83
|
+
|
|
84
|
+
async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
85
|
+
logger.info(f"Result: {context.output}")
|
|
86
|
+
return None # Pass through unchanged
|
|
50
87
|
"""
|
|
51
88
|
|
|
52
|
-
|
|
89
|
+
@property
|
|
90
|
+
def enabled(self) -> bool:
|
|
91
|
+
"""Check if this middleware is enabled.
|
|
92
|
+
|
|
93
|
+
Returns:
|
|
94
|
+
True if the middleware should be applied, False otherwise.
|
|
95
|
+
Default implementation always returns True.
|
|
96
|
+
"""
|
|
97
|
+
return True
|
|
98
|
+
|
|
99
|
+
async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
100
|
+
"""Pre-invocation hook called before the function is invoked.
|
|
101
|
+
|
|
102
|
+
Args:
|
|
103
|
+
context: Invocation context containing function metadata and args
|
|
104
|
+
|
|
105
|
+
Returns:
|
|
106
|
+
InvocationContext if modified, or None to pass through unchanged.
|
|
107
|
+
Default implementation does nothing.
|
|
108
|
+
"""
|
|
109
|
+
del context # Unused by default implementation
|
|
110
|
+
return None
|
|
111
|
+
|
|
112
|
+
async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
113
|
+
"""Post-invocation hook called after the function returns.
|
|
114
|
+
|
|
115
|
+
Args:
|
|
116
|
+
context: Invocation context containing function metadata, args, and output
|
|
117
|
+
|
|
118
|
+
Returns:
|
|
119
|
+
InvocationContext if modified, or None to pass through unchanged.
|
|
120
|
+
Default implementation does nothing.
|
|
121
|
+
"""
|
|
122
|
+
del context # Unused by default implementation
|
|
123
|
+
return None
|
|
124
|
+
|
|
125
|
+
# ==================== Middleware Delegation ====================
|
|
126
|
+
async def middleware_invoke(self,
|
|
127
|
+
*args: Any,
|
|
128
|
+
call_next: CallNext,
|
|
129
|
+
context: FunctionMiddlewareContext,
|
|
130
|
+
**kwargs: Any) -> Any:
|
|
53
131
|
"""Delegate to function_middleware_invoke for function-specific handling."""
|
|
54
|
-
return await self.function_middleware_invoke(
|
|
132
|
+
return await self.function_middleware_invoke(*args, call_next=call_next, context=context, **kwargs)
|
|
55
133
|
|
|
56
|
-
async def middleware_stream(self,
|
|
57
|
-
|
|
134
|
+
async def middleware_stream(self,
|
|
135
|
+
*args: Any,
|
|
136
|
+
call_next: CallNextStream,
|
|
137
|
+
context: FunctionMiddlewareContext,
|
|
138
|
+
**kwargs: Any) -> AsyncIterator[Any]:
|
|
58
139
|
"""Delegate to function_middleware_stream for function-specific handling."""
|
|
59
|
-
async for chunk in self.function_middleware_stream(
|
|
140
|
+
async for chunk in self.function_middleware_stream(*args, call_next=call_next, context=context, **kwargs):
|
|
60
141
|
yield chunk
|
|
61
142
|
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
143
|
+
# ==================== Orchestration ====================
|
|
144
|
+
|
|
145
|
+
async def function_middleware_invoke(
|
|
146
|
+
self,
|
|
147
|
+
*args: Any,
|
|
148
|
+
call_next: CallNext,
|
|
149
|
+
context: FunctionMiddlewareContext,
|
|
150
|
+
**kwargs: Any,
|
|
151
|
+
) -> Any:
|
|
152
|
+
"""Execute middleware hooks around function call.
|
|
153
|
+
|
|
154
|
+
Default implementation orchestrates: pre_invoke → call_next → post_invoke
|
|
155
|
+
|
|
156
|
+
Override for full control over execution flow (e.g., caching,
|
|
157
|
+
retry logic, conditional execution).
|
|
158
|
+
|
|
159
|
+
Note: Framework checks ``enabled`` before calling this method.
|
|
160
|
+
You do NOT need to check ``enabled`` yourself.
|
|
65
161
|
|
|
66
162
|
Args:
|
|
67
|
-
|
|
68
|
-
call_next: Callable to invoke
|
|
69
|
-
context:
|
|
163
|
+
args: Positional arguments for the function (first arg is typically the input value).
|
|
164
|
+
call_next: Callable to invoke next middleware or target function.
|
|
165
|
+
context: Static function metadata.
|
|
166
|
+
kwargs: Keyword arguments for the function.
|
|
70
167
|
|
|
71
168
|
Returns:
|
|
72
|
-
The (potentially
|
|
73
|
-
|
|
74
|
-
The default implementation simply delegates to ``call_next``. Override this
|
|
75
|
-
in subclasses to add function-specific preprocessing and postprocessing.
|
|
169
|
+
The (potentially transformed) function output.
|
|
76
170
|
"""
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
171
|
+
# Build invocation context with frozen originals + mutable current
|
|
172
|
+
# output starts as None (pre-invoke phase)
|
|
173
|
+
ctx = InvocationContext(
|
|
174
|
+
function_context=context,
|
|
175
|
+
original_args=args,
|
|
176
|
+
original_kwargs=dict(kwargs),
|
|
177
|
+
modified_args=args,
|
|
178
|
+
modified_kwargs=dict(kwargs),
|
|
179
|
+
output=None,
|
|
180
|
+
)
|
|
181
|
+
|
|
182
|
+
# Pre-invoke transformation (output is None at this phase)
|
|
183
|
+
result = await self.pre_invoke(ctx)
|
|
184
|
+
if result is not None:
|
|
185
|
+
ctx = result
|
|
186
|
+
|
|
187
|
+
# Execute function with (potentially modified) args/kwargs
|
|
188
|
+
ctx.output = await call_next(*ctx.modified_args, **ctx.modified_kwargs)
|
|
189
|
+
|
|
190
|
+
# Post-invoke transformation (output now has the result)
|
|
191
|
+
result = await self.post_invoke(ctx)
|
|
192
|
+
if result is not None:
|
|
193
|
+
ctx = result
|
|
194
|
+
|
|
195
|
+
return ctx.output
|
|
196
|
+
|
|
197
|
+
async def function_middleware_stream(
|
|
198
|
+
self,
|
|
199
|
+
*args: Any,
|
|
200
|
+
call_next: CallNextStream,
|
|
201
|
+
context: FunctionMiddlewareContext,
|
|
202
|
+
**kwargs: Any,
|
|
203
|
+
) -> AsyncIterator[Any]:
|
|
204
|
+
"""Execute middleware hooks around streaming function call.
|
|
205
|
+
|
|
206
|
+
Pre-invoke runs once before streaming starts.
|
|
207
|
+
Post-invoke runs per-chunk as they stream through.
|
|
208
|
+
|
|
209
|
+
Override for custom streaming behavior (e.g., buffering,
|
|
210
|
+
aggregation, chunk filtering).
|
|
211
|
+
|
|
212
|
+
Note: Framework checks ``enabled`` before calling this method.
|
|
213
|
+
You do NOT need to check ``enabled`` yourself.
|
|
84
214
|
|
|
85
215
|
Args:
|
|
86
|
-
|
|
87
|
-
call_next: Callable to invoke
|
|
88
|
-
context:
|
|
216
|
+
args: Positional arguments for the function (first arg is typically the input value).
|
|
217
|
+
call_next: Callable to invoke next middleware or target stream.
|
|
218
|
+
context: Static function metadata.
|
|
219
|
+
kwargs: Keyword arguments for the function.
|
|
89
220
|
|
|
90
221
|
Yields:
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
The default implementation forwards to ``call_next`` untouched. Override this
|
|
94
|
-
in subclasses to add function-specific preprocessing and chunk transformations.
|
|
222
|
+
Stream chunks (potentially transformed by post_invoke).
|
|
95
223
|
"""
|
|
96
|
-
|
|
97
|
-
|
|
224
|
+
# Build invocation context with frozen originals + mutable current
|
|
225
|
+
# output starts as None (pre-invoke phase)
|
|
226
|
+
ctx = InvocationContext(
|
|
227
|
+
function_context=context,
|
|
228
|
+
original_args=args,
|
|
229
|
+
original_kwargs=dict(kwargs),
|
|
230
|
+
modified_args=args,
|
|
231
|
+
modified_kwargs=dict(kwargs),
|
|
232
|
+
output=None,
|
|
233
|
+
)
|
|
234
|
+
|
|
235
|
+
# Pre-invoke transformation (once before streaming)
|
|
236
|
+
result = await self.pre_invoke(ctx)
|
|
237
|
+
if result is not None:
|
|
238
|
+
ctx = result
|
|
239
|
+
|
|
240
|
+
# Stream with per-chunk post-invoke
|
|
241
|
+
async for chunk in call_next(*ctx.modified_args, **ctx.modified_kwargs):
|
|
242
|
+
# Set output for this chunk
|
|
243
|
+
ctx.output = chunk
|
|
244
|
+
|
|
245
|
+
# Post-invoke transformation per chunk
|
|
246
|
+
result = await self.post_invoke(ctx)
|
|
247
|
+
if result is not None:
|
|
248
|
+
ctx = result
|
|
249
|
+
|
|
250
|
+
yield ctx.output
|
|
98
251
|
|
|
99
252
|
|
|
100
253
|
class FunctionMiddlewareChain:
|
|
101
|
-
"""
|
|
254
|
+
"""Composes middleware into an execution chain.
|
|
255
|
+
|
|
256
|
+
The chain builder checks each middleware's ``enabled`` property.
|
|
257
|
+
Disabled middleware is skipped entirely—no methods are called.
|
|
102
258
|
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
259
|
+
Execution order:
|
|
260
|
+
- Pre-invoke: first middleware → last middleware → function
|
|
261
|
+
- Post-invoke: function → last middleware → first middleware
|
|
262
|
+
|
|
263
|
+
Context:
|
|
264
|
+
- FunctionMiddlewareContext contains only static function metadata
|
|
265
|
+
- Original args/kwargs are captured by the orchestration layer
|
|
266
|
+
- Middleware receives InvocationContext with frozen originals and mutable args/output
|
|
106
267
|
"""
|
|
107
268
|
|
|
108
|
-
def __init__(self, *, middleware: Sequence[
|
|
269
|
+
def __init__(self, *, middleware: Sequence[FunctionMiddleware], context: FunctionMiddlewareContext) -> None:
|
|
270
|
+
"""Initialize the middleware chain.
|
|
271
|
+
|
|
272
|
+
Args:
|
|
273
|
+
middleware: Sequence of middleware to chain (order matters)
|
|
274
|
+
context: Static function metadata
|
|
275
|
+
"""
|
|
109
276
|
self._middleware = tuple(middleware)
|
|
110
277
|
self._context = context
|
|
111
278
|
|
|
112
279
|
def build_single(self, final_call: CallNext) -> CallNext:
|
|
113
280
|
"""Build the middleware chain for single-output invocations.
|
|
114
281
|
|
|
282
|
+
Disabled middleware (enabled=False) is skipped entirely.
|
|
283
|
+
|
|
115
284
|
Args:
|
|
116
285
|
final_call: The final function to call (the actual function implementation)
|
|
117
286
|
|
|
@@ -121,18 +290,28 @@ class FunctionMiddlewareChain:
|
|
|
121
290
|
call = final_call
|
|
122
291
|
|
|
123
292
|
for mw in reversed(self._middleware):
|
|
293
|
+
# Framework-enforced: skip disabled middleware
|
|
294
|
+
if not mw.enabled:
|
|
295
|
+
continue
|
|
296
|
+
|
|
124
297
|
call_next = call
|
|
125
298
|
|
|
126
|
-
async def wrapped(
|
|
127
|
-
|
|
299
|
+
async def wrapped(*args: Any,
|
|
300
|
+
_middleware: FunctionMiddleware = mw,
|
|
301
|
+
_call_next: CallNext = call_next,
|
|
302
|
+
_context: FunctionMiddlewareContext = self._context,
|
|
303
|
+
**kwargs: Any) -> Any:
|
|
304
|
+
return await _middleware.middleware_invoke(*args, call_next=_call_next, context=_context, **kwargs)
|
|
128
305
|
|
|
129
|
-
call = wrapped
|
|
306
|
+
call = wrapped # type: ignore[assignment]
|
|
130
307
|
|
|
131
308
|
return call
|
|
132
309
|
|
|
133
310
|
def build_stream(self, final_call: CallNextStream) -> CallNextStream:
|
|
134
311
|
"""Build the middleware chain for streaming invocations.
|
|
135
312
|
|
|
313
|
+
Disabled middleware (enabled=False) is skipped entirely.
|
|
314
|
+
|
|
136
315
|
Args:
|
|
137
316
|
final_call: The final function to call (the actual function implementation)
|
|
138
317
|
|
|
@@ -142,16 +321,21 @@ class FunctionMiddlewareChain:
|
|
|
142
321
|
call = final_call
|
|
143
322
|
|
|
144
323
|
for mw in reversed(self._middleware):
|
|
324
|
+
if not mw.enabled:
|
|
325
|
+
continue
|
|
326
|
+
|
|
145
327
|
call_next = call
|
|
146
328
|
|
|
147
|
-
async def wrapped(
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
329
|
+
async def wrapped(*args: Any,
|
|
330
|
+
_middleware: FunctionMiddleware = mw,
|
|
331
|
+
_call_next: CallNextStream = call_next,
|
|
332
|
+
_context: FunctionMiddlewareContext = self._context,
|
|
333
|
+
**kwargs: Any) -> AsyncIterator[Any]:
|
|
334
|
+
stream = _middleware.middleware_stream(*args, call_next=_call_next, context=_context, **kwargs)
|
|
335
|
+
async for chunk in stream:
|
|
152
336
|
yield chunk
|
|
153
337
|
|
|
154
|
-
call = wrapped
|
|
338
|
+
call = wrapped # type: ignore[assignment]
|
|
155
339
|
|
|
156
340
|
return call
|
|
157
341
|
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import logging
|
|
17
|
+
|
|
18
|
+
from nat.builder.builder import Builder
|
|
19
|
+
from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware
|
|
20
|
+
from nat.middleware.logging.logging_middleware_config import LoggingMiddlewareConfig
|
|
21
|
+
from nat.middleware.middleware import InvocationContext
|
|
22
|
+
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class LoggingMiddleware(DynamicFunctionMiddleware):
|
|
27
|
+
"""Middleware that logs intercepted function inputs/outputs.
|
|
28
|
+
|
|
29
|
+
This middleware extends DynamicFunctionMiddleware to get automatic chain
|
|
30
|
+
orchestration and dynamic discovery features. Custom logic is implemented
|
|
31
|
+
through the pre_invoke and post_invoke hooks.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(self, config: LoggingMiddlewareConfig, builder: Builder):
|
|
35
|
+
"""Initialize logging middleware.
|
|
36
|
+
|
|
37
|
+
Args:
|
|
38
|
+
config: Logging middleware configuration
|
|
39
|
+
builder: Workflow builder
|
|
40
|
+
"""
|
|
41
|
+
super().__init__(config=config, builder=builder)
|
|
42
|
+
self._config: LoggingMiddlewareConfig = config
|
|
43
|
+
|
|
44
|
+
async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
45
|
+
"""Log inputs before function execution.
|
|
46
|
+
|
|
47
|
+
Returns:
|
|
48
|
+
None to pass through unchanged (logging only, no modification)
|
|
49
|
+
"""
|
|
50
|
+
log_level = getattr(logging, self._config.log_level.upper(), logging.INFO)
|
|
51
|
+
logger.log(log_level, f"Calling {context.function_context.name} with args: {context.modified_args}")
|
|
52
|
+
|
|
53
|
+
# Log if args were modified by prior middleware
|
|
54
|
+
if context.modified_args != context.original_args:
|
|
55
|
+
logger.log(log_level, f" (original args were: {context.original_args})")
|
|
56
|
+
|
|
57
|
+
return None # Pass through unchanged
|
|
58
|
+
|
|
59
|
+
async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
|
|
60
|
+
"""Log outputs after function execution.
|
|
61
|
+
|
|
62
|
+
Returns:
|
|
63
|
+
None to pass through unchanged (logging only, no modification)
|
|
64
|
+
"""
|
|
65
|
+
log_level = getattr(logging, self._config.log_level.upper(), logging.INFO)
|
|
66
|
+
logger.log(log_level, f"Function {context.function_context.name} returned: {context.output}")
|
|
67
|
+
return None # Pass through unchanged
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from pydantic import Field
|
|
17
|
+
|
|
18
|
+
from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class LoggingMiddlewareConfig(DynamicMiddlewareConfig, name="logging_middleware"):
|
|
22
|
+
"""Configuration for logging middleware.
|
|
23
|
+
|
|
24
|
+
Extends DynamicMiddlewareConfig to inherit all dynamic discovery features
|
|
25
|
+
(register_llms, register_workflow_functions, and so on) and the enabled toggle.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
log_level: str = Field(default="INFO", description="Logging level (DEBUG, INFO, WARNING, ERROR)")
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from nat.builder.builder import Builder
|
|
17
|
+
from nat.cli.register_workflow import register_middleware
|
|
18
|
+
from nat.middleware.logging.logging_middleware import LoggingMiddleware
|
|
19
|
+
from nat.middleware.logging.logging_middleware_config import LoggingMiddlewareConfig
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@register_middleware(config_type=LoggingMiddlewareConfig)
|
|
23
|
+
async def logging_middleware(config: LoggingMiddlewareConfig, builder: Builder):
|
|
24
|
+
"""Build a logging middleware from configuration.
|
|
25
|
+
|
|
26
|
+
Args:
|
|
27
|
+
config: The logging middleware configuration
|
|
28
|
+
builder: The workflow builder
|
|
29
|
+
|
|
30
|
+
Yields:
|
|
31
|
+
A configured logging middleware instance
|
|
32
|
+
"""
|
|
33
|
+
yield LoggingMiddleware(config=config, builder=builder)
|