nvidia-nat 1.2.0rc5__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/agent/__init__.py +0 -0
- aiq/agent/base.py +239 -0
- aiq/agent/dual_node.py +67 -0
- aiq/agent/react_agent/__init__.py +0 -0
- aiq/agent/react_agent/agent.py +355 -0
- aiq/agent/react_agent/output_parser.py +104 -0
- aiq/agent/react_agent/prompt.py +41 -0
- aiq/agent/react_agent/register.py +149 -0
- aiq/agent/reasoning_agent/__init__.py +0 -0
- aiq/agent/reasoning_agent/reasoning_agent.py +225 -0
- aiq/agent/register.py +23 -0
- aiq/agent/rewoo_agent/__init__.py +0 -0
- aiq/agent/rewoo_agent/agent.py +411 -0
- aiq/agent/rewoo_agent/prompt.py +108 -0
- aiq/agent/rewoo_agent/register.py +158 -0
- aiq/agent/tool_calling_agent/__init__.py +0 -0
- aiq/agent/tool_calling_agent/agent.py +119 -0
- aiq/agent/tool_calling_agent/register.py +106 -0
- aiq/authentication/__init__.py +14 -0
- aiq/authentication/api_key/__init__.py +14 -0
- aiq/authentication/api_key/api_key_auth_provider.py +96 -0
- aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
- aiq/authentication/api_key/register.py +26 -0
- aiq/authentication/exceptions/__init__.py +14 -0
- aiq/authentication/exceptions/api_key_exceptions.py +38 -0
- aiq/authentication/http_basic_auth/__init__.py +0 -0
- aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
- aiq/authentication/http_basic_auth/register.py +30 -0
- aiq/authentication/interfaces.py +93 -0
- aiq/authentication/oauth2/__init__.py +14 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
- aiq/authentication/oauth2/register.py +25 -0
- aiq/authentication/register.py +21 -0
- aiq/builder/__init__.py +0 -0
- aiq/builder/builder.py +285 -0
- aiq/builder/component_utils.py +316 -0
- aiq/builder/context.py +264 -0
- aiq/builder/embedder.py +24 -0
- aiq/builder/eval_builder.py +161 -0
- aiq/builder/evaluator.py +29 -0
- aiq/builder/framework_enum.py +24 -0
- aiq/builder/front_end.py +73 -0
- aiq/builder/function.py +344 -0
- aiq/builder/function_base.py +380 -0
- aiq/builder/function_info.py +627 -0
- aiq/builder/intermediate_step_manager.py +174 -0
- aiq/builder/llm.py +25 -0
- aiq/builder/retriever.py +25 -0
- aiq/builder/user_interaction_manager.py +74 -0
- aiq/builder/workflow.py +148 -0
- aiq/builder/workflow_builder.py +1117 -0
- aiq/cli/__init__.py +14 -0
- aiq/cli/cli_utils/__init__.py +0 -0
- aiq/cli/cli_utils/config_override.py +231 -0
- aiq/cli/cli_utils/validation.py +37 -0
- aiq/cli/commands/__init__.py +0 -0
- aiq/cli/commands/configure/__init__.py +0 -0
- aiq/cli/commands/configure/channel/__init__.py +0 -0
- aiq/cli/commands/configure/channel/add.py +28 -0
- aiq/cli/commands/configure/channel/channel.py +36 -0
- aiq/cli/commands/configure/channel/remove.py +30 -0
- aiq/cli/commands/configure/channel/update.py +30 -0
- aiq/cli/commands/configure/configure.py +33 -0
- aiq/cli/commands/evaluate.py +139 -0
- aiq/cli/commands/info/__init__.py +14 -0
- aiq/cli/commands/info/info.py +39 -0
- aiq/cli/commands/info/list_channels.py +32 -0
- aiq/cli/commands/info/list_components.py +129 -0
- aiq/cli/commands/info/list_mcp.py +213 -0
- aiq/cli/commands/registry/__init__.py +14 -0
- aiq/cli/commands/registry/publish.py +88 -0
- aiq/cli/commands/registry/pull.py +118 -0
- aiq/cli/commands/registry/registry.py +38 -0
- aiq/cli/commands/registry/remove.py +108 -0
- aiq/cli/commands/registry/search.py +155 -0
- aiq/cli/commands/sizing/__init__.py +14 -0
- aiq/cli/commands/sizing/calc.py +297 -0
- aiq/cli/commands/sizing/sizing.py +27 -0
- aiq/cli/commands/start.py +246 -0
- aiq/cli/commands/uninstall.py +81 -0
- aiq/cli/commands/validate.py +47 -0
- aiq/cli/commands/workflow/__init__.py +14 -0
- aiq/cli/commands/workflow/templates/__init__.py.j2 +0 -0
- aiq/cli/commands/workflow/templates/config.yml.j2 +16 -0
- aiq/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
- aiq/cli/commands/workflow/templates/register.py.j2 +5 -0
- aiq/cli/commands/workflow/templates/workflow.py.j2 +36 -0
- aiq/cli/commands/workflow/workflow.py +37 -0
- aiq/cli/commands/workflow/workflow_commands.py +313 -0
- aiq/cli/entrypoint.py +135 -0
- aiq/cli/main.py +44 -0
- aiq/cli/register_workflow.py +488 -0
- aiq/cli/type_registry.py +1000 -0
- aiq/data_models/__init__.py +14 -0
- aiq/data_models/api_server.py +694 -0
- aiq/data_models/authentication.py +231 -0
- aiq/data_models/common.py +171 -0
- aiq/data_models/component.py +54 -0
- aiq/data_models/component_ref.py +168 -0
- aiq/data_models/config.py +406 -0
- aiq/data_models/dataset_handler.py +123 -0
- aiq/data_models/discovery_metadata.py +335 -0
- aiq/data_models/embedder.py +27 -0
- aiq/data_models/evaluate.py +127 -0
- aiq/data_models/evaluator.py +26 -0
- aiq/data_models/front_end.py +26 -0
- aiq/data_models/function.py +30 -0
- aiq/data_models/function_dependencies.py +72 -0
- aiq/data_models/interactive.py +246 -0
- aiq/data_models/intermediate_step.py +302 -0
- aiq/data_models/invocation_node.py +38 -0
- aiq/data_models/llm.py +27 -0
- aiq/data_models/logging.py +26 -0
- aiq/data_models/memory.py +27 -0
- aiq/data_models/object_store.py +44 -0
- aiq/data_models/profiler.py +54 -0
- aiq/data_models/registry_handler.py +26 -0
- aiq/data_models/retriever.py +30 -0
- aiq/data_models/retry_mixin.py +35 -0
- aiq/data_models/span.py +187 -0
- aiq/data_models/step_adaptor.py +64 -0
- aiq/data_models/streaming.py +33 -0
- aiq/data_models/swe_bench_model.py +54 -0
- aiq/data_models/telemetry_exporter.py +26 -0
- aiq/data_models/ttc_strategy.py +30 -0
- aiq/embedder/__init__.py +0 -0
- aiq/embedder/langchain_client.py +41 -0
- aiq/embedder/nim_embedder.py +59 -0
- aiq/embedder/openai_embedder.py +43 -0
- aiq/embedder/register.py +24 -0
- aiq/eval/__init__.py +14 -0
- aiq/eval/config.py +60 -0
- aiq/eval/dataset_handler/__init__.py +0 -0
- aiq/eval/dataset_handler/dataset_downloader.py +106 -0
- aiq/eval/dataset_handler/dataset_filter.py +52 -0
- aiq/eval/dataset_handler/dataset_handler.py +254 -0
- aiq/eval/evaluate.py +506 -0
- aiq/eval/evaluator/__init__.py +14 -0
- aiq/eval/evaluator/base_evaluator.py +73 -0
- aiq/eval/evaluator/evaluator_model.py +45 -0
- aiq/eval/intermediate_step_adapter.py +99 -0
- aiq/eval/rag_evaluator/__init__.py +0 -0
- aiq/eval/rag_evaluator/evaluate.py +178 -0
- aiq/eval/rag_evaluator/register.py +143 -0
- aiq/eval/register.py +23 -0
- aiq/eval/remote_workflow.py +133 -0
- aiq/eval/runners/__init__.py +14 -0
- aiq/eval/runners/config.py +39 -0
- aiq/eval/runners/multi_eval_runner.py +54 -0
- aiq/eval/runtime_event_subscriber.py +52 -0
- aiq/eval/swe_bench_evaluator/__init__.py +0 -0
- aiq/eval/swe_bench_evaluator/evaluate.py +215 -0
- aiq/eval/swe_bench_evaluator/register.py +36 -0
- aiq/eval/trajectory_evaluator/__init__.py +0 -0
- aiq/eval/trajectory_evaluator/evaluate.py +75 -0
- aiq/eval/trajectory_evaluator/register.py +40 -0
- aiq/eval/tunable_rag_evaluator/__init__.py +0 -0
- aiq/eval/tunable_rag_evaluator/evaluate.py +245 -0
- aiq/eval/tunable_rag_evaluator/register.py +52 -0
- aiq/eval/usage_stats.py +41 -0
- aiq/eval/utils/__init__.py +0 -0
- aiq/eval/utils/output_uploader.py +140 -0
- aiq/eval/utils/tqdm_position_registry.py +40 -0
- aiq/eval/utils/weave_eval.py +184 -0
- aiq/experimental/__init__.py +0 -0
- aiq/experimental/decorators/__init__.py +0 -0
- aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
- aiq/experimental/test_time_compute/__init__.py +0 -0
- aiq/experimental/test_time_compute/editing/__init__.py +0 -0
- aiq/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
- aiq/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
- aiq/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
- aiq/experimental/test_time_compute/functions/__init__.py +0 -0
- aiq/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
- aiq/experimental/test_time_compute/functions/its_tool_orchestration_function.py +205 -0
- aiq/experimental/test_time_compute/functions/its_tool_wrapper_function.py +146 -0
- aiq/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
- aiq/experimental/test_time_compute/models/__init__.py +0 -0
- aiq/experimental/test_time_compute/models/editor_config.py +132 -0
- aiq/experimental/test_time_compute/models/scoring_config.py +112 -0
- aiq/experimental/test_time_compute/models/search_config.py +120 -0
- aiq/experimental/test_time_compute/models/selection_config.py +154 -0
- aiq/experimental/test_time_compute/models/stage_enums.py +43 -0
- aiq/experimental/test_time_compute/models/strategy_base.py +66 -0
- aiq/experimental/test_time_compute/models/tool_use_config.py +41 -0
- aiq/experimental/test_time_compute/models/ttc_item.py +48 -0
- aiq/experimental/test_time_compute/register.py +36 -0
- aiq/experimental/test_time_compute/scoring/__init__.py +0 -0
- aiq/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
- aiq/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
- aiq/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
- aiq/experimental/test_time_compute/search/__init__.py +0 -0
- aiq/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
- aiq/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
- aiq/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
- aiq/experimental/test_time_compute/selection/__init__.py +0 -0
- aiq/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
- aiq/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
- aiq/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
- aiq/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
- aiq/experimental/test_time_compute/selection/threshold_selector.py +58 -0
- aiq/front_ends/__init__.py +14 -0
- aiq/front_ends/console/__init__.py +14 -0
- aiq/front_ends/console/authentication_flow_handler.py +233 -0
- aiq/front_ends/console/console_front_end_config.py +32 -0
- aiq/front_ends/console/console_front_end_plugin.py +96 -0
- aiq/front_ends/console/register.py +25 -0
- aiq/front_ends/cron/__init__.py +14 -0
- aiq/front_ends/fastapi/__init__.py +14 -0
- aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
- aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
- aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
- aiq/front_ends/fastapi/fastapi_front_end_config.py +234 -0
- aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
- aiq/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
- aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1092 -0
- aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
- aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
- aiq/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
- aiq/front_ends/fastapi/job_store.py +183 -0
- aiq/front_ends/fastapi/main.py +72 -0
- aiq/front_ends/fastapi/message_handler.py +298 -0
- aiq/front_ends/fastapi/message_validator.py +345 -0
- aiq/front_ends/fastapi/register.py +25 -0
- aiq/front_ends/fastapi/response_helpers.py +195 -0
- aiq/front_ends/fastapi/step_adaptor.py +321 -0
- aiq/front_ends/mcp/__init__.py +14 -0
- aiq/front_ends/mcp/mcp_front_end_config.py +32 -0
- aiq/front_ends/mcp/mcp_front_end_plugin.py +93 -0
- aiq/front_ends/mcp/register.py +27 -0
- aiq/front_ends/mcp/tool_converter.py +242 -0
- aiq/front_ends/register.py +22 -0
- aiq/front_ends/simple_base/__init__.py +14 -0
- aiq/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
- aiq/llm/__init__.py +0 -0
- aiq/llm/aws_bedrock_llm.py +57 -0
- aiq/llm/nim_llm.py +46 -0
- aiq/llm/openai_llm.py +46 -0
- aiq/llm/register.py +23 -0
- aiq/llm/utils/__init__.py +14 -0
- aiq/llm/utils/env_config_value.py +94 -0
- aiq/llm/utils/error.py +17 -0
- aiq/memory/__init__.py +20 -0
- aiq/memory/interfaces.py +183 -0
- aiq/memory/models.py +112 -0
- aiq/meta/module_to_distro.json +3 -0
- aiq/meta/pypi.md +58 -0
- aiq/object_store/__init__.py +20 -0
- aiq/object_store/in_memory_object_store.py +76 -0
- aiq/object_store/interfaces.py +84 -0
- aiq/object_store/models.py +36 -0
- aiq/object_store/register.py +20 -0
- aiq/observability/__init__.py +14 -0
- aiq/observability/exporter/__init__.py +14 -0
- aiq/observability/exporter/base_exporter.py +449 -0
- aiq/observability/exporter/exporter.py +78 -0
- aiq/observability/exporter/file_exporter.py +33 -0
- aiq/observability/exporter/processing_exporter.py +322 -0
- aiq/observability/exporter/raw_exporter.py +52 -0
- aiq/observability/exporter/span_exporter.py +265 -0
- aiq/observability/exporter_manager.py +335 -0
- aiq/observability/mixin/__init__.py +14 -0
- aiq/observability/mixin/batch_config_mixin.py +26 -0
- aiq/observability/mixin/collector_config_mixin.py +23 -0
- aiq/observability/mixin/file_mixin.py +288 -0
- aiq/observability/mixin/file_mode.py +23 -0
- aiq/observability/mixin/resource_conflict_mixin.py +134 -0
- aiq/observability/mixin/serialize_mixin.py +61 -0
- aiq/observability/mixin/type_introspection_mixin.py +183 -0
- aiq/observability/processor/__init__.py +14 -0
- aiq/observability/processor/batching_processor.py +310 -0
- aiq/observability/processor/callback_processor.py +42 -0
- aiq/observability/processor/intermediate_step_serializer.py +28 -0
- aiq/observability/processor/processor.py +71 -0
- aiq/observability/register.py +96 -0
- aiq/observability/utils/__init__.py +14 -0
- aiq/observability/utils/dict_utils.py +236 -0
- aiq/observability/utils/time_utils.py +31 -0
- aiq/plugins/.namespace +1 -0
- aiq/profiler/__init__.py +0 -0
- aiq/profiler/calc/__init__.py +14 -0
- aiq/profiler/calc/calc_runner.py +627 -0
- aiq/profiler/calc/calculations.py +288 -0
- aiq/profiler/calc/data_models.py +188 -0
- aiq/profiler/calc/plot.py +345 -0
- aiq/profiler/callbacks/__init__.py +0 -0
- aiq/profiler/callbacks/agno_callback_handler.py +295 -0
- aiq/profiler/callbacks/base_callback_class.py +20 -0
- aiq/profiler/callbacks/langchain_callback_handler.py +290 -0
- aiq/profiler/callbacks/llama_index_callback_handler.py +205 -0
- aiq/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
- aiq/profiler/callbacks/token_usage_base_model.py +27 -0
- aiq/profiler/data_frame_row.py +51 -0
- aiq/profiler/data_models.py +24 -0
- aiq/profiler/decorators/__init__.py +0 -0
- aiq/profiler/decorators/framework_wrapper.py +131 -0
- aiq/profiler/decorators/function_tracking.py +254 -0
- aiq/profiler/forecasting/__init__.py +0 -0
- aiq/profiler/forecasting/config.py +18 -0
- aiq/profiler/forecasting/model_trainer.py +75 -0
- aiq/profiler/forecasting/models/__init__.py +22 -0
- aiq/profiler/forecasting/models/forecasting_base_model.py +40 -0
- aiq/profiler/forecasting/models/linear_model.py +196 -0
- aiq/profiler/forecasting/models/random_forest_regressor.py +268 -0
- aiq/profiler/inference_metrics_model.py +28 -0
- aiq/profiler/inference_optimization/__init__.py +0 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
- aiq/profiler/inference_optimization/data_models.py +386 -0
- aiq/profiler/inference_optimization/experimental/__init__.py +0 -0
- aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
- aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
- aiq/profiler/inference_optimization/llm_metrics.py +212 -0
- aiq/profiler/inference_optimization/prompt_caching.py +163 -0
- aiq/profiler/inference_optimization/token_uniqueness.py +107 -0
- aiq/profiler/inference_optimization/workflow_runtimes.py +72 -0
- aiq/profiler/intermediate_property_adapter.py +102 -0
- aiq/profiler/profile_runner.py +473 -0
- aiq/profiler/utils.py +184 -0
- aiq/registry_handlers/__init__.py +0 -0
- aiq/registry_handlers/local/__init__.py +0 -0
- aiq/registry_handlers/local/local_handler.py +176 -0
- aiq/registry_handlers/local/register_local.py +37 -0
- aiq/registry_handlers/metadata_factory.py +60 -0
- aiq/registry_handlers/package_utils.py +567 -0
- aiq/registry_handlers/pypi/__init__.py +0 -0
- aiq/registry_handlers/pypi/pypi_handler.py +251 -0
- aiq/registry_handlers/pypi/register_pypi.py +40 -0
- aiq/registry_handlers/register.py +21 -0
- aiq/registry_handlers/registry_handler_base.py +157 -0
- aiq/registry_handlers/rest/__init__.py +0 -0
- aiq/registry_handlers/rest/register_rest.py +56 -0
- aiq/registry_handlers/rest/rest_handler.py +237 -0
- aiq/registry_handlers/schemas/__init__.py +0 -0
- aiq/registry_handlers/schemas/headers.py +42 -0
- aiq/registry_handlers/schemas/package.py +68 -0
- aiq/registry_handlers/schemas/publish.py +63 -0
- aiq/registry_handlers/schemas/pull.py +82 -0
- aiq/registry_handlers/schemas/remove.py +36 -0
- aiq/registry_handlers/schemas/search.py +91 -0
- aiq/registry_handlers/schemas/status.py +47 -0
- aiq/retriever/__init__.py +0 -0
- aiq/retriever/interface.py +37 -0
- aiq/retriever/milvus/__init__.py +14 -0
- aiq/retriever/milvus/register.py +81 -0
- aiq/retriever/milvus/retriever.py +228 -0
- aiq/retriever/models.py +74 -0
- aiq/retriever/nemo_retriever/__init__.py +14 -0
- aiq/retriever/nemo_retriever/register.py +60 -0
- aiq/retriever/nemo_retriever/retriever.py +190 -0
- aiq/retriever/register.py +22 -0
- aiq/runtime/__init__.py +14 -0
- aiq/runtime/loader.py +215 -0
- aiq/runtime/runner.py +190 -0
- aiq/runtime/session.py +158 -0
- aiq/runtime/user_metadata.py +130 -0
- aiq/settings/__init__.py +0 -0
- aiq/settings/global_settings.py +318 -0
- aiq/test/.namespace +1 -0
- aiq/tool/__init__.py +0 -0
- aiq/tool/chat_completion.py +74 -0
- aiq/tool/code_execution/README.md +151 -0
- aiq/tool/code_execution/__init__.py +0 -0
- aiq/tool/code_execution/code_sandbox.py +267 -0
- aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
- aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
- aiq/tool/code_execution/local_sandbox/__init__.py +13 -0
- aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
- aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
- aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
- aiq/tool/code_execution/register.py +74 -0
- aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
- aiq/tool/code_execution/utils.py +100 -0
- aiq/tool/datetime_tools.py +42 -0
- aiq/tool/document_search.py +141 -0
- aiq/tool/github_tools/__init__.py +0 -0
- aiq/tool/github_tools/create_github_commit.py +133 -0
- aiq/tool/github_tools/create_github_issue.py +87 -0
- aiq/tool/github_tools/create_github_pr.py +106 -0
- aiq/tool/github_tools/get_github_file.py +106 -0
- aiq/tool/github_tools/get_github_issue.py +166 -0
- aiq/tool/github_tools/get_github_pr.py +256 -0
- aiq/tool/github_tools/update_github_issue.py +100 -0
- aiq/tool/mcp/__init__.py +14 -0
- aiq/tool/mcp/exceptions.py +142 -0
- aiq/tool/mcp/mcp_client.py +255 -0
- aiq/tool/mcp/mcp_tool.py +96 -0
- aiq/tool/memory_tools/__init__.py +0 -0
- aiq/tool/memory_tools/add_memory_tool.py +79 -0
- aiq/tool/memory_tools/delete_memory_tool.py +67 -0
- aiq/tool/memory_tools/get_memory_tool.py +72 -0
- aiq/tool/nvidia_rag.py +95 -0
- aiq/tool/register.py +38 -0
- aiq/tool/retriever.py +89 -0
- aiq/tool/server_tools.py +66 -0
- aiq/utils/__init__.py +0 -0
- aiq/utils/data_models/__init__.py +0 -0
- aiq/utils/data_models/schema_validator.py +58 -0
- aiq/utils/debugging_utils.py +43 -0
- aiq/utils/dump_distro_mapping.py +32 -0
- aiq/utils/exception_handlers/__init__.py +0 -0
- aiq/utils/exception_handlers/automatic_retries.py +289 -0
- aiq/utils/exception_handlers/mcp.py +211 -0
- aiq/utils/exception_handlers/schemas.py +114 -0
- aiq/utils/io/__init__.py +0 -0
- aiq/utils/io/model_processing.py +28 -0
- aiq/utils/io/yaml_tools.py +119 -0
- aiq/utils/log_utils.py +37 -0
- aiq/utils/metadata_utils.py +74 -0
- aiq/utils/optional_imports.py +142 -0
- aiq/utils/producer_consumer_queue.py +178 -0
- aiq/utils/reactive/__init__.py +0 -0
- aiq/utils/reactive/base/__init__.py +0 -0
- aiq/utils/reactive/base/observable_base.py +65 -0
- aiq/utils/reactive/base/observer_base.py +55 -0
- aiq/utils/reactive/base/subject_base.py +79 -0
- aiq/utils/reactive/observable.py +59 -0
- aiq/utils/reactive/observer.py +76 -0
- aiq/utils/reactive/subject.py +131 -0
- aiq/utils/reactive/subscription.py +49 -0
- aiq/utils/settings/__init__.py +0 -0
- aiq/utils/settings/global_settings.py +197 -0
- aiq/utils/string_utils.py +38 -0
- aiq/utils/type_converter.py +290 -0
- aiq/utils/type_utils.py +484 -0
- aiq/utils/url_utils.py +27 -0
- nvidia_nat-1.2.0rc5.dist-info/METADATA +363 -0
- nvidia_nat-1.2.0rc5.dist-info/RECORD +435 -0
- nvidia_nat-1.2.0rc5.dist-info/WHEEL +5 -0
- nvidia_nat-1.2.0rc5.dist-info/entry_points.txt +20 -0
- nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE-3rd-party.txt +3686 -0
- nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE.md +201 -0
- nvidia_nat-1.2.0rc5.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,411 @@
|
|
|
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 json
|
|
17
|
+
# pylint: disable=R0917
|
|
18
|
+
import logging
|
|
19
|
+
from json import JSONDecodeError
|
|
20
|
+
|
|
21
|
+
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
22
|
+
from langchain_core.language_models import BaseChatModel
|
|
23
|
+
from langchain_core.messages.ai import AIMessage
|
|
24
|
+
from langchain_core.messages.human import HumanMessage
|
|
25
|
+
from langchain_core.messages.tool import ToolMessage
|
|
26
|
+
from langchain_core.prompts.chat import ChatPromptTemplate
|
|
27
|
+
from langchain_core.runnables.config import RunnableConfig
|
|
28
|
+
from langchain_core.tools import BaseTool
|
|
29
|
+
from langgraph.graph import StateGraph
|
|
30
|
+
from pydantic import BaseModel
|
|
31
|
+
from pydantic import Field
|
|
32
|
+
|
|
33
|
+
from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
|
|
34
|
+
from aiq.agent.base import AGENT_LOG_PREFIX
|
|
35
|
+
from aiq.agent.base import INPUT_SCHEMA_MESSAGE
|
|
36
|
+
from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
|
|
37
|
+
from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
|
|
38
|
+
from aiq.agent.base import AgentDecision
|
|
39
|
+
from aiq.agent.base import BaseAgent
|
|
40
|
+
|
|
41
|
+
logger = logging.getLogger(__name__)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class ReWOOGraphState(BaseModel):
|
|
45
|
+
"""State schema for the ReWOO Agent Graph"""
|
|
46
|
+
task: HumanMessage = Field(default_factory=lambda: HumanMessage(content="")) # the task provided by user
|
|
47
|
+
plan: AIMessage = Field(
|
|
48
|
+
default_factory=lambda: AIMessage(content="")) # the plan generated by the planner to solve the task
|
|
49
|
+
steps: AIMessage = Field(
|
|
50
|
+
default_factory=lambda: AIMessage(content="")) # the steps to solve the task, parsed from the plan
|
|
51
|
+
intermediate_results: dict[str, ToolMessage] = Field(default_factory=dict) # the intermediate results of each step
|
|
52
|
+
result: AIMessage = Field(
|
|
53
|
+
default_factory=lambda: AIMessage(content="")) # the final result of the task, generated by the solver
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class ReWOOAgentGraph(BaseAgent):
|
|
57
|
+
"""Configurable LangGraph ReWOO Agent. A ReWOO Agent performs reasoning by interacting with other objects or tools
|
|
58
|
+
and utilizes their outputs to make decisions. Supports retrying on output parsing errors. Argument
|
|
59
|
+
"detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
|
|
60
|
+
|
|
61
|
+
def __init__(self,
|
|
62
|
+
llm: BaseChatModel,
|
|
63
|
+
planner_prompt: ChatPromptTemplate,
|
|
64
|
+
solver_prompt: ChatPromptTemplate,
|
|
65
|
+
tools: list[BaseTool],
|
|
66
|
+
use_tool_schema: bool = True,
|
|
67
|
+
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
68
|
+
detailed_logs: bool = False):
|
|
69
|
+
super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
|
|
70
|
+
|
|
71
|
+
logger.debug(
|
|
72
|
+
"%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
|
|
73
|
+
AGENT_LOG_PREFIX)
|
|
74
|
+
tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name # prevent trailing ","
|
|
75
|
+
if not use_tool_schema:
|
|
76
|
+
tool_names_and_descriptions = "\n".join(
|
|
77
|
+
[f"{tool.name}: {tool.description}"
|
|
78
|
+
for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}" # prevent trailing "\n"
|
|
79
|
+
else:
|
|
80
|
+
logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
|
|
81
|
+
tool_names_and_descriptions = "\n".join([
|
|
82
|
+
f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
|
|
83
|
+
for tool in tools[:-1]
|
|
84
|
+
]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
|
|
85
|
+
f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
|
|
86
|
+
|
|
87
|
+
self.planner_prompt = planner_prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
|
|
88
|
+
self.solver_prompt = solver_prompt
|
|
89
|
+
self.tools_dict = {tool.name: tool for tool in tools}
|
|
90
|
+
|
|
91
|
+
logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
|
|
92
|
+
|
|
93
|
+
def _get_tool(self, tool_name: str):
|
|
94
|
+
try:
|
|
95
|
+
return self.tools_dict.get(tool_name)
|
|
96
|
+
except Exception as ex:
|
|
97
|
+
logger.exception("%s Unable to find tool with the name %s\n%s",
|
|
98
|
+
AGENT_LOG_PREFIX,
|
|
99
|
+
tool_name,
|
|
100
|
+
ex,
|
|
101
|
+
exc_info=True)
|
|
102
|
+
raise ex
|
|
103
|
+
|
|
104
|
+
@staticmethod
|
|
105
|
+
def _get_current_step(state: ReWOOGraphState) -> int:
|
|
106
|
+
steps = state.steps.content
|
|
107
|
+
if len(steps) == 0:
|
|
108
|
+
raise RuntimeError('No steps received in ReWOOGraphState')
|
|
109
|
+
|
|
110
|
+
if len(state.intermediate_results) == len(steps):
|
|
111
|
+
# all steps are done
|
|
112
|
+
return -1
|
|
113
|
+
|
|
114
|
+
return len(state.intermediate_results)
|
|
115
|
+
|
|
116
|
+
@staticmethod
|
|
117
|
+
def _parse_planner_output(planner_output: str) -> AIMessage:
|
|
118
|
+
|
|
119
|
+
try:
|
|
120
|
+
steps = json.loads(planner_output)
|
|
121
|
+
except json.JSONDecodeError as ex:
|
|
122
|
+
raise ValueError(f"The output of planner is invalid JSON format: {planner_output}") from ex
|
|
123
|
+
|
|
124
|
+
return AIMessage(content=steps)
|
|
125
|
+
|
|
126
|
+
@staticmethod
|
|
127
|
+
def _replace_placeholder(placeholder: str, tool_input: str | dict, tool_output: str | dict) -> str | dict:
|
|
128
|
+
|
|
129
|
+
# Replace the placeholders in the tool input with the previous tool output
|
|
130
|
+
if isinstance(tool_input, dict):
|
|
131
|
+
for key, value in tool_input.items():
|
|
132
|
+
if value is not None:
|
|
133
|
+
if value == placeholder:
|
|
134
|
+
tool_input[key] = tool_output
|
|
135
|
+
elif placeholder in value:
|
|
136
|
+
# If the placeholder is part of the value, replace it with the stringified output
|
|
137
|
+
tool_input[key] = value.replace(placeholder, str(tool_output))
|
|
138
|
+
|
|
139
|
+
elif isinstance(tool_input, str):
|
|
140
|
+
tool_input = tool_input.replace(placeholder, str(tool_output))
|
|
141
|
+
|
|
142
|
+
else:
|
|
143
|
+
assert False, f"Unexpected type for tool_input: {type(tool_input)}"
|
|
144
|
+
return tool_input
|
|
145
|
+
|
|
146
|
+
@staticmethod
|
|
147
|
+
def _parse_tool_input(tool_input: str | dict):
|
|
148
|
+
|
|
149
|
+
# If the input is already a dictionary, return it as is
|
|
150
|
+
if isinstance(tool_input, dict):
|
|
151
|
+
logger.debug("%s Tool input is already a dictionary. Use the tool input as is.", AGENT_LOG_PREFIX)
|
|
152
|
+
return tool_input
|
|
153
|
+
|
|
154
|
+
# If the input is a string, attempt to parse it as JSON
|
|
155
|
+
try:
|
|
156
|
+
tool_input = tool_input.strip()
|
|
157
|
+
# If the input is already a valid JSON string, load it
|
|
158
|
+
tool_input_parsed = json.loads(tool_input)
|
|
159
|
+
logger.debug("%s Successfully parsed structured tool input", AGENT_LOG_PREFIX)
|
|
160
|
+
|
|
161
|
+
except JSONDecodeError:
|
|
162
|
+
try:
|
|
163
|
+
# Replace single quotes with double quotes and attempt parsing again
|
|
164
|
+
tool_input_fixed = tool_input.replace("'", '"')
|
|
165
|
+
tool_input_parsed = json.loads(tool_input_fixed)
|
|
166
|
+
logger.debug(
|
|
167
|
+
"%s Successfully parsed structured tool input after replacing single quotes with double quotes",
|
|
168
|
+
AGENT_LOG_PREFIX)
|
|
169
|
+
|
|
170
|
+
except JSONDecodeError:
|
|
171
|
+
# If it still fails, fall back to using the input as a raw string
|
|
172
|
+
tool_input_parsed = tool_input
|
|
173
|
+
logger.debug("%s Unable to parse structured tool input. Using raw tool input as is.", AGENT_LOG_PREFIX)
|
|
174
|
+
|
|
175
|
+
return tool_input_parsed
|
|
176
|
+
|
|
177
|
+
async def planner_node(self, state: ReWOOGraphState):
|
|
178
|
+
try:
|
|
179
|
+
logger.debug("%s Starting the ReWOO Planner Node", AGENT_LOG_PREFIX)
|
|
180
|
+
|
|
181
|
+
planner = self.planner_prompt | self.llm
|
|
182
|
+
task = str(state.task.content)
|
|
183
|
+
if not task:
|
|
184
|
+
logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
|
|
185
|
+
return {"result": NO_INPUT_ERROR_MESSAGE}
|
|
186
|
+
|
|
187
|
+
plan = await self._stream_llm(
|
|
188
|
+
planner,
|
|
189
|
+
{"task": task},
|
|
190
|
+
RunnableConfig(callbacks=self.callbacks) # type: ignore
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
steps = self._parse_planner_output(str(plan.content))
|
|
194
|
+
|
|
195
|
+
if self.detailed_logs:
|
|
196
|
+
agent_response_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(plan.content))
|
|
197
|
+
logger.info("ReWOO agent planner output: %s", agent_response_log_message)
|
|
198
|
+
|
|
199
|
+
return {"plan": plan, "steps": steps}
|
|
200
|
+
|
|
201
|
+
except Exception as ex:
|
|
202
|
+
logger.exception("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
|
|
203
|
+
raise ex
|
|
204
|
+
|
|
205
|
+
async def executor_node(self, state: ReWOOGraphState):
|
|
206
|
+
try:
|
|
207
|
+
logger.debug("%s Starting the ReWOO Executor Node", AGENT_LOG_PREFIX)
|
|
208
|
+
|
|
209
|
+
current_step = self._get_current_step(state)
|
|
210
|
+
# The executor node should not be invoked after all steps are finished
|
|
211
|
+
if current_step < 0:
|
|
212
|
+
logger.error("%s ReWOO Executor is invoked with an invalid step number: %s",
|
|
213
|
+
AGENT_LOG_PREFIX,
|
|
214
|
+
current_step)
|
|
215
|
+
raise RuntimeError(f"ReWOO Executor is invoked with an invalid step number: {current_step}")
|
|
216
|
+
|
|
217
|
+
steps_content = state.steps.content
|
|
218
|
+
if isinstance(steps_content, list) and current_step < len(steps_content):
|
|
219
|
+
step = steps_content[current_step]
|
|
220
|
+
if isinstance(step, dict) and "evidence" in step:
|
|
221
|
+
step_info = step["evidence"]
|
|
222
|
+
placeholder = step_info.get("placeholder", "")
|
|
223
|
+
tool = step_info.get("tool", "")
|
|
224
|
+
tool_input = step_info.get("tool_input", "")
|
|
225
|
+
else:
|
|
226
|
+
logger.error("%s Invalid step format at index %s", AGENT_LOG_PREFIX, current_step)
|
|
227
|
+
return {"intermediate_results": state.intermediate_results}
|
|
228
|
+
else:
|
|
229
|
+
logger.error("%s Invalid steps content or index %s", AGENT_LOG_PREFIX, current_step)
|
|
230
|
+
return {"intermediate_results": state.intermediate_results}
|
|
231
|
+
|
|
232
|
+
intermediate_results = state.intermediate_results
|
|
233
|
+
|
|
234
|
+
# Replace the placeholder in the tool input with the previous tool output
|
|
235
|
+
for _placeholder, _tool_output in intermediate_results.items():
|
|
236
|
+
_tool_output = _tool_output.content
|
|
237
|
+
# If the content is a list, get the first element which should be a dict
|
|
238
|
+
if isinstance(_tool_output, list):
|
|
239
|
+
_tool_output = _tool_output[0]
|
|
240
|
+
assert isinstance(_tool_output, dict)
|
|
241
|
+
|
|
242
|
+
tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
|
|
243
|
+
|
|
244
|
+
requested_tool = self._get_tool(tool)
|
|
245
|
+
if not requested_tool:
|
|
246
|
+
configured_tool_names = list(self.tools_dict.keys())
|
|
247
|
+
logger.warning(
|
|
248
|
+
"%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
|
|
249
|
+
"there is no tool with that name: %s",
|
|
250
|
+
AGENT_LOG_PREFIX,
|
|
251
|
+
tool,
|
|
252
|
+
configured_tool_names)
|
|
253
|
+
|
|
254
|
+
intermediate_results[placeholder] = ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(
|
|
255
|
+
tool_name=tool, tools=configured_tool_names),
|
|
256
|
+
tool_call_id=tool)
|
|
257
|
+
return {"intermediate_results": intermediate_results}
|
|
258
|
+
|
|
259
|
+
if self.detailed_logs:
|
|
260
|
+
logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
|
|
261
|
+
|
|
262
|
+
# Run the tool. Try to use structured input, if possible
|
|
263
|
+
tool_input_parsed = self._parse_tool_input(tool_input)
|
|
264
|
+
tool_response = await self._call_tool(requested_tool,
|
|
265
|
+
tool_input_parsed,
|
|
266
|
+
RunnableConfig(callbacks=self.callbacks),
|
|
267
|
+
max_retries=3)
|
|
268
|
+
|
|
269
|
+
# ToolMessage only accepts str or list[str | dict] as content.
|
|
270
|
+
# Convert into list if the response is a dict.
|
|
271
|
+
if isinstance(tool_response, dict):
|
|
272
|
+
tool_response = [tool_response]
|
|
273
|
+
|
|
274
|
+
tool_response_message = ToolMessage(name=tool, tool_call_id=tool, content=tool_response)
|
|
275
|
+
|
|
276
|
+
if self.detailed_logs:
|
|
277
|
+
self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
|
|
278
|
+
|
|
279
|
+
intermediate_results[placeholder] = tool_response_message
|
|
280
|
+
return {"intermediate_results": intermediate_results}
|
|
281
|
+
|
|
282
|
+
except Exception as ex:
|
|
283
|
+
logger.exception("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
|
|
284
|
+
raise ex
|
|
285
|
+
|
|
286
|
+
async def solver_node(self, state: ReWOOGraphState):
|
|
287
|
+
try:
|
|
288
|
+
logger.debug("%s Starting the ReWOO Solver Node", AGENT_LOG_PREFIX)
|
|
289
|
+
|
|
290
|
+
plan = ""
|
|
291
|
+
# Add the tool outputs of each step to the plan
|
|
292
|
+
for step in state.steps.content:
|
|
293
|
+
step_info = step["evidence"]
|
|
294
|
+
placeholder = step_info.get("placeholder", "")
|
|
295
|
+
tool_input = step_info.get("tool_input", "")
|
|
296
|
+
|
|
297
|
+
intermediate_results = state.intermediate_results
|
|
298
|
+
for _placeholder, _tool_output in intermediate_results.items():
|
|
299
|
+
_tool_output = _tool_output.content
|
|
300
|
+
# If the content is a list, get the first element which should be a dict
|
|
301
|
+
if isinstance(_tool_output, list):
|
|
302
|
+
_tool_output = _tool_output[0]
|
|
303
|
+
assert isinstance(_tool_output, dict)
|
|
304
|
+
|
|
305
|
+
tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
|
|
306
|
+
|
|
307
|
+
placeholder = placeholder.replace(_placeholder, str(_tool_output))
|
|
308
|
+
|
|
309
|
+
_plan = step.get("plan")
|
|
310
|
+
tool = step_info.get("tool")
|
|
311
|
+
plan += f"Plan: {_plan}\n{placeholder} = {tool}[{tool_input}]"
|
|
312
|
+
|
|
313
|
+
task = str(state.task.content)
|
|
314
|
+
solver_prompt = self.solver_prompt.partial(plan=plan)
|
|
315
|
+
solver = solver_prompt | self.llm
|
|
316
|
+
|
|
317
|
+
output_message = await self._stream_llm(solver, {"task": task},
|
|
318
|
+
RunnableConfig(callbacks=self.callbacks)) # type: ignore
|
|
319
|
+
|
|
320
|
+
if self.detailed_logs:
|
|
321
|
+
solver_output_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(output_message.content))
|
|
322
|
+
logger.info("ReWOO agent solver output: %s", solver_output_log_message)
|
|
323
|
+
|
|
324
|
+
return {"result": output_message}
|
|
325
|
+
|
|
326
|
+
except Exception as ex:
|
|
327
|
+
logger.exception("%s Failed to call solver_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
|
|
328
|
+
raise ex
|
|
329
|
+
|
|
330
|
+
async def conditional_edge(self, state: ReWOOGraphState):
|
|
331
|
+
try:
|
|
332
|
+
logger.debug("%s Starting the ReWOO Conditional Edge", AGENT_LOG_PREFIX)
|
|
333
|
+
|
|
334
|
+
current_step = self._get_current_step(state)
|
|
335
|
+
if current_step == -1:
|
|
336
|
+
logger.debug("%s The ReWOO Executor has finished its task", AGENT_LOG_PREFIX)
|
|
337
|
+
return AgentDecision.END
|
|
338
|
+
|
|
339
|
+
logger.debug("%s The ReWOO Executor is still working on the task", AGENT_LOG_PREFIX)
|
|
340
|
+
return AgentDecision.TOOL
|
|
341
|
+
|
|
342
|
+
except Exception as ex:
|
|
343
|
+
logger.exception("%s Failed to determine whether agent is calling a tool: %s",
|
|
344
|
+
AGENT_LOG_PREFIX,
|
|
345
|
+
ex,
|
|
346
|
+
exc_info=True)
|
|
347
|
+
logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
|
|
348
|
+
return AgentDecision.END
|
|
349
|
+
|
|
350
|
+
async def _build_graph(self, state_schema):
|
|
351
|
+
try:
|
|
352
|
+
logger.debug("%s Building and compiling the ReWOO Graph", AGENT_LOG_PREFIX)
|
|
353
|
+
|
|
354
|
+
graph = StateGraph(state_schema)
|
|
355
|
+
graph.add_node("planner", self.planner_node)
|
|
356
|
+
graph.add_node("executor", self.executor_node)
|
|
357
|
+
graph.add_node("solver", self.solver_node)
|
|
358
|
+
|
|
359
|
+
graph.add_edge("planner", "executor")
|
|
360
|
+
conditional_edge_possible_outputs = {AgentDecision.TOOL: "executor", AgentDecision.END: "solver"}
|
|
361
|
+
graph.add_conditional_edges("executor", self.conditional_edge, conditional_edge_possible_outputs)
|
|
362
|
+
|
|
363
|
+
graph.set_entry_point("planner")
|
|
364
|
+
graph.set_finish_point("solver")
|
|
365
|
+
|
|
366
|
+
self.graph = graph.compile()
|
|
367
|
+
logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
368
|
+
|
|
369
|
+
return self.graph
|
|
370
|
+
|
|
371
|
+
except Exception as ex:
|
|
372
|
+
logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
|
|
373
|
+
raise ex
|
|
374
|
+
|
|
375
|
+
async def build_graph(self):
|
|
376
|
+
try:
|
|
377
|
+
await self._build_graph(state_schema=ReWOOGraphState)
|
|
378
|
+
logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
379
|
+
return self.graph
|
|
380
|
+
except Exception as ex:
|
|
381
|
+
logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
|
|
382
|
+
raise ex
|
|
383
|
+
|
|
384
|
+
@staticmethod
|
|
385
|
+
def validate_planner_prompt(planner_prompt: str) -> bool:
|
|
386
|
+
errors = []
|
|
387
|
+
if not planner_prompt:
|
|
388
|
+
errors.append("The planner prompt cannot be empty.")
|
|
389
|
+
required_prompt_variables = {
|
|
390
|
+
"{tools}": "The planner prompt must contain {tools} so the planner agent knows about configured tools.",
|
|
391
|
+
"{tool_names}": "The planner prompt must contain {tool_names} so the planner agent knows tool names."
|
|
392
|
+
}
|
|
393
|
+
for variable_name, error_message in required_prompt_variables.items():
|
|
394
|
+
if variable_name not in planner_prompt:
|
|
395
|
+
errors.append(error_message)
|
|
396
|
+
if errors:
|
|
397
|
+
error_text = "\n".join(errors)
|
|
398
|
+
logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
|
|
399
|
+
raise ValueError(error_text)
|
|
400
|
+
return True
|
|
401
|
+
|
|
402
|
+
@staticmethod
|
|
403
|
+
def validate_solver_prompt(solver_prompt: str) -> bool:
|
|
404
|
+
errors = []
|
|
405
|
+
if not solver_prompt:
|
|
406
|
+
errors.append("The solver prompt cannot be empty.")
|
|
407
|
+
if errors:
|
|
408
|
+
error_text = "\n".join(errors)
|
|
409
|
+
logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
|
|
410
|
+
raise ValueError(error_text)
|
|
411
|
+
return True
|
|
@@ -0,0 +1,108 @@
|
|
|
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
|
+
# flake8: noqa
|
|
17
|
+
from langchain_core.prompts.chat import ChatPromptTemplate
|
|
18
|
+
|
|
19
|
+
PLANNER_SYSTEM_PROMPT = """
|
|
20
|
+
For the following task, make plans that can solve the problem step by step. For each plan, indicate \
|
|
21
|
+
which external tool together with tool input to retrieve evidence. You can store the evidence into a \
|
|
22
|
+
placeholder #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)
|
|
23
|
+
|
|
24
|
+
You may ask the human to the following tools:
|
|
25
|
+
|
|
26
|
+
{tools}
|
|
27
|
+
|
|
28
|
+
The tools should be one of the following: [{tool_names}]
|
|
29
|
+
|
|
30
|
+
You are not required to use all the tools listed. Choose only the ones that best fit the needs of each plan step.
|
|
31
|
+
|
|
32
|
+
Your output must be a JSON array where each element represents one planning step. Each step must be an object with exactly two keys:
|
|
33
|
+
|
|
34
|
+
1. "plan": A string that describes in detail the action or reasoning for that step.
|
|
35
|
+
|
|
36
|
+
2. "evidence": An object representing the external tool call associated with that plan step. This object must have the following keys:
|
|
37
|
+
|
|
38
|
+
-"placeholder": A string that identifies the evidence placeholder (e.g., "#E1", "#E2", etc.). The numbering should be sequential based on the order of steps.
|
|
39
|
+
|
|
40
|
+
-"tool": A string specifying the name of the external tool used.
|
|
41
|
+
|
|
42
|
+
-"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the tool.
|
|
43
|
+
|
|
44
|
+
Do not include any additional keys or characters in your output, and do not wrap your response with markdown formatting. Your output must be strictly valid JSON.
|
|
45
|
+
|
|
46
|
+
Important instructions:
|
|
47
|
+
|
|
48
|
+
Do not output any additional text, comments, or markdown formatting.
|
|
49
|
+
|
|
50
|
+
Do not include any explanation or reasoning text outside of the JSON array.
|
|
51
|
+
|
|
52
|
+
The output must be a valid JSON array that can be parsed directly.
|
|
53
|
+
|
|
54
|
+
Here is an example of how a valid JSON output should look:
|
|
55
|
+
|
|
56
|
+
[
|
|
57
|
+
\'{{
|
|
58
|
+
"plan": "Calculate the result of 2023 minus 25.",
|
|
59
|
+
"evidence": \'{{
|
|
60
|
+
"placeholder": "#E1",
|
|
61
|
+
"tool": "calculator_subtract",
|
|
62
|
+
"tool_input": [2023, 25]
|
|
63
|
+
}}\'
|
|
64
|
+
}}\',
|
|
65
|
+
\'{{
|
|
66
|
+
"plan": "Retrieve the year represented by the result stored in #E1.",
|
|
67
|
+
"evidence": \'{{
|
|
68
|
+
"placeholder": "#E2",
|
|
69
|
+
"tool": "haystack_chitchat_agent",
|
|
70
|
+
"tool_input": "Response with the result number contained in #E1"
|
|
71
|
+
}}\'
|
|
72
|
+
}}\',
|
|
73
|
+
\'{{
|
|
74
|
+
"plan": "Search for the CEO of Golden State Warriors in the year stored in #E2.",
|
|
75
|
+
"evidence": \'{{
|
|
76
|
+
"placeholder": "#E3",
|
|
77
|
+
"tool": "internet_search",
|
|
78
|
+
"tool_input": "Who was the CEO of Golden State Warriors in the year #E2?"
|
|
79
|
+
}}\'
|
|
80
|
+
}}\'
|
|
81
|
+
]
|
|
82
|
+
|
|
83
|
+
Begin!
|
|
84
|
+
"""
|
|
85
|
+
|
|
86
|
+
PLANNER_USER_PROMPT = """
|
|
87
|
+
task: {task}
|
|
88
|
+
"""
|
|
89
|
+
|
|
90
|
+
rewoo_planner_prompt = ChatPromptTemplate([("system", PLANNER_SYSTEM_PROMPT), ("user", PLANNER_USER_PROMPT)])
|
|
91
|
+
|
|
92
|
+
SOLVER_SYSTEM_PROMPT = """
|
|
93
|
+
Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
|
|
94
|
+
retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
|
|
95
|
+
contain irrelevant information.
|
|
96
|
+
|
|
97
|
+
Now solve the question or task according to provided Evidence above. Respond with the answer
|
|
98
|
+
directly with no extra words.
|
|
99
|
+
|
|
100
|
+
"""
|
|
101
|
+
SOLVER_USER_PROMPT = """
|
|
102
|
+
plan: {plan}
|
|
103
|
+
task: {task}
|
|
104
|
+
|
|
105
|
+
Response:
|
|
106
|
+
"""
|
|
107
|
+
|
|
108
|
+
rewoo_solver_prompt = ChatPromptTemplate([("system", SOLVER_SYSTEM_PROMPT), ("user", SOLVER_USER_PROMPT)])
|
|
@@ -0,0 +1,158 @@
|
|
|
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
|
+
|
|
18
|
+
from pydantic import Field
|
|
19
|
+
|
|
20
|
+
from aiq.builder.builder import Builder
|
|
21
|
+
from aiq.builder.framework_enum import LLMFrameworkEnum
|
|
22
|
+
from aiq.builder.function_info import FunctionInfo
|
|
23
|
+
from aiq.cli.register_workflow import register_function
|
|
24
|
+
from aiq.data_models.api_server import AIQChatRequest
|
|
25
|
+
from aiq.data_models.api_server import AIQChatResponse
|
|
26
|
+
from aiq.data_models.component_ref import FunctionRef
|
|
27
|
+
from aiq.data_models.component_ref import LLMRef
|
|
28
|
+
from aiq.data_models.function import FunctionBaseConfig
|
|
29
|
+
from aiq.utils.type_converter import GlobalTypeConverter
|
|
30
|
+
|
|
31
|
+
logger = logging.getLogger(__name__)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
|
|
35
|
+
"""
|
|
36
|
+
Defines an AIQ Toolkit function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
|
|
37
|
+
tool names and descriptions to select the optimal tool.
|
|
38
|
+
"""
|
|
39
|
+
|
|
40
|
+
tool_names: list[FunctionRef] = Field(default_factory=list,
|
|
41
|
+
description="The list of tools to provide to the rewoo agent.")
|
|
42
|
+
llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
|
|
43
|
+
verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
|
|
44
|
+
include_tool_input_schema_in_tool_description: bool = Field(
|
|
45
|
+
default=True, description="Specify inclusion of tool input schemas in the prompt.")
|
|
46
|
+
description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
|
|
47
|
+
planner_prompt: str | None = Field(
|
|
48
|
+
default=None,
|
|
49
|
+
description="Provides the PLANNER_PROMPT to use with the agent") # defaults to PLANNER_PROMPT in prompt.py
|
|
50
|
+
solver_prompt: str | None = Field(
|
|
51
|
+
default=None,
|
|
52
|
+
description="Provides the SOLVER_PROMPT to use with the agent") # defaults to SOLVER_PROMPT in prompt.py
|
|
53
|
+
max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
|
|
54
|
+
use_openai_api: bool = Field(default=False,
|
|
55
|
+
description=("Use OpenAI API for the input/output types to the function. "
|
|
56
|
+
"If False, strings will be used."))
|
|
57
|
+
additional_instructions: str | None = Field(
|
|
58
|
+
default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
|
|
62
|
+
async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
|
|
63
|
+
from langchain.schema import BaseMessage
|
|
64
|
+
from langchain_core.messages import trim_messages
|
|
65
|
+
from langchain_core.messages.human import HumanMessage
|
|
66
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
67
|
+
from langgraph.graph.graph import CompiledGraph
|
|
68
|
+
|
|
69
|
+
from aiq.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
|
|
70
|
+
from aiq.agent.rewoo_agent.prompt import SOLVER_USER_PROMPT
|
|
71
|
+
|
|
72
|
+
from .agent import ReWOOAgentGraph
|
|
73
|
+
from .agent import ReWOOGraphState
|
|
74
|
+
from .prompt import rewoo_planner_prompt
|
|
75
|
+
from .prompt import rewoo_solver_prompt
|
|
76
|
+
|
|
77
|
+
# the ReWOO Agent prompt comes from prompt.py, and can be customized there or via config option planner_prompt and
|
|
78
|
+
# solver_prompt.
|
|
79
|
+
if config.planner_prompt:
|
|
80
|
+
planner_prompt = config.planner_prompt
|
|
81
|
+
if config.additional_instructions:
|
|
82
|
+
planner_prompt += f"{config.additional_instructions}"
|
|
83
|
+
valid = ReWOOAgentGraph.validate_planner_prompt(config.planner_prompt)
|
|
84
|
+
if not valid:
|
|
85
|
+
logger.exception("Invalid planner_prompt")
|
|
86
|
+
raise ValueError("Invalid planner_prompt")
|
|
87
|
+
planner_prompt = ChatPromptTemplate([("system", config.planner_prompt), ("user", PLANNER_USER_PROMPT)])
|
|
88
|
+
else:
|
|
89
|
+
planner_prompt = rewoo_planner_prompt
|
|
90
|
+
|
|
91
|
+
if config.solver_prompt:
|
|
92
|
+
solver_prompt = config.solver_prompt
|
|
93
|
+
if config.additional_instructions:
|
|
94
|
+
solver_prompt += f"{config.additional_instructions}"
|
|
95
|
+
valid = ReWOOAgentGraph.validate_solver_prompt(config.solver_prompt)
|
|
96
|
+
if not valid:
|
|
97
|
+
logger.exception("Invalid solver_prompt")
|
|
98
|
+
raise ValueError("Invalid solver_prompt")
|
|
99
|
+
solver_prompt = ChatPromptTemplate([("system", config.solver_prompt), ("user", SOLVER_USER_PROMPT)])
|
|
100
|
+
else:
|
|
101
|
+
solver_prompt = rewoo_solver_prompt
|
|
102
|
+
|
|
103
|
+
# we can choose an LLM for the ReWOO agent in the config file
|
|
104
|
+
llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
105
|
+
|
|
106
|
+
# the agent can run any installed tool, simply install the tool and add it to the config file
|
|
107
|
+
# the sample tool provided can easily be copied or changed
|
|
108
|
+
tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
109
|
+
if not tools:
|
|
110
|
+
raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
|
|
111
|
+
|
|
112
|
+
# construct the ReWOO Agent Graph from the configured llm, prompt, and tools
|
|
113
|
+
graph: CompiledGraph = await ReWOOAgentGraph(llm=llm,
|
|
114
|
+
planner_prompt=planner_prompt,
|
|
115
|
+
solver_prompt=solver_prompt,
|
|
116
|
+
tools=tools,
|
|
117
|
+
use_tool_schema=config.include_tool_input_schema_in_tool_description,
|
|
118
|
+
detailed_logs=config.verbose).build_graph()
|
|
119
|
+
|
|
120
|
+
async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
|
|
121
|
+
try:
|
|
122
|
+
# initialize the starting state with the user query
|
|
123
|
+
messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
|
|
124
|
+
max_tokens=config.max_history,
|
|
125
|
+
strategy="last",
|
|
126
|
+
token_counter=len,
|
|
127
|
+
start_on="human",
|
|
128
|
+
include_system=True)
|
|
129
|
+
task = HumanMessage(content=messages[0].content)
|
|
130
|
+
state = ReWOOGraphState(task=task)
|
|
131
|
+
|
|
132
|
+
# run the ReWOO Agent Graph
|
|
133
|
+
state = await graph.ainvoke(state)
|
|
134
|
+
|
|
135
|
+
# get and return the output from the state
|
|
136
|
+
state = ReWOOGraphState(**state)
|
|
137
|
+
output_message = state.result.content # pylint: disable=E1101
|
|
138
|
+
return AIQChatResponse.from_string(output_message)
|
|
139
|
+
|
|
140
|
+
except Exception as ex:
|
|
141
|
+
logger.exception("ReWOO Agent failed with exception: %s", ex, exc_info=ex)
|
|
142
|
+
# here, we can implement custom error messages
|
|
143
|
+
if config.verbose:
|
|
144
|
+
return AIQChatResponse.from_string(str(ex))
|
|
145
|
+
return AIQChatResponse.from_string("I seem to be having a problem.")
|
|
146
|
+
|
|
147
|
+
if (config.use_openai_api):
|
|
148
|
+
yield FunctionInfo.from_fn(_response_fn, description=config.description)
|
|
149
|
+
|
|
150
|
+
else:
|
|
151
|
+
|
|
152
|
+
async def _str_api_fn(input_message: str) -> str:
|
|
153
|
+
oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
|
|
154
|
+
oai_output = await _response_fn(oai_input)
|
|
155
|
+
|
|
156
|
+
return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
|
|
157
|
+
|
|
158
|
+
yield FunctionInfo.from_fn(_str_api_fn, description=config.description)
|
|
File without changes
|