nvidia-nat 1.2.0__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 +66 -0
- nat/agent/__init__.py +0 -0
- nat/agent/base.py +256 -0
- nat/agent/dual_node.py +67 -0
- nat/agent/react_agent/__init__.py +0 -0
- nat/agent/react_agent/agent.py +363 -0
- nat/agent/react_agent/output_parser.py +104 -0
- nat/agent/react_agent/prompt.py +44 -0
- nat/agent/react_agent/register.py +149 -0
- nat/agent/reasoning_agent/__init__.py +0 -0
- nat/agent/reasoning_agent/reasoning_agent.py +225 -0
- nat/agent/register.py +23 -0
- nat/agent/rewoo_agent/__init__.py +0 -0
- nat/agent/rewoo_agent/agent.py +415 -0
- nat/agent/rewoo_agent/prompt.py +110 -0
- nat/agent/rewoo_agent/register.py +157 -0
- nat/agent/tool_calling_agent/__init__.py +0 -0
- nat/agent/tool_calling_agent/agent.py +119 -0
- nat/agent/tool_calling_agent/register.py +106 -0
- nat/authentication/__init__.py +14 -0
- nat/authentication/api_key/__init__.py +14 -0
- nat/authentication/api_key/api_key_auth_provider.py +96 -0
- nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
- nat/authentication/api_key/register.py +26 -0
- nat/authentication/exceptions/__init__.py +14 -0
- nat/authentication/exceptions/api_key_exceptions.py +38 -0
- nat/authentication/http_basic_auth/__init__.py +0 -0
- nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
- nat/authentication/http_basic_auth/register.py +30 -0
- nat/authentication/interfaces.py +93 -0
- nat/authentication/oauth2/__init__.py +14 -0
- nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
- nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
- nat/authentication/oauth2/register.py +25 -0
- nat/authentication/register.py +21 -0
- nat/builder/__init__.py +0 -0
- nat/builder/builder.py +285 -0
- nat/builder/component_utils.py +316 -0
- nat/builder/context.py +270 -0
- nat/builder/embedder.py +24 -0
- nat/builder/eval_builder.py +161 -0
- nat/builder/evaluator.py +29 -0
- nat/builder/framework_enum.py +24 -0
- nat/builder/front_end.py +73 -0
- nat/builder/function.py +344 -0
- nat/builder/function_base.py +380 -0
- nat/builder/function_info.py +627 -0
- nat/builder/intermediate_step_manager.py +174 -0
- nat/builder/llm.py +25 -0
- nat/builder/retriever.py +25 -0
- nat/builder/user_interaction_manager.py +78 -0
- nat/builder/workflow.py +148 -0
- nat/builder/workflow_builder.py +1117 -0
- nat/cli/__init__.py +14 -0
- nat/cli/cli_utils/__init__.py +0 -0
- nat/cli/cli_utils/config_override.py +231 -0
- nat/cli/cli_utils/validation.py +37 -0
- nat/cli/commands/__init__.py +0 -0
- nat/cli/commands/configure/__init__.py +0 -0
- nat/cli/commands/configure/channel/__init__.py +0 -0
- nat/cli/commands/configure/channel/add.py +28 -0
- nat/cli/commands/configure/channel/channel.py +34 -0
- nat/cli/commands/configure/channel/remove.py +30 -0
- nat/cli/commands/configure/channel/update.py +30 -0
- nat/cli/commands/configure/configure.py +33 -0
- nat/cli/commands/evaluate.py +139 -0
- nat/cli/commands/info/__init__.py +14 -0
- nat/cli/commands/info/info.py +37 -0
- nat/cli/commands/info/list_channels.py +32 -0
- nat/cli/commands/info/list_components.py +129 -0
- nat/cli/commands/info/list_mcp.py +304 -0
- nat/cli/commands/registry/__init__.py +14 -0
- nat/cli/commands/registry/publish.py +88 -0
- nat/cli/commands/registry/pull.py +118 -0
- nat/cli/commands/registry/registry.py +36 -0
- nat/cli/commands/registry/remove.py +108 -0
- nat/cli/commands/registry/search.py +155 -0
- nat/cli/commands/sizing/__init__.py +14 -0
- nat/cli/commands/sizing/calc.py +297 -0
- nat/cli/commands/sizing/sizing.py +27 -0
- nat/cli/commands/start.py +246 -0
- nat/cli/commands/uninstall.py +81 -0
- nat/cli/commands/validate.py +47 -0
- nat/cli/commands/workflow/__init__.py +14 -0
- nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
- nat/cli/commands/workflow/templates/config.yml.j2 +16 -0
- nat/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
- nat/cli/commands/workflow/templates/register.py.j2 +5 -0
- nat/cli/commands/workflow/templates/workflow.py.j2 +36 -0
- nat/cli/commands/workflow/workflow.py +37 -0
- nat/cli/commands/workflow/workflow_commands.py +317 -0
- nat/cli/entrypoint.py +135 -0
- nat/cli/main.py +57 -0
- nat/cli/register_workflow.py +488 -0
- nat/cli/type_registry.py +1000 -0
- nat/data_models/__init__.py +14 -0
- nat/data_models/api_server.py +716 -0
- nat/data_models/authentication.py +231 -0
- nat/data_models/common.py +171 -0
- nat/data_models/component.py +58 -0
- nat/data_models/component_ref.py +168 -0
- nat/data_models/config.py +410 -0
- nat/data_models/dataset_handler.py +169 -0
- nat/data_models/discovery_metadata.py +305 -0
- nat/data_models/embedder.py +27 -0
- nat/data_models/evaluate.py +127 -0
- nat/data_models/evaluator.py +26 -0
- nat/data_models/front_end.py +26 -0
- nat/data_models/function.py +30 -0
- nat/data_models/function_dependencies.py +72 -0
- nat/data_models/interactive.py +246 -0
- nat/data_models/intermediate_step.py +302 -0
- nat/data_models/invocation_node.py +38 -0
- nat/data_models/llm.py +27 -0
- nat/data_models/logging.py +26 -0
- nat/data_models/memory.py +27 -0
- nat/data_models/object_store.py +44 -0
- nat/data_models/profiler.py +54 -0
- nat/data_models/registry_handler.py +26 -0
- nat/data_models/retriever.py +30 -0
- nat/data_models/retry_mixin.py +35 -0
- nat/data_models/span.py +190 -0
- nat/data_models/step_adaptor.py +64 -0
- nat/data_models/streaming.py +33 -0
- nat/data_models/swe_bench_model.py +54 -0
- nat/data_models/telemetry_exporter.py +26 -0
- nat/data_models/ttc_strategy.py +30 -0
- nat/embedder/__init__.py +0 -0
- nat/embedder/nim_embedder.py +59 -0
- nat/embedder/openai_embedder.py +43 -0
- nat/embedder/register.py +22 -0
- nat/eval/__init__.py +14 -0
- nat/eval/config.py +60 -0
- nat/eval/dataset_handler/__init__.py +0 -0
- nat/eval/dataset_handler/dataset_downloader.py +106 -0
- nat/eval/dataset_handler/dataset_filter.py +52 -0
- nat/eval/dataset_handler/dataset_handler.py +367 -0
- nat/eval/evaluate.py +510 -0
- nat/eval/evaluator/__init__.py +14 -0
- nat/eval/evaluator/base_evaluator.py +77 -0
- nat/eval/evaluator/evaluator_model.py +45 -0
- nat/eval/intermediate_step_adapter.py +99 -0
- nat/eval/rag_evaluator/__init__.py +0 -0
- nat/eval/rag_evaluator/evaluate.py +178 -0
- nat/eval/rag_evaluator/register.py +143 -0
- nat/eval/register.py +23 -0
- nat/eval/remote_workflow.py +133 -0
- nat/eval/runners/__init__.py +14 -0
- nat/eval/runners/config.py +39 -0
- nat/eval/runners/multi_eval_runner.py +54 -0
- nat/eval/runtime_event_subscriber.py +52 -0
- nat/eval/swe_bench_evaluator/__init__.py +0 -0
- nat/eval/swe_bench_evaluator/evaluate.py +215 -0
- nat/eval/swe_bench_evaluator/register.py +36 -0
- nat/eval/trajectory_evaluator/__init__.py +0 -0
- nat/eval/trajectory_evaluator/evaluate.py +75 -0
- nat/eval/trajectory_evaluator/register.py +40 -0
- nat/eval/tunable_rag_evaluator/__init__.py +0 -0
- nat/eval/tunable_rag_evaluator/evaluate.py +245 -0
- nat/eval/tunable_rag_evaluator/register.py +52 -0
- nat/eval/usage_stats.py +41 -0
- nat/eval/utils/__init__.py +0 -0
- nat/eval/utils/output_uploader.py +140 -0
- nat/eval/utils/tqdm_position_registry.py +40 -0
- nat/eval/utils/weave_eval.py +184 -0
- nat/experimental/__init__.py +0 -0
- nat/experimental/decorators/__init__.py +0 -0
- nat/experimental/decorators/experimental_warning_decorator.py +134 -0
- nat/experimental/test_time_compute/__init__.py +0 -0
- nat/experimental/test_time_compute/editing/__init__.py +0 -0
- nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
- nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
- nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
- nat/experimental/test_time_compute/functions/__init__.py +0 -0
- nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
- nat/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
- nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
- nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
- nat/experimental/test_time_compute/models/__init__.py +0 -0
- nat/experimental/test_time_compute/models/editor_config.py +132 -0
- nat/experimental/test_time_compute/models/scoring_config.py +112 -0
- nat/experimental/test_time_compute/models/search_config.py +120 -0
- nat/experimental/test_time_compute/models/selection_config.py +154 -0
- nat/experimental/test_time_compute/models/stage_enums.py +43 -0
- nat/experimental/test_time_compute/models/strategy_base.py +66 -0
- nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
- nat/experimental/test_time_compute/models/ttc_item.py +48 -0
- nat/experimental/test_time_compute/register.py +36 -0
- nat/experimental/test_time_compute/scoring/__init__.py +0 -0
- nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
- nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
- nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
- nat/experimental/test_time_compute/search/__init__.py +0 -0
- nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
- nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
- nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
- nat/experimental/test_time_compute/selection/__init__.py +0 -0
- nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
- nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
- nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
- nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
- nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
- nat/front_ends/__init__.py +14 -0
- nat/front_ends/console/__init__.py +14 -0
- nat/front_ends/console/authentication_flow_handler.py +233 -0
- nat/front_ends/console/console_front_end_config.py +32 -0
- nat/front_ends/console/console_front_end_plugin.py +96 -0
- nat/front_ends/console/register.py +25 -0
- nat/front_ends/cron/__init__.py +14 -0
- nat/front_ends/fastapi/__init__.py +14 -0
- nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
- nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
- nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
- nat/front_ends/fastapi/fastapi_front_end_config.py +241 -0
- nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
- nat/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
- nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1087 -0
- nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
- nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
- nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
- nat/front_ends/fastapi/job_store.py +183 -0
- nat/front_ends/fastapi/main.py +72 -0
- nat/front_ends/fastapi/message_handler.py +320 -0
- nat/front_ends/fastapi/message_validator.py +352 -0
- nat/front_ends/fastapi/register.py +25 -0
- nat/front_ends/fastapi/response_helpers.py +195 -0
- nat/front_ends/fastapi/step_adaptor.py +319 -0
- nat/front_ends/mcp/__init__.py +14 -0
- nat/front_ends/mcp/mcp_front_end_config.py +36 -0
- nat/front_ends/mcp/mcp_front_end_plugin.py +81 -0
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +143 -0
- nat/front_ends/mcp/register.py +27 -0
- nat/front_ends/mcp/tool_converter.py +241 -0
- nat/front_ends/register.py +22 -0
- nat/front_ends/simple_base/__init__.py +14 -0
- nat/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
- nat/llm/__init__.py +0 -0
- nat/llm/aws_bedrock_llm.py +57 -0
- nat/llm/nim_llm.py +46 -0
- nat/llm/openai_llm.py +46 -0
- nat/llm/register.py +23 -0
- nat/llm/utils/__init__.py +14 -0
- nat/llm/utils/env_config_value.py +94 -0
- nat/llm/utils/error.py +17 -0
- nat/memory/__init__.py +20 -0
- nat/memory/interfaces.py +183 -0
- nat/memory/models.py +112 -0
- nat/meta/pypi.md +58 -0
- nat/object_store/__init__.py +20 -0
- nat/object_store/in_memory_object_store.py +76 -0
- nat/object_store/interfaces.py +84 -0
- nat/object_store/models.py +38 -0
- nat/object_store/register.py +20 -0
- nat/observability/__init__.py +14 -0
- nat/observability/exporter/__init__.py +14 -0
- nat/observability/exporter/base_exporter.py +449 -0
- nat/observability/exporter/exporter.py +78 -0
- nat/observability/exporter/file_exporter.py +33 -0
- nat/observability/exporter/processing_exporter.py +322 -0
- nat/observability/exporter/raw_exporter.py +52 -0
- nat/observability/exporter/span_exporter.py +288 -0
- nat/observability/exporter_manager.py +335 -0
- nat/observability/mixin/__init__.py +14 -0
- nat/observability/mixin/batch_config_mixin.py +26 -0
- nat/observability/mixin/collector_config_mixin.py +23 -0
- nat/observability/mixin/file_mixin.py +288 -0
- nat/observability/mixin/file_mode.py +23 -0
- nat/observability/mixin/resource_conflict_mixin.py +134 -0
- nat/observability/mixin/serialize_mixin.py +61 -0
- nat/observability/mixin/type_introspection_mixin.py +183 -0
- nat/observability/processor/__init__.py +14 -0
- nat/observability/processor/batching_processor.py +310 -0
- nat/observability/processor/callback_processor.py +42 -0
- nat/observability/processor/intermediate_step_serializer.py +28 -0
- nat/observability/processor/processor.py +71 -0
- nat/observability/register.py +96 -0
- nat/observability/utils/__init__.py +14 -0
- nat/observability/utils/dict_utils.py +236 -0
- nat/observability/utils/time_utils.py +31 -0
- nat/plugins/.namespace +1 -0
- nat/profiler/__init__.py +0 -0
- nat/profiler/calc/__init__.py +14 -0
- nat/profiler/calc/calc_runner.py +627 -0
- nat/profiler/calc/calculations.py +288 -0
- nat/profiler/calc/data_models.py +188 -0
- nat/profiler/calc/plot.py +345 -0
- nat/profiler/callbacks/__init__.py +0 -0
- nat/profiler/callbacks/agno_callback_handler.py +295 -0
- nat/profiler/callbacks/base_callback_class.py +20 -0
- nat/profiler/callbacks/langchain_callback_handler.py +290 -0
- nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
- nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
- nat/profiler/callbacks/token_usage_base_model.py +27 -0
- nat/profiler/data_frame_row.py +51 -0
- nat/profiler/data_models.py +24 -0
- nat/profiler/decorators/__init__.py +0 -0
- nat/profiler/decorators/framework_wrapper.py +131 -0
- nat/profiler/decorators/function_tracking.py +254 -0
- nat/profiler/forecasting/__init__.py +0 -0
- nat/profiler/forecasting/config.py +18 -0
- nat/profiler/forecasting/model_trainer.py +75 -0
- nat/profiler/forecasting/models/__init__.py +22 -0
- nat/profiler/forecasting/models/forecasting_base_model.py +40 -0
- nat/profiler/forecasting/models/linear_model.py +197 -0
- nat/profiler/forecasting/models/random_forest_regressor.py +269 -0
- nat/profiler/inference_metrics_model.py +28 -0
- nat/profiler/inference_optimization/__init__.py +0 -0
- nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
- nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
- nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
- nat/profiler/inference_optimization/data_models.py +386 -0
- nat/profiler/inference_optimization/experimental/__init__.py +0 -0
- nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
- nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
- nat/profiler/inference_optimization/llm_metrics.py +212 -0
- nat/profiler/inference_optimization/prompt_caching.py +163 -0
- nat/profiler/inference_optimization/token_uniqueness.py +107 -0
- nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
- nat/profiler/intermediate_property_adapter.py +102 -0
- nat/profiler/profile_runner.py +473 -0
- nat/profiler/utils.py +184 -0
- nat/registry_handlers/__init__.py +0 -0
- nat/registry_handlers/local/__init__.py +0 -0
- nat/registry_handlers/local/local_handler.py +176 -0
- nat/registry_handlers/local/register_local.py +37 -0
- nat/registry_handlers/metadata_factory.py +60 -0
- nat/registry_handlers/package_utils.py +571 -0
- nat/registry_handlers/pypi/__init__.py +0 -0
- nat/registry_handlers/pypi/pypi_handler.py +251 -0
- nat/registry_handlers/pypi/register_pypi.py +40 -0
- nat/registry_handlers/register.py +21 -0
- nat/registry_handlers/registry_handler_base.py +157 -0
- nat/registry_handlers/rest/__init__.py +0 -0
- nat/registry_handlers/rest/register_rest.py +56 -0
- nat/registry_handlers/rest/rest_handler.py +237 -0
- nat/registry_handlers/schemas/__init__.py +0 -0
- nat/registry_handlers/schemas/headers.py +42 -0
- nat/registry_handlers/schemas/package.py +68 -0
- nat/registry_handlers/schemas/publish.py +68 -0
- nat/registry_handlers/schemas/pull.py +82 -0
- nat/registry_handlers/schemas/remove.py +36 -0
- nat/registry_handlers/schemas/search.py +91 -0
- nat/registry_handlers/schemas/status.py +47 -0
- nat/retriever/__init__.py +0 -0
- nat/retriever/interface.py +41 -0
- nat/retriever/milvus/__init__.py +14 -0
- nat/retriever/milvus/register.py +81 -0
- nat/retriever/milvus/retriever.py +228 -0
- nat/retriever/models.py +77 -0
- nat/retriever/nemo_retriever/__init__.py +14 -0
- nat/retriever/nemo_retriever/register.py +60 -0
- nat/retriever/nemo_retriever/retriever.py +190 -0
- nat/retriever/register.py +22 -0
- nat/runtime/__init__.py +14 -0
- nat/runtime/loader.py +220 -0
- nat/runtime/runner.py +195 -0
- nat/runtime/session.py +162 -0
- nat/runtime/user_metadata.py +130 -0
- nat/settings/__init__.py +0 -0
- nat/settings/global_settings.py +318 -0
- nat/test/.namespace +1 -0
- nat/tool/__init__.py +0 -0
- nat/tool/chat_completion.py +74 -0
- nat/tool/code_execution/README.md +151 -0
- nat/tool/code_execution/__init__.py +0 -0
- nat/tool/code_execution/code_sandbox.py +267 -0
- nat/tool/code_execution/local_sandbox/.gitignore +1 -0
- nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
- nat/tool/code_execution/local_sandbox/__init__.py +13 -0
- nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
- nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
- nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
- nat/tool/code_execution/register.py +74 -0
- nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
- nat/tool/code_execution/utils.py +100 -0
- nat/tool/datetime_tools.py +42 -0
- nat/tool/document_search.py +141 -0
- nat/tool/github_tools/__init__.py +0 -0
- nat/tool/github_tools/create_github_commit.py +133 -0
- nat/tool/github_tools/create_github_issue.py +87 -0
- nat/tool/github_tools/create_github_pr.py +106 -0
- nat/tool/github_tools/get_github_file.py +106 -0
- nat/tool/github_tools/get_github_issue.py +166 -0
- nat/tool/github_tools/get_github_pr.py +256 -0
- nat/tool/github_tools/update_github_issue.py +100 -0
- nat/tool/mcp/__init__.py +14 -0
- nat/tool/mcp/exceptions.py +142 -0
- nat/tool/mcp/mcp_client.py +255 -0
- nat/tool/mcp/mcp_tool.py +96 -0
- nat/tool/memory_tools/__init__.py +0 -0
- nat/tool/memory_tools/add_memory_tool.py +79 -0
- nat/tool/memory_tools/delete_memory_tool.py +67 -0
- nat/tool/memory_tools/get_memory_tool.py +72 -0
- nat/tool/nvidia_rag.py +95 -0
- nat/tool/register.py +38 -0
- nat/tool/retriever.py +94 -0
- nat/tool/server_tools.py +66 -0
- nat/utils/__init__.py +0 -0
- nat/utils/data_models/__init__.py +0 -0
- nat/utils/data_models/schema_validator.py +58 -0
- nat/utils/debugging_utils.py +43 -0
- nat/utils/dump_distro_mapping.py +32 -0
- nat/utils/exception_handlers/__init__.py +0 -0
- nat/utils/exception_handlers/automatic_retries.py +289 -0
- nat/utils/exception_handlers/mcp.py +211 -0
- nat/utils/exception_handlers/schemas.py +114 -0
- nat/utils/io/__init__.py +0 -0
- nat/utils/io/model_processing.py +28 -0
- nat/utils/io/yaml_tools.py +119 -0
- nat/utils/log_utils.py +37 -0
- nat/utils/metadata_utils.py +74 -0
- nat/utils/optional_imports.py +142 -0
- nat/utils/producer_consumer_queue.py +178 -0
- nat/utils/reactive/__init__.py +0 -0
- nat/utils/reactive/base/__init__.py +0 -0
- nat/utils/reactive/base/observable_base.py +65 -0
- nat/utils/reactive/base/observer_base.py +55 -0
- nat/utils/reactive/base/subject_base.py +79 -0
- nat/utils/reactive/observable.py +59 -0
- nat/utils/reactive/observer.py +76 -0
- nat/utils/reactive/subject.py +131 -0
- nat/utils/reactive/subscription.py +49 -0
- nat/utils/settings/__init__.py +0 -0
- nat/utils/settings/global_settings.py +197 -0
- nat/utils/string_utils.py +38 -0
- nat/utils/type_converter.py +290 -0
- nat/utils/type_utils.py +484 -0
- nat/utils/url_utils.py +27 -0
- nvidia_nat-1.2.0.dist-info/METADATA +365 -0
- nvidia_nat-1.2.0.dist-info/RECORD +435 -0
- nvidia_nat-1.2.0.dist-info/WHEEL +5 -0
- nvidia_nat-1.2.0.dist-info/entry_points.txt +21 -0
- nvidia_nat-1.2.0.dist-info/licenses/LICENSE-3rd-party.txt +5478 -0
- nvidia_nat-1.2.0.dist-info/licenses/LICENSE.md +201 -0
- nvidia_nat-1.2.0.dist-info/top_level.txt +2 -0
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, 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.cli.register_workflow import register_front_end
|
|
17
|
+
from nat.data_models.config import Config
|
|
18
|
+
from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@register_front_end(config_type=FastApiFrontEndConfig)
|
|
22
|
+
async def register_fastapi_front_end(config: FastApiFrontEndConfig, full_config: Config):
|
|
23
|
+
from nat.front_ends.fastapi.fastapi_front_end_plugin import FastApiFrontEndPlugin
|
|
24
|
+
|
|
25
|
+
yield FastApiFrontEndPlugin(full_config=full_config)
|
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, 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 asyncio
|
|
17
|
+
import typing
|
|
18
|
+
from collections.abc import AsyncGenerator
|
|
19
|
+
|
|
20
|
+
from nat.data_models.api_server import ResponseIntermediateStep
|
|
21
|
+
from nat.data_models.api_server import ResponsePayloadOutput
|
|
22
|
+
from nat.data_models.api_server import ResponseSerializable
|
|
23
|
+
from nat.data_models.step_adaptor import StepAdaptorConfig
|
|
24
|
+
from nat.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
|
|
25
|
+
from nat.front_ends.fastapi.step_adaptor import StepAdaptor
|
|
26
|
+
from nat.runtime.session import SessionManager
|
|
27
|
+
from nat.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
async def generate_streaming_response_as_str(payload: typing.Any,
|
|
31
|
+
*,
|
|
32
|
+
session_manager: SessionManager,
|
|
33
|
+
streaming: bool,
|
|
34
|
+
step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
|
|
35
|
+
result_type: type | None = None,
|
|
36
|
+
output_type: type | None = None) -> AsyncGenerator[str]:
|
|
37
|
+
|
|
38
|
+
async for item in generate_streaming_response(payload,
|
|
39
|
+
session_manager=session_manager,
|
|
40
|
+
streaming=streaming,
|
|
41
|
+
step_adaptor=step_adaptor,
|
|
42
|
+
result_type=result_type,
|
|
43
|
+
output_type=output_type):
|
|
44
|
+
|
|
45
|
+
if (isinstance(item, ResponseSerializable)):
|
|
46
|
+
yield item.get_stream_data()
|
|
47
|
+
else:
|
|
48
|
+
raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
|
|
49
|
+
str(type(item)))
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
async def generate_streaming_response(payload: typing.Any,
|
|
53
|
+
*,
|
|
54
|
+
session_manager: SessionManager,
|
|
55
|
+
streaming: bool,
|
|
56
|
+
step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
|
|
57
|
+
result_type: type | None = None,
|
|
58
|
+
output_type: type | None = None) -> AsyncGenerator[ResponseSerializable]:
|
|
59
|
+
|
|
60
|
+
async with session_manager.run(payload) as runner:
|
|
61
|
+
|
|
62
|
+
q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
|
|
63
|
+
|
|
64
|
+
# Start the intermediate stream
|
|
65
|
+
intermediate_complete = await pull_intermediate(q, step_adaptor)
|
|
66
|
+
|
|
67
|
+
async def pull_result():
|
|
68
|
+
if session_manager.workflow.has_streaming_output and streaming:
|
|
69
|
+
async for chunk in runner.result_stream(to_type=output_type):
|
|
70
|
+
await q.put(chunk)
|
|
71
|
+
else:
|
|
72
|
+
result = await runner.result(to_type=result_type)
|
|
73
|
+
await q.put(runner.convert(result, output_type))
|
|
74
|
+
|
|
75
|
+
# Wait until the intermediate subscription is done before closing q
|
|
76
|
+
# But we have no direct "intermediate_done" reference here
|
|
77
|
+
# because it's encapsulated in pull_intermediate. So we can do:
|
|
78
|
+
# await some_event.wait()
|
|
79
|
+
# If needed. Alternatively, you can skip that if the intermediate
|
|
80
|
+
# subscriber won't block the main flow.
|
|
81
|
+
#
|
|
82
|
+
# For example, if you *need* to guarantee the subscriber is done before
|
|
83
|
+
# closing the queue, you can structure the code to store or return
|
|
84
|
+
# the 'intermediate_done' event from pull_intermediate.
|
|
85
|
+
#
|
|
86
|
+
|
|
87
|
+
await intermediate_complete.wait()
|
|
88
|
+
|
|
89
|
+
await q.close()
|
|
90
|
+
|
|
91
|
+
try:
|
|
92
|
+
# Start the result stream
|
|
93
|
+
asyncio.create_task(pull_result())
|
|
94
|
+
|
|
95
|
+
async for item in q:
|
|
96
|
+
|
|
97
|
+
if (isinstance(item, ResponseSerializable)):
|
|
98
|
+
yield item
|
|
99
|
+
else:
|
|
100
|
+
yield ResponsePayloadOutput(payload=item)
|
|
101
|
+
except Exception as e:
|
|
102
|
+
# Handle exceptions here
|
|
103
|
+
raise e
|
|
104
|
+
finally:
|
|
105
|
+
await q.close()
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
async def generate_single_response(
|
|
109
|
+
payload: typing.Any,
|
|
110
|
+
session_manager: SessionManager,
|
|
111
|
+
result_type: type | None = None,
|
|
112
|
+
) -> typing.Any:
|
|
113
|
+
if (not session_manager.workflow.has_single_output):
|
|
114
|
+
raise ValueError("Cannot get a single output value for streaming workflows")
|
|
115
|
+
|
|
116
|
+
async with session_manager.run(payload) as runner:
|
|
117
|
+
return await runner.result(to_type=result_type)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
async def generate_streaming_response_full(payload: typing.Any,
|
|
121
|
+
*,
|
|
122
|
+
session_manager: SessionManager,
|
|
123
|
+
streaming: bool,
|
|
124
|
+
result_type: type | None = None,
|
|
125
|
+
output_type: type | None = None,
|
|
126
|
+
filter_steps: str | None = None) -> AsyncGenerator[ResponseSerializable]:
|
|
127
|
+
"""
|
|
128
|
+
Similar to generate_streaming_response but provides raw ResponseIntermediateStep objects
|
|
129
|
+
without any step adaptor translations.
|
|
130
|
+
"""
|
|
131
|
+
# Parse filter_steps into a set of allowed types if provided
|
|
132
|
+
# Special case: if filter_steps is "none", suppress all steps
|
|
133
|
+
allowed_types = None
|
|
134
|
+
if filter_steps:
|
|
135
|
+
if filter_steps.lower() == "none":
|
|
136
|
+
allowed_types = set() # Empty set means no steps allowed
|
|
137
|
+
else:
|
|
138
|
+
allowed_types = set(filter_steps.split(','))
|
|
139
|
+
|
|
140
|
+
async with session_manager.run(payload) as runner:
|
|
141
|
+
q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
|
|
142
|
+
|
|
143
|
+
# Start the intermediate stream without step adaptor
|
|
144
|
+
intermediate_complete = await pull_intermediate(q, None)
|
|
145
|
+
|
|
146
|
+
async def pull_result():
|
|
147
|
+
if session_manager.workflow.has_streaming_output and streaming:
|
|
148
|
+
async for chunk in runner.result_stream(to_type=output_type):
|
|
149
|
+
await q.put(chunk)
|
|
150
|
+
else:
|
|
151
|
+
result = await runner.result(to_type=result_type)
|
|
152
|
+
await q.put(runner.convert(result, output_type))
|
|
153
|
+
|
|
154
|
+
await intermediate_complete.wait()
|
|
155
|
+
await q.close()
|
|
156
|
+
|
|
157
|
+
try:
|
|
158
|
+
# Start the result stream
|
|
159
|
+
asyncio.create_task(pull_result())
|
|
160
|
+
|
|
161
|
+
async for item in q:
|
|
162
|
+
if (isinstance(item, ResponseIntermediateStep)):
|
|
163
|
+
# Filter intermediate steps if filter_steps is provided
|
|
164
|
+
if allowed_types is None or item.type in allowed_types:
|
|
165
|
+
yield item
|
|
166
|
+
else:
|
|
167
|
+
yield ResponsePayloadOutput(payload=item)
|
|
168
|
+
except Exception as e:
|
|
169
|
+
# Handle exceptions here
|
|
170
|
+
raise e
|
|
171
|
+
finally:
|
|
172
|
+
await q.close()
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
async def generate_streaming_response_full_as_str(payload: typing.Any,
|
|
176
|
+
*,
|
|
177
|
+
session_manager: SessionManager,
|
|
178
|
+
streaming: bool,
|
|
179
|
+
result_type: type | None = None,
|
|
180
|
+
output_type: type | None = None,
|
|
181
|
+
filter_steps: str | None = None) -> AsyncGenerator[str]:
|
|
182
|
+
"""
|
|
183
|
+
Similar to generate_streaming_response but converts the response to a string format.
|
|
184
|
+
"""
|
|
185
|
+
async for item in generate_streaming_response_full(payload,
|
|
186
|
+
session_manager=session_manager,
|
|
187
|
+
streaming=streaming,
|
|
188
|
+
result_type=result_type,
|
|
189
|
+
output_type=output_type,
|
|
190
|
+
filter_steps=filter_steps):
|
|
191
|
+
if (isinstance(item, ResponseIntermediateStep) or isinstance(item, ResponsePayloadOutput)):
|
|
192
|
+
yield item.get_stream_data()
|
|
193
|
+
else:
|
|
194
|
+
raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
|
|
195
|
+
str(type(item)))
|
|
@@ -0,0 +1,319 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, 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 html
|
|
17
|
+
import logging
|
|
18
|
+
from functools import reduce
|
|
19
|
+
from textwrap import dedent
|
|
20
|
+
|
|
21
|
+
from nat.data_models.api_server import ResponseIntermediateStep
|
|
22
|
+
from nat.data_models.api_server import ResponseSerializable
|
|
23
|
+
from nat.data_models.intermediate_step import IntermediateStep
|
|
24
|
+
from nat.data_models.intermediate_step import IntermediateStepCategory
|
|
25
|
+
from nat.data_models.intermediate_step import IntermediateStepPayload
|
|
26
|
+
from nat.data_models.intermediate_step import IntermediateStepType
|
|
27
|
+
from nat.data_models.invocation_node import InvocationNode
|
|
28
|
+
from nat.data_models.step_adaptor import StepAdaptorConfig
|
|
29
|
+
from nat.data_models.step_adaptor import StepAdaptorMode
|
|
30
|
+
from nat.utils.type_utils import is_valid_json
|
|
31
|
+
|
|
32
|
+
logger = logging.getLogger(__name__)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class StepAdaptor:
|
|
36
|
+
|
|
37
|
+
def __init__(self, config: StepAdaptorConfig):
|
|
38
|
+
|
|
39
|
+
self._history: list[IntermediateStep] = []
|
|
40
|
+
self.config = config
|
|
41
|
+
|
|
42
|
+
def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool:
|
|
43
|
+
"""
|
|
44
|
+
Returns True if this intermediate step should be included (based on the config.mode).
|
|
45
|
+
"""
|
|
46
|
+
|
|
47
|
+
if config.mode == StepAdaptorMode.OFF:
|
|
48
|
+
return False
|
|
49
|
+
|
|
50
|
+
if config.mode == StepAdaptorMode.DEFAULT:
|
|
51
|
+
# default existing behavior: show LLM events + TOOL_END + FUNCTION events
|
|
52
|
+
if step.event_category == IntermediateStepCategory.LLM:
|
|
53
|
+
return True
|
|
54
|
+
if step.event_category == IntermediateStepCategory.TOOL:
|
|
55
|
+
return True
|
|
56
|
+
if step.event_category == IntermediateStepCategory.FUNCTION:
|
|
57
|
+
return True
|
|
58
|
+
return False
|
|
59
|
+
|
|
60
|
+
if config.mode == StepAdaptorMode.CUSTOM:
|
|
61
|
+
# pass only what the user explicitly listed
|
|
62
|
+
return step.event_type in config.custom_event_types
|
|
63
|
+
|
|
64
|
+
return False
|
|
65
|
+
|
|
66
|
+
def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
|
|
67
|
+
input_str: str | None = None
|
|
68
|
+
output_str: str | None = None
|
|
69
|
+
|
|
70
|
+
# Find the start in the history with matching run_id
|
|
71
|
+
start_step = next(
|
|
72
|
+
(x for x in self._history if x.event_type == IntermediateStepType.LLM_START and x.UUID == step.UUID), None)
|
|
73
|
+
|
|
74
|
+
if not start_step:
|
|
75
|
+
# If we don't have a start step, we can't do anything
|
|
76
|
+
return None
|
|
77
|
+
|
|
78
|
+
input_str = str(start_step.data.input)
|
|
79
|
+
|
|
80
|
+
if step.event_type == IntermediateStepType.LLM_NEW_TOKEN:
|
|
81
|
+
|
|
82
|
+
# Find all of the previous LLM chunks and concatenate them
|
|
83
|
+
output_str = reduce(
|
|
84
|
+
lambda x, y: x + y,
|
|
85
|
+
(str(x.data.chunk)
|
|
86
|
+
for x in self._history if x.event_type == IntermediateStepType.LLM_NEW_TOKEN and x.UUID == step.UUID),
|
|
87
|
+
"")
|
|
88
|
+
|
|
89
|
+
elif step.event_type == IntermediateStepType.LLM_END:
|
|
90
|
+
output_str = str(step.data.output)
|
|
91
|
+
|
|
92
|
+
if not input_str and not output_str:
|
|
93
|
+
return None
|
|
94
|
+
|
|
95
|
+
escaped_input = html.escape(input_str, quote=False)
|
|
96
|
+
|
|
97
|
+
# Dont use f-strings here because the payload is markdown and screws up the dedent
|
|
98
|
+
payload = dedent("""
|
|
99
|
+
**Input:**
|
|
100
|
+
```python
|
|
101
|
+
{input_value}
|
|
102
|
+
```
|
|
103
|
+
""").strip("\n").format(input_value=escaped_input)
|
|
104
|
+
|
|
105
|
+
if (output_str):
|
|
106
|
+
escaped_output = html.escape(output_str, quote=False) if output_str else ""
|
|
107
|
+
|
|
108
|
+
# Dont use f-strings here because the payload is markdown and screws up the dedent
|
|
109
|
+
payload = dedent("""
|
|
110
|
+
{payload}
|
|
111
|
+
|
|
112
|
+
**Output:**
|
|
113
|
+
{output_value}
|
|
114
|
+
""").strip("\n").format(payload=payload, output_value=escaped_output)
|
|
115
|
+
|
|
116
|
+
event = ResponseIntermediateStep(id=step.UUID,
|
|
117
|
+
name=step.name or "",
|
|
118
|
+
payload=payload,
|
|
119
|
+
parent_id=ancestry.function_id)
|
|
120
|
+
|
|
121
|
+
return event
|
|
122
|
+
|
|
123
|
+
def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
|
|
124
|
+
"""
|
|
125
|
+
Handles both TOOL_START and TOOL_END events
|
|
126
|
+
"""
|
|
127
|
+
input_str: str | None = None
|
|
128
|
+
output_str: str | None = None
|
|
129
|
+
|
|
130
|
+
# Find the start in the history with matching run_id
|
|
131
|
+
start_step = next(
|
|
132
|
+
(x for x in self._history if x.event_type == IntermediateStepType.TOOL_START and x.UUID == step.UUID), None)
|
|
133
|
+
|
|
134
|
+
if not start_step:
|
|
135
|
+
# If we don't have a start step, we can't do anything
|
|
136
|
+
return None
|
|
137
|
+
|
|
138
|
+
input_str = str(start_step.data.input)
|
|
139
|
+
|
|
140
|
+
if step.event_type == IntermediateStepType.TOOL_END:
|
|
141
|
+
output_str = str(step.data.output)
|
|
142
|
+
|
|
143
|
+
if not input_str and not output_str:
|
|
144
|
+
return None
|
|
145
|
+
|
|
146
|
+
escaped_input = html.escape(input_str, quote=False)
|
|
147
|
+
format_input_type = "json" if is_valid_json(escaped_input) else "python"
|
|
148
|
+
|
|
149
|
+
# Dont use f-strings here because the payload is markdown and screws up the dedent
|
|
150
|
+
payload = dedent("""
|
|
151
|
+
**Input:**
|
|
152
|
+
```{format_input_type}
|
|
153
|
+
{input_value}
|
|
154
|
+
```
|
|
155
|
+
""").strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
|
|
156
|
+
|
|
157
|
+
if output_str:
|
|
158
|
+
escaped_output = html.escape(output_str, quote=False)
|
|
159
|
+
format_output_type = "json" if is_valid_json(escaped_output) else "python"
|
|
160
|
+
|
|
161
|
+
# Dont use f-strings here because the payload is markdown and screws up the dedent
|
|
162
|
+
payload = dedent("""
|
|
163
|
+
{payload}
|
|
164
|
+
|
|
165
|
+
**Output:**
|
|
166
|
+
```{format_output_type}
|
|
167
|
+
{output_value}
|
|
168
|
+
```
|
|
169
|
+
""").strip("\n").format(payload=payload, output_value=escaped_output, format_output_type=format_output_type)
|
|
170
|
+
|
|
171
|
+
event = ResponseIntermediateStep(id=step.UUID,
|
|
172
|
+
name=f"Tool: {step.name}",
|
|
173
|
+
payload=payload,
|
|
174
|
+
parent_id=ancestry.function_id)
|
|
175
|
+
|
|
176
|
+
return event
|
|
177
|
+
|
|
178
|
+
def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
|
|
179
|
+
"""
|
|
180
|
+
Handles the FUNCTION_START and FUNCTION_END events
|
|
181
|
+
"""
|
|
182
|
+
input_str: str | None = None
|
|
183
|
+
output_str: str | None = None
|
|
184
|
+
|
|
185
|
+
if step.event_type == IntermediateStepType.FUNCTION_START:
|
|
186
|
+
# For function start events, display input data
|
|
187
|
+
if step.data and hasattr(step.data, 'input'):
|
|
188
|
+
input_str = str(step.data.input)
|
|
189
|
+
elif step.data:
|
|
190
|
+
input_str = str(step.data)
|
|
191
|
+
|
|
192
|
+
if not input_str:
|
|
193
|
+
return None
|
|
194
|
+
|
|
195
|
+
escaped_input = html.escape(input_str, quote=False)
|
|
196
|
+
format_input_type = "json" if is_valid_json(escaped_input) else "python"
|
|
197
|
+
|
|
198
|
+
# Create payload for function start
|
|
199
|
+
payload_str = dedent("""
|
|
200
|
+
**Function Input:**
|
|
201
|
+
```{format_input_type}
|
|
202
|
+
{input_value}
|
|
203
|
+
```
|
|
204
|
+
""").strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
|
|
205
|
+
|
|
206
|
+
event = ResponseIntermediateStep(id=step.UUID,
|
|
207
|
+
name=f"Function Start: {step.name}",
|
|
208
|
+
payload=payload_str,
|
|
209
|
+
parent_id=ancestry.parent_id)
|
|
210
|
+
return event
|
|
211
|
+
|
|
212
|
+
if step.event_type == IntermediateStepType.FUNCTION_END:
|
|
213
|
+
# Find the start event with matching UUID
|
|
214
|
+
start_step = next(
|
|
215
|
+
(x
|
|
216
|
+
for x in self._history if x.event_type == IntermediateStepType.FUNCTION_START and x.UUID == step.UUID),
|
|
217
|
+
None)
|
|
218
|
+
|
|
219
|
+
# For function end events, display output data
|
|
220
|
+
if step.data and hasattr(step.data, 'output'):
|
|
221
|
+
output_str = str(step.data.output)
|
|
222
|
+
elif step.data:
|
|
223
|
+
output_str = str(step.data)
|
|
224
|
+
|
|
225
|
+
if not output_str:
|
|
226
|
+
return None
|
|
227
|
+
|
|
228
|
+
escaped_output = html.escape(output_str, quote=False)
|
|
229
|
+
format_output_type = "json" if is_valid_json(escaped_output) else "python"
|
|
230
|
+
|
|
231
|
+
# Get input from start step if available
|
|
232
|
+
input_payload = ""
|
|
233
|
+
if start_step and start_step.data:
|
|
234
|
+
if hasattr(start_step.data, 'input'):
|
|
235
|
+
input_str = str(start_step.data.input)
|
|
236
|
+
else:
|
|
237
|
+
input_str = str(start_step.data)
|
|
238
|
+
|
|
239
|
+
if input_str:
|
|
240
|
+
escaped_input = html.escape(input_str, quote=False)
|
|
241
|
+
format_input_type = "json" if is_valid_json(escaped_input) else "python"
|
|
242
|
+
input_payload = dedent("""
|
|
243
|
+
**Function Input:**
|
|
244
|
+
```{format_input_type}
|
|
245
|
+
{input_value}
|
|
246
|
+
```
|
|
247
|
+
""").strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
|
|
248
|
+
|
|
249
|
+
# Create payload for function end
|
|
250
|
+
payload_str = dedent("""
|
|
251
|
+
{input_payload}**Function Output:**
|
|
252
|
+
```{format_output_type}
|
|
253
|
+
{output_value}
|
|
254
|
+
```
|
|
255
|
+
""").strip("\n").format(input_payload=input_payload,
|
|
256
|
+
output_value=escaped_output,
|
|
257
|
+
format_output_type=format_output_type)
|
|
258
|
+
|
|
259
|
+
event = ResponseIntermediateStep(id=step.UUID,
|
|
260
|
+
name=f"Function Complete: {step.name}",
|
|
261
|
+
payload=payload_str,
|
|
262
|
+
parent_id=ancestry.parent_id)
|
|
263
|
+
return event
|
|
264
|
+
|
|
265
|
+
return None
|
|
266
|
+
|
|
267
|
+
def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
|
|
268
|
+
"""
|
|
269
|
+
Handles the CUSTOM event
|
|
270
|
+
"""
|
|
271
|
+
escaped_payload = html.escape(str(payload), quote=False)
|
|
272
|
+
escaped_payload = escaped_payload.replace("\n", "")
|
|
273
|
+
|
|
274
|
+
# Attempt to determine type
|
|
275
|
+
format_type = "json" if is_valid_json(escaped_payload) else "python"
|
|
276
|
+
|
|
277
|
+
# Don't use f-strings here because the payload is markdown and screws up the dedent
|
|
278
|
+
payload_str = dedent("""
|
|
279
|
+
```{format_type}
|
|
280
|
+
{payload}
|
|
281
|
+
```
|
|
282
|
+
""").strip("\n").format(payload=escaped_payload, format_type=format_type)
|
|
283
|
+
|
|
284
|
+
# Return the event
|
|
285
|
+
event = ResponseIntermediateStep(id=payload.UUID,
|
|
286
|
+
name=f"{payload.event_type}",
|
|
287
|
+
payload=payload_str,
|
|
288
|
+
parent_id=ancestry.function_id)
|
|
289
|
+
|
|
290
|
+
return event
|
|
291
|
+
|
|
292
|
+
def process(self, step: IntermediateStep) -> ResponseSerializable | None: # pylint: disable=R1710
|
|
293
|
+
|
|
294
|
+
# Track the chunk
|
|
295
|
+
self._history.append(step)
|
|
296
|
+
payload = step.payload
|
|
297
|
+
ancestry = step.function_ancestry
|
|
298
|
+
|
|
299
|
+
if not self._step_matches_filter(step, self.config):
|
|
300
|
+
return None
|
|
301
|
+
|
|
302
|
+
try:
|
|
303
|
+
|
|
304
|
+
if step.event_category == IntermediateStepCategory.LLM:
|
|
305
|
+
return self._handle_llm(payload, ancestry)
|
|
306
|
+
|
|
307
|
+
if step.event_category == IntermediateStepCategory.TOOL:
|
|
308
|
+
return self._handle_tool(payload, ancestry)
|
|
309
|
+
|
|
310
|
+
if step.event_category == IntermediateStepCategory.FUNCTION:
|
|
311
|
+
return self._handle_function(payload, ancestry)
|
|
312
|
+
|
|
313
|
+
if step.event_category == IntermediateStepCategory.CUSTOM:
|
|
314
|
+
return self._handle_custom(payload, ancestry)
|
|
315
|
+
|
|
316
|
+
except Exception as e:
|
|
317
|
+
logger.error("Error processing intermediate step: %s", e, exc_info=True)
|
|
318
|
+
|
|
319
|
+
return None
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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,36 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, 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.data_models.front_end import FrontEndBaseConfig
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
|
|
22
|
+
"""MCP front end configuration.
|
|
23
|
+
|
|
24
|
+
A simple MCP (Modular Communication Protocol) front end for NeMo Agent toolkit.
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
name: str = Field(default="NeMo Agent Toolkit MCP",
|
|
28
|
+
description="Name of the MCP server (default: NeMo Agent Toolkit MCP)")
|
|
29
|
+
host: str = Field(default="localhost", description="Host to bind the server to (default: localhost)")
|
|
30
|
+
port: int = Field(default=9901, description="Port to bind the server to (default: 9901)", ge=0, le=65535)
|
|
31
|
+
debug: bool = Field(default=False, description="Enable debug mode (default: False)")
|
|
32
|
+
log_level: str = Field(default="INFO", description="Log level for the MCP server (default: INFO)")
|
|
33
|
+
tool_names: list[str] = Field(default_factory=list,
|
|
34
|
+
description="The list of tools MCP server will expose (default: all tools)")
|
|
35
|
+
runner_class: str | None = Field(
|
|
36
|
+
default=None, description="Custom worker class for handling MCP routes (default: built-in worker)")
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, 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
|
+
import typing
|
|
18
|
+
|
|
19
|
+
from nat.builder.front_end import FrontEndBase
|
|
20
|
+
from nat.builder.workflow_builder import WorkflowBuilder
|
|
21
|
+
from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
|
|
22
|
+
from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorkerBase
|
|
23
|
+
|
|
24
|
+
logger = logging.getLogger(__name__)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
|
|
28
|
+
"""MCP front end plugin implementation."""
|
|
29
|
+
|
|
30
|
+
def get_worker_class(self) -> type[MCPFrontEndPluginWorkerBase]:
|
|
31
|
+
"""Get the worker class for handling MCP routes."""
|
|
32
|
+
from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorker
|
|
33
|
+
|
|
34
|
+
return MCPFrontEndPluginWorker
|
|
35
|
+
|
|
36
|
+
@typing.final
|
|
37
|
+
def get_worker_class_name(self) -> str:
|
|
38
|
+
"""Get the worker class name from configuration or default."""
|
|
39
|
+
if self.front_end_config.runner_class:
|
|
40
|
+
return self.front_end_config.runner_class
|
|
41
|
+
|
|
42
|
+
worker_class = self.get_worker_class()
|
|
43
|
+
return f"{worker_class.__module__}.{worker_class.__qualname__}"
|
|
44
|
+
|
|
45
|
+
def _get_worker_instance(self) -> MCPFrontEndPluginWorkerBase:
|
|
46
|
+
"""Get an instance of the worker class."""
|
|
47
|
+
# Import the worker class dynamically if specified in config
|
|
48
|
+
if self.front_end_config.runner_class:
|
|
49
|
+
module_name, class_name = self.front_end_config.runner_class.rsplit(".", 1)
|
|
50
|
+
import importlib
|
|
51
|
+
module = importlib.import_module(module_name)
|
|
52
|
+
worker_class = getattr(module, class_name)
|
|
53
|
+
else:
|
|
54
|
+
worker_class = self.get_worker_class()
|
|
55
|
+
|
|
56
|
+
return worker_class(self.full_config)
|
|
57
|
+
|
|
58
|
+
async def run(self) -> None:
|
|
59
|
+
"""Run the MCP server."""
|
|
60
|
+
# Import FastMCP
|
|
61
|
+
from mcp.server.fastmcp import FastMCP
|
|
62
|
+
|
|
63
|
+
# Create an MCP server with the configured parameters
|
|
64
|
+
mcp = FastMCP(
|
|
65
|
+
self.front_end_config.name,
|
|
66
|
+
host=self.front_end_config.host,
|
|
67
|
+
port=self.front_end_config.port,
|
|
68
|
+
debug=self.front_end_config.debug,
|
|
69
|
+
log_level=self.front_end_config.log_level,
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
# Get the worker instance and set up routes
|
|
73
|
+
worker = self._get_worker_instance()
|
|
74
|
+
|
|
75
|
+
# Build the workflow and add routes using the worker
|
|
76
|
+
async with WorkflowBuilder.from_config(config=self.full_config) as builder:
|
|
77
|
+
# Add routes through the worker (includes health endpoint and function registration)
|
|
78
|
+
await worker.add_routes(mcp, builder)
|
|
79
|
+
|
|
80
|
+
# Start the MCP server
|
|
81
|
+
await mcp.run_sse_async()
|