PyPI - aiqtoolkit - Versions diffs - 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl - Mend

aiqtoolkit 1.2.0.dev0py3-none-any.whl → 1.2.0rc2py3-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.

Potentially problematic release.

This version of aiqtoolkit might be problematic. Click here for more details.

Files changed (220) hide show

aiq/agent/base.py +170 -8
aiq/agent/dual_node.py +1 -1
aiq/agent/react_agent/agent.py +146 -112
aiq/agent/react_agent/prompt.py +1 -6
aiq/agent/react_agent/register.py +36 -35
aiq/agent/rewoo_agent/agent.py +36 -35
aiq/agent/rewoo_agent/register.py +2 -2
aiq/agent/tool_calling_agent/agent.py +3 -7
aiq/agent/tool_calling_agent/register.py +1 -1
aiq/authentication/__init__.py +14 -0
aiq/authentication/api_key/__init__.py +14 -0
aiq/authentication/api_key/api_key_auth_provider.py +92 -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/exceptions/auth_code_grant_exceptions.py +86 -0
aiq/authentication/exceptions/call_back_exceptions.py +38 -0
aiq/authentication/exceptions/request_exceptions.py +54 -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/builder.py +64 -2
aiq/builder/component_utils.py +16 -3
aiq/builder/context.py +37 -0
aiq/builder/eval_builder.py +43 -2
aiq/builder/function.py +44 -12
aiq/builder/function_base.py +1 -1
aiq/builder/intermediate_step_manager.py +6 -8
aiq/builder/user_interaction_manager.py +3 -0
aiq/builder/workflow.py +23 -18
aiq/builder/workflow_builder.py +421 -61
aiq/cli/commands/info/list_mcp.py +103 -16
aiq/cli/commands/sizing/__init__.py +14 -0
aiq/cli/commands/sizing/calc.py +294 -0
aiq/cli/commands/sizing/sizing.py +27 -0
aiq/cli/commands/start.py +2 -1
aiq/cli/entrypoint.py +2 -0
aiq/cli/register_workflow.py +80 -0
aiq/cli/type_registry.py +151 -30
aiq/data_models/api_server.py +124 -12
aiq/data_models/authentication.py +231 -0
aiq/data_models/common.py +35 -7
aiq/data_models/component.py +17 -9
aiq/data_models/component_ref.py +33 -0
aiq/data_models/config.py +60 -3
aiq/data_models/dataset_handler.py +2 -1
aiq/data_models/embedder.py +1 -0
aiq/data_models/evaluate.py +23 -0
aiq/data_models/function_dependencies.py +8 -0
aiq/data_models/interactive.py +10 -1
aiq/data_models/intermediate_step.py +38 -5
aiq/data_models/its_strategy.py +30 -0
aiq/data_models/llm.py +1 -0
aiq/data_models/memory.py +1 -0
aiq/data_models/object_store.py +44 -0
aiq/data_models/profiler.py +1 -0
aiq/data_models/retry_mixin.py +35 -0
aiq/data_models/span.py +187 -0
aiq/data_models/telemetry_exporter.py +2 -2
aiq/embedder/nim_embedder.py +2 -1
aiq/embedder/openai_embedder.py +2 -1
aiq/eval/config.py +19 -1
aiq/eval/dataset_handler/dataset_handler.py +87 -2
aiq/eval/evaluate.py +208 -27
aiq/eval/evaluator/base_evaluator.py +73 -0
aiq/eval/evaluator/evaluator_model.py +1 -0
aiq/eval/intermediate_step_adapter.py +11 -5
aiq/eval/rag_evaluator/evaluate.py +55 -15
aiq/eval/rag_evaluator/register.py +6 -1
aiq/eval/remote_workflow.py +7 -2
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/trajectory_evaluator/evaluate.py +22 -65
aiq/eval/tunable_rag_evaluator/evaluate.py +150 -168
aiq/eval/tunable_rag_evaluator/register.py +2 -0
aiq/eval/usage_stats.py +41 -0
aiq/eval/utils/output_uploader.py +10 -1
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/inference_time_scaling/__init__.py +0 -0
aiq/experimental/inference_time_scaling/editing/__init__.py +0 -0
aiq/experimental/inference_time_scaling/editing/iterative_plan_refinement_editor.py +147 -0
aiq/experimental/inference_time_scaling/editing/llm_as_a_judge_editor.py +204 -0
aiq/experimental/inference_time_scaling/editing/motivation_aware_summarization.py +107 -0
aiq/experimental/inference_time_scaling/functions/__init__.py +0 -0
aiq/experimental/inference_time_scaling/functions/execute_score_select_function.py +105 -0
aiq/experimental/inference_time_scaling/functions/its_tool_orchestration_function.py +205 -0
aiq/experimental/inference_time_scaling/functions/its_tool_wrapper_function.py +146 -0
aiq/experimental/inference_time_scaling/functions/plan_select_execute_function.py +224 -0
aiq/experimental/inference_time_scaling/models/__init__.py +0 -0
aiq/experimental/inference_time_scaling/models/editor_config.py +132 -0
aiq/experimental/inference_time_scaling/models/its_item.py +48 -0
aiq/experimental/inference_time_scaling/models/scoring_config.py +112 -0
aiq/experimental/inference_time_scaling/models/search_config.py +120 -0
aiq/experimental/inference_time_scaling/models/selection_config.py +154 -0
aiq/experimental/inference_time_scaling/models/stage_enums.py +43 -0
aiq/experimental/inference_time_scaling/models/strategy_base.py +66 -0
aiq/experimental/inference_time_scaling/models/tool_use_config.py +41 -0
aiq/experimental/inference_time_scaling/register.py +36 -0
aiq/experimental/inference_time_scaling/scoring/__init__.py +0 -0
aiq/experimental/inference_time_scaling/scoring/llm_based_agent_scorer.py +168 -0
aiq/experimental/inference_time_scaling/scoring/llm_based_plan_scorer.py +168 -0
aiq/experimental/inference_time_scaling/scoring/motivation_aware_scorer.py +111 -0
aiq/experimental/inference_time_scaling/search/__init__.py +0 -0
aiq/experimental/inference_time_scaling/search/multi_llm_planner.py +128 -0
aiq/experimental/inference_time_scaling/search/multi_query_retrieval_search.py +122 -0
aiq/experimental/inference_time_scaling/search/single_shot_multi_plan_planner.py +128 -0
aiq/experimental/inference_time_scaling/selection/__init__.py +0 -0
aiq/experimental/inference_time_scaling/selection/best_of_n_selector.py +63 -0
aiq/experimental/inference_time_scaling/selection/llm_based_agent_output_selector.py +131 -0
aiq/experimental/inference_time_scaling/selection/llm_based_output_merging_selector.py +159 -0
aiq/experimental/inference_time_scaling/selection/llm_based_plan_selector.py +128 -0
aiq/experimental/inference_time_scaling/selection/threshold_selector.py +58 -0
aiq/front_ends/console/authentication_flow_handler.py +233 -0
aiq/front_ends/console/console_front_end_plugin.py +11 -2
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 +93 -9
aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
aiq/front_ends/fastapi/fastapi_front_end_plugin.py +14 -1
aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +537 -52
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/job_store.py +47 -25
aiq/front_ends/fastapi/main.py +2 -0
aiq/front_ends/fastapi/message_handler.py +108 -89
aiq/front_ends/fastapi/step_adaptor.py +2 -1
aiq/llm/aws_bedrock_llm.py +57 -0
aiq/llm/nim_llm.py +2 -1
aiq/llm/openai_llm.py +3 -2
aiq/llm/register.py +1 -0
aiq/meta/pypi.md +12 -12
aiq/object_store/__init__.py +20 -0
aiq/object_store/in_memory_object_store.py +74 -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 +269 -0
aiq/observability/exporter/raw_exporter.py +52 -0
aiq/observability/exporter/span_exporter.py +264 -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 +316 -0
aiq/observability/processor/intermediate_step_serializer.py +28 -0
aiq/observability/processor/processor.py +68 -0
aiq/observability/register.py +36 -39
aiq/observability/utils/__init__.py +14 -0
aiq/observability/utils/dict_utils.py +236 -0
aiq/observability/utils/time_utils.py +31 -0
aiq/profiler/calc/__init__.py +14 -0
aiq/profiler/calc/calc_runner.py +623 -0
aiq/profiler/calc/calculations.py +288 -0
aiq/profiler/calc/data_models.py +176 -0
aiq/profiler/calc/plot.py +345 -0
aiq/profiler/callbacks/langchain_callback_handler.py +22 -10
aiq/profiler/data_models.py +24 -0
aiq/profiler/inference_metrics_model.py +3 -0
aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +8 -0
aiq/profiler/inference_optimization/data_models.py +2 -2
aiq/profiler/inference_optimization/llm_metrics.py +2 -2
aiq/profiler/profile_runner.py +61 -21
aiq/runtime/loader.py +9 -3
aiq/runtime/runner.py +23 -9
aiq/runtime/session.py +25 -7
aiq/runtime/user_metadata.py +2 -3
aiq/tool/chat_completion.py +74 -0
aiq/tool/code_execution/README.md +152 -0
aiq/tool/code_execution/code_sandbox.py +151 -72
aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +139 -24
aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +3 -1
aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +27 -2
aiq/tool/code_execution/register.py +7 -3
aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
aiq/tool/mcp/exceptions.py +142 -0
aiq/tool/mcp/mcp_client.py +41 -6
aiq/tool/mcp/mcp_tool.py +3 -2
aiq/tool/register.py +1 -0
aiq/tool/server_tools.py +6 -3
aiq/utils/exception_handlers/automatic_retries.py +289 -0
aiq/utils/exception_handlers/mcp.py +211 -0
aiq/utils/io/model_processing.py +28 -0
aiq/utils/log_utils.py +37 -0
aiq/utils/string_utils.py +38 -0
aiq/utils/type_converter.py +18 -2
aiq/utils/type_utils.py +87 -0
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/METADATA +53 -21
aiqtoolkit-1.2.0rc2.dist-info/RECORD +436 -0
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/WHEEL +1 -1
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/entry_points.txt +3 -0
aiq/front_ends/fastapi/websocket.py +0 -148
aiq/observability/async_otel_listener.py +0 -429
aiqtoolkit-1.2.0.dev0.dist-info/RECORD +0 -316
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/licenses/LICENSE.md +0 -0
{aiqtoolkit-1.2.0.dev0.dist-info → aiqtoolkit-1.2.0rc2.dist-info}/top_level.txt +0 -0

