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,335 @@
|
|
|
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 importlib.metadata
|
|
17
|
+
import inspect
|
|
18
|
+
import json
|
|
19
|
+
import logging
|
|
20
|
+
import typing
|
|
21
|
+
from enum import Enum
|
|
22
|
+
from functools import lru_cache
|
|
23
|
+
from pathlib import Path
|
|
24
|
+
from types import ModuleType
|
|
25
|
+
from typing import TYPE_CHECKING
|
|
26
|
+
|
|
27
|
+
from pydantic import BaseModel
|
|
28
|
+
from pydantic import field_validator
|
|
29
|
+
|
|
30
|
+
from aiq.builder.framework_enum import LLMFrameworkEnum
|
|
31
|
+
from aiq.data_models.component import AIQComponentEnum
|
|
32
|
+
from aiq.utils.metadata_utils import generate_config_type_docs
|
|
33
|
+
|
|
34
|
+
if TYPE_CHECKING:
|
|
35
|
+
from aiq.cli.type_registry import ToolWrapperBuildCallableT
|
|
36
|
+
from aiq.data_models.common import TypedBaseModelT
|
|
37
|
+
|
|
38
|
+
logger = logging.getLogger(__name__)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class DiscoveryStatusEnum(str, Enum):
|
|
42
|
+
SUCCESS = "success"
|
|
43
|
+
FAILURE = "failure"
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class DiscoveryContractFieldsEnum(str, Enum):
|
|
47
|
+
PACKAGE = "package"
|
|
48
|
+
VERSION = "version"
|
|
49
|
+
COMPONENT_TYPE = "component_type"
|
|
50
|
+
COMPONENT_NAME = "component_name"
|
|
51
|
+
DESCRIPTION = "description"
|
|
52
|
+
DEVELOPER_NOTES = "developer_notes"
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class DiscoveryMetadata(BaseModel):
|
|
56
|
+
"""A data model representing metadata about each registered component to faciliate its discovery.
|
|
57
|
+
|
|
58
|
+
Args:
|
|
59
|
+
package (str): The name of the package containing the AIQ Toolkit component.
|
|
60
|
+
version (str): The version number of the package containing the AIQ Toolkit component.
|
|
61
|
+
component_type (AIQComponentEnum): The type of AIQ Toolkit component this metadata represents.
|
|
62
|
+
component_name (str): The registered name of the AIQ Toolkit component.
|
|
63
|
+
description (str): Description of the AIQ Toolkit component pulled from its config objects docstrings.
|
|
64
|
+
developer_notes (str): Other notes to a developers to aid in the use of the component.
|
|
65
|
+
status (DiscoveryStatusEnum): Provides the status of the metadata discovery process.
|
|
66
|
+
"""
|
|
67
|
+
|
|
68
|
+
package: str = ""
|
|
69
|
+
version: str = ""
|
|
70
|
+
component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED
|
|
71
|
+
component_name: str = ""
|
|
72
|
+
description: str = ""
|
|
73
|
+
developer_notes: str = ""
|
|
74
|
+
status: DiscoveryStatusEnum = DiscoveryStatusEnum.SUCCESS
|
|
75
|
+
|
|
76
|
+
@field_validator("description", mode="before")
|
|
77
|
+
@classmethod
|
|
78
|
+
def ensure_description_string(cls, v: typing.Any):
|
|
79
|
+
if not isinstance(v, str):
|
|
80
|
+
return ""
|
|
81
|
+
return v
|
|
82
|
+
|
|
83
|
+
@staticmethod
|
|
84
|
+
def get_preferred_item(items: list, preferred: str) -> str:
|
|
85
|
+
return preferred if preferred in items else items[0]
|
|
86
|
+
|
|
87
|
+
@staticmethod
|
|
88
|
+
@lru_cache
|
|
89
|
+
def get_distribution_name_from_metadata(root_package_name: str) -> str | None:
|
|
90
|
+
"""
|
|
91
|
+
This is not performant and is only present to be used (not used
|
|
92
|
+
currently) as a fallback when the distro name doesn't match the
|
|
93
|
+
module name and private_data is not available to map it.
|
|
94
|
+
"""
|
|
95
|
+
mapping = importlib.metadata.packages_distributions()
|
|
96
|
+
try:
|
|
97
|
+
distro_names = mapping.get(root_package_name, [None])
|
|
98
|
+
distro_name = DiscoveryMetadata.get_preferred_item(distro_names, "aiqtoolkit")
|
|
99
|
+
except KeyError:
|
|
100
|
+
return root_package_name
|
|
101
|
+
|
|
102
|
+
return distro_name if distro_name else root_package_name
|
|
103
|
+
|
|
104
|
+
@staticmethod
|
|
105
|
+
@lru_cache
|
|
106
|
+
def get_distribution_name_from_private_data(root_package: str) -> str | None:
|
|
107
|
+
# Locate distibution mapping stored in the packages private data
|
|
108
|
+
module = __import__(root_package)
|
|
109
|
+
for path in module.__path__:
|
|
110
|
+
package_dir = Path(path).resolve()
|
|
111
|
+
distinfo_path = package_dir / "meta" / "module_to_distro.json"
|
|
112
|
+
|
|
113
|
+
if distinfo_path.exists():
|
|
114
|
+
with distinfo_path.open("r") as f:
|
|
115
|
+
data = json.load(f)
|
|
116
|
+
return data.get(root_package, None)
|
|
117
|
+
return None
|
|
118
|
+
|
|
119
|
+
@staticmethod
|
|
120
|
+
@lru_cache
|
|
121
|
+
def get_distribution_name_from_module(module: ModuleType | None) -> str:
|
|
122
|
+
"""Get the distribution name from the config type using the mapping of module names to distro names.
|
|
123
|
+
|
|
124
|
+
Args:
|
|
125
|
+
module (ModuleType): A registered component's module.
|
|
126
|
+
|
|
127
|
+
Returns:
|
|
128
|
+
str: The distribution name of the AIQ Toolkit component.
|
|
129
|
+
"""
|
|
130
|
+
from aiq.runtime.loader import get_all_aiq_entrypoints_distro_mapping
|
|
131
|
+
|
|
132
|
+
if module is None:
|
|
133
|
+
return "aiqtoolkit"
|
|
134
|
+
|
|
135
|
+
# Get the mapping of module names to distro names
|
|
136
|
+
mapping = get_all_aiq_entrypoints_distro_mapping()
|
|
137
|
+
module_package = module.__package__
|
|
138
|
+
|
|
139
|
+
if module_package is None:
|
|
140
|
+
return "aiqtoolkit"
|
|
141
|
+
|
|
142
|
+
# Traverse the module package parts in reverse order to find the distro name
|
|
143
|
+
# This is because the module package is the root package for the AIQ Toolkit component
|
|
144
|
+
# and the distro name is the name of the package that contains the component
|
|
145
|
+
module_package_parts = module_package.split(".")
|
|
146
|
+
for part_idx in range(len(module_package_parts), 0, -1):
|
|
147
|
+
candidate_module_name = ".".join(module_package_parts[0:part_idx])
|
|
148
|
+
candidate_distro_name = mapping.get(candidate_module_name, None)
|
|
149
|
+
if candidate_distro_name is not None:
|
|
150
|
+
return candidate_distro_name
|
|
151
|
+
|
|
152
|
+
return "aiqtoolkit"
|
|
153
|
+
|
|
154
|
+
@staticmethod
|
|
155
|
+
@lru_cache
|
|
156
|
+
def get_distribution_name_from_config_type(config_type: type["TypedBaseModelT"]) -> str:
|
|
157
|
+
"""Get the distribution name from the config type using the mapping of module names to distro names.
|
|
158
|
+
|
|
159
|
+
Args:
|
|
160
|
+
config_type (type[TypedBaseModelT]): A registered component's configuration object.
|
|
161
|
+
|
|
162
|
+
Returns:
|
|
163
|
+
str: The distribution name of the AIQ Toolkit component.
|
|
164
|
+
"""
|
|
165
|
+
module = inspect.getmodule(config_type)
|
|
166
|
+
return DiscoveryMetadata.get_distribution_name_from_module(module)
|
|
167
|
+
|
|
168
|
+
@staticmethod
|
|
169
|
+
@lru_cache
|
|
170
|
+
def get_distribution_name(root_package: str) -> str:
|
|
171
|
+
"""
|
|
172
|
+
The aiq library packages use a distro name 'aiqtoolkit[]' and
|
|
173
|
+
root package name 'aiq'. They provide mapping in a metadata file
|
|
174
|
+
for optimized installation.
|
|
175
|
+
"""
|
|
176
|
+
|
|
177
|
+
distro_name = DiscoveryMetadata.get_distribution_name_from_private_data(root_package)
|
|
178
|
+
return distro_name if distro_name else root_package
|
|
179
|
+
|
|
180
|
+
@staticmethod
|
|
181
|
+
def from_config_type(config_type: type["TypedBaseModelT"],
|
|
182
|
+
component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
|
|
183
|
+
"""Generates discovery metadata from an AIQ Toolkit config object.
|
|
184
|
+
|
|
185
|
+
Args:
|
|
186
|
+
config_type (type[TypedBaseModelT]): A registered component's configuration object.
|
|
187
|
+
component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
|
|
188
|
+
AIQComponentEnum.UNDEFINED.
|
|
189
|
+
|
|
190
|
+
Returns:
|
|
191
|
+
DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
|
|
192
|
+
"""
|
|
193
|
+
|
|
194
|
+
try:
|
|
195
|
+
module = inspect.getmodule(config_type)
|
|
196
|
+
distro_name = DiscoveryMetadata.get_distribution_name_from_config_type(config_type)
|
|
197
|
+
|
|
198
|
+
if not distro_name:
|
|
199
|
+
# raise an exception
|
|
200
|
+
logger.error("Encountered issue getting distro_name for module %s", module.__name__)
|
|
201
|
+
return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
|
|
202
|
+
|
|
203
|
+
try:
|
|
204
|
+
version = importlib.metadata.version(distro_name) if distro_name != "" else ""
|
|
205
|
+
except importlib.metadata.PackageNotFoundError:
|
|
206
|
+
logger.warning("Package metadata not found for %s", distro_name)
|
|
207
|
+
version = ""
|
|
208
|
+
except Exception as e:
|
|
209
|
+
logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
|
|
210
|
+
return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
|
|
211
|
+
|
|
212
|
+
description = generate_config_type_docs(config_type=config_type)
|
|
213
|
+
|
|
214
|
+
return DiscoveryMetadata(package=distro_name,
|
|
215
|
+
version=version,
|
|
216
|
+
component_type=component_type,
|
|
217
|
+
component_name=config_type.static_type(),
|
|
218
|
+
description=description)
|
|
219
|
+
|
|
220
|
+
@staticmethod
|
|
221
|
+
def from_fn_wrapper(fn: "ToolWrapperBuildCallableT",
|
|
222
|
+
wrapper_type: LLMFrameworkEnum | str,
|
|
223
|
+
component_type: AIQComponentEnum = AIQComponentEnum.TOOL_WRAPPER) -> "DiscoveryMetadata":
|
|
224
|
+
"""Generates discovery metadata from function with specified wrapper type.
|
|
225
|
+
|
|
226
|
+
Args:
|
|
227
|
+
fn (ToolWrapperBuildCallableT): A tool wrapper callable to source component metadata.
|
|
228
|
+
wrapper_type (LLMFrameworkEnum): The wrapper to apply to the callable to faciliate inter-framwork
|
|
229
|
+
interoperability.
|
|
230
|
+
|
|
231
|
+
component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
|
|
232
|
+
AIQComponentEnum.TOOL_WRAPPER.
|
|
233
|
+
|
|
234
|
+
Returns:
|
|
235
|
+
DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
|
|
236
|
+
"""
|
|
237
|
+
|
|
238
|
+
try:
|
|
239
|
+
module = inspect.getmodule(fn)
|
|
240
|
+
distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
|
|
241
|
+
|
|
242
|
+
try:
|
|
243
|
+
# version = importlib.metadata.version(root_package) if root_package != "" else ""
|
|
244
|
+
version = importlib.metadata.version(distro_name) if distro_name != "" else ""
|
|
245
|
+
except importlib.metadata.PackageNotFoundError:
|
|
246
|
+
logger.warning("Package metadata not found for %s", distro_name)
|
|
247
|
+
version = ""
|
|
248
|
+
except Exception as e:
|
|
249
|
+
logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
|
|
250
|
+
return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
|
|
251
|
+
|
|
252
|
+
if isinstance(wrapper_type, LLMFrameworkEnum):
|
|
253
|
+
wrapper_type = wrapper_type.value
|
|
254
|
+
|
|
255
|
+
return DiscoveryMetadata(package=distro_name,
|
|
256
|
+
version=version,
|
|
257
|
+
component_type=component_type,
|
|
258
|
+
component_name=wrapper_type,
|
|
259
|
+
description=fn.__doc__ or "")
|
|
260
|
+
|
|
261
|
+
@staticmethod
|
|
262
|
+
def from_package_name(package_name: str, package_version: str | None) -> "DiscoveryMetadata":
|
|
263
|
+
"""Generates discovery metadata from an installed package name.
|
|
264
|
+
|
|
265
|
+
Args:
|
|
266
|
+
package_name (str): The name of the AIQ Toolkit plugin package containing registered components.
|
|
267
|
+
package_version (str, optional): The version of the package, Defaults to None.
|
|
268
|
+
|
|
269
|
+
Returns:
|
|
270
|
+
DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
|
|
271
|
+
"""
|
|
272
|
+
|
|
273
|
+
try:
|
|
274
|
+
try:
|
|
275
|
+
metadata = importlib.metadata.metadata(package_name)
|
|
276
|
+
description = metadata.get("Summary", "")
|
|
277
|
+
if (package_version is None):
|
|
278
|
+
package_version = importlib.metadata.version(package_name)
|
|
279
|
+
except importlib.metadata.PackageNotFoundError:
|
|
280
|
+
logger.warning("Package metadata not found for %s", package_name)
|
|
281
|
+
description = ""
|
|
282
|
+
package_version = package_version or ""
|
|
283
|
+
except Exception as e:
|
|
284
|
+
logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e, exc_info=True)
|
|
285
|
+
return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
|
|
286
|
+
|
|
287
|
+
return DiscoveryMetadata(package=package_name,
|
|
288
|
+
version=package_version,
|
|
289
|
+
component_type=AIQComponentEnum.PACKAGE,
|
|
290
|
+
component_name=package_name,
|
|
291
|
+
description=description)
|
|
292
|
+
|
|
293
|
+
@staticmethod
|
|
294
|
+
def from_provider_framework_map(
|
|
295
|
+
config_type: type["TypedBaseModelT"],
|
|
296
|
+
wrapper_type: LLMFrameworkEnum | str | None,
|
|
297
|
+
provider_type: AIQComponentEnum,
|
|
298
|
+
component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
|
|
299
|
+
"""Generates discovery metadata from provider and framework mapping information.
|
|
300
|
+
|
|
301
|
+
Args:
|
|
302
|
+
config_type (type[TypedBaseModelT]): A registered component's configuration object.
|
|
303
|
+
wrapper_type (LLMFrameworkEnum | str): The wrapper to apply to the callable to faciliate inter-framwork
|
|
304
|
+
interoperability.
|
|
305
|
+
|
|
306
|
+
provider_type (AIQComponentEnum): The type of provider the registered component supports.
|
|
307
|
+
component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
|
|
308
|
+
AIQComponentEnum.UNDEFINED.
|
|
309
|
+
|
|
310
|
+
Returns:
|
|
311
|
+
DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
|
|
312
|
+
"""
|
|
313
|
+
|
|
314
|
+
try:
|
|
315
|
+
module = inspect.getmodule(config_type)
|
|
316
|
+
distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
|
|
317
|
+
try:
|
|
318
|
+
version = importlib.metadata.version(distro_name) if distro_name != "" else ""
|
|
319
|
+
except importlib.metadata.PackageNotFoundError:
|
|
320
|
+
logger.warning("Package metadata not found for %s", distro_name)
|
|
321
|
+
version = ""
|
|
322
|
+
except Exception as e:
|
|
323
|
+
logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
|
|
324
|
+
return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
|
|
325
|
+
|
|
326
|
+
wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type
|
|
327
|
+
component_name = f"{config_type.static_type()} ({provider_type.value}) - {wrapper_type}"
|
|
328
|
+
|
|
329
|
+
description = generate_config_type_docs(config_type=config_type)
|
|
330
|
+
|
|
331
|
+
return DiscoveryMetadata(package=distro_name,
|
|
332
|
+
version=version,
|
|
333
|
+
component_type=component_type,
|
|
334
|
+
component_name=component_name,
|
|
335
|
+
description=description)
|
|
@@ -0,0 +1,27 @@
|
|
|
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
|
+
|
|
18
|
+
from .common import BaseModelRegistryTag
|
|
19
|
+
from .common import TypedBaseModel
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class EmbedderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
|
|
23
|
+
""" Base configuration for embedding model providers. """
|
|
24
|
+
pass
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
EmbedderBaseConfigT = typing.TypeVar("EmbedderBaseConfigT", bound=EmbedderBaseConfig)
|
|
@@ -0,0 +1,127 @@
|
|
|
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 enum import Enum
|
|
18
|
+
from pathlib import Path
|
|
19
|
+
|
|
20
|
+
from pydantic import BaseModel
|
|
21
|
+
from pydantic import Discriminator
|
|
22
|
+
from pydantic import model_validator
|
|
23
|
+
|
|
24
|
+
from aiq.data_models.common import TypedBaseModel
|
|
25
|
+
from aiq.data_models.dataset_handler import EvalDatasetConfig
|
|
26
|
+
from aiq.data_models.dataset_handler import EvalS3Config
|
|
27
|
+
from aiq.data_models.evaluator import EvaluatorBaseConfig
|
|
28
|
+
from aiq.data_models.intermediate_step import IntermediateStepType
|
|
29
|
+
from aiq.data_models.profiler import ProfilerConfig
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class JobEvictionPolicy(str, Enum):
|
|
33
|
+
"""Policy for evicting old jobs when max_jobs is exceeded."""
|
|
34
|
+
TIME_CREATED = "time_created"
|
|
35
|
+
TIME_MODIFIED = "time_modified"
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class EvalCustomScriptConfig(BaseModel):
|
|
39
|
+
# Path to the script to run
|
|
40
|
+
script: Path
|
|
41
|
+
# Keyword arguments to pass to the script
|
|
42
|
+
kwargs: dict[str, str] = {}
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class JobManagementConfig(BaseModel):
|
|
46
|
+
# Whether to append a unique job ID to the output directory for each run
|
|
47
|
+
append_job_id_to_output_dir: bool = False
|
|
48
|
+
# Maximum number of jobs to keep in the output directory. Oldest jobs will be evicted.
|
|
49
|
+
# A value of 0 means no limit.
|
|
50
|
+
max_jobs: int = 0
|
|
51
|
+
# Policy for evicting old jobs. Defaults to using time_created.
|
|
52
|
+
eviction_policy: JobEvictionPolicy = JobEvictionPolicy.TIME_CREATED
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class EvalOutputConfig(BaseModel):
|
|
56
|
+
# Output directory for the workflow and evaluation results
|
|
57
|
+
dir: Path = Path("/tmp/aiq/examples/default/")
|
|
58
|
+
# S3 prefix for the workflow and evaluation results
|
|
59
|
+
remote_dir: str | None = None
|
|
60
|
+
# Custom scripts to run after the workflow and evaluation results are saved
|
|
61
|
+
custom_scripts: dict[str, EvalCustomScriptConfig] = {}
|
|
62
|
+
# S3 config for uploading the contents of the output directory
|
|
63
|
+
s3: EvalS3Config | None = None
|
|
64
|
+
# Whether to cleanup the output directory before running the workflow
|
|
65
|
+
cleanup: bool = True
|
|
66
|
+
# Job management configuration (job id, eviction, etc.)
|
|
67
|
+
job_management: JobManagementConfig = JobManagementConfig()
|
|
68
|
+
# Filter for the workflow output steps
|
|
69
|
+
workflow_output_step_filter: list[IntermediateStepType] | None = None
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class EvalGeneralConfig(BaseModel):
|
|
73
|
+
max_concurrency: int = 8
|
|
74
|
+
|
|
75
|
+
# Workflow alias for displaying in evaluation UI, if not provided,
|
|
76
|
+
# the workflow type will be used
|
|
77
|
+
workflow_alias: str | None = None
|
|
78
|
+
|
|
79
|
+
# Output directory for the workflow and evaluation results
|
|
80
|
+
output_dir: Path = Path("/tmp/aiq/examples/default/")
|
|
81
|
+
|
|
82
|
+
# If present overrides output_dir
|
|
83
|
+
output: EvalOutputConfig | None = None
|
|
84
|
+
|
|
85
|
+
# Dataset for running the workflow and evaluating
|
|
86
|
+
dataset: EvalDatasetConfig | None = None
|
|
87
|
+
|
|
88
|
+
# Inference profiler
|
|
89
|
+
profiler: ProfilerConfig | None = None
|
|
90
|
+
|
|
91
|
+
# overwrite the output_dir with the output config if present
|
|
92
|
+
@model_validator(mode="before")
|
|
93
|
+
@classmethod
|
|
94
|
+
def override_output_dir(cls, values):
|
|
95
|
+
if values.get("output") and values["output"].get("dir"):
|
|
96
|
+
values["output_dir"] = values["output"]["dir"]
|
|
97
|
+
return values
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
class EvalConfig(BaseModel):
|
|
101
|
+
|
|
102
|
+
# General Evaluation Options
|
|
103
|
+
general: EvalGeneralConfig = EvalGeneralConfig()
|
|
104
|
+
|
|
105
|
+
# Evaluators
|
|
106
|
+
evaluators: dict[str, EvaluatorBaseConfig] = {}
|
|
107
|
+
|
|
108
|
+
@classmethod
|
|
109
|
+
def rebuild_annotations(cls):
|
|
110
|
+
|
|
111
|
+
from aiq.cli.type_registry import GlobalTypeRegistry # pylint: disable=cyclic-import
|
|
112
|
+
|
|
113
|
+
type_registry = GlobalTypeRegistry.get()
|
|
114
|
+
|
|
115
|
+
EvaluatorsAnnotation = dict[str,
|
|
116
|
+
typing.Annotated[type_registry.compute_annotation(EvaluatorBaseConfig),
|
|
117
|
+
Discriminator(TypedBaseModel.discriminator)]]
|
|
118
|
+
|
|
119
|
+
should_rebuild = False
|
|
120
|
+
|
|
121
|
+
evaluators_field = cls.model_fields.get("evaluators")
|
|
122
|
+
if evaluators_field is not None and evaluators_field.annotation != EvaluatorsAnnotation:
|
|
123
|
+
evaluators_field.annotation = EvaluatorsAnnotation
|
|
124
|
+
should_rebuild = True
|
|
125
|
+
|
|
126
|
+
if (should_rebuild):
|
|
127
|
+
cls.model_rebuild(force=True)
|
|
@@ -0,0 +1,26 @@
|
|
|
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 typing
|
|
17
|
+
|
|
18
|
+
from .common import BaseModelRegistryTag
|
|
19
|
+
from .common import TypedBaseModel
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class EvaluatorBaseConfig(TypedBaseModel, BaseModelRegistryTag):
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
EvaluatorBaseConfigT = typing.TypeVar("EvaluatorBaseConfigT", bound=EvaluatorBaseConfig)
|
|
@@ -0,0 +1,26 @@
|
|
|
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
|
+
|
|
18
|
+
from .common import BaseModelRegistryTag
|
|
19
|
+
from .common import TypedBaseModel
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class FrontEndBaseConfig(TypedBaseModel, BaseModelRegistryTag):
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
FrontEndConfigT = typing.TypeVar("FrontEndConfigT", bound=FrontEndBaseConfig)
|
|
@@ -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
|
+
import typing
|
|
17
|
+
|
|
18
|
+
from .common import BaseModelRegistryTag
|
|
19
|
+
from .common import TypedBaseModel
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
|
|
27
|
+
pass
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
from pydantic import BaseModel
|
|
17
|
+
from pydantic import Field
|
|
18
|
+
from pydantic import field_serializer
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class FunctionDependencies(BaseModel):
|
|
22
|
+
"""
|
|
23
|
+
A class to represent the dependencies of a function.
|
|
24
|
+
"""
|
|
25
|
+
functions: set[str] = Field(default_factory=set)
|
|
26
|
+
llms: set[str] = Field(default_factory=set)
|
|
27
|
+
embedders: set[str] = Field(default_factory=set)
|
|
28
|
+
memory_clients: set[str] = Field(default_factory=set)
|
|
29
|
+
object_stores: set[str] = Field(default_factory=set)
|
|
30
|
+
retrievers: set[str] = Field(default_factory=set)
|
|
31
|
+
|
|
32
|
+
@field_serializer("functions", when_used="json")
|
|
33
|
+
def serialize_functions(self, v: set[str]) -> list[str]:
|
|
34
|
+
return list(v)
|
|
35
|
+
|
|
36
|
+
@field_serializer("llms", when_used="json")
|
|
37
|
+
def serialize_llms(self, v: set[str]) -> list[str]:
|
|
38
|
+
return list(v)
|
|
39
|
+
|
|
40
|
+
@field_serializer("embedders", when_used="json")
|
|
41
|
+
def serialize_embedders(self, v: set[str]) -> list[str]:
|
|
42
|
+
return list(v)
|
|
43
|
+
|
|
44
|
+
@field_serializer("memory_clients", when_used="json")
|
|
45
|
+
def serialize_memory_clients(self, v: set[str]) -> list[str]:
|
|
46
|
+
return list(v)
|
|
47
|
+
|
|
48
|
+
@field_serializer("object_stores", when_used="json")
|
|
49
|
+
def serialize_object_stores(self, v: set[str]) -> list[str]:
|
|
50
|
+
return list(v)
|
|
51
|
+
|
|
52
|
+
@field_serializer("retrievers", when_used="json")
|
|
53
|
+
def serialize_retrievers(self, v: set[str]) -> list[str]:
|
|
54
|
+
return list(v)
|
|
55
|
+
|
|
56
|
+
def add_function(self, function: str):
|
|
57
|
+
self.functions.add(function) # pylint: disable=no-member
|
|
58
|
+
|
|
59
|
+
def add_llm(self, llm: str):
|
|
60
|
+
self.llms.add(llm) # pylint: disable=no-member
|
|
61
|
+
|
|
62
|
+
def add_embedder(self, embedder: str):
|
|
63
|
+
self.embedders.add(embedder) # pylint: disable=no-member
|
|
64
|
+
|
|
65
|
+
def add_memory_client(self, memory_client: str):
|
|
66
|
+
self.memory_clients.add(memory_client) # pylint: disable=no-member
|
|
67
|
+
|
|
68
|
+
def add_object_store(self, object_store: str):
|
|
69
|
+
self.object_stores.add(object_store) # pylint: disable=no-member
|
|
70
|
+
|
|
71
|
+
def add_retriever(self, retriever: str):
|
|
72
|
+
self.retrievers.add(retriever) # pylint: disable=no-member
|