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,14 @@
|
|
|
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.
|
|
@@ -0,0 +1,35 @@
|
|
|
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
|
+
AUTH_REDIRECT_SUCCESS_HTML = """
|
|
17
|
+
<!DOCTYPE html>
|
|
18
|
+
<html>
|
|
19
|
+
<head>
|
|
20
|
+
<title>Authentication Complete</title>
|
|
21
|
+
<script>
|
|
22
|
+
(function () {
|
|
23
|
+
window.history.replaceState(null, "", window.location.pathname);
|
|
24
|
+
|
|
25
|
+
window.opener?.postMessage({ type: 'AUTH_SUCCESS' }, '*');
|
|
26
|
+
|
|
27
|
+
window.close();
|
|
28
|
+
})();
|
|
29
|
+
</script>
|
|
30
|
+
</head>
|
|
31
|
+
<body>
|
|
32
|
+
<p>Authentication complete. You may now close this window.</p>
|
|
33
|
+
</body>
|
|
34
|
+
</html>
|
|
35
|
+
"""
|
|
@@ -0,0 +1,80 @@
|
|
|
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 logging
|
|
18
|
+
|
|
19
|
+
from aiq.builder.context import AIQContext
|
|
20
|
+
from aiq.data_models.api_server import AIQResponseIntermediateStep
|
|
21
|
+
from aiq.data_models.intermediate_step import IntermediateStep
|
|
22
|
+
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
async def pull_intermediate(_q, adapter):
|
|
27
|
+
"""
|
|
28
|
+
Subscribes to the runner's event stream (which is now a simplified Observable)
|
|
29
|
+
using direct callbacks. Processes each event with the adapter and enqueues
|
|
30
|
+
results to `_q`.
|
|
31
|
+
"""
|
|
32
|
+
intermediate_done = asyncio.Event()
|
|
33
|
+
context = AIQContext.get()
|
|
34
|
+
loop = asyncio.get_running_loop()
|
|
35
|
+
|
|
36
|
+
async def set_intermediate_done():
|
|
37
|
+
intermediate_done.set()
|
|
38
|
+
|
|
39
|
+
def on_next_cb(item: IntermediateStep):
|
|
40
|
+
"""
|
|
41
|
+
Synchronously called whenever the runner publishes an event.
|
|
42
|
+
We process it, then place it into the async queue (via a small async task).
|
|
43
|
+
If adapter is None, convert the raw IntermediateStep into the complete
|
|
44
|
+
AIQResponseIntermediateStep and place it into the queue.
|
|
45
|
+
"""
|
|
46
|
+
if adapter is None:
|
|
47
|
+
adapted = AIQResponseIntermediateStep(id=item.UUID,
|
|
48
|
+
type=item.event_type,
|
|
49
|
+
name=item.name or "",
|
|
50
|
+
parent_id=item.parent_id,
|
|
51
|
+
payload=item.payload.model_dump_json())
|
|
52
|
+
else:
|
|
53
|
+
adapted = adapter.process(item)
|
|
54
|
+
|
|
55
|
+
if adapted is not None:
|
|
56
|
+
loop.create_task(_q.put(adapted))
|
|
57
|
+
|
|
58
|
+
def on_error_cb(exc: Exception):
|
|
59
|
+
"""
|
|
60
|
+
Called if the runner signals an error. We log it and unblock our wait.
|
|
61
|
+
"""
|
|
62
|
+
logger.error("Hit on_error: %s", exc)
|
|
63
|
+
|
|
64
|
+
loop.create_task(set_intermediate_done())
|
|
65
|
+
|
|
66
|
+
def on_complete_cb():
|
|
67
|
+
"""
|
|
68
|
+
Called once the runner signals no more items. We unblock our wait.
|
|
69
|
+
"""
|
|
70
|
+
logger.debug("Completed reading intermediate steps")
|
|
71
|
+
|
|
72
|
+
loop.create_task(set_intermediate_done())
|
|
73
|
+
|
|
74
|
+
# Subscribe to the runner's "reactive_event_stream" (now a simple Observable)
|
|
75
|
+
_ = context.intermediate_step_manager.subscribe(on_next=on_next_cb,
|
|
76
|
+
on_error=on_error_cb,
|
|
77
|
+
on_complete=on_complete_cb)
|
|
78
|
+
|
|
79
|
+
# Wait until on_complete or on_error sets intermediate_done
|
|
80
|
+
return intermediate_done
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import logging
|
|
17
|
+
import os
|
|
18
|
+
import shutil
|
|
19
|
+
import threading
|
|
20
|
+
from datetime import UTC
|
|
21
|
+
from datetime import datetime
|
|
22
|
+
from datetime import timedelta
|
|
23
|
+
from enum import Enum
|
|
24
|
+
from uuid import uuid4
|
|
25
|
+
|
|
26
|
+
from pydantic import BaseModel
|
|
27
|
+
|
|
28
|
+
logger = logging.getLogger(__name__)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class JobStatus(str, Enum):
|
|
32
|
+
SUBMITTED = "submitted"
|
|
33
|
+
RUNNING = "running"
|
|
34
|
+
SUCCESS = "success"
|
|
35
|
+
FAILURE = "failure"
|
|
36
|
+
INTERRUPTED = "interrupted"
|
|
37
|
+
NOT_FOUND = "not_found"
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
# pydantic model for the job status
|
|
41
|
+
class JobInfo(BaseModel):
|
|
42
|
+
job_id: str
|
|
43
|
+
status: JobStatus
|
|
44
|
+
config_file: str | None
|
|
45
|
+
error: str | None
|
|
46
|
+
output_path: str | None
|
|
47
|
+
created_at: datetime
|
|
48
|
+
updated_at: datetime
|
|
49
|
+
expiry_seconds: int
|
|
50
|
+
output: BaseModel | None = None
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
class JobStore:
|
|
54
|
+
|
|
55
|
+
MIN_EXPIRY = 600 # 10 minutes
|
|
56
|
+
MAX_EXPIRY = 86400 # 24 hours
|
|
57
|
+
DEFAULT_EXPIRY = 3600 # 1 hour
|
|
58
|
+
|
|
59
|
+
# active jobs are exempt from expiry
|
|
60
|
+
ACTIVE_STATUS = {"running", "submitted"}
|
|
61
|
+
|
|
62
|
+
def __init__(self):
|
|
63
|
+
self._jobs = {}
|
|
64
|
+
self._lock = threading.Lock() # Ensure thread safety for job operations
|
|
65
|
+
|
|
66
|
+
def create_job(self,
|
|
67
|
+
config_file: str | None = None,
|
|
68
|
+
job_id: str | None = None,
|
|
69
|
+
expiry_seconds: int = DEFAULT_EXPIRY) -> str:
|
|
70
|
+
if job_id is None:
|
|
71
|
+
job_id = str(uuid4())
|
|
72
|
+
|
|
73
|
+
clamped_expiry = max(self.MIN_EXPIRY, min(expiry_seconds, self.MAX_EXPIRY))
|
|
74
|
+
if expiry_seconds != clamped_expiry:
|
|
75
|
+
logger.info("Clamped expiry_seconds from %d to %d for job %s", expiry_seconds, clamped_expiry, job_id)
|
|
76
|
+
|
|
77
|
+
job = JobInfo(job_id=job_id,
|
|
78
|
+
status=JobStatus.SUBMITTED,
|
|
79
|
+
config_file=config_file,
|
|
80
|
+
created_at=datetime.now(UTC),
|
|
81
|
+
updated_at=datetime.now(UTC),
|
|
82
|
+
error=None,
|
|
83
|
+
output_path=None,
|
|
84
|
+
expiry_seconds=clamped_expiry)
|
|
85
|
+
|
|
86
|
+
with self._lock:
|
|
87
|
+
self._jobs[job_id] = job
|
|
88
|
+
|
|
89
|
+
logger.info("Created new job %s with config %s", job_id, config_file)
|
|
90
|
+
return job_id
|
|
91
|
+
|
|
92
|
+
def update_status(self,
|
|
93
|
+
job_id: str,
|
|
94
|
+
status: str,
|
|
95
|
+
error: str | None = None,
|
|
96
|
+
output_path: str | None = None,
|
|
97
|
+
output: BaseModel | None = None):
|
|
98
|
+
if job_id not in self._jobs:
|
|
99
|
+
raise ValueError(f"Job {job_id} not found")
|
|
100
|
+
|
|
101
|
+
with self._lock:
|
|
102
|
+
job = self._jobs[job_id]
|
|
103
|
+
job.status = status
|
|
104
|
+
job.error = error
|
|
105
|
+
job.output_path = output_path
|
|
106
|
+
job.updated_at = datetime.now(UTC)
|
|
107
|
+
job.output = output
|
|
108
|
+
|
|
109
|
+
def get_status(self, job_id: str) -> JobInfo | None:
|
|
110
|
+
with self._lock:
|
|
111
|
+
return self._jobs.get(job_id)
|
|
112
|
+
|
|
113
|
+
def list_jobs(self):
|
|
114
|
+
with self._lock:
|
|
115
|
+
return self._jobs
|
|
116
|
+
|
|
117
|
+
def get_job(self, job_id: str) -> JobInfo | None:
|
|
118
|
+
"""Get a job by its ID."""
|
|
119
|
+
with self._lock:
|
|
120
|
+
return self._jobs.get(job_id)
|
|
121
|
+
|
|
122
|
+
def get_last_job(self) -> JobInfo | None:
|
|
123
|
+
"""Get the last created job."""
|
|
124
|
+
with self._lock:
|
|
125
|
+
if not self._jobs:
|
|
126
|
+
logger.info("No jobs found in job store")
|
|
127
|
+
return None
|
|
128
|
+
last_job = max(self._jobs.values(), key=lambda job: job.created_at)
|
|
129
|
+
logger.info("Retrieved last job %s created at %s", last_job.job_id, last_job.created_at)
|
|
130
|
+
return last_job
|
|
131
|
+
|
|
132
|
+
def get_jobs_by_status(self, status: str) -> list[JobInfo]:
|
|
133
|
+
"""Get all jobs with the specified status."""
|
|
134
|
+
with self._lock:
|
|
135
|
+
return [job for job in self._jobs.values() if job.status == status]
|
|
136
|
+
|
|
137
|
+
def get_all_jobs(self) -> list[JobInfo]:
|
|
138
|
+
"""Get all jobs in the store."""
|
|
139
|
+
with self._lock:
|
|
140
|
+
return list(self._jobs.values())
|
|
141
|
+
|
|
142
|
+
def get_expires_at(self, job: JobInfo) -> datetime | None:
|
|
143
|
+
"""Get the time for a job to expire."""
|
|
144
|
+
if job.status in self.ACTIVE_STATUS:
|
|
145
|
+
return None
|
|
146
|
+
return job.updated_at + timedelta(seconds=job.expiry_seconds)
|
|
147
|
+
|
|
148
|
+
def cleanup_expired_jobs(self):
|
|
149
|
+
"""
|
|
150
|
+
Cleanup expired jobs, keeping the most recent one.
|
|
151
|
+
Updated_at is used instead of created_at to determine the most recent job.
|
|
152
|
+
This is because jobs may not be processed in the order they are created.
|
|
153
|
+
"""
|
|
154
|
+
now = datetime.now(UTC)
|
|
155
|
+
|
|
156
|
+
# Filter out active jobs
|
|
157
|
+
with self._lock:
|
|
158
|
+
finished_jobs = {job_id: job for job_id, job in self._jobs.items() if job.status not in self.ACTIVE_STATUS}
|
|
159
|
+
|
|
160
|
+
# Sort finished jobs by updated_at descending
|
|
161
|
+
sorted_finished = sorted(finished_jobs.items(), key=lambda item: item[1].updated_at, reverse=True)
|
|
162
|
+
|
|
163
|
+
# Always keep the most recent finished job
|
|
164
|
+
jobs_to_check = sorted_finished[1:]
|
|
165
|
+
|
|
166
|
+
expired_ids = []
|
|
167
|
+
for job_id, job in jobs_to_check:
|
|
168
|
+
expires_at = self.get_expires_at(job)
|
|
169
|
+
if expires_at and now > expires_at:
|
|
170
|
+
expired_ids.append(job_id)
|
|
171
|
+
# cleanup output dir if present
|
|
172
|
+
if job.output_path:
|
|
173
|
+
logger.info("Cleaning up output directory for job %s at %s", job_id, job.output_path)
|
|
174
|
+
# If it is a file remove it
|
|
175
|
+
if os.path.isfile(job.output_path):
|
|
176
|
+
os.remove(job.output_path)
|
|
177
|
+
# If it is a directory remove it
|
|
178
|
+
elif os.path.isdir(job.output_path):
|
|
179
|
+
shutil.rmtree(job.output_path)
|
|
180
|
+
|
|
181
|
+
with self._lock:
|
|
182
|
+
for job_id in expired_ids:
|
|
183
|
+
del self._jobs[job_id]
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import importlib
|
|
17
|
+
import logging
|
|
18
|
+
import os
|
|
19
|
+
|
|
20
|
+
from aiq.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorkerBase
|
|
21
|
+
from aiq.runtime.loader import load_config
|
|
22
|
+
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def get_app():
|
|
27
|
+
|
|
28
|
+
config_file_path = os.getenv("AIQ_CONFIG_FILE")
|
|
29
|
+
front_end_worker_full_name = os.getenv("AIQ_FRONT_END_WORKER")
|
|
30
|
+
|
|
31
|
+
if (not config_file_path):
|
|
32
|
+
raise ValueError("Config file not found in environment variable AIQ_CONFIG_FILE.")
|
|
33
|
+
|
|
34
|
+
if (not front_end_worker_full_name):
|
|
35
|
+
raise ValueError("Front end worker not found in environment variable AIQ_FRONT_END_WORKER.")
|
|
36
|
+
|
|
37
|
+
# Try to import the front end worker class
|
|
38
|
+
try:
|
|
39
|
+
# Split the package from the class
|
|
40
|
+
front_end_worker_parts = front_end_worker_full_name.split(".")
|
|
41
|
+
|
|
42
|
+
front_end_worker_module_name = ".".join(front_end_worker_parts[:-1])
|
|
43
|
+
front_end_worker_class_name = front_end_worker_parts[-1]
|
|
44
|
+
|
|
45
|
+
front_end_worker_module = importlib.import_module(front_end_worker_module_name)
|
|
46
|
+
|
|
47
|
+
if not hasattr(front_end_worker_module, front_end_worker_class_name):
|
|
48
|
+
raise ValueError(f"Front end worker {front_end_worker_full_name} not found.")
|
|
49
|
+
|
|
50
|
+
front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = getattr(front_end_worker_module,
|
|
51
|
+
front_end_worker_class_name)
|
|
52
|
+
|
|
53
|
+
if (not issubclass(front_end_worker_class, FastApiFrontEndPluginWorkerBase)):
|
|
54
|
+
raise ValueError(
|
|
55
|
+
f"Front end worker {front_end_worker_full_name} is not a subclass of FastApiFrontEndPluginWorker.")
|
|
56
|
+
|
|
57
|
+
# Load the config
|
|
58
|
+
abs_config_file_path = os.path.abspath(config_file_path)
|
|
59
|
+
|
|
60
|
+
config = load_config(abs_config_file_path)
|
|
61
|
+
|
|
62
|
+
# Create an instance of the front end worker class
|
|
63
|
+
front_end_worker = front_end_worker_class(config)
|
|
64
|
+
|
|
65
|
+
aiq_app = front_end_worker.build_app()
|
|
66
|
+
|
|
67
|
+
return aiq_app
|
|
68
|
+
|
|
69
|
+
except ImportError as e:
|
|
70
|
+
raise ValueError(f"Front end worker {front_end_worker_full_name} not found.") from e
|
|
71
|
+
except Exception as e:
|
|
72
|
+
raise ValueError(f"Error loading front end worker {front_end_worker_full_name}: {e}") from e
|
|
@@ -0,0 +1,298 @@
|
|
|
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 logging
|
|
18
|
+
import typing
|
|
19
|
+
import uuid
|
|
20
|
+
from typing import Any
|
|
21
|
+
|
|
22
|
+
from fastapi import WebSocket
|
|
23
|
+
from pydantic import BaseModel
|
|
24
|
+
from pydantic import ValidationError
|
|
25
|
+
from starlette.websockets import WebSocketDisconnect
|
|
26
|
+
|
|
27
|
+
from aiq.authentication.interfaces import FlowHandlerBase
|
|
28
|
+
from aiq.data_models.api_server import AIQChatResponse
|
|
29
|
+
from aiq.data_models.api_server import AIQResponsePayloadOutput
|
|
30
|
+
from aiq.data_models.api_server import AIQResponseSerializable
|
|
31
|
+
from aiq.data_models.api_server import Error
|
|
32
|
+
from aiq.data_models.api_server import ErrorTypes
|
|
33
|
+
from aiq.data_models.api_server import SystemResponseContent
|
|
34
|
+
from aiq.data_models.api_server import TextContent
|
|
35
|
+
from aiq.data_models.api_server import WebSocketMessageStatus
|
|
36
|
+
from aiq.data_models.api_server import WebSocketMessageType
|
|
37
|
+
from aiq.data_models.api_server import WebSocketSystemInteractionMessage
|
|
38
|
+
from aiq.data_models.api_server import WebSocketSystemIntermediateStepMessage
|
|
39
|
+
from aiq.data_models.api_server import WebSocketSystemResponseTokenMessage
|
|
40
|
+
from aiq.data_models.api_server import WebSocketUserInteractionResponseMessage
|
|
41
|
+
from aiq.data_models.api_server import WebSocketUserMessage
|
|
42
|
+
from aiq.data_models.interactive import HumanPromptNotification
|
|
43
|
+
from aiq.data_models.interactive import HumanResponse
|
|
44
|
+
from aiq.data_models.interactive import HumanResponseNotification
|
|
45
|
+
from aiq.data_models.interactive import InteractionPrompt
|
|
46
|
+
from aiq.front_ends.fastapi.message_validator import MessageValidator
|
|
47
|
+
from aiq.front_ends.fastapi.response_helpers import generate_streaming_response
|
|
48
|
+
from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
|
|
49
|
+
from aiq.runtime.session import AIQSessionManager
|
|
50
|
+
|
|
51
|
+
logger = logging.getLogger(__name__)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
class WebSocketMessageHandler:
|
|
55
|
+
|
|
56
|
+
def __init__(self, socket: WebSocket, session_manager: AIQSessionManager, step_adaptor: StepAdaptor):
|
|
57
|
+
self._socket: WebSocket = socket
|
|
58
|
+
self._session_manager: AIQSessionManager = session_manager
|
|
59
|
+
self._step_adaptor: StepAdaptor = step_adaptor
|
|
60
|
+
|
|
61
|
+
self._message_validator: MessageValidator = MessageValidator()
|
|
62
|
+
self._running_workflow_task: asyncio.Task | None = None
|
|
63
|
+
self._message_parent_id: str = "default_id"
|
|
64
|
+
self._workflow_schema_type: str = None
|
|
65
|
+
self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
|
|
66
|
+
|
|
67
|
+
self._flow_handler: FlowHandlerBase | None = None
|
|
68
|
+
|
|
69
|
+
def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
|
|
70
|
+
self._flow_handler = flow_handler
|
|
71
|
+
|
|
72
|
+
async def __aenter__(self) -> "WebSocketMessageHandler":
|
|
73
|
+
await self._socket.accept()
|
|
74
|
+
|
|
75
|
+
return self
|
|
76
|
+
|
|
77
|
+
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
|
78
|
+
|
|
79
|
+
# TODO: Handle the exit
|
|
80
|
+
pass
|
|
81
|
+
|
|
82
|
+
async def run(self) -> None:
|
|
83
|
+
"""
|
|
84
|
+
Processes received messages from websocket and routes them appropriately.
|
|
85
|
+
"""
|
|
86
|
+
while True:
|
|
87
|
+
|
|
88
|
+
try:
|
|
89
|
+
|
|
90
|
+
message: dict[str, Any] = await self._socket.receive_json()
|
|
91
|
+
|
|
92
|
+
validated_message: BaseModel = await self._message_validator.validate_message(message)
|
|
93
|
+
|
|
94
|
+
# Received a request to start a workflow
|
|
95
|
+
if (isinstance(validated_message, WebSocketUserMessage)):
|
|
96
|
+
await self.process_workflow_request(validated_message)
|
|
97
|
+
|
|
98
|
+
elif isinstance(
|
|
99
|
+
validated_message,
|
|
100
|
+
( # noqa: E131
|
|
101
|
+
WebSocketSystemResponseTokenMessage,
|
|
102
|
+
WebSocketSystemIntermediateStepMessage,
|
|
103
|
+
WebSocketSystemInteractionMessage)):
|
|
104
|
+
# These messages are already handled by self.create_websocket_message(data_model=value, …)
|
|
105
|
+
# No further processing is needed here.
|
|
106
|
+
pass
|
|
107
|
+
|
|
108
|
+
elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
|
|
109
|
+
user_content = await self.process_user_message_content(validated_message)
|
|
110
|
+
self._user_interaction_response.set_result(user_content)
|
|
111
|
+
except (asyncio.CancelledError, WebSocketDisconnect):
|
|
112
|
+
# TODO: Handle the disconnect
|
|
113
|
+
break
|
|
114
|
+
|
|
115
|
+
return None
|
|
116
|
+
|
|
117
|
+
async def process_user_message_content(
|
|
118
|
+
self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
|
|
119
|
+
"""
|
|
120
|
+
Processes the contents of a user message.
|
|
121
|
+
|
|
122
|
+
:param user_content: Incoming content data model.
|
|
123
|
+
:return: A validated Pydantic user content model or None if not found.
|
|
124
|
+
"""
|
|
125
|
+
|
|
126
|
+
for user_message in user_content.content.messages[::-1]:
|
|
127
|
+
if (user_message.role == "user"):
|
|
128
|
+
|
|
129
|
+
for attachment in user_message.content:
|
|
130
|
+
|
|
131
|
+
if isinstance(attachment, TextContent):
|
|
132
|
+
return attachment
|
|
133
|
+
|
|
134
|
+
return None
|
|
135
|
+
|
|
136
|
+
async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
|
|
137
|
+
"""
|
|
138
|
+
Process user messages and routes them appropriately.
|
|
139
|
+
|
|
140
|
+
:param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
|
|
141
|
+
"""
|
|
142
|
+
|
|
143
|
+
try:
|
|
144
|
+
self._message_parent_id = user_message_as_validated_type.id
|
|
145
|
+
self._workflow_schema_type = user_message_as_validated_type.schema_type
|
|
146
|
+
conversation_id: str = user_message_as_validated_type.conversation_id
|
|
147
|
+
|
|
148
|
+
content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
|
|
149
|
+
|
|
150
|
+
if content is None:
|
|
151
|
+
raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
|
|
152
|
+
|
|
153
|
+
if isinstance(content, TextContent) and (self._running_workflow_task is None):
|
|
154
|
+
|
|
155
|
+
def _done_callback(task: asyncio.Task):
|
|
156
|
+
self._running_workflow_task = None
|
|
157
|
+
|
|
158
|
+
# await self._process_response()
|
|
159
|
+
self._running_workflow_task = asyncio.create_task(
|
|
160
|
+
self._run_workflow(content.text, conversation_id,
|
|
161
|
+
result_type=AIQChatResponse)).add_done_callback(_done_callback)
|
|
162
|
+
|
|
163
|
+
except ValueError as e:
|
|
164
|
+
logger.error("User message content not found: %s", str(e), exc_info=True)
|
|
165
|
+
await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT,
|
|
166
|
+
message="User message content could not be found",
|
|
167
|
+
details=str(e)),
|
|
168
|
+
message_type=WebSocketMessageType.ERROR_MESSAGE,
|
|
169
|
+
status=WebSocketMessageStatus.IN_PROGRESS)
|
|
170
|
+
|
|
171
|
+
async def create_websocket_message(self,
|
|
172
|
+
data_model: BaseModel,
|
|
173
|
+
message_type: str | None = None,
|
|
174
|
+
status: str = WebSocketMessageStatus.IN_PROGRESS) -> None:
|
|
175
|
+
"""
|
|
176
|
+
Creates a websocket message that will be ready for routing based on message type or data model.
|
|
177
|
+
|
|
178
|
+
:param data_model: Message content model.
|
|
179
|
+
:param message_type: Message content model.
|
|
180
|
+
:param status: Message content model.
|
|
181
|
+
"""
|
|
182
|
+
try:
|
|
183
|
+
message: BaseModel | None = None
|
|
184
|
+
|
|
185
|
+
if message_type is None:
|
|
186
|
+
message_type = await self._message_validator.resolve_message_type_by_data(data_model)
|
|
187
|
+
|
|
188
|
+
message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type)
|
|
189
|
+
|
|
190
|
+
if 'id' in data_model.model_fields:
|
|
191
|
+
message_id: str = data_model.id
|
|
192
|
+
else:
|
|
193
|
+
message_id = str(uuid.uuid4())
|
|
194
|
+
|
|
195
|
+
content: BaseModel = await self._message_validator.convert_data_to_message_content(data_model)
|
|
196
|
+
|
|
197
|
+
if issubclass(message_schema, WebSocketSystemResponseTokenMessage):
|
|
198
|
+
message = await self._message_validator.create_system_response_token_message(
|
|
199
|
+
message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
|
|
200
|
+
|
|
201
|
+
elif issubclass(message_schema, WebSocketSystemIntermediateStepMessage):
|
|
202
|
+
message = await self._message_validator.create_system_intermediate_step_message(
|
|
203
|
+
message_id=message_id,
|
|
204
|
+
parent_id=await self._message_validator.get_intermediate_step_parent_id(data_model),
|
|
205
|
+
content=content,
|
|
206
|
+
status=status)
|
|
207
|
+
|
|
208
|
+
elif issubclass(message_schema, WebSocketSystemInteractionMessage):
|
|
209
|
+
message = await self._message_validator.create_system_interaction_message(
|
|
210
|
+
message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
|
|
211
|
+
|
|
212
|
+
elif isinstance(content, Error):
|
|
213
|
+
raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
|
|
214
|
+
|
|
215
|
+
elif issubclass(message_schema, Error):
|
|
216
|
+
raise TypeError(f"Invalid message type: {message_type}")
|
|
217
|
+
|
|
218
|
+
elif (message is None):
|
|
219
|
+
raise ValueError(
|
|
220
|
+
f"Message type could not be resolved by input data model: {data_model.model_dump_json()}")
|
|
221
|
+
|
|
222
|
+
except (ValidationError, TypeError, ValueError) as e:
|
|
223
|
+
logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
|
|
224
|
+
message = await self._message_validator.create_system_response_token_message(
|
|
225
|
+
message_type=WebSocketMessageType.ERROR_MESSAGE,
|
|
226
|
+
content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
|
|
227
|
+
|
|
228
|
+
finally:
|
|
229
|
+
if (message is not None):
|
|
230
|
+
await self._socket.send_json(message.model_dump())
|
|
231
|
+
|
|
232
|
+
async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse:
|
|
233
|
+
"""
|
|
234
|
+
Registered human interaction callback that processes human interactions and returns
|
|
235
|
+
responses from websocket connection.
|
|
236
|
+
|
|
237
|
+
:param prompt: Incoming interaction content data model.
|
|
238
|
+
:return: A Text Content Base Pydantic model.
|
|
239
|
+
"""
|
|
240
|
+
|
|
241
|
+
# First create a future from the loop for the human response
|
|
242
|
+
human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
|
|
243
|
+
|
|
244
|
+
# Then add the future to the outstanding human prompts dictionary
|
|
245
|
+
self._user_interaction_response = human_response_future
|
|
246
|
+
|
|
247
|
+
try:
|
|
248
|
+
|
|
249
|
+
await self.create_websocket_message(data_model=prompt.content,
|
|
250
|
+
message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
|
|
251
|
+
status=WebSocketMessageStatus.IN_PROGRESS)
|
|
252
|
+
|
|
253
|
+
if (isinstance(prompt.content, HumanPromptNotification)):
|
|
254
|
+
|
|
255
|
+
return HumanResponseNotification()
|
|
256
|
+
|
|
257
|
+
# Wait for the human response future to complete
|
|
258
|
+
interaction_response: HumanResponse = await human_response_future
|
|
259
|
+
|
|
260
|
+
interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
|
|
261
|
+
interaction_response, prompt.content)
|
|
262
|
+
|
|
263
|
+
return interaction_response
|
|
264
|
+
|
|
265
|
+
finally:
|
|
266
|
+
# Delete the future from the outstanding human prompts dictionary
|
|
267
|
+
self._user_interaction_response = None
|
|
268
|
+
|
|
269
|
+
async def _run_workflow(self,
|
|
270
|
+
payload: typing.Any,
|
|
271
|
+
conversation_id: str | None = None,
|
|
272
|
+
result_type: type | None = None,
|
|
273
|
+
output_type: type | None = None) -> None:
|
|
274
|
+
|
|
275
|
+
try:
|
|
276
|
+
async with self._session_manager.session(
|
|
277
|
+
conversation_id=conversation_id,
|
|
278
|
+
request=self._socket,
|
|
279
|
+
user_input_callback=self.human_interaction_callback,
|
|
280
|
+
user_authentication_callback=(self._flow_handler.authenticate
|
|
281
|
+
if self._flow_handler else None)) as session:
|
|
282
|
+
|
|
283
|
+
async for value in generate_streaming_response(payload,
|
|
284
|
+
session_manager=session,
|
|
285
|
+
streaming=True,
|
|
286
|
+
step_adaptor=self._step_adaptor,
|
|
287
|
+
result_type=result_type,
|
|
288
|
+
output_type=output_type):
|
|
289
|
+
|
|
290
|
+
if not isinstance(value, AIQResponseSerializable):
|
|
291
|
+
value = AIQResponsePayloadOutput(payload=value)
|
|
292
|
+
|
|
293
|
+
await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
|
|
294
|
+
|
|
295
|
+
finally:
|
|
296
|
+
await self.create_websocket_message(data_model=SystemResponseContent(),
|
|
297
|
+
message_type=WebSocketMessageType.RESPONSE_MESSAGE,
|
|
298
|
+
status=WebSocketMessageStatus.COMPLETE)
|