aiq/front_ends/fastapi/html_snippets/__init__.py ADDED Viewed

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py ADDED Viewed

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+AUTH_REDIRECT_SUCCESS_HTML = """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Authentication Complete</title>
+    <script>
+        (function () {
+            window.history.replaceState(null, "", window.location.pathname);
+            window.opener?.postMessage({ type: 'AUTH_SUCCESS' }, '*');
+            window.close();
+        })();
+    </script>
+</head>
+<body>
+    <p>Authentication complete. You may now close this window.</p>
+</body>
+</html>
+"""

aiq/front_ends/fastapi/job_store.py CHANGED Viewed

@@ -16,6 +16,7 @@
 import logging
 import os
 import shutil
+import threading
 from datetime import UTC
 from datetime import datetime
 from datetime import timedelta
@@ -40,12 +41,13 @@ class JobStatus(str, Enum):
 class JobInfo(BaseModel):
     job_id: str
     status: JobStatus
-    config_file: str
+    config_file: str | None
     error: str | None
     output_path: str | None
     created_at: datetime
     updated_at: datetime
     expiry_seconds: int
+    output: BaseModel | None = None
 class JobStore:
@@ -59,8 +61,12 @@ class JobStore:
     def __init__(self):
         self._jobs = {}
+        self._lock = threading.Lock()  # Ensure thread safety for job operations
-    def create_job(self, config_file: str, job_id: str | None = None, expiry_seconds: int = DEFAULT_EXPIRY) -> str:
+    def create_job(self,
+                   config_file: str | None = None,
+                   job_id: str | None = None,
+                   expiry_seconds: int = DEFAULT_EXPIRY) -> str:
         if job_id is None:
             job_id = str(uuid4())
