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,30 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from aiq.builder.builder import Builder
|
|
17
|
+
from aiq.cli.register_workflow import register_auth_provider
|
|
18
|
+
from aiq.data_models.authentication import AuthProviderBaseConfig
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
|
|
22
|
+
pass
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
|
|
26
|
+
async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
|
|
27
|
+
|
|
28
|
+
from aiq.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
|
|
29
|
+
|
|
30
|
+
yield HTTPBasicAuthProvider(config)
|
|
@@ -0,0 +1,93 @@
|
|
|
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 typing
|
|
17
|
+
from abc import ABC
|
|
18
|
+
from abc import abstractmethod
|
|
19
|
+
|
|
20
|
+
from aiq.data_models.authentication import AuthenticatedContext
|
|
21
|
+
from aiq.data_models.authentication import AuthFlowType
|
|
22
|
+
from aiq.data_models.authentication import AuthProviderBaseConfig
|
|
23
|
+
from aiq.data_models.authentication import AuthProviderBaseConfigT
|
|
24
|
+
from aiq.data_models.authentication import AuthResult
|
|
25
|
+
|
|
26
|
+
AUTHORIZATION_HEADER = "Authorization"
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
|
|
30
|
+
"""
|
|
31
|
+
Base class for authenticating to API services.
|
|
32
|
+
This class provides an interface for authenticating to API services.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
def __init__(self, config: AuthProviderBaseConfigT):
|
|
36
|
+
"""
|
|
37
|
+
Initialize the AuthProviderBase with the given configuration.
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
config (AuthProviderBaseConfig): Configuration items for authentication.
|
|
41
|
+
"""
|
|
42
|
+
self._config = config
|
|
43
|
+
|
|
44
|
+
@property
|
|
45
|
+
def config(self) -> AuthProviderBaseConfigT:
|
|
46
|
+
"""
|
|
47
|
+
Returns the auth provider configuration object.
|
|
48
|
+
|
|
49
|
+
Returns
|
|
50
|
+
-------
|
|
51
|
+
AuthProviderBaseConfigT
|
|
52
|
+
The auth provider configuration object.
|
|
53
|
+
"""
|
|
54
|
+
return self._config
|
|
55
|
+
|
|
56
|
+
@abstractmethod
|
|
57
|
+
async def authenticate(self, user_id: str | None = None) -> AuthResult:
|
|
58
|
+
"""
|
|
59
|
+
Perform the authentication process for the client.
|
|
60
|
+
|
|
61
|
+
This method handles the necessary steps to authenticate the client with the
|
|
62
|
+
target API service, which may include obtaining tokens, refreshing credentials,
|
|
63
|
+
or completing multi-step authentication flows.
|
|
64
|
+
|
|
65
|
+
Raises:
|
|
66
|
+
NotImplementedError: Must be implemented by subclasses.
|
|
67
|
+
"""
|
|
68
|
+
# This method will call the frontend FlowHandlerBase `authenticate` method
|
|
69
|
+
pass
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class FlowHandlerBase(ABC):
|
|
73
|
+
"""
|
|
74
|
+
Handles front-end specifc flows for authentication clients.
|
|
75
|
+
|
|
76
|
+
Each front end will define a FlowHandler that will implement the authenticate method.
|
|
77
|
+
|
|
78
|
+
The `authenticate` method will be stored as the callback in the AIQContextState.user_auth_callback
|
|
79
|
+
"""
|
|
80
|
+
|
|
81
|
+
@abstractmethod
|
|
82
|
+
async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
|
|
83
|
+
"""
|
|
84
|
+
Perform the authentication process for the client.
|
|
85
|
+
|
|
86
|
+
This method handles the necessary steps to authenticate the client with the
|
|
87
|
+
target API service, which may include obtaining tokens, refreshing credentials,
|
|
88
|
+
or completing multistep authentication flows.
|
|
89
|
+
|
|
90
|
+
Raises:
|
|
91
|
+
NotImplementedError: Must be implemented by subclasses.
|
|
92
|
+
"""
|
|
93
|
+
pass
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from datetime import datetime
|
|
17
|
+
from datetime import timezone
|
|
18
|
+
|
|
19
|
+
from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
|
|
20
|
+
from pydantic import SecretStr
|
|
21
|
+
|
|
22
|
+
from aiq.authentication.interfaces import AuthProviderBase
|
|
23
|
+
from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
|
|
24
|
+
from aiq.builder.context import AIQContext
|
|
25
|
+
from aiq.data_models.authentication import AuthFlowType
|
|
26
|
+
from aiq.data_models.authentication import AuthResult
|
|
27
|
+
from aiq.data_models.authentication import BearerTokenCred
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
|
|
31
|
+
|
|
32
|
+
def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
|
|
33
|
+
super().__init__(config)
|
|
34
|
+
self._authenticated_tokens: dict[str, AuthResult] = {}
|
|
35
|
+
self._context = AIQContext.get()
|
|
36
|
+
|
|
37
|
+
async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
|
|
38
|
+
refresh_token = auth_result.raw.get("refresh_token")
|
|
39
|
+
if not isinstance(refresh_token, str):
|
|
40
|
+
return None
|
|
41
|
+
|
|
42
|
+
with AuthlibOAuth2Client(
|
|
43
|
+
client_id=self.config.client_id,
|
|
44
|
+
client_secret=self.config.client_secret,
|
|
45
|
+
) as client:
|
|
46
|
+
try:
|
|
47
|
+
new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
|
|
48
|
+
except Exception:
|
|
49
|
+
# On any failure, we'll fall back to the full auth flow.
|
|
50
|
+
return None
|
|
51
|
+
|
|
52
|
+
expires_at_ts = new_token_data.get("expires_at")
|
|
53
|
+
new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
|
|
54
|
+
|
|
55
|
+
new_auth_result = AuthResult(
|
|
56
|
+
credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
|
|
57
|
+
token_expires_at=new_expires_at,
|
|
58
|
+
raw=new_token_data,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
self._authenticated_tokens[user_id] = new_auth_result
|
|
62
|
+
|
|
63
|
+
return new_auth_result
|
|
64
|
+
|
|
65
|
+
async def authenticate(self, user_id: str | None = None) -> AuthResult:
|
|
66
|
+
if user_id is None and hasattr(AIQContext.get(), "metadata") and hasattr(
|
|
67
|
+
AIQContext.get().metadata, "cookies") and AIQContext.get().metadata.cookies is not None:
|
|
68
|
+
session_id = AIQContext.get().metadata.cookies.get("aiqtoolkit-session", None)
|
|
69
|
+
if not session_id:
|
|
70
|
+
raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
|
|
71
|
+
|
|
72
|
+
user_id = session_id
|
|
73
|
+
|
|
74
|
+
if user_id and user_id in self._authenticated_tokens:
|
|
75
|
+
auth_result = self._authenticated_tokens[user_id]
|
|
76
|
+
if not auth_result.is_expired():
|
|
77
|
+
return auth_result
|
|
78
|
+
|
|
79
|
+
refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
|
|
80
|
+
if refreshed_auth_result:
|
|
81
|
+
return refreshed_auth_result
|
|
82
|
+
|
|
83
|
+
auth_callback = self._context.user_auth_callback
|
|
84
|
+
if not auth_callback:
|
|
85
|
+
raise RuntimeError("Authentication callback not set on AIQContext.")
|
|
86
|
+
|
|
87
|
+
try:
|
|
88
|
+
authenticated_context = await auth_callback(self.config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
|
|
89
|
+
except Exception as e:
|
|
90
|
+
raise RuntimeError(f"Authentication callback failed: {e}") from e
|
|
91
|
+
|
|
92
|
+
auth_header = authenticated_context.headers.get("Authorization", "")
|
|
93
|
+
if not auth_header.startswith("Bearer "):
|
|
94
|
+
raise RuntimeError("Invalid Authorization header")
|
|
95
|
+
|
|
96
|
+
token = auth_header.split(" ")[1]
|
|
97
|
+
|
|
98
|
+
auth_result = AuthResult(
|
|
99
|
+
credentials=[BearerTokenCred(token=SecretStr(token))],
|
|
100
|
+
token_expires_at=authenticated_context.metadata.get("expires_at"),
|
|
101
|
+
raw=authenticated_context.metadata.get("raw_token"),
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
if user_id:
|
|
105
|
+
self._authenticated_tokens[user_id] = auth_result
|
|
106
|
+
|
|
107
|
+
return auth_result
|
|
@@ -0,0 +1,39 @@
|
|
|
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
|
+
from pydantic import Field
|
|
17
|
+
|
|
18
|
+
from aiq.data_models.authentication import AuthProviderBaseConfig
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"):
|
|
22
|
+
|
|
23
|
+
client_id: str = Field(description="The client ID for OAuth 2.0 authentication.")
|
|
24
|
+
client_secret: str = Field(description="The secret associated with the client_id.")
|
|
25
|
+
authorization_url: str = Field(description="The authorization URL for OAuth 2.0 authentication.")
|
|
26
|
+
token_url: str = Field(description="The token URL for OAuth 2.0 authentication.")
|
|
27
|
+
token_endpoint_auth_method: str | None = Field(
|
|
28
|
+
description=("The authentication method for the token endpoint. "
|
|
29
|
+
"Usually one of `client_secret_post` or `client_secret_basic`."),
|
|
30
|
+
default=None)
|
|
31
|
+
redirect_uri: str = Field(description="The redirect URI for OAuth 2.0 authentication. Must match the registered "
|
|
32
|
+
"redirect URI with the OAuth provider.")
|
|
33
|
+
scopes: list[str] = Field(description="The scopes for OAuth 2.0 authentication.", default_factory=list)
|
|
34
|
+
use_pkce: bool = Field(default=False,
|
|
35
|
+
description="Whether to use PKCE (Proof Key for Code Exchange) in the OAuth 2.0 flow.")
|
|
36
|
+
|
|
37
|
+
authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
|
|
38
|
+
"authorization request."),
|
|
39
|
+
default=None)
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
|
|
17
|
+
from aiq.builder.builder import Builder
|
|
18
|
+
from aiq.cli.register_workflow import register_auth_provider
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig)
|
|
22
|
+
async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder):
|
|
23
|
+
from aiq.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
|
|
24
|
+
|
|
25
|
+
yield OAuth2AuthCodeFlowProvider(authentication_provider)
|
|
@@ -0,0 +1,21 @@
|
|
|
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
|
+
# pylint: disable=unused-import
|
|
17
|
+
# flake8: noqa
|
|
18
|
+
|
|
19
|
+
from aiq.authentication.api_key import register as register_api_key
|
|
20
|
+
from aiq.authentication.http_basic_auth import register as register_http_basic_auth
|
|
21
|
+
from aiq.authentication.oauth2 import register as register_oauth2
|
aiq/builder/__init__.py
ADDED
|
File without changes
|
aiq/builder/builder.py
ADDED
|
@@ -0,0 +1,285 @@
|
|
|
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 asyncio
|
|
17
|
+
import typing
|
|
18
|
+
from abc import ABC
|
|
19
|
+
from abc import abstractmethod
|
|
20
|
+
from collections.abc import Sequence
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
|
|
23
|
+
from aiq.authentication.interfaces import AuthProviderBase
|
|
24
|
+
from aiq.builder.context import AIQContext
|
|
25
|
+
from aiq.builder.framework_enum import LLMFrameworkEnum
|
|
26
|
+
from aiq.builder.function import Function
|
|
27
|
+
from aiq.data_models.authentication import AuthProviderBaseConfig
|
|
28
|
+
from aiq.data_models.component_ref import AuthenticationRef
|
|
29
|
+
from aiq.data_models.component_ref import EmbedderRef
|
|
30
|
+
from aiq.data_models.component_ref import FunctionRef
|
|
31
|
+
from aiq.data_models.component_ref import LLMRef
|
|
32
|
+
from aiq.data_models.component_ref import MemoryRef
|
|
33
|
+
from aiq.data_models.component_ref import ObjectStoreRef
|
|
34
|
+
from aiq.data_models.component_ref import RetrieverRef
|
|
35
|
+
from aiq.data_models.component_ref import TTCStrategyRef
|
|
36
|
+
from aiq.data_models.embedder import EmbedderBaseConfig
|
|
37
|
+
from aiq.data_models.evaluator import EvaluatorBaseConfig
|
|
38
|
+
from aiq.data_models.function import FunctionBaseConfig
|
|
39
|
+
from aiq.data_models.function_dependencies import FunctionDependencies
|
|
40
|
+
from aiq.data_models.llm import LLMBaseConfig
|
|
41
|
+
from aiq.data_models.memory import MemoryBaseConfig
|
|
42
|
+
from aiq.data_models.object_store import ObjectStoreBaseConfig
|
|
43
|
+
from aiq.data_models.retriever import RetrieverBaseConfig
|
|
44
|
+
from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
|
|
45
|
+
from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
|
|
46
|
+
from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
|
|
47
|
+
from aiq.memory.interfaces import MemoryEditor
|
|
48
|
+
from aiq.object_store.interfaces import ObjectStore
|
|
49
|
+
from aiq.retriever.interface import AIQRetriever
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
class UserManagerHolder():
|
|
53
|
+
|
|
54
|
+
def __init__(self, context: AIQContext) -> None:
|
|
55
|
+
self._context = context
|
|
56
|
+
|
|
57
|
+
def get_id(self):
|
|
58
|
+
return self._context.user_manager.get_id()
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
class Builder(ABC): # pylint: disable=too-many-public-methods
|
|
62
|
+
|
|
63
|
+
@abstractmethod
|
|
64
|
+
async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
|
|
65
|
+
pass
|
|
66
|
+
|
|
67
|
+
@abstractmethod
|
|
68
|
+
def get_function(self, name: str | FunctionRef) -> Function:
|
|
69
|
+
pass
|
|
70
|
+
|
|
71
|
+
def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
|
|
72
|
+
|
|
73
|
+
return [self.get_function(name) for name in function_names]
|
|
74
|
+
|
|
75
|
+
@abstractmethod
|
|
76
|
+
def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
|
|
77
|
+
pass
|
|
78
|
+
|
|
79
|
+
@abstractmethod
|
|
80
|
+
async def set_workflow(self, config: FunctionBaseConfig) -> Function:
|
|
81
|
+
pass
|
|
82
|
+
|
|
83
|
+
@abstractmethod
|
|
84
|
+
def get_workflow(self) -> Function:
|
|
85
|
+
pass
|
|
86
|
+
|
|
87
|
+
@abstractmethod
|
|
88
|
+
def get_workflow_config(self) -> FunctionBaseConfig:
|
|
89
|
+
pass
|
|
90
|
+
|
|
91
|
+
def get_tools(self, tool_names: Sequence[str | FunctionRef],
|
|
92
|
+
wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
|
|
93
|
+
|
|
94
|
+
return [self.get_tool(fn_name=n, wrapper_type=wrapper_type) for n in tool_names]
|
|
95
|
+
|
|
96
|
+
@abstractmethod
|
|
97
|
+
def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
|
|
98
|
+
pass
|
|
99
|
+
|
|
100
|
+
@abstractmethod
|
|
101
|
+
async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
|
|
102
|
+
pass
|
|
103
|
+
|
|
104
|
+
@abstractmethod
|
|
105
|
+
async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
|
|
106
|
+
pass
|
|
107
|
+
|
|
108
|
+
async def get_llms(self, llm_names: Sequence[str | LLMRef],
|
|
109
|
+
wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
|
|
110
|
+
|
|
111
|
+
coros = [self.get_llm(llm_name=n, wrapper_type=wrapper_type) for n in llm_names]
|
|
112
|
+
|
|
113
|
+
llms = await asyncio.gather(*coros, return_exceptions=False)
|
|
114
|
+
|
|
115
|
+
return list(llms)
|
|
116
|
+
|
|
117
|
+
@abstractmethod
|
|
118
|
+
def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
|
|
119
|
+
pass
|
|
120
|
+
|
|
121
|
+
@abstractmethod
|
|
122
|
+
async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
|
|
123
|
+
pass
|
|
124
|
+
|
|
125
|
+
@abstractmethod
|
|
126
|
+
async def get_auth_provider(self, auth_provider_name: str | AuthenticationRef) -> AuthProviderBase:
|
|
127
|
+
pass
|
|
128
|
+
|
|
129
|
+
async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]):
|
|
130
|
+
|
|
131
|
+
coros = [self.get_auth_provider(auth_provider_name=n) for n in auth_provider_names]
|
|
132
|
+
|
|
133
|
+
auth_providers = await asyncio.gather(*coros, return_exceptions=False)
|
|
134
|
+
|
|
135
|
+
return list(auth_providers)
|
|
136
|
+
|
|
137
|
+
@abstractmethod
|
|
138
|
+
async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
|
|
139
|
+
pass
|
|
140
|
+
|
|
141
|
+
async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
|
|
142
|
+
"""
|
|
143
|
+
Return a list of all object store clients.
|
|
144
|
+
"""
|
|
145
|
+
return list(await asyncio.gather(*[self.get_object_store_client(name) for name in object_store_names]))
|
|
146
|
+
|
|
147
|
+
@abstractmethod
|
|
148
|
+
async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
|
|
149
|
+
pass
|
|
150
|
+
|
|
151
|
+
@abstractmethod
|
|
152
|
+
def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
|
|
153
|
+
pass
|
|
154
|
+
|
|
155
|
+
@abstractmethod
|
|
156
|
+
async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
|
|
157
|
+
pass
|
|
158
|
+
|
|
159
|
+
async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
|
|
160
|
+
wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
|
|
161
|
+
|
|
162
|
+
coros = [self.get_embedder(embedder_name=n, wrapper_type=wrapper_type) for n in embedder_names]
|
|
163
|
+
|
|
164
|
+
embedders = await asyncio.gather(*coros, return_exceptions=False)
|
|
165
|
+
|
|
166
|
+
return list(embedders)
|
|
167
|
+
|
|
168
|
+
@abstractmethod
|
|
169
|
+
async def get_embedder(self, embedder_name: str | EmbedderRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
|
|
170
|
+
pass
|
|
171
|
+
|
|
172
|
+
@abstractmethod
|
|
173
|
+
def get_embedder_config(self, embedder_name: str | EmbedderRef) -> EmbedderBaseConfig:
|
|
174
|
+
pass
|
|
175
|
+
|
|
176
|
+
@abstractmethod
|
|
177
|
+
async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig):
|
|
178
|
+
pass
|
|
179
|
+
|
|
180
|
+
def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
|
|
181
|
+
"""
|
|
182
|
+
Return a list of memory clients for the specified names.
|
|
183
|
+
"""
|
|
184
|
+
return [self.get_memory_client(n) for n in memory_names]
|
|
185
|
+
|
|
186
|
+
@abstractmethod
|
|
187
|
+
def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
|
|
188
|
+
"""
|
|
189
|
+
Return the instantiated memory client for the given name.
|
|
190
|
+
"""
|
|
191
|
+
pass
|
|
192
|
+
|
|
193
|
+
@abstractmethod
|
|
194
|
+
def get_memory_client_config(self, memory_name: str | MemoryRef) -> MemoryBaseConfig:
|
|
195
|
+
pass
|
|
196
|
+
|
|
197
|
+
@abstractmethod
|
|
198
|
+
async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
|
|
199
|
+
pass
|
|
200
|
+
|
|
201
|
+
async def get_retrievers(self,
|
|
202
|
+
retriever_names: Sequence[str | RetrieverRef],
|
|
203
|
+
wrapper_type: LLMFrameworkEnum | str | None = None):
|
|
204
|
+
|
|
205
|
+
tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
|
|
206
|
+
|
|
207
|
+
retrievers = await asyncio.gather(*tasks, return_exceptions=False)
|
|
208
|
+
|
|
209
|
+
return list(retrievers)
|
|
210
|
+
|
|
211
|
+
@typing.overload
|
|
212
|
+
async def get_retriever(self, retriever_name: str | RetrieverRef,
|
|
213
|
+
wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
|
|
214
|
+
...
|
|
215
|
+
|
|
216
|
+
@typing.overload
|
|
217
|
+
async def get_retriever(self, retriever_name: str | RetrieverRef, wrapper_type: None) -> AIQRetriever:
|
|
218
|
+
...
|
|
219
|
+
|
|
220
|
+
@typing.overload
|
|
221
|
+
async def get_retriever(self, retriever_name: str | RetrieverRef) -> AIQRetriever:
|
|
222
|
+
...
|
|
223
|
+
|
|
224
|
+
@abstractmethod
|
|
225
|
+
async def get_retriever(self,
|
|
226
|
+
retriever_name: str | RetrieverRef,
|
|
227
|
+
wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:
|
|
228
|
+
pass
|
|
229
|
+
|
|
230
|
+
@abstractmethod
|
|
231
|
+
async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
|
|
232
|
+
pass
|
|
233
|
+
|
|
234
|
+
@abstractmethod
|
|
235
|
+
async def add_ttc_strategy(self, name: str | str, config: TTCStrategyBaseConfig):
|
|
236
|
+
pass
|
|
237
|
+
|
|
238
|
+
@abstractmethod
|
|
239
|
+
async def get_ttc_strategy(self,
|
|
240
|
+
strategy_name: str | TTCStrategyRef,
|
|
241
|
+
pipeline_type: PipelineTypeEnum,
|
|
242
|
+
stage_type: StageTypeEnum):
|
|
243
|
+
pass
|
|
244
|
+
|
|
245
|
+
@abstractmethod
|
|
246
|
+
async def get_ttc_strategy_config(self,
|
|
247
|
+
strategy_name: str | TTCStrategyRef,
|
|
248
|
+
pipeline_type: PipelineTypeEnum,
|
|
249
|
+
stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
|
|
250
|
+
pass
|
|
251
|
+
|
|
252
|
+
@abstractmethod
|
|
253
|
+
def get_user_manager(self) -> UserManagerHolder:
|
|
254
|
+
pass
|
|
255
|
+
|
|
256
|
+
@abstractmethod
|
|
257
|
+
def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
|
|
258
|
+
pass
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
class EvalBuilder(Builder):
|
|
262
|
+
|
|
263
|
+
@abstractmethod
|
|
264
|
+
async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
|
|
265
|
+
pass
|
|
266
|
+
|
|
267
|
+
@abstractmethod
|
|
268
|
+
def get_evaluator(self, evaluator_name: str) -> typing.Any:
|
|
269
|
+
pass
|
|
270
|
+
|
|
271
|
+
@abstractmethod
|
|
272
|
+
def get_evaluator_config(self, evaluator_name: str) -> EvaluatorBaseConfig:
|
|
273
|
+
pass
|
|
274
|
+
|
|
275
|
+
@abstractmethod
|
|
276
|
+
def get_max_concurrency(self) -> int:
|
|
277
|
+
pass
|
|
278
|
+
|
|
279
|
+
@abstractmethod
|
|
280
|
+
def get_output_dir(self) -> Path:
|
|
281
|
+
pass
|
|
282
|
+
|
|
283
|
+
@abstractmethod
|
|
284
|
+
def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
|
|
285
|
+
pass
|