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
aiq/agent/__init__.py
ADDED
|
File without changes
|
aiq/agent/base.py
ADDED
|
@@ -0,0 +1,239 @@
|
|
|
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 json
|
|
18
|
+
import logging
|
|
19
|
+
from abc import ABC
|
|
20
|
+
from abc import abstractmethod
|
|
21
|
+
from enum import Enum
|
|
22
|
+
from typing import Any
|
|
23
|
+
|
|
24
|
+
from colorama import Fore
|
|
25
|
+
from langchain_core.callbacks import AsyncCallbackHandler
|
|
26
|
+
from langchain_core.language_models import BaseChatModel
|
|
27
|
+
from langchain_core.messages import AIMessage
|
|
28
|
+
from langchain_core.messages import BaseMessage
|
|
29
|
+
from langchain_core.messages import ToolMessage
|
|
30
|
+
from langchain_core.runnables import RunnableConfig
|
|
31
|
+
from langchain_core.tools import BaseTool
|
|
32
|
+
from langgraph.graph.graph import CompiledGraph
|
|
33
|
+
|
|
34
|
+
logger = logging.getLogger(__name__)
|
|
35
|
+
|
|
36
|
+
TOOL_NOT_FOUND_ERROR_MESSAGE = "There is no tool named {tool_name}. Tool must be one of {tools}."
|
|
37
|
+
INPUT_SCHEMA_MESSAGE = ". Arguments must be provided as a valid JSON object following this format: {schema}"
|
|
38
|
+
NO_INPUT_ERROR_MESSAGE = "No human input received to the agent, Please ask a valid question."
|
|
39
|
+
|
|
40
|
+
AGENT_LOG_PREFIX = "[AGENT]"
|
|
41
|
+
AGENT_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
|
|
42
|
+
AGENT_LOG_PREFIX + "\n" + \
|
|
43
|
+
Fore.YELLOW + \
|
|
44
|
+
"Agent input: %s\n" + \
|
|
45
|
+
Fore.CYAN + \
|
|
46
|
+
"Agent's thoughts: \n%s" + \
|
|
47
|
+
Fore.RESET + \
|
|
48
|
+
f"\n{'-' * 30}"
|
|
49
|
+
|
|
50
|
+
TOOL_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
|
|
51
|
+
AGENT_LOG_PREFIX + "\n" + \
|
|
52
|
+
Fore.WHITE + \
|
|
53
|
+
"Calling tools: %s\n" + \
|
|
54
|
+
Fore.YELLOW + \
|
|
55
|
+
"Tool's input: %s\n" + \
|
|
56
|
+
Fore.CYAN + \
|
|
57
|
+
"Tool's response: \n%s" + \
|
|
58
|
+
Fore.RESET + \
|
|
59
|
+
f"\n{'-' * 30}"
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
class AgentDecision(Enum):
|
|
63
|
+
TOOL = "tool"
|
|
64
|
+
END = "finished"
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
class BaseAgent(ABC):
|
|
68
|
+
|
|
69
|
+
def __init__(self,
|
|
70
|
+
llm: BaseChatModel,
|
|
71
|
+
tools: list[BaseTool],
|
|
72
|
+
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
73
|
+
detailed_logs: bool = False) -> None:
|
|
74
|
+
logger.debug("Initializing Agent Graph")
|
|
75
|
+
self.llm = llm
|
|
76
|
+
self.tools = tools
|
|
77
|
+
self.callbacks = callbacks or []
|
|
78
|
+
self.detailed_logs = detailed_logs
|
|
79
|
+
self.graph = None
|
|
80
|
+
|
|
81
|
+
async def _stream_llm(self,
|
|
82
|
+
runnable: Any,
|
|
83
|
+
inputs: dict[str, Any],
|
|
84
|
+
config: RunnableConfig | None = None) -> AIMessage:
|
|
85
|
+
"""
|
|
86
|
+
Stream from LLM runnable. Retry logic is handled automatically by the underlying LLM client.
|
|
87
|
+
|
|
88
|
+
Parameters
|
|
89
|
+
----------
|
|
90
|
+
runnable : Any
|
|
91
|
+
The LLM runnable (prompt | llm or similar)
|
|
92
|
+
inputs : Dict[str, Any]
|
|
93
|
+
The inputs to pass to the runnable
|
|
94
|
+
config : RunnableConfig | None
|
|
95
|
+
The config to pass to the runnable (should include callbacks)
|
|
96
|
+
|
|
97
|
+
Returns
|
|
98
|
+
-------
|
|
99
|
+
AIMessage
|
|
100
|
+
The LLM response
|
|
101
|
+
"""
|
|
102
|
+
output_message = ""
|
|
103
|
+
async for event in runnable.astream(inputs, config=config):
|
|
104
|
+
output_message += event.content
|
|
105
|
+
|
|
106
|
+
return AIMessage(content=output_message)
|
|
107
|
+
|
|
108
|
+
async def _call_llm(self, messages: list[BaseMessage]) -> AIMessage:
|
|
109
|
+
"""
|
|
110
|
+
Call the LLM directly. Retry logic is handled automatically by the underlying LLM client.
|
|
111
|
+
|
|
112
|
+
Parameters
|
|
113
|
+
----------
|
|
114
|
+
messages : list[BaseMessage]
|
|
115
|
+
The messages to send to the LLM
|
|
116
|
+
|
|
117
|
+
Returns
|
|
118
|
+
-------
|
|
119
|
+
AIMessage
|
|
120
|
+
The LLM response
|
|
121
|
+
"""
|
|
122
|
+
response = await self.llm.ainvoke(messages)
|
|
123
|
+
return AIMessage(content=str(response.content))
|
|
124
|
+
|
|
125
|
+
async def _call_tool(self,
|
|
126
|
+
tool: BaseTool,
|
|
127
|
+
tool_input: dict[str, Any] | str,
|
|
128
|
+
config: RunnableConfig | None = None,
|
|
129
|
+
max_retries: int = 3) -> ToolMessage:
|
|
130
|
+
"""
|
|
131
|
+
Call a tool with retry logic and error handling.
|
|
132
|
+
|
|
133
|
+
Parameters
|
|
134
|
+
----------
|
|
135
|
+
tool : BaseTool
|
|
136
|
+
The tool to call
|
|
137
|
+
tool_input : Union[Dict[str, Any], str]
|
|
138
|
+
The input to pass to the tool
|
|
139
|
+
config : RunnableConfig | None
|
|
140
|
+
The config to pass to the tool
|
|
141
|
+
max_retries : int
|
|
142
|
+
Maximum number of retry attempts (default: 3)
|
|
143
|
+
|
|
144
|
+
Returns
|
|
145
|
+
-------
|
|
146
|
+
ToolMessage
|
|
147
|
+
The tool response
|
|
148
|
+
"""
|
|
149
|
+
last_exception = None
|
|
150
|
+
|
|
151
|
+
for attempt in range(1, max_retries + 1):
|
|
152
|
+
try:
|
|
153
|
+
response = await tool.ainvoke(tool_input, config=config)
|
|
154
|
+
|
|
155
|
+
# Handle empty responses
|
|
156
|
+
if response is None or (isinstance(response, str) and response == ""):
|
|
157
|
+
return ToolMessage(name=tool.name,
|
|
158
|
+
tool_call_id=tool.name,
|
|
159
|
+
content=f"The tool {tool.name} provided an empty response.")
|
|
160
|
+
|
|
161
|
+
return ToolMessage(name=tool.name, tool_call_id=tool.name, content=response)
|
|
162
|
+
|
|
163
|
+
except Exception as e:
|
|
164
|
+
last_exception = e
|
|
165
|
+
|
|
166
|
+
# If this was the last attempt, don't sleep
|
|
167
|
+
if attempt == max_retries:
|
|
168
|
+
break
|
|
169
|
+
|
|
170
|
+
logger.warning("%s Tool call attempt %d/%d failed for tool %s: %s",
|
|
171
|
+
AGENT_LOG_PREFIX,
|
|
172
|
+
attempt,
|
|
173
|
+
max_retries,
|
|
174
|
+
tool.name,
|
|
175
|
+
str(e))
|
|
176
|
+
|
|
177
|
+
# Exponential backoff: 2^attempt seconds
|
|
178
|
+
sleep_time = 2**attempt
|
|
179
|
+
logger.debug("%s Retrying tool call for %s in %d seconds...", AGENT_LOG_PREFIX, tool.name, sleep_time)
|
|
180
|
+
await asyncio.sleep(sleep_time)
|
|
181
|
+
|
|
182
|
+
# All retries exhausted, return error message
|
|
183
|
+
error_content = "Tool call failed after all retry attempts. Last error: %s" % str(last_exception)
|
|
184
|
+
logger.error("%s %s", AGENT_LOG_PREFIX, error_content)
|
|
185
|
+
return ToolMessage(name=tool.name, tool_call_id=tool.name, content=error_content, status="error")
|
|
186
|
+
|
|
187
|
+
def _log_tool_response(self, tool_name: str, tool_input: Any, tool_response: str, max_chars: int = 1000) -> None:
|
|
188
|
+
"""
|
|
189
|
+
Log tool response with consistent formatting and length limits.
|
|
190
|
+
|
|
191
|
+
Parameters
|
|
192
|
+
----------
|
|
193
|
+
tool_name : str
|
|
194
|
+
The name of the tool that was called
|
|
195
|
+
tool_input : Any
|
|
196
|
+
The input that was passed to the tool
|
|
197
|
+
tool_response : str
|
|
198
|
+
The response from the tool
|
|
199
|
+
max_chars : int
|
|
200
|
+
Maximum number of characters to log (default: 1000)
|
|
201
|
+
"""
|
|
202
|
+
if self.detailed_logs:
|
|
203
|
+
# Truncate tool response if too long
|
|
204
|
+
display_response = tool_response[:max_chars] + "...(rest of response truncated)" if len(
|
|
205
|
+
tool_response) > max_chars else tool_response
|
|
206
|
+
|
|
207
|
+
# Format the tool input for display
|
|
208
|
+
tool_input_str = str(tool_input)
|
|
209
|
+
|
|
210
|
+
tool_response_log_message = TOOL_CALL_LOG_MESSAGE % (tool_name, tool_input_str, display_response)
|
|
211
|
+
logger.info(tool_response_log_message)
|
|
212
|
+
|
|
213
|
+
def _parse_json(self, json_string: str) -> dict[str, Any]:
|
|
214
|
+
"""
|
|
215
|
+
Safely parse JSON with graceful error handling.
|
|
216
|
+
If JSON parsing fails, returns an empty dict or error info.
|
|
217
|
+
|
|
218
|
+
Parameters
|
|
219
|
+
----------
|
|
220
|
+
json_string : str
|
|
221
|
+
The JSON string to parse
|
|
222
|
+
|
|
223
|
+
Returns
|
|
224
|
+
-------
|
|
225
|
+
Dict[str, Any]
|
|
226
|
+
The parsed JSON or error information
|
|
227
|
+
"""
|
|
228
|
+
try:
|
|
229
|
+
return json.loads(json_string)
|
|
230
|
+
except json.JSONDecodeError as e:
|
|
231
|
+
logger.warning("%s JSON parsing failed, returning the original string: %s", AGENT_LOG_PREFIX, str(e))
|
|
232
|
+
return {"error": f"JSON parsing failed: {str(e)}", "original_string": json_string}
|
|
233
|
+
except Exception as e:
|
|
234
|
+
logger.warning("%s Unexpected error during JSON parsing: %s", AGENT_LOG_PREFIX, str(e))
|
|
235
|
+
return {"error": f"Unexpected parsing error: {str(e)}", "original_string": json_string}
|
|
236
|
+
|
|
237
|
+
@abstractmethod
|
|
238
|
+
async def _build_graph(self, state_schema: type) -> CompiledGraph:
|
|
239
|
+
pass
|
aiq/agent/dual_node.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
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
|
+
from abc import abstractmethod
|
|
18
|
+
|
|
19
|
+
from langchain_core.callbacks import AsyncCallbackHandler
|
|
20
|
+
from langchain_core.language_models import BaseChatModel
|
|
21
|
+
from langchain_core.tools import BaseTool
|
|
22
|
+
from langgraph.graph import StateGraph
|
|
23
|
+
from langgraph.graph.graph import CompiledGraph
|
|
24
|
+
from pydantic import BaseModel
|
|
25
|
+
|
|
26
|
+
from .base import AgentDecision
|
|
27
|
+
from .base import BaseAgent
|
|
28
|
+
|
|
29
|
+
log = logging.getLogger(__name__)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class DualNodeAgent(BaseAgent):
|
|
33
|
+
|
|
34
|
+
def __init__(self,
|
|
35
|
+
llm: BaseChatModel,
|
|
36
|
+
tools: list[BaseTool],
|
|
37
|
+
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
38
|
+
detailed_logs: bool = False):
|
|
39
|
+
super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
|
|
40
|
+
|
|
41
|
+
@abstractmethod
|
|
42
|
+
async def agent_node(self, state: BaseModel) -> BaseModel:
|
|
43
|
+
pass
|
|
44
|
+
|
|
45
|
+
@abstractmethod
|
|
46
|
+
async def tool_node(self, state: BaseModel) -> BaseModel:
|
|
47
|
+
pass
|
|
48
|
+
|
|
49
|
+
@abstractmethod
|
|
50
|
+
async def conditional_edge(self, state: BaseModel) -> str:
|
|
51
|
+
pass
|
|
52
|
+
|
|
53
|
+
async def _build_graph(self, state_schema) -> CompiledGraph:
|
|
54
|
+
log.debug("Building and compiling the Agent Graph")
|
|
55
|
+
|
|
56
|
+
graph = StateGraph(state_schema)
|
|
57
|
+
graph.add_node("agent", self.agent_node)
|
|
58
|
+
graph.add_node("tool", self.tool_node)
|
|
59
|
+
graph.add_edge("tool", "agent")
|
|
60
|
+
|
|
61
|
+
conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
|
|
62
|
+
graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
|
|
63
|
+
|
|
64
|
+
graph.set_entry_point("agent")
|
|
65
|
+
self.graph = graph.compile()
|
|
66
|
+
|
|
67
|
+
return self.graph
|
|
File without changes
|
|
@@ -0,0 +1,355 @@
|
|
|
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.agents import AgentAction
|
|
22
|
+
from langchain_core.agents import AgentFinish
|
|
23
|
+
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
24
|
+
from langchain_core.language_models import BaseChatModel
|
|
25
|
+
from langchain_core.messages.ai import AIMessage
|
|
26
|
+
from langchain_core.messages.base import BaseMessage
|
|
27
|
+
from langchain_core.messages.human import HumanMessage
|
|
28
|
+
from langchain_core.messages.tool import ToolMessage
|
|
29
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
30
|
+
from langchain_core.prompts import MessagesPlaceholder
|
|
31
|
+
from langchain_core.runnables.config import RunnableConfig
|
|
32
|
+
from langchain_core.tools import BaseTool
|
|
33
|
+
from pydantic import BaseModel
|
|
34
|
+
from pydantic import Field
|
|
35
|
+
|
|
36
|
+
from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
|
|
37
|
+
from aiq.agent.base import AGENT_LOG_PREFIX
|
|
38
|
+
from aiq.agent.base import INPUT_SCHEMA_MESSAGE
|
|
39
|
+
from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
|
|
40
|
+
from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
|
|
41
|
+
from aiq.agent.base import AgentDecision
|
|
42
|
+
from aiq.agent.dual_node import DualNodeAgent
|
|
43
|
+
from aiq.agent.react_agent.output_parser import ReActOutputParser
|
|
44
|
+
from aiq.agent.react_agent.output_parser import ReActOutputParserException
|
|
45
|
+
from aiq.agent.react_agent.prompt import SYSTEM_PROMPT
|
|
46
|
+
from aiq.agent.react_agent.prompt import USER_PROMPT
|
|
47
|
+
from aiq.agent.react_agent.register import ReActAgentWorkflowConfig
|
|
48
|
+
|
|
49
|
+
logger = logging.getLogger(__name__)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
class ReActGraphState(BaseModel):
|
|
53
|
+
"""State schema for the ReAct Agent Graph"""
|
|
54
|
+
messages: list[BaseMessage] = Field(default_factory=list) # input and output of the ReAct Agent
|
|
55
|
+
agent_scratchpad: list[AgentAction] = Field(default_factory=list) # agent thoughts / intermediate steps
|
|
56
|
+
tool_responses: list[BaseMessage] = Field(default_factory=list) # the responses from any tool calls
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class ReActAgentGraph(DualNodeAgent):
|
|
60
|
+
"""Configurable LangGraph ReAct Agent. A ReAct Agent performs reasoning inbetween tool calls, and utilizes the tool
|
|
61
|
+
names and descriptions to select the optimal tool. Supports retrying on output parsing errors. Argument
|
|
62
|
+
"detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
|
|
63
|
+
|
|
64
|
+
def __init__(self,
|
|
65
|
+
llm: BaseChatModel,
|
|
66
|
+
prompt: ChatPromptTemplate,
|
|
67
|
+
tools: list[BaseTool],
|
|
68
|
+
use_tool_schema: bool = True,
|
|
69
|
+
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
70
|
+
detailed_logs: bool = False,
|
|
71
|
+
retry_agent_response_parsing_errors: bool = True,
|
|
72
|
+
parse_agent_response_max_retries: int = 1,
|
|
73
|
+
tool_call_max_retries: int = 1,
|
|
74
|
+
pass_tool_call_errors_to_agent: bool = True):
|
|
75
|
+
super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
|
|
76
|
+
self.parse_agent_response_max_retries = (parse_agent_response_max_retries
|
|
77
|
+
if retry_agent_response_parsing_errors else 1)
|
|
78
|
+
self.tool_call_max_retries = tool_call_max_retries
|
|
79
|
+
self.pass_tool_call_errors_to_agent = pass_tool_call_errors_to_agent
|
|
80
|
+
logger.debug(
|
|
81
|
+
"%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
|
|
82
|
+
AGENT_LOG_PREFIX)
|
|
83
|
+
tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name # prevent trailing ","
|
|
84
|
+
if not use_tool_schema:
|
|
85
|
+
tool_names_and_descriptions = "\n".join(
|
|
86
|
+
[f"{tool.name}: {tool.description}"
|
|
87
|
+
for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}" # prevent trailing "\n"
|
|
88
|
+
else:
|
|
89
|
+
logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
|
|
90
|
+
tool_names_and_descriptions = "\n".join([
|
|
91
|
+
f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
|
|
92
|
+
for tool in tools[:-1]
|
|
93
|
+
]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
|
|
94
|
+
f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
|
|
95
|
+
prompt = prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
|
|
96
|
+
# construct the ReAct Agent
|
|
97
|
+
bound_llm = llm.bind(stop=["Observation:"]) # type: ignore
|
|
98
|
+
self.agent = prompt | bound_llm
|
|
99
|
+
self.tools_dict = {tool.name: tool for tool in tools}
|
|
100
|
+
logger.debug("%s Initialized ReAct Agent Graph", AGENT_LOG_PREFIX)
|
|
101
|
+
|
|
102
|
+
def _get_tool(self, tool_name: str):
|
|
103
|
+
try:
|
|
104
|
+
return self.tools_dict.get(tool_name)
|
|
105
|
+
except Exception as ex:
|
|
106
|
+
logger.exception("%s Unable to find tool with the name %s\n%s",
|
|
107
|
+
AGENT_LOG_PREFIX,
|
|
108
|
+
tool_name,
|
|
109
|
+
ex,
|
|
110
|
+
exc_info=True)
|
|
111
|
+
raise ex
|
|
112
|
+
|
|
113
|
+
async def agent_node(self, state: ReActGraphState):
|
|
114
|
+
try:
|
|
115
|
+
logger.debug("%s Starting the ReAct Agent Node", AGENT_LOG_PREFIX)
|
|
116
|
+
# keeping a working state allows us to resolve parsing errors without polluting the agent scratchpad
|
|
117
|
+
# the agent "forgets" about the parsing error after solving it - prevents hallucinations in next cycles
|
|
118
|
+
working_state = []
|
|
119
|
+
# Starting from attempt 1 instead of 0 for logging
|
|
120
|
+
for attempt in range(1, self.parse_agent_response_max_retries + 1):
|
|
121
|
+
# the first time we are invoking the ReAct Agent, it won't have any intermediate steps / agent thoughts
|
|
122
|
+
if len(state.agent_scratchpad) == 0 and len(working_state) == 0:
|
|
123
|
+
# the user input comes from the "messages" state channel
|
|
124
|
+
if len(state.messages) == 0:
|
|
125
|
+
raise RuntimeError('No input received in state: "messages"')
|
|
126
|
+
# to check is any human input passed or not, if no input passed Agent will return the state
|
|
127
|
+
content = str(state.messages[0].content)
|
|
128
|
+
if content.strip() == "":
|
|
129
|
+
logger.error("%s No human input passed to the agent.", AGENT_LOG_PREFIX)
|
|
130
|
+
state.messages += [AIMessage(content=NO_INPUT_ERROR_MESSAGE)]
|
|
131
|
+
return state
|
|
132
|
+
question = content
|
|
133
|
+
logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
|
|
134
|
+
|
|
135
|
+
output_message = await self._stream_llm(
|
|
136
|
+
self.agent,
|
|
137
|
+
{"question": question},
|
|
138
|
+
RunnableConfig(callbacks=self.callbacks) # type: ignore
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
if self.detailed_logs:
|
|
142
|
+
logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
|
|
143
|
+
else:
|
|
144
|
+
# ReAct Agents require agentic cycles
|
|
145
|
+
# in an agentic cycle, preserve the agent's thoughts from the previous cycles,
|
|
146
|
+
# and give the agent the response from the tool it called
|
|
147
|
+
agent_scratchpad = []
|
|
148
|
+
for index, intermediate_step in enumerate(state.agent_scratchpad):
|
|
149
|
+
agent_thoughts = AIMessage(content=intermediate_step.log)
|
|
150
|
+
agent_scratchpad.append(agent_thoughts)
|
|
151
|
+
tool_response_content = str(state.tool_responses[index].content)
|
|
152
|
+
tool_response = HumanMessage(content=tool_response_content)
|
|
153
|
+
agent_scratchpad.append(tool_response)
|
|
154
|
+
agent_scratchpad += working_state
|
|
155
|
+
question = str(state.messages[0].content)
|
|
156
|
+
logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
|
|
157
|
+
|
|
158
|
+
output_message = await self._stream_llm(self.agent, {
|
|
159
|
+
"question": question, "agent_scratchpad": agent_scratchpad
|
|
160
|
+
},
|
|
161
|
+
RunnableConfig(callbacks=self.callbacks))
|
|
162
|
+
|
|
163
|
+
if self.detailed_logs:
|
|
164
|
+
logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
|
|
165
|
+
logger.debug("%s The agent's scratchpad (with tool result) was:\n%s",
|
|
166
|
+
AGENT_LOG_PREFIX,
|
|
167
|
+
agent_scratchpad)
|
|
168
|
+
try:
|
|
169
|
+
# check if the agent has the final answer yet
|
|
170
|
+
logger.debug("%s Successfully obtained agent response. Parsing agent's response", AGENT_LOG_PREFIX)
|
|
171
|
+
agent_output = await ReActOutputParser().aparse(output_message.content)
|
|
172
|
+
logger.debug("%s Successfully parsed agent response after %s attempts", AGENT_LOG_PREFIX, attempt)
|
|
173
|
+
if isinstance(agent_output, AgentFinish):
|
|
174
|
+
final_answer = agent_output.return_values.get('output', output_message.content)
|
|
175
|
+
logger.debug("%s The agent has finished, and has the final answer", AGENT_LOG_PREFIX)
|
|
176
|
+
# this is where we handle the final output of the Agent, we can clean-up/format/postprocess here
|
|
177
|
+
# the final answer goes in the "messages" state channel
|
|
178
|
+
state.messages += [AIMessage(content=final_answer)]
|
|
179
|
+
else:
|
|
180
|
+
# the agent wants to call a tool, ensure the thoughts are preserved for the next agentic cycle
|
|
181
|
+
agent_output.log = output_message.content
|
|
182
|
+
logger.debug("%s The agent wants to call a tool: %s", AGENT_LOG_PREFIX, agent_output.tool)
|
|
183
|
+
state.agent_scratchpad += [agent_output]
|
|
184
|
+
|
|
185
|
+
return state
|
|
186
|
+
except ReActOutputParserException as ex:
|
|
187
|
+
# the agent output did not meet the expected ReAct output format. This can happen for a few reasons:
|
|
188
|
+
# the agent mentioned a tool, but already has the final answer, this can happen with Llama models
|
|
189
|
+
# - the ReAct Agent already has the answer, and is reflecting on how it obtained the answer
|
|
190
|
+
# the agent might have also missed Action or Action Input in its output
|
|
191
|
+
logger.debug("%s Error parsing agent output\nObservation:%s\nAgent Output:\n%s",
|
|
192
|
+
AGENT_LOG_PREFIX,
|
|
193
|
+
ex.observation,
|
|
194
|
+
output_message.content)
|
|
195
|
+
if attempt == self.parse_agent_response_max_retries:
|
|
196
|
+
logger.warning(
|
|
197
|
+
"%s Failed to parse agent output after %d attempts, consider enabling or "
|
|
198
|
+
"increasing parse_agent_response_max_retries",
|
|
199
|
+
AGENT_LOG_PREFIX,
|
|
200
|
+
attempt)
|
|
201
|
+
# the final answer goes in the "messages" state channel
|
|
202
|
+
combined_content = str(ex.observation) + '\n' + str(output_message.content)
|
|
203
|
+
output_message.content = combined_content
|
|
204
|
+
state.messages += [output_message]
|
|
205
|
+
return state
|
|
206
|
+
# retry parsing errors, if configured
|
|
207
|
+
logger.info("%s Retrying ReAct Agent, including output parsing Observation", AGENT_LOG_PREFIX)
|
|
208
|
+
working_state.append(output_message)
|
|
209
|
+
working_state.append(HumanMessage(content=str(ex.observation)))
|
|
210
|
+
except Exception as ex:
|
|
211
|
+
logger.exception("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
|
|
212
|
+
raise ex
|
|
213
|
+
|
|
214
|
+
async def conditional_edge(self, state: ReActGraphState):
|
|
215
|
+
try:
|
|
216
|
+
logger.debug("%s Starting the ReAct Conditional Edge", AGENT_LOG_PREFIX)
|
|
217
|
+
if len(state.messages) > 1:
|
|
218
|
+
# the ReAct Agent has finished executing, the last agent output was AgentFinish
|
|
219
|
+
last_message_content = str(state.messages[-1].content)
|
|
220
|
+
logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, last_message_content)
|
|
221
|
+
return AgentDecision.END
|
|
222
|
+
# else the agent wants to call a tool
|
|
223
|
+
agent_output = state.agent_scratchpad[-1]
|
|
224
|
+
logger.debug("%s The agent wants to call: %s with input: %s",
|
|
225
|
+
AGENT_LOG_PREFIX,
|
|
226
|
+
agent_output.tool,
|
|
227
|
+
agent_output.tool_input)
|
|
228
|
+
return AgentDecision.TOOL
|
|
229
|
+
except Exception as ex:
|
|
230
|
+
logger.exception("Failed to determine whether agent is calling a tool: %s", ex, exc_info=True)
|
|
231
|
+
logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
|
|
232
|
+
return AgentDecision.END
|
|
233
|
+
|
|
234
|
+
async def tool_node(self, state: ReActGraphState):
|
|
235
|
+
|
|
236
|
+
logger.debug("%s Starting the Tool Call Node", AGENT_LOG_PREFIX)
|
|
237
|
+
if len(state.agent_scratchpad) == 0:
|
|
238
|
+
raise RuntimeError('No tool input received in state: "agent_scratchpad"')
|
|
239
|
+
agent_thoughts = state.agent_scratchpad[-1]
|
|
240
|
+
# the agent can run any installed tool, simply install the tool and add it to the config file
|
|
241
|
+
requested_tool = self._get_tool(agent_thoughts.tool)
|
|
242
|
+
if not requested_tool:
|
|
243
|
+
configured_tool_names = list(self.tools_dict.keys())
|
|
244
|
+
logger.warning(
|
|
245
|
+
"%s ReAct Agent wants to call tool %s. In the ReAct Agent's configuration within the config file,"
|
|
246
|
+
"there is no tool with that name: %s",
|
|
247
|
+
AGENT_LOG_PREFIX,
|
|
248
|
+
agent_thoughts.tool,
|
|
249
|
+
configured_tool_names)
|
|
250
|
+
tool_response = ToolMessage(name='agent_error',
|
|
251
|
+
tool_call_id='agent_error',
|
|
252
|
+
content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(tool_name=agent_thoughts.tool,
|
|
253
|
+
tools=configured_tool_names))
|
|
254
|
+
state.tool_responses += [tool_response]
|
|
255
|
+
return state
|
|
256
|
+
|
|
257
|
+
logger.debug("%s Calling tool %s with input: %s",
|
|
258
|
+
AGENT_LOG_PREFIX,
|
|
259
|
+
requested_tool.name,
|
|
260
|
+
agent_thoughts.tool_input)
|
|
261
|
+
|
|
262
|
+
# Run the tool. Try to use structured input, if possible.
|
|
263
|
+
try:
|
|
264
|
+
tool_input_str = str(agent_thoughts.tool_input).strip().replace("'", '"')
|
|
265
|
+
tool_input_dict = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
|
|
266
|
+
logger.debug("%s Successfully parsed structured tool input from Action Input", AGENT_LOG_PREFIX)
|
|
267
|
+
|
|
268
|
+
tool_response = await self._call_tool(requested_tool,
|
|
269
|
+
tool_input_dict,
|
|
270
|
+
RunnableConfig(callbacks=self.callbacks),
|
|
271
|
+
max_retries=self.tool_call_max_retries)
|
|
272
|
+
|
|
273
|
+
if self.detailed_logs:
|
|
274
|
+
self._log_tool_response(requested_tool.name, tool_input_dict, str(tool_response.content))
|
|
275
|
+
|
|
276
|
+
except JSONDecodeError as ex:
|
|
277
|
+
logger.debug(
|
|
278
|
+
"%s Unable to parse structured tool input from Action Input. Using Action Input as is."
|
|
279
|
+
"\nParsing error: %s",
|
|
280
|
+
AGENT_LOG_PREFIX,
|
|
281
|
+
ex,
|
|
282
|
+
exc_info=True)
|
|
283
|
+
tool_input_str = str(agent_thoughts.tool_input)
|
|
284
|
+
|
|
285
|
+
tool_response = await self._call_tool(requested_tool,
|
|
286
|
+
tool_input_str,
|
|
287
|
+
RunnableConfig(callbacks=self.callbacks),
|
|
288
|
+
max_retries=self.tool_call_max_retries)
|
|
289
|
+
|
|
290
|
+
if self.detailed_logs:
|
|
291
|
+
self._log_tool_response(requested_tool.name, tool_input_str, str(tool_response.content))
|
|
292
|
+
|
|
293
|
+
if not self.pass_tool_call_errors_to_agent:
|
|
294
|
+
if tool_response.status == "error":
|
|
295
|
+
logger.error("%s Tool %s failed: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_response.content)
|
|
296
|
+
raise RuntimeError("Tool call failed: " + str(tool_response.content))
|
|
297
|
+
|
|
298
|
+
state.tool_responses += [tool_response]
|
|
299
|
+
return state
|
|
300
|
+
|
|
301
|
+
async def build_graph(self):
|
|
302
|
+
try:
|
|
303
|
+
await super()._build_graph(state_schema=ReActGraphState)
|
|
304
|
+
logger.debug("%s ReAct Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
305
|
+
return self.graph
|
|
306
|
+
except Exception as ex:
|
|
307
|
+
logger.exception("%s Failed to build ReAct Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
|
|
308
|
+
raise ex
|
|
309
|
+
|
|
310
|
+
@staticmethod
|
|
311
|
+
def validate_system_prompt(system_prompt: str) -> bool:
|
|
312
|
+
errors = []
|
|
313
|
+
if not system_prompt:
|
|
314
|
+
errors.append("The system prompt cannot be empty.")
|
|
315
|
+
required_prompt_variables = {
|
|
316
|
+
"{tools}": "The system prompt must contain {tools} so the agent knows about configured tools.",
|
|
317
|
+
"{tool_names}": "The system prompt must contain {tool_names} so the agent knows tool names."
|
|
318
|
+
}
|
|
319
|
+
for variable_name, error_message in required_prompt_variables.items():
|
|
320
|
+
if variable_name not in system_prompt:
|
|
321
|
+
errors.append(error_message)
|
|
322
|
+
if errors:
|
|
323
|
+
error_text = "\n".join(errors)
|
|
324
|
+
logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
|
|
325
|
+
raise ValueError(error_text)
|
|
326
|
+
return True
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTemplate:
|
|
330
|
+
"""
|
|
331
|
+
Create a ReAct Agent prompt from the config.
|
|
332
|
+
|
|
333
|
+
Args:
|
|
334
|
+
config (ReActAgentWorkflowConfig): The config to use for the prompt.
|
|
335
|
+
|
|
336
|
+
Returns:
|
|
337
|
+
ChatPromptTemplate: The ReAct Agent prompt.
|
|
338
|
+
"""
|
|
339
|
+
# the ReAct Agent prompt can be customized via config option system_prompt and additional_instructions.
|
|
340
|
+
|
|
341
|
+
if config.system_prompt:
|
|
342
|
+
prompt_str = config.system_prompt
|
|
343
|
+
else:
|
|
344
|
+
prompt_str = SYSTEM_PROMPT
|
|
345
|
+
|
|
346
|
+
if config.additional_instructions:
|
|
347
|
+
prompt_str += f" {config.additional_instructions}"
|
|
348
|
+
|
|
349
|
+
valid_prompt = ReActAgentGraph.validate_system_prompt(prompt_str)
|
|
350
|
+
if not valid_prompt:
|
|
351
|
+
logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
|
|
352
|
+
raise ValueError("Invalid system_prompt")
|
|
353
|
+
prompt = ChatPromptTemplate([("system", prompt_str), ("user", USER_PROMPT),
|
|
354
|
+
MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
|
|
355
|
+
return prompt
|