@@ -76,46 +82,62 @@ class JobStore:
                       error=None,
                       output_path=None,
                       expiry_seconds=clamped_expiry)
-        self._jobs[job_id] = job
+        with self._lock:
+            self._jobs[job_id] = job
         logger.info("Created new job %s with config %s", job_id, config_file)
         return job_id
-    def update_status(self, job_id: str, status: str, error: str | None = None, output_path: str | None = None):
+    def update_status(self,
+                      job_id: str,
+                      status: str,
+                      error: str | None = None,
+                      output_path: str | None = None,
+                      output: BaseModel | None = None):
         if job_id not in self._jobs:
             raise ValueError(f"Job {job_id} not found")
-        job = self._jobs[job_id]
-        job.status = status
-        job.error = error
-        job.output_path = output_path
-        job.updated_at = datetime.now(UTC)
+        with self._lock:
+            job = self._jobs[job_id]
+            job.status = status
+            job.error = error
+            job.output_path = output_path
+            job.updated_at = datetime.now(UTC)
+            job.output = output
     def get_status(self, job_id: str) -> JobInfo | None:
-        return self._jobs.get(job_id)
+        with self._lock:
+            return self._jobs.get(job_id)
     def list_jobs(self):
-        return self._jobs
+        with self._lock:
+            return self._jobs
     def get_job(self, job_id: str) -> JobInfo | None:
         """Get a job by its ID."""
