nvidia-nat 1.2.0rc5__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aiq/agent/__init__.py +0 -0
- aiq/agent/base.py +239 -0
- aiq/agent/dual_node.py +67 -0
- aiq/agent/react_agent/__init__.py +0 -0
- aiq/agent/react_agent/agent.py +355 -0
- aiq/agent/react_agent/output_parser.py +104 -0
- aiq/agent/react_agent/prompt.py +41 -0
- aiq/agent/react_agent/register.py +149 -0
- aiq/agent/reasoning_agent/__init__.py +0 -0
- aiq/agent/reasoning_agent/reasoning_agent.py +225 -0
- aiq/agent/register.py +23 -0
- aiq/agent/rewoo_agent/__init__.py +0 -0
- aiq/agent/rewoo_agent/agent.py +411 -0
- aiq/agent/rewoo_agent/prompt.py +108 -0
- aiq/agent/rewoo_agent/register.py +158 -0
- aiq/agent/tool_calling_agent/__init__.py +0 -0
- aiq/agent/tool_calling_agent/agent.py +119 -0
- aiq/agent/tool_calling_agent/register.py +106 -0
- aiq/authentication/__init__.py +14 -0
- aiq/authentication/api_key/__init__.py +14 -0
- aiq/authentication/api_key/api_key_auth_provider.py +96 -0
- aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
- aiq/authentication/api_key/register.py +26 -0
- aiq/authentication/exceptions/__init__.py +14 -0
- aiq/authentication/exceptions/api_key_exceptions.py +38 -0
- aiq/authentication/http_basic_auth/__init__.py +0 -0
- aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
- aiq/authentication/http_basic_auth/register.py +30 -0
- aiq/authentication/interfaces.py +93 -0
- aiq/authentication/oauth2/__init__.py +14 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
- aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
- aiq/authentication/oauth2/register.py +25 -0
- aiq/authentication/register.py +21 -0
- aiq/builder/__init__.py +0 -0
- aiq/builder/builder.py +285 -0
- aiq/builder/component_utils.py +316 -0
- aiq/builder/context.py +264 -0
- aiq/builder/embedder.py +24 -0
- aiq/builder/eval_builder.py +161 -0
- aiq/builder/evaluator.py +29 -0
- aiq/builder/framework_enum.py +24 -0
- aiq/builder/front_end.py +73 -0
- aiq/builder/function.py +344 -0
- aiq/builder/function_base.py +380 -0
- aiq/builder/function_info.py +627 -0
- aiq/builder/intermediate_step_manager.py +174 -0
- aiq/builder/llm.py +25 -0
- aiq/builder/retriever.py +25 -0
- aiq/builder/user_interaction_manager.py +74 -0
- aiq/builder/workflow.py +148 -0
- aiq/builder/workflow_builder.py +1117 -0
- aiq/cli/__init__.py +14 -0
- aiq/cli/cli_utils/__init__.py +0 -0
- aiq/cli/cli_utils/config_override.py +231 -0
- aiq/cli/cli_utils/validation.py +37 -0
- aiq/cli/commands/__init__.py +0 -0
- aiq/cli/commands/configure/__init__.py +0 -0
- aiq/cli/commands/configure/channel/__init__.py +0 -0
- aiq/cli/commands/configure/channel/add.py +28 -0
- aiq/cli/commands/configure/channel/channel.py +36 -0
- aiq/cli/commands/configure/channel/remove.py +30 -0
- aiq/cli/commands/configure/channel/update.py +30 -0
- aiq/cli/commands/configure/configure.py +33 -0
- aiq/cli/commands/evaluate.py +139 -0
- aiq/cli/commands/info/__init__.py +14 -0
- aiq/cli/commands/info/info.py +39 -0
- aiq/cli/commands/info/list_channels.py +32 -0
- aiq/cli/commands/info/list_components.py +129 -0
- aiq/cli/commands/info/list_mcp.py +213 -0
- aiq/cli/commands/registry/__init__.py +14 -0
- aiq/cli/commands/registry/publish.py +88 -0
- aiq/cli/commands/registry/pull.py +118 -0
- aiq/cli/commands/registry/registry.py +38 -0
- aiq/cli/commands/registry/remove.py +108 -0
- aiq/cli/commands/registry/search.py +155 -0
- aiq/cli/commands/sizing/__init__.py +14 -0
- aiq/cli/commands/sizing/calc.py +297 -0
- aiq/cli/commands/sizing/sizing.py +27 -0
- aiq/cli/commands/start.py +246 -0
- aiq/cli/commands/uninstall.py +81 -0
- aiq/cli/commands/validate.py +47 -0
- aiq/cli/commands/workflow/__init__.py +14 -0
- aiq/cli/commands/workflow/templates/__init__.py.j2 +0 -0
- aiq/cli/commands/workflow/templates/config.yml.j2 +16 -0
- aiq/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
- aiq/cli/commands/workflow/templates/register.py.j2 +5 -0
- aiq/cli/commands/workflow/templates/workflow.py.j2 +36 -0
- aiq/cli/commands/workflow/workflow.py +37 -0
- aiq/cli/commands/workflow/workflow_commands.py +313 -0
- aiq/cli/entrypoint.py +135 -0
- aiq/cli/main.py +44 -0
- aiq/cli/register_workflow.py +488 -0
- aiq/cli/type_registry.py +1000 -0
- aiq/data_models/__init__.py +14 -0
- aiq/data_models/api_server.py +694 -0
- aiq/data_models/authentication.py +231 -0
- aiq/data_models/common.py +171 -0
- aiq/data_models/component.py +54 -0
- aiq/data_models/component_ref.py +168 -0
- aiq/data_models/config.py +406 -0
- aiq/data_models/dataset_handler.py +123 -0
- aiq/data_models/discovery_metadata.py +335 -0
- aiq/data_models/embedder.py +27 -0
- aiq/data_models/evaluate.py +127 -0
- aiq/data_models/evaluator.py +26 -0
- aiq/data_models/front_end.py +26 -0
- aiq/data_models/function.py +30 -0
- aiq/data_models/function_dependencies.py +72 -0
- aiq/data_models/interactive.py +246 -0
- aiq/data_models/intermediate_step.py +302 -0
- aiq/data_models/invocation_node.py +38 -0
- aiq/data_models/llm.py +27 -0
- aiq/data_models/logging.py +26 -0
- aiq/data_models/memory.py +27 -0
- aiq/data_models/object_store.py +44 -0
- aiq/data_models/profiler.py +54 -0
- aiq/data_models/registry_handler.py +26 -0
- aiq/data_models/retriever.py +30 -0
- aiq/data_models/retry_mixin.py +35 -0
- aiq/data_models/span.py +187 -0
- aiq/data_models/step_adaptor.py +64 -0
- aiq/data_models/streaming.py +33 -0
- aiq/data_models/swe_bench_model.py +54 -0
- aiq/data_models/telemetry_exporter.py +26 -0
- aiq/data_models/ttc_strategy.py +30 -0
- aiq/embedder/__init__.py +0 -0
- aiq/embedder/langchain_client.py +41 -0
- aiq/embedder/nim_embedder.py +59 -0
- aiq/embedder/openai_embedder.py +43 -0
- aiq/embedder/register.py +24 -0
- aiq/eval/__init__.py +14 -0
- aiq/eval/config.py +60 -0
- aiq/eval/dataset_handler/__init__.py +0 -0
- aiq/eval/dataset_handler/dataset_downloader.py +106 -0
- aiq/eval/dataset_handler/dataset_filter.py +52 -0
- aiq/eval/dataset_handler/dataset_handler.py +254 -0
- aiq/eval/evaluate.py +506 -0
- aiq/eval/evaluator/__init__.py +14 -0
- aiq/eval/evaluator/base_evaluator.py +73 -0
- aiq/eval/evaluator/evaluator_model.py +45 -0
- aiq/eval/intermediate_step_adapter.py +99 -0
- aiq/eval/rag_evaluator/__init__.py +0 -0
- aiq/eval/rag_evaluator/evaluate.py +178 -0
- aiq/eval/rag_evaluator/register.py +143 -0
- aiq/eval/register.py +23 -0
- aiq/eval/remote_workflow.py +133 -0
- aiq/eval/runners/__init__.py +14 -0
- aiq/eval/runners/config.py +39 -0
- aiq/eval/runners/multi_eval_runner.py +54 -0
- aiq/eval/runtime_event_subscriber.py +52 -0
- aiq/eval/swe_bench_evaluator/__init__.py +0 -0
- aiq/eval/swe_bench_evaluator/evaluate.py +215 -0
- aiq/eval/swe_bench_evaluator/register.py +36 -0
- aiq/eval/trajectory_evaluator/__init__.py +0 -0
- aiq/eval/trajectory_evaluator/evaluate.py +75 -0
- aiq/eval/trajectory_evaluator/register.py +40 -0
- aiq/eval/tunable_rag_evaluator/__init__.py +0 -0
- aiq/eval/tunable_rag_evaluator/evaluate.py +245 -0
- aiq/eval/tunable_rag_evaluator/register.py +52 -0
- aiq/eval/usage_stats.py +41 -0
- aiq/eval/utils/__init__.py +0 -0
- aiq/eval/utils/output_uploader.py +140 -0
- aiq/eval/utils/tqdm_position_registry.py +40 -0
- aiq/eval/utils/weave_eval.py +184 -0
- aiq/experimental/__init__.py +0 -0
- aiq/experimental/decorators/__init__.py +0 -0
- aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
- aiq/experimental/test_time_compute/__init__.py +0 -0
- aiq/experimental/test_time_compute/editing/__init__.py +0 -0
- aiq/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
- aiq/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
- aiq/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
- aiq/experimental/test_time_compute/functions/__init__.py +0 -0
- aiq/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
- aiq/experimental/test_time_compute/functions/its_tool_orchestration_function.py +205 -0
- aiq/experimental/test_time_compute/functions/its_tool_wrapper_function.py +146 -0
- aiq/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
- aiq/experimental/test_time_compute/models/__init__.py +0 -0
- aiq/experimental/test_time_compute/models/editor_config.py +132 -0
- aiq/experimental/test_time_compute/models/scoring_config.py +112 -0
- aiq/experimental/test_time_compute/models/search_config.py +120 -0
- aiq/experimental/test_time_compute/models/selection_config.py +154 -0
- aiq/experimental/test_time_compute/models/stage_enums.py +43 -0
- aiq/experimental/test_time_compute/models/strategy_base.py +66 -0
- aiq/experimental/test_time_compute/models/tool_use_config.py +41 -0
- aiq/experimental/test_time_compute/models/ttc_item.py +48 -0
- aiq/experimental/test_time_compute/register.py +36 -0
- aiq/experimental/test_time_compute/scoring/__init__.py +0 -0
- aiq/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
- aiq/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
- aiq/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
- aiq/experimental/test_time_compute/search/__init__.py +0 -0
- aiq/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
- aiq/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
- aiq/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
- aiq/experimental/test_time_compute/selection/__init__.py +0 -0
- aiq/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
- aiq/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
- aiq/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
- aiq/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
- aiq/experimental/test_time_compute/selection/threshold_selector.py +58 -0
- aiq/front_ends/__init__.py +14 -0
- aiq/front_ends/console/__init__.py +14 -0
- aiq/front_ends/console/authentication_flow_handler.py +233 -0
- aiq/front_ends/console/console_front_end_config.py +32 -0
- aiq/front_ends/console/console_front_end_plugin.py +96 -0
- aiq/front_ends/console/register.py +25 -0
- aiq/front_ends/cron/__init__.py +14 -0
- aiq/front_ends/fastapi/__init__.py +14 -0
- aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
- aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
- aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
- aiq/front_ends/fastapi/fastapi_front_end_config.py +234 -0
- aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
- aiq/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
- aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1092 -0
- aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
- aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
- aiq/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
- aiq/front_ends/fastapi/job_store.py +183 -0
- aiq/front_ends/fastapi/main.py +72 -0
- aiq/front_ends/fastapi/message_handler.py +298 -0
- aiq/front_ends/fastapi/message_validator.py +345 -0
- aiq/front_ends/fastapi/register.py +25 -0
- aiq/front_ends/fastapi/response_helpers.py +195 -0
- aiq/front_ends/fastapi/step_adaptor.py +321 -0
- aiq/front_ends/mcp/__init__.py +14 -0
- aiq/front_ends/mcp/mcp_front_end_config.py +32 -0
- aiq/front_ends/mcp/mcp_front_end_plugin.py +93 -0
- aiq/front_ends/mcp/register.py +27 -0
- aiq/front_ends/mcp/tool_converter.py +242 -0
- aiq/front_ends/register.py +22 -0
- aiq/front_ends/simple_base/__init__.py +14 -0
- aiq/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
- aiq/llm/__init__.py +0 -0
- aiq/llm/aws_bedrock_llm.py +57 -0
- aiq/llm/nim_llm.py +46 -0
- aiq/llm/openai_llm.py +46 -0
- aiq/llm/register.py +23 -0
- aiq/llm/utils/__init__.py +14 -0
- aiq/llm/utils/env_config_value.py +94 -0
- aiq/llm/utils/error.py +17 -0
- aiq/memory/__init__.py +20 -0
- aiq/memory/interfaces.py +183 -0
- aiq/memory/models.py +112 -0
- aiq/meta/module_to_distro.json +3 -0
- aiq/meta/pypi.md +58 -0
- aiq/object_store/__init__.py +20 -0
- aiq/object_store/in_memory_object_store.py +76 -0
- aiq/object_store/interfaces.py +84 -0
- aiq/object_store/models.py +36 -0
- aiq/object_store/register.py +20 -0
- aiq/observability/__init__.py +14 -0
- aiq/observability/exporter/__init__.py +14 -0
- aiq/observability/exporter/base_exporter.py +449 -0
- aiq/observability/exporter/exporter.py +78 -0
- aiq/observability/exporter/file_exporter.py +33 -0
- aiq/observability/exporter/processing_exporter.py +322 -0
- aiq/observability/exporter/raw_exporter.py +52 -0
- aiq/observability/exporter/span_exporter.py +265 -0
- aiq/observability/exporter_manager.py +335 -0
- aiq/observability/mixin/__init__.py +14 -0
- aiq/observability/mixin/batch_config_mixin.py +26 -0
- aiq/observability/mixin/collector_config_mixin.py +23 -0
- aiq/observability/mixin/file_mixin.py +288 -0
- aiq/observability/mixin/file_mode.py +23 -0
- aiq/observability/mixin/resource_conflict_mixin.py +134 -0
- aiq/observability/mixin/serialize_mixin.py +61 -0
- aiq/observability/mixin/type_introspection_mixin.py +183 -0
- aiq/observability/processor/__init__.py +14 -0
- aiq/observability/processor/batching_processor.py +310 -0
- aiq/observability/processor/callback_processor.py +42 -0
- aiq/observability/processor/intermediate_step_serializer.py +28 -0
- aiq/observability/processor/processor.py +71 -0
- aiq/observability/register.py +96 -0
- aiq/observability/utils/__init__.py +14 -0
- aiq/observability/utils/dict_utils.py +236 -0
- aiq/observability/utils/time_utils.py +31 -0
- aiq/plugins/.namespace +1 -0
- aiq/profiler/__init__.py +0 -0
- aiq/profiler/calc/__init__.py +14 -0
- aiq/profiler/calc/calc_runner.py +627 -0
- aiq/profiler/calc/calculations.py +288 -0
- aiq/profiler/calc/data_models.py +188 -0
- aiq/profiler/calc/plot.py +345 -0
- aiq/profiler/callbacks/__init__.py +0 -0
- aiq/profiler/callbacks/agno_callback_handler.py +295 -0
- aiq/profiler/callbacks/base_callback_class.py +20 -0
- aiq/profiler/callbacks/langchain_callback_handler.py +290 -0
- aiq/profiler/callbacks/llama_index_callback_handler.py +205 -0
- aiq/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
- aiq/profiler/callbacks/token_usage_base_model.py +27 -0
- aiq/profiler/data_frame_row.py +51 -0
- aiq/profiler/data_models.py +24 -0
- aiq/profiler/decorators/__init__.py +0 -0
- aiq/profiler/decorators/framework_wrapper.py +131 -0
- aiq/profiler/decorators/function_tracking.py +254 -0
- aiq/profiler/forecasting/__init__.py +0 -0
- aiq/profiler/forecasting/config.py +18 -0
- aiq/profiler/forecasting/model_trainer.py +75 -0
- aiq/profiler/forecasting/models/__init__.py +22 -0
- aiq/profiler/forecasting/models/forecasting_base_model.py +40 -0
- aiq/profiler/forecasting/models/linear_model.py +196 -0
- aiq/profiler/forecasting/models/random_forest_regressor.py +268 -0
- aiq/profiler/inference_metrics_model.py +28 -0
- aiq/profiler/inference_optimization/__init__.py +0 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
- aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
- aiq/profiler/inference_optimization/data_models.py +386 -0
- aiq/profiler/inference_optimization/experimental/__init__.py +0 -0
- aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
- aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
- aiq/profiler/inference_optimization/llm_metrics.py +212 -0
- aiq/profiler/inference_optimization/prompt_caching.py +163 -0
- aiq/profiler/inference_optimization/token_uniqueness.py +107 -0
- aiq/profiler/inference_optimization/workflow_runtimes.py +72 -0
- aiq/profiler/intermediate_property_adapter.py +102 -0
- aiq/profiler/profile_runner.py +473 -0
- aiq/profiler/utils.py +184 -0
- aiq/registry_handlers/__init__.py +0 -0
- aiq/registry_handlers/local/__init__.py +0 -0
- aiq/registry_handlers/local/local_handler.py +176 -0
- aiq/registry_handlers/local/register_local.py +37 -0
- aiq/registry_handlers/metadata_factory.py +60 -0
- aiq/registry_handlers/package_utils.py +567 -0
- aiq/registry_handlers/pypi/__init__.py +0 -0
- aiq/registry_handlers/pypi/pypi_handler.py +251 -0
- aiq/registry_handlers/pypi/register_pypi.py +40 -0
- aiq/registry_handlers/register.py +21 -0
- aiq/registry_handlers/registry_handler_base.py +157 -0
- aiq/registry_handlers/rest/__init__.py +0 -0
- aiq/registry_handlers/rest/register_rest.py +56 -0
- aiq/registry_handlers/rest/rest_handler.py +237 -0
- aiq/registry_handlers/schemas/__init__.py +0 -0
- aiq/registry_handlers/schemas/headers.py +42 -0
- aiq/registry_handlers/schemas/package.py +68 -0
- aiq/registry_handlers/schemas/publish.py +63 -0
- aiq/registry_handlers/schemas/pull.py +82 -0
- aiq/registry_handlers/schemas/remove.py +36 -0
- aiq/registry_handlers/schemas/search.py +91 -0
- aiq/registry_handlers/schemas/status.py +47 -0
- aiq/retriever/__init__.py +0 -0
- aiq/retriever/interface.py +37 -0
- aiq/retriever/milvus/__init__.py +14 -0
- aiq/retriever/milvus/register.py +81 -0
- aiq/retriever/milvus/retriever.py +228 -0
- aiq/retriever/models.py +74 -0
- aiq/retriever/nemo_retriever/__init__.py +14 -0
- aiq/retriever/nemo_retriever/register.py +60 -0
- aiq/retriever/nemo_retriever/retriever.py +190 -0
- aiq/retriever/register.py +22 -0
- aiq/runtime/__init__.py +14 -0
- aiq/runtime/loader.py +215 -0
- aiq/runtime/runner.py +190 -0
- aiq/runtime/session.py +158 -0
- aiq/runtime/user_metadata.py +130 -0
- aiq/settings/__init__.py +0 -0
- aiq/settings/global_settings.py +318 -0
- aiq/test/.namespace +1 -0
- aiq/tool/__init__.py +0 -0
- aiq/tool/chat_completion.py +74 -0
- aiq/tool/code_execution/README.md +151 -0
- aiq/tool/code_execution/__init__.py +0 -0
- aiq/tool/code_execution/code_sandbox.py +267 -0
- aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
- aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
- aiq/tool/code_execution/local_sandbox/__init__.py +13 -0
- aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
- aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
- aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
- aiq/tool/code_execution/register.py +74 -0
- aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
- aiq/tool/code_execution/utils.py +100 -0
- aiq/tool/datetime_tools.py +42 -0
- aiq/tool/document_search.py +141 -0
- aiq/tool/github_tools/__init__.py +0 -0
- aiq/tool/github_tools/create_github_commit.py +133 -0
- aiq/tool/github_tools/create_github_issue.py +87 -0
- aiq/tool/github_tools/create_github_pr.py +106 -0
- aiq/tool/github_tools/get_github_file.py +106 -0
- aiq/tool/github_tools/get_github_issue.py +166 -0
- aiq/tool/github_tools/get_github_pr.py +256 -0
- aiq/tool/github_tools/update_github_issue.py +100 -0
- aiq/tool/mcp/__init__.py +14 -0
- aiq/tool/mcp/exceptions.py +142 -0
- aiq/tool/mcp/mcp_client.py +255 -0
- aiq/tool/mcp/mcp_tool.py +96 -0
- aiq/tool/memory_tools/__init__.py +0 -0
- aiq/tool/memory_tools/add_memory_tool.py +79 -0
- aiq/tool/memory_tools/delete_memory_tool.py +67 -0
- aiq/tool/memory_tools/get_memory_tool.py +72 -0
- aiq/tool/nvidia_rag.py +95 -0
- aiq/tool/register.py +38 -0
- aiq/tool/retriever.py +89 -0
- aiq/tool/server_tools.py +66 -0
- aiq/utils/__init__.py +0 -0
- aiq/utils/data_models/__init__.py +0 -0
- aiq/utils/data_models/schema_validator.py +58 -0
- aiq/utils/debugging_utils.py +43 -0
- aiq/utils/dump_distro_mapping.py +32 -0
- aiq/utils/exception_handlers/__init__.py +0 -0
- aiq/utils/exception_handlers/automatic_retries.py +289 -0
- aiq/utils/exception_handlers/mcp.py +211 -0
- aiq/utils/exception_handlers/schemas.py +114 -0
- aiq/utils/io/__init__.py +0 -0
- aiq/utils/io/model_processing.py +28 -0
- aiq/utils/io/yaml_tools.py +119 -0
- aiq/utils/log_utils.py +37 -0
- aiq/utils/metadata_utils.py +74 -0
- aiq/utils/optional_imports.py +142 -0
- aiq/utils/producer_consumer_queue.py +178 -0
- aiq/utils/reactive/__init__.py +0 -0
- aiq/utils/reactive/base/__init__.py +0 -0
- aiq/utils/reactive/base/observable_base.py +65 -0
- aiq/utils/reactive/base/observer_base.py +55 -0
- aiq/utils/reactive/base/subject_base.py +79 -0
- aiq/utils/reactive/observable.py +59 -0
- aiq/utils/reactive/observer.py +76 -0
- aiq/utils/reactive/subject.py +131 -0
- aiq/utils/reactive/subscription.py +49 -0
- aiq/utils/settings/__init__.py +0 -0
- aiq/utils/settings/global_settings.py +197 -0
- aiq/utils/string_utils.py +38 -0
- aiq/utils/type_converter.py +290 -0
- aiq/utils/type_utils.py +484 -0
- aiq/utils/url_utils.py +27 -0
- nvidia_nat-1.2.0rc5.dist-info/METADATA +363 -0
- nvidia_nat-1.2.0rc5.dist-info/RECORD +435 -0
- nvidia_nat-1.2.0rc5.dist-info/WHEEL +5 -0
- nvidia_nat-1.2.0rc5.dist-info/entry_points.txt +20 -0
- nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE-3rd-party.txt +3686 -0
- nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE.md +201 -0
- nvidia_nat-1.2.0rc5.dist-info/top_level.txt +1 -0
aiq/memory/interfaces.py
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
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 abc import ABC
|
|
17
|
+
from abc import abstractmethod
|
|
18
|
+
from collections.abc import Callable
|
|
19
|
+
|
|
20
|
+
from .models import MemoryItem
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class MemoryEditor(ABC):
|
|
24
|
+
"""
|
|
25
|
+
Abstract interface for editing and
|
|
26
|
+
retrieving memory items.
|
|
27
|
+
|
|
28
|
+
A MemoryEditor is responsible for adding, searching, and
|
|
29
|
+
removing MemoryItems.
|
|
30
|
+
|
|
31
|
+
Implementations may integrate with
|
|
32
|
+
vector stores or other indexing backends.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
@abstractmethod
|
|
36
|
+
async def add_items(self, items: list[MemoryItem]) -> None:
|
|
37
|
+
"""
|
|
38
|
+
Insert multiple MemoryItems into the memory.
|
|
39
|
+
|
|
40
|
+
Args:
|
|
41
|
+
items (list[MemoryItem]): The items to be added.
|
|
42
|
+
"""
|
|
43
|
+
raise NotImplementedError
|
|
44
|
+
|
|
45
|
+
@abstractmethod
|
|
46
|
+
async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
|
|
47
|
+
"""
|
|
48
|
+
Retrieve items relevant to the given query.
|
|
49
|
+
Relevance criteria depend on implementation.
|
|
50
|
+
|
|
51
|
+
Args:
|
|
52
|
+
query (str): The query string to match.
|
|
53
|
+
top_k (int): Maximum number of items to return.
|
|
54
|
+
kwargs (dict): Keyword arguments to pass to the search method.
|
|
55
|
+
|
|
56
|
+
Returns:
|
|
57
|
+
list[MemoryItem]: The most relevant MemoryItems.
|
|
58
|
+
"""
|
|
59
|
+
raise NotImplementedError
|
|
60
|
+
|
|
61
|
+
@abstractmethod
|
|
62
|
+
async def remove_items(self, **kwargs) -> None:
|
|
63
|
+
"""
|
|
64
|
+
Remove items. Additional parameters
|
|
65
|
+
needed for deletion can be specified in keyword arguments.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
kwargs (dict): Keyword arguments to pass to the remove-items method.
|
|
69
|
+
"""
|
|
70
|
+
raise NotImplementedError
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class MemoryIOBase(ABC):
|
|
74
|
+
"""
|
|
75
|
+
Base abstract class for I/O operations
|
|
76
|
+
on memory, providing a common interface for
|
|
77
|
+
|
|
78
|
+
MemoryReader and MemoryWriter to interact
|
|
79
|
+
with a MemoryEditor.
|
|
80
|
+
|
|
81
|
+
Concrete subclasses should hold a
|
|
82
|
+
reference to a MemoryEditor instance.
|
|
83
|
+
"""
|
|
84
|
+
|
|
85
|
+
def __init__(self, editor: MemoryEditor) -> None:
|
|
86
|
+
self._editor = editor
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
class MemoryReader(MemoryIOBase):
|
|
90
|
+
"""
|
|
91
|
+
Responsible for retrieving MemoryItems
|
|
92
|
+
from the MemoryEditor based on context or queries.
|
|
93
|
+
"""
|
|
94
|
+
|
|
95
|
+
@abstractmethod
|
|
96
|
+
async def retrieve(self, context: str, top_k: int = 5) -> list[MemoryItem]:
|
|
97
|
+
"""
|
|
98
|
+
Retrieve a subset of
|
|
99
|
+
MemoryItems relevant to the provided context.
|
|
100
|
+
|
|
101
|
+
Args:
|
|
102
|
+
context (str): A string representing
|
|
103
|
+
the current user context or query.
|
|
104
|
+
top_k (int): Maximum number of items to return.
|
|
105
|
+
|
|
106
|
+
Returns:
|
|
107
|
+
list[MemoryItem]: Relevant MemoryItems.
|
|
108
|
+
"""
|
|
109
|
+
raise NotImplementedError
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
class MemoryWriter(MemoryIOBase):
|
|
113
|
+
"""
|
|
114
|
+
Responsible for converting new observations
|
|
115
|
+
(textual inputs) into MemoryItems andstoring
|
|
116
|
+
them via the MemoryEditor.
|
|
117
|
+
"""
|
|
118
|
+
|
|
119
|
+
@abstractmethod
|
|
120
|
+
async def write(self, observation: str, context: str | None = None) -> list[MemoryItem]:
|
|
121
|
+
"""
|
|
122
|
+
Process the given observation and store the resulting MemoryItems.
|
|
123
|
+
|
|
124
|
+
Args:
|
|
125
|
+
observation (str): The new textual input to record.
|
|
126
|
+
context (Optional[str]): Additional
|
|
127
|
+
context that might influence how the observation is stored.
|
|
128
|
+
|
|
129
|
+
Returns:
|
|
130
|
+
list[MemoryItem]: The newly created MemoryItems.
|
|
131
|
+
"""
|
|
132
|
+
raise NotImplementedError
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
class MemoryManager(ABC):
|
|
136
|
+
"""
|
|
137
|
+
Manages the lifecycle of the stored
|
|
138
|
+
memory by applying policies such as summarization,
|
|
139
|
+
reflection, forgetting, and mergingn
|
|
140
|
+
to ensure long-term coherence and relevance.
|
|
141
|
+
"""
|
|
142
|
+
|
|
143
|
+
@abstractmethod
|
|
144
|
+
async def summarize(self) -> None:
|
|
145
|
+
"""
|
|
146
|
+
Summarize long or numerous MemoryItems into a more compact form.
|
|
147
|
+
This may remove the original items and store a new summary item.
|
|
148
|
+
"""
|
|
149
|
+
raise NotImplementedError
|
|
150
|
+
|
|
151
|
+
@abstractmethod
|
|
152
|
+
async def reflect(self) -> None:
|
|
153
|
+
"""
|
|
154
|
+
Generate higher-level insights or
|
|
155
|
+
abstractions from existing MemoryItems.
|
|
156
|
+
This may call out to an LLM or other
|
|
157
|
+
logic to produce conceptual memory.
|
|
158
|
+
"""
|
|
159
|
+
raise NotImplementedError
|
|
160
|
+
|
|
161
|
+
@abstractmethod
|
|
162
|
+
async def forget(self, criteria: Callable[[MemoryItem], bool]) -> None:
|
|
163
|
+
"""
|
|
164
|
+
Remove MemoryItems that are no
|
|
165
|
+
longer relevant or have low importance.
|
|
166
|
+
|
|
167
|
+
Args:
|
|
168
|
+
criteria (Callable[[MemoryItem], bool]): A function that
|
|
169
|
+
returns True for items to forget.
|
|
170
|
+
"""
|
|
171
|
+
raise NotImplementedError
|
|
172
|
+
|
|
173
|
+
@abstractmethod
|
|
174
|
+
async def merge(self, criteria: Callable[[MemoryItem, MemoryItem], bool]) -> None:
|
|
175
|
+
"""
|
|
176
|
+
Merge similar or redundant MemoryItems
|
|
177
|
+
into a smaller set of more concise items.
|
|
178
|
+
|
|
179
|
+
Args:
|
|
180
|
+
criteria (Callable[[MemoryItem, MemoryItem], bool]): A function
|
|
181
|
+
that determines which items can be merged.
|
|
182
|
+
"""
|
|
183
|
+
raise NotImplementedError
|
aiq/memory/models.py
ADDED
|
@@ -0,0 +1,112 @@
|
|
|
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 pydantic import BaseModel
|
|
19
|
+
from pydantic import ConfigDict
|
|
20
|
+
from pydantic import Field
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class MemoryItem(BaseModel):
|
|
24
|
+
"""
|
|
25
|
+
Represents a single memory item consisting of structured content and associated metadata.
|
|
26
|
+
|
|
27
|
+
Attributes
|
|
28
|
+
----------
|
|
29
|
+
conversation : list[dict[str, str]]
|
|
30
|
+
A list of dictionaries, each containing string key-value pairs.
|
|
31
|
+
user_id : str
|
|
32
|
+
Unique identifier for this MemoryItem's user.
|
|
33
|
+
tags : list[str]
|
|
34
|
+
A list of strings representing tags attached to the item.
|
|
35
|
+
metadata : dict[str, typing.Any]
|
|
36
|
+
Metadata providing context and utility for management operations.
|
|
37
|
+
memory : str or None
|
|
38
|
+
Optional memory string. Helpful when returning a memory.
|
|
39
|
+
"""
|
|
40
|
+
# yapf: disable
|
|
41
|
+
model_config = ConfigDict(
|
|
42
|
+
json_schema_extra={
|
|
43
|
+
"examples": [
|
|
44
|
+
{
|
|
45
|
+
"conversation": [
|
|
46
|
+
{
|
|
47
|
+
"role": "user",
|
|
48
|
+
"content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts."
|
|
49
|
+
},
|
|
50
|
+
{
|
|
51
|
+
"role": "assistant",
|
|
52
|
+
"content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy."
|
|
53
|
+
}
|
|
54
|
+
],
|
|
55
|
+
"user_id": "user_abc",
|
|
56
|
+
"tags": ["diet", "allergy"],
|
|
57
|
+
"metadata": {
|
|
58
|
+
"key_value_pairs": {
|
|
59
|
+
"type": "profile",
|
|
60
|
+
"relevance": "high"
|
|
61
|
+
}
|
|
62
|
+
}
|
|
63
|
+
},
|
|
64
|
+
{
|
|
65
|
+
"memory": "User prefers expensive hotels and is vegan.",
|
|
66
|
+
"user_id": "user_abc",
|
|
67
|
+
"tags": ["hotel", "restaurant"]
|
|
68
|
+
}
|
|
69
|
+
]
|
|
70
|
+
},
|
|
71
|
+
# Allow population of models from arbitrary types (e.g., ORM objects)
|
|
72
|
+
arbitrary_types_allowed=True,
|
|
73
|
+
# Enable aliasing if needed
|
|
74
|
+
populate_by_name=True
|
|
75
|
+
)
|
|
76
|
+
# yapf: enable
|
|
77
|
+
conversation: list[dict[str, str]] | None = Field(
|
|
78
|
+
description="List of conversation messages. Each message must have a \"role\" "
|
|
79
|
+
"key (user or assistant. It must also have a \"content\" key.",
|
|
80
|
+
default=None)
|
|
81
|
+
tags: list[str] = Field(default_factory=list, description="List of tags applied to the item.")
|
|
82
|
+
metadata: dict[str, typing.Any] = Field(description="Metadata about the memory item.", default={})
|
|
83
|
+
user_id: str = Field(description="The user's ID.")
|
|
84
|
+
memory: str | None = Field(default=None)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class SearchMemoryInput(BaseModel):
|
|
88
|
+
"""
|
|
89
|
+
Represents a search memory input structure.
|
|
90
|
+
"""
|
|
91
|
+
model_config = ConfigDict(json_schema_extra={
|
|
92
|
+
"example": {
|
|
93
|
+
"query": "What is the user's preferred programming language?",
|
|
94
|
+
"top_k": 1,
|
|
95
|
+
"user_id": "user_abc",
|
|
96
|
+
}
|
|
97
|
+
})
|
|
98
|
+
|
|
99
|
+
query: str = Field(description="Search query for which to retrieve memory.") # noqa: E501
|
|
100
|
+
top_k: int = Field(description="Maximum number of memories to return")
|
|
101
|
+
user_id: str = Field(description="ID of the user to search for.")
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class DeleteMemoryInput(BaseModel):
|
|
105
|
+
"""
|
|
106
|
+
Represents a delete memory input structure.
|
|
107
|
+
"""
|
|
108
|
+
model_config = ConfigDict(json_schema_extra={"example": {"user_id": "user_abc", }})
|
|
109
|
+
|
|
110
|
+
user_id: str = Field(description="ID of the user to delete memory for. Careful when using "
|
|
111
|
+
"this tool; make sure you use the "
|
|
112
|
+
"username present in the conversation.")
|
aiq/meta/pypi.md
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
<!--
|
|
2
|
+
SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
3
|
+
SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
6
|
+
you may not use this file except in compliance with the License.
|
|
7
|
+
You may obtain a copy of the License at
|
|
8
|
+
|
|
9
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
|
|
11
|
+
Unless required by applicable law or agreed to in writing, software
|
|
12
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
13
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
14
|
+
See the License for the specific language governing permissions and
|
|
15
|
+
limitations under the License.
|
|
16
|
+
-->
|
|
17
|
+
|
|
18
|
+

|
|
19
|
+
|
|
20
|
+
# NVIDIA NeMo Agent Toolkit
|
|
21
|
+
|
|
22
|
+
AIQ toolkit is a flexible library designed to seamlessly integrate your enterprise agents—regardless of framework—with various data sources and tools. By treating agents, tools, and agentic workflows as simple function calls, AIQ toolkit enables true composability: build once and reuse anywhere.
|
|
23
|
+
|
|
24
|
+
## Key Features
|
|
25
|
+
|
|
26
|
+
- [**Framework Agnostic:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
|
|
27
|
+
- [**Reusability:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
|
|
28
|
+
- [**Rapid Development:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
|
|
29
|
+
- [**Profiling:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
|
|
30
|
+
- [**Observability:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/observe/observe-workflow-with-weave.html).
|
|
31
|
+
- [**Evaluation System:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
|
|
32
|
+
- [**User Interface:**](https://docs.nvidia.com/aiqtoolkit/1.2.0/quick-start/launching-ui.html) Use the AIQ toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
|
|
33
|
+
- [**MCP Compatibility**](https://docs.nvidia.com/aiqtoolkit/1.2.0/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as AIQ toolkit functions.
|
|
34
|
+
|
|
35
|
+
With AIQ toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
|
|
36
|
+
|
|
37
|
+
## Links
|
|
38
|
+
* [Documentation](https://docs.nvidia.com/aiqtoolkit/1.2.0/index.html): Explore the full documentation for AIQ toolkit.
|
|
39
|
+
|
|
40
|
+
## First time user?
|
|
41
|
+
If this is your first time using AIQ toolkit, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/NeMo-Agent-Toolkit?tab=readme-ov-file#quick-start) on GitHub. This package is intended for users who are familiar with AIQ toolkit applications and need to add AIQ toolkit as a dependency to their project.
|
|
42
|
+
|
|
43
|
+
## Feedback
|
|
44
|
+
|
|
45
|
+
We would love to hear from you! Please file an issue on [GitHub](https://github.com/NVIDIA/NeMo-Agent-Toolkit/issues) if you have any feedback or feature requests.
|
|
46
|
+
|
|
47
|
+
## Acknowledgements
|
|
48
|
+
|
|
49
|
+
We would like to thank the following open source projects that made AIQ toolkit possible:
|
|
50
|
+
|
|
51
|
+
- [CrewAI](https://github.com/crewAIInc/crewAI)
|
|
52
|
+
- [FastAPI](https://github.com/tiangolo/fastapi)
|
|
53
|
+
- [LangChain](https://github.com/langchain-ai/langchain)
|
|
54
|
+
- [Llama-Index](https://github.com/run-llama/llama_index)
|
|
55
|
+
- [Mem0ai](https://github.com/mem0ai/mem0)
|
|
56
|
+
- [Ragas](https://github.com/explodinggradients/ragas)
|
|
57
|
+
- [Semantic Kernel](https://github.com/microsoft/semantic-kernel)
|
|
58
|
+
- [uv](https://github.com/astral-sh/uv)
|
|
@@ -0,0 +1,20 @@
|
|
|
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
|
+
AIQ Toolkit Object Store Module
|
|
17
|
+
|
|
18
|
+
This package provides foundational classes and interfaces
|
|
19
|
+
for managing object storage in AIQ Toolkit's LLM-based agents.
|
|
20
|
+
"""
|
|
@@ -0,0 +1,76 @@
|
|
|
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
|
+
|
|
18
|
+
from aiq.builder.builder import Builder
|
|
19
|
+
from aiq.cli.register_workflow import register_object_store
|
|
20
|
+
from aiq.data_models.object_store import KeyAlreadyExistsError
|
|
21
|
+
from aiq.data_models.object_store import NoSuchKeyError
|
|
22
|
+
from aiq.data_models.object_store import ObjectStoreBaseConfig
|
|
23
|
+
from aiq.utils.type_utils import override
|
|
24
|
+
|
|
25
|
+
from .interfaces import ObjectStore
|
|
26
|
+
from .models import ObjectStoreItem
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class InMemoryObjectStoreConfig(ObjectStoreBaseConfig, name="in_memory"):
|
|
30
|
+
"""
|
|
31
|
+
Object store that stores objects in memory. Objects are not persisted when the process shuts down.
|
|
32
|
+
"""
|
|
33
|
+
pass
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class InMemoryObjectStore(ObjectStore):
|
|
37
|
+
"""
|
|
38
|
+
Implementation of ObjectStore that stores objects in memory. Objects are not persisted when the process shuts down.
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
def __init__(self) -> None:
|
|
42
|
+
self._lock = asyncio.Lock()
|
|
43
|
+
self._store: dict[str, ObjectStoreItem] = {}
|
|
44
|
+
|
|
45
|
+
@override
|
|
46
|
+
async def put_object(self, key: str, item: ObjectStoreItem) -> None:
|
|
47
|
+
async with self._lock:
|
|
48
|
+
if key in self._store:
|
|
49
|
+
raise KeyAlreadyExistsError(key)
|
|
50
|
+
self._store[key] = item
|
|
51
|
+
|
|
52
|
+
@override
|
|
53
|
+
async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
|
|
54
|
+
async with self._lock:
|
|
55
|
+
self._store[key] = item
|
|
56
|
+
|
|
57
|
+
@override
|
|
58
|
+
async def get_object(self, key: str) -> ObjectStoreItem:
|
|
59
|
+
async with self._lock:
|
|
60
|
+
value = self._store.get(key)
|
|
61
|
+
if value is None:
|
|
62
|
+
raise NoSuchKeyError(key)
|
|
63
|
+
return value
|
|
64
|
+
|
|
65
|
+
@override
|
|
66
|
+
async def delete_object(self, key: str) -> None:
|
|
67
|
+
try:
|
|
68
|
+
async with self._lock:
|
|
69
|
+
self._store.pop(key)
|
|
70
|
+
except KeyError:
|
|
71
|
+
raise NoSuchKeyError(key)
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
@register_object_store(config_type=InMemoryObjectStoreConfig)
|
|
75
|
+
async def in_memory_object_store(config: InMemoryObjectStoreConfig, builder: Builder):
|
|
76
|
+
yield InMemoryObjectStore()
|
|
@@ -0,0 +1,84 @@
|
|
|
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 abc import ABC
|
|
17
|
+
from abc import abstractmethod
|
|
18
|
+
|
|
19
|
+
from .models import ObjectStoreItem
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class ObjectStore(ABC):
|
|
23
|
+
"""
|
|
24
|
+
Abstract interface for an object store.
|
|
25
|
+
|
|
26
|
+
Implementations may integrate with various object stores,
|
|
27
|
+
such as S3, MySQL, etc.
|
|
28
|
+
"""
|
|
29
|
+
|
|
30
|
+
@abstractmethod
|
|
31
|
+
async def put_object(self, key: str, item: ObjectStoreItem) -> None:
|
|
32
|
+
"""
|
|
33
|
+
Save an ObjectStoreItem in the object store with the given key.
|
|
34
|
+
If the key already exists, raise an error.
|
|
35
|
+
|
|
36
|
+
Args:
|
|
37
|
+
key (str): The key to save the item under.
|
|
38
|
+
item (ObjectStoreItem): The item to save.
|
|
39
|
+
|
|
40
|
+
Raises:
|
|
41
|
+
KeyAlreadyExistsError: If the key already exists.
|
|
42
|
+
"""
|
|
43
|
+
pass
|
|
44
|
+
|
|
45
|
+
@abstractmethod
|
|
46
|
+
async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
|
|
47
|
+
"""
|
|
48
|
+
Save an ObjectStoreItem in the object store with the given key.
|
|
49
|
+
If the key already exists, update the item.
|
|
50
|
+
|
|
51
|
+
Args:
|
|
52
|
+
key (str): The key to save the item under.
|
|
53
|
+
item (ObjectStoreItem): The item to save.
|
|
54
|
+
"""
|
|
55
|
+
pass
|
|
56
|
+
|
|
57
|
+
@abstractmethod
|
|
58
|
+
async def get_object(self, key: str) -> ObjectStoreItem:
|
|
59
|
+
"""
|
|
60
|
+
Get an ObjectStoreItem from the object store by key.
|
|
61
|
+
|
|
62
|
+
Args:
|
|
63
|
+
key (str): The key to get the item from.
|
|
64
|
+
|
|
65
|
+
Returns:
|
|
66
|
+
ObjectStoreItem: The item retrieved from the object store.
|
|
67
|
+
|
|
68
|
+
Raises:
|
|
69
|
+
NoSuchKeyError: If the item does not exist.
|
|
70
|
+
"""
|
|
71
|
+
pass
|
|
72
|
+
|
|
73
|
+
@abstractmethod
|
|
74
|
+
async def delete_object(self, key: str) -> None:
|
|
75
|
+
"""
|
|
76
|
+
Delete an ObjectStoreItem from the object store by key.
|
|
77
|
+
|
|
78
|
+
Args:
|
|
79
|
+
key (str): The key to delete the item from.
|
|
80
|
+
|
|
81
|
+
Raises:
|
|
82
|
+
NoSuchKeyError: If the item does not exist.
|
|
83
|
+
"""
|
|
84
|
+
pass
|
|
@@ -0,0 +1,36 @@
|
|
|
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 BaseModel
|
|
17
|
+
from pydantic import Field
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class ObjectStoreItem(BaseModel):
|
|
21
|
+
"""
|
|
22
|
+
Represents an object store item consisting of bytes and associated metadata.
|
|
23
|
+
|
|
24
|
+
Attributes
|
|
25
|
+
----------
|
|
26
|
+
data : bytes
|
|
27
|
+
The data to store in the object store.
|
|
28
|
+
content_type : str | None
|
|
29
|
+
The content type of the data.
|
|
30
|
+
metadata : dict[str, str] | None
|
|
31
|
+
Metadata providing context and utility for management operations.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
data: bytes = Field(description="The data to store in the object store.")
|
|
35
|
+
content_type: str | None = Field(description="The content type of the data.", default=None)
|
|
36
|
+
metadata: dict[str, str] | None = Field(description="The metadata of the data.", default=None)
|
|
@@ -0,0 +1,20 @@
|
|
|
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
|
+
# pylint: disable=unused-import
|
|
17
|
+
# flake8: noqa
|
|
18
|
+
# isort:skip_file
|
|
19
|
+
|
|
20
|
+
from . import in_memory_object_store
|
|
@@ -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,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.
|