-        return self._jobs.get(job_id)
+        with self._lock:
+            return self._jobs.get(job_id)
     def get_last_job(self) -> JobInfo | None:
         """Get the last created job."""
-        if not self._jobs:
-            logger.info("No jobs found in job store")
-            return None
-        last_job = max(self._jobs.values(), key=lambda job: job.created_at)
-        logger.info("Retrieved last job %s created at %s", last_job.job_id, last_job.created_at)
-        return last_job
+        with self._lock:
+            if not self._jobs:
+                logger.info("No jobs found in job store")
+                return None
+            last_job = max(self._jobs.values(), key=lambda job: job.created_at)
+            logger.info("Retrieved last job %s created at %s", last_job.job_id, last_job.created_at)
+            return last_job
     def get_jobs_by_status(self, status: str) -> list[JobInfo]:
         """Get all jobs with the specified status."""
-        return [job for job in self._jobs.values() if job.status == status]
+        with self._lock:
+            return [job for job in self._jobs.values() if job.status == status]
     def get_all_jobs(self) -> list[JobInfo]:
         """Get all jobs in the store."""
-        return list(self._jobs.values())
+        with self._lock:
+            return list(self._jobs.values())
     def get_expires_at(self, job: JobInfo) -> datetime | None:
         """Get the time for a job to expire."""
@@ -132,7 +154,8 @@ class JobStore:
         now = datetime.now(UTC)
         # Filter out active jobs
-        finished_jobs = {job_id: job for job_id, job in self._jobs.items() if job.status not in self.ACTIVE_STATUS}
+        with self._lock:
+            finished_jobs = {job_id: job for job_id, job in self._jobs.items() if job.status not in self.ACTIVE_STATUS}
         # Sort finished jobs by updated_at descending
         sorted_finished = sorted(finished_jobs.items(), key=lambda item: item[1].updated_at, reverse=True)
@@ -155,7 +178,6 @@ class JobStore:
                     elif os.path.isdir(job.output_path):
                         shutil.rmtree(job.output_path)
-        for job_id in expired_ids:
-            # cleanup output dir if present
-            del self._jobs[job_id]
+        with self._lock:
+            for job_id in expired_ids:
+                del self._jobs[job_id]

aiq/front_ends/fastapi/main.py CHANGED Viewed

@@ -68,3 +68,5 @@ def get_app():
     except ImportError as e:
         raise ValueError(f"Front end worker {front_end_worker_full_name} not found.") from e
+    except Exception as e:
+        raise ValueError(f"Error loading front end worker {front_end_worker_full_name}: {e}") from e

aiq/front_ends/fastapi/message_handler.py CHANGED Viewed

@@ -15,14 +15,19 @@
 import asyncio
 import logging
+import typing
 import uuid
 from typing import Any
 from fastapi import WebSocket
 from pydantic import BaseModel
 from pydantic import ValidationError
-from starlette.endpoints import WebSocketEndpoint
+from starlette.websockets import WebSocketDisconnect
+from aiq.authentication.interfaces import FlowHandlerBase
+from aiq.data_models.api_server import AIQChatResponse
+from aiq.data_models.api_server import AIQResponsePayloadOutput
+from aiq.data_models.api_server import AIQResponseSerializable
 from aiq.data_models.api_server import Error
 from aiq.data_models.api_server import ErrorTypes
 from aiq.data_models.api_server import SystemResponseContent
@@ -39,74 +44,72 @@ from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import HumanResponseNotification
 from aiq.data_models.interactive import InteractionPrompt
 from aiq.front_ends.fastapi.message_validator import MessageValidator
+from aiq.front_ends.fastapi.response_helpers import generate_streaming_response
+from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
+from aiq.runtime.session import AIQSessionManager
 logger = logging.getLogger(__name__)
-class MessageHandler:
+class WebSocketMessageHandler:
+    def __init__(self, socket: WebSocket, session_manager: AIQSessionManager, step_adaptor: StepAdaptor):
+        self._socket: WebSocket = socket
+        self._session_manager: AIQSessionManager = session_manager
+        self._step_adaptor: StepAdaptor = step_adaptor
-    def __init__(self, websocket_reference: WebSocketEndpoint):
-        self._websocket_reference: WebSocketEndpoint = websocket_reference
         self._message_validator: MessageValidator = MessageValidator()
-        self._messages_queue: asyncio.Queue[dict[str, str]] = asyncio.Queue()
-        self._out_going_messages_queue: asyncio.Queue[dict] = asyncio.Queue()
-        self._process_messages_task: asyncio.Task | None = None
-        self._process_out_going_messages_task: asyncio.Task = None
-        self._background_task: asyncio.Task = None
+        self._running_workflow_task: asyncio.Task | None = None
         self._message_parent_id: str = "default_id"
         self._workflow_schema_type: str = None
-        self._user_interaction_response: asyncio.Future[TextContent] = asyncio.Future()
+        self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
-    @property
-    def messages_queue(self) -> asyncio.Queue[dict[str, str]]:
-        return self._messages_queue
+        self._flow_handler: FlowHandlerBase | None = None
-    @property
-    def background_task(self) -> asyncio.Task:
-        return self._background_task
+    def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
+        self._flow_handler = flow_handler
-    @property
-    def process_messages_task(self) -> asyncio.Task | None:
-        return self._process_messages_task
+    async def __aenter__(self) -> "WebSocketMessageHandler":
+        await self._socket.accept()
-    @process_messages_task.setter
-    def process_messages_task(self, process_messages_task) -> None:
-        self._process_messages_task = process_messages_task
+        return self
-    @property
-    def process_out_going_messages_task(self) -> asyncio.Task:
-        return self._process_out_going_messages_task
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
-    @process_out_going_messages_task.setter
-    def process_out_going_messages_task(self, process_out_going_messages_task) -> None:
-        self._process_out_going_messages_task = process_out_going_messages_task
+        # TODO: Handle the exit
+        pass
-    async def process_messages(self) -> None:
+    async def run(self) -> None:
         """
         Processes received messages from websocket and routes them appropriately.
         """
         while True:
             try:
-                message: dict[str, Any] = await self._messages_queue.get()
+                message: dict[str, Any] = await self._socket.receive_json()
                 validated_message: BaseModel = await self._message_validator.validate_message(message)
+                # Received a request to start a workflow
                 if (isinstance(validated_message, WebSocketUserMessage)):
-                    await self.process_user_message(validated_message)
+                    await self.process_workflow_request(validated_message)
-                if isinstance(
+                elif isinstance(
                         validated_message,
                     (  # noqa: E131
                         WebSocketSystemResponseTokenMessage,
                         WebSocketSystemIntermediateStepMessage,
                         WebSocketSystemInteractionMessage)):
-                    await self._out_going_messages_queue.put(validated_message.model_dump())
+                    # These messages are already handled by self.create_websocket_message(data_model=value, …)
+                    # No further processing is needed here.
+                    pass
-                if (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
+                elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
-            except (asyncio.CancelledError):
+            except (asyncio.CancelledError, WebSocketDisconnect):
+                # TODO: Handle the disconnect
                 break
         return None
@@ -130,29 +133,32 @@ class MessageHandler:
         return None
-    async def process_user_message(self, message_as_validated_type: WebSocketUserMessage) -> None:
+    async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
         """
         Process user messages and routes them appropriately.
-        :param message_as_validated_type: A WebSocketUserMessage Data Model instance.
+        :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
         """
         try:
-            self._message_parent_id = message_as_validated_type.id
-            self._workflow_schema_type = message_as_validated_type.schema_type
+            self._message_parent_id = user_message_as_validated_type.id
+            self._workflow_schema_type = user_message_as_validated_type.schema_type
+            conversation_id: str = user_message_as_validated_type.conversation_id
-            content: BaseModel | None = await self.process_user_message_content(message_as_validated_type)
+            content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
             if content is None:
-                raise ValueError(f"User message content could not be found: {message_as_validated_type}")
+                raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
-            if isinstance(content, TextContent) and (self._background_task is None):
+            if isinstance(content, TextContent) and (self._running_workflow_task is None):
-                await self._process_response()
-                self._background_task = asyncio.create_task(
-                    self._websocket_reference.workflow_schema_type.get(self._workflow_schema_type)(
-                        content.text)).add_done_callback(
-                            lambda task: asyncio.create_task(self._on_process_stream_task_done(task)))
+                def _done_callback(task: asyncio.Task):
+                    self._running_workflow_task = None
+                # await self._process_response()
+                self._running_workflow_task = asyncio.create_task(
+                    self._run_workflow(content.text, conversation_id,
+                                       result_type=AIQChatResponse)).add_done_callback(_done_callback)
         except ValueError as e:
             logger.error("User message content not found: %s", str(e), exc_info=True)
@@ -220,60 +226,73 @@ class MessageHandler:
                 content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
         finally:
-            await self._messages_queue.put(message.model_dump())
-    async def _on_process_stream_task_done(self, task: asyncio.Task) -> None:
-        await self.create_websocket_message(data_model=SystemResponseContent(),
-                                            message_type=WebSocketMessageType.RESPONSE_MESSAGE,
-                                            status=WebSocketMessageStatus.COMPLETE)
+            if (message is not None):
+                await self._socket.send_json(message.model_dump())
-        return None
-    async def process_out_going_messages(self, websocket: WebSocket) -> None:
+    async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse:
         """
-        Spawns out going message processing task.
+        Registered human interaction callback that processes human interactions and returns
+        responses from websocket connection.
-        :param websocket: Websocket instance.
+        :param prompt: Incoming interaction content data model.
+        :return: A Text Content Base Pydantic model.
         """
-        while True:
-            try:
-                out_going_message = await self._out_going_messages_queue.get()
-                await self._websocket_reference.on_send(websocket, out_going_message)
-            except (asyncio.CancelledError, ValidationError):
-                break
+        # First create a future from the loop for the human response
+        human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
-        return None
+        # Then add the future to the outstanding human prompts dictionary
+        self._user_interaction_response = human_response_future
+        try:
-    async def _process_response(self):
-        self._websocket_reference.process_response_event.set()
+            await self.create_websocket_message(data_model=prompt.content,
+                                                message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
+                                                status=WebSocketMessageStatus.IN_PROGRESS)
-    async def _pause_response(self):
-        self._websocket_reference.process_response_event.clear()
+            if (isinstance(prompt.content, HumanPromptNotification)):
-    async def __reset_user_interaction_response(self):
-        self._user_interaction_response = asyncio.Future()
+                return HumanResponseNotification()
-    async def human_interaction(self, prompt: InteractionPrompt) -> HumanResponse:
-        """
-        Registered human interaction callback that processes human interactions and returns
-        responses from websocket connection.
+            # Wait for the human response future to complete
+            interaction_response: HumanResponse = await human_response_future
-        :param prompt: Incoming interaction content data model.
-        :return: A Text Content Base Pydantic model.
-        """
-        await self.create_websocket_message(data_model=prompt.content,
-                                            message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
-                                            status=WebSocketMessageStatus.IN_PROGRESS)
+            interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
+                interaction_response, prompt.content)
-        if (isinstance(prompt.content, HumanPromptNotification)):
-            return HumanResponseNotification()
+            return interaction_response
-        user_message_repsonse_content: TextContent = await self._user_interaction_response
-        interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
-            user_message_repsonse_content, prompt.content)
+        finally:
+            # Delete the future from the outstanding human prompts dictionary
+            self._user_interaction_response = None
+    async def _run_workflow(self,
+                            payload: typing.Any,
+                            conversation_id: str | None = None,
+                            result_type: type | None = None,
+                            output_type: type | None = None) -> None:
+        try:
+            async with self._session_manager.session(
+                    conversation_id=conversation_id,
+                    request=self._socket,
+                    user_input_callback=self.human_interaction_callback,
+                    user_authentication_callback=(self._flow_handler.authenticate
+                                                  if self._flow_handler else None)) as session:
-        await self.__reset_user_interaction_response()
-        await self._process_response()
+                async for value in generate_streaming_response(payload,
+                                                               session_manager=session,
+                                                               streaming=True,
+                                                               step_adaptor=self._step_adaptor,
+                                                               result_type=result_type,
+                                                               output_type=output_type):
-        return interaction_response
+                    if not isinstance(value, AIQResponseSerializable):
+                        value = AIQResponsePayloadOutput(payload=value)
+                    await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
+        finally:
+            await self.create_websocket_message(data_model=SystemResponseContent(),
+                                                message_type=WebSocketMessageType.RESPONSE_MESSAGE,
+                                                status=WebSocketMessageStatus.COMPLETE)

aiq/front_ends/fastapi/step_adaptor.py CHANGED Viewed

@@ -291,7 +291,8 @@ class StepAdaptor:
         return event
-    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:
+    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:  # pylint: disable=R1710
         # Track the chunk
         self._history.append(step)
         payload = step.payload

aiq/llm/aws_bedrock_llm.py ADDED Viewed

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+from aiq.builder.builder import Builder
+from aiq.builder.llm import LLMProviderInfo
+from aiq.cli.register_workflow import register_llm_provider
+from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
+    """An AWS Bedrock llm provider to be used with an LLM client."""
+    model_config = ConfigDict(protected_namespaces=())
+    # Completion parameters
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The model name for the hosted AWS Bedrock.")
+    temperature: float = Field(default=0.0, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")
+    max_tokens: int | None = Field(default=1024,
+                                   gt=0,
+                                   description="Maximum number of tokens to generate."
+                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
+    context_size: int | None = Field(default=1024,
+                                     gt=0,
+                                     description="Maximum number of tokens to generate."
+                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+    # Client parameters
+    region_name: str | None = Field(default="None", description="AWS region to use.")
+    base_url: str | None = Field(
+        default=None, description="Bedrock endpoint to use. Needed if you don't want to default to us-east-1 endpoint.")
+    credentials_profile_name: str | None = Field(
+        default=None, description="The name of the profile in the ~/.aws/credentials or ~/.aws/config files.")
+@register_llm_provider(config_type=AWSBedrockModelConfig)
+async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, builder: Builder):
+    yield LLMProviderInfo(config=llm_config, description="A AWS Bedrock model for use with an LLM client.")

aiq/llm/nim_llm.py CHANGED Viewed

@@ -22,9 +22,10 @@ from aiq.builder.builder import Builder
 from aiq.builder.llm import LLMProviderInfo
 from aiq.cli.register_workflow import register_llm_provider
 from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
-class NIMModelConfig(LLMBaseConfig, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
     model_config = ConfigDict(protected_namespaces=())

aiq/llm/openai_llm.py CHANGED Viewed

@@ -21,12 +21,13 @@ from aiq.builder.builder import Builder
 from aiq.builder.llm import LLMProviderInfo
 from aiq.cli.register_workflow import register_llm_provider
 from aiq.data_models.llm import LLMBaseConfig
+from aiq.data_models.retry_mixin import RetryMixin
-class OpenAIModelConfig(LLMBaseConfig, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
     api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
     base_url: str | None = Field(default=None, description="Base url to the hosted model.")

aiq/llm/register.py CHANGED Viewed

@@ -18,5 +18,6 @@
 # isort:skip_file
 # Import any providers which need to be automatically registered here
+from . import aws_bedrock_llm
 from . import nim_llm
 from . import openai_llm

aiq/meta/pypi.md CHANGED Viewed

@@ -15,7 +15,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 -->
-![NVIDIA Agent Intelligence Toolkit](https://media.githubusercontent.com/media/NVIDIA/AIQToolkit/refs/heads/main/docs/source/_static/aiqtoolkit_banner.png "AIQ toolkit banner image")
+![NVIDIA Agent Intelligence Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/aiqtoolkit_banner.png "AIQ toolkit banner image")
 # NVIDIA Agent Intelligence Toolkit
@@ -23,26 +23,26 @@ AIQ toolkit is a flexible library designed to seamlessly integrate your enterpri
 ## Key Features
-- [**Framework Agnostic:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
-- [**Reusability:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
-- [**Rapid Development:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
-- [**Profiling:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
-- [**Observability:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/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/v1.2.0-dev/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/workflows/observe/observe-workflow-with-weave.html).
-- [**Evaluation System:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
-- [**User Interface:**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/quick-start/launching-ui.html) Use the AIQ toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
-- [**MCP Compatibility**](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as AIQ toolkit functions.
+- [**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.
+- [**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.
+- [**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.
+- [**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.
+- [**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).
+- [**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.
+- [**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.
+- [**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.
 With AIQ toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
 ## Links
- * [Documentation](https://docs.nvidia.com/aiqtoolkit/v1.2.0-dev/index.html): Explore the full documentation for AIQ toolkit.
+ * [Documentation](https://docs.nvidia.com/aiqtoolkit/1.2.0/index.html): Explore the full documentation for AIQ toolkit.
 ## First time user?
- 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/AIQToolkit?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.
+ 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.
 ## Feedback
-We would love to hear from you! Please file an issue on [GitHub](https://github.com/NVIDIA/AIQToolkit/issues) if you have any feedback or feature requests.
+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.
 ## Acknowledgements

aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

Potentially problematic release.

aiqtoolkit 1.2.0.dev0py3-none-any.whl → 1.2.0rc2py3-none-any.whl