aiqtoolkit 1.1.0rc6__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiq/front_ends/mcp/tool_converter.py

@@ -1,242 +0,0 @@
-# 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.
-
-import json
-import logging
-from inspect import Parameter
-from inspect import Signature
-
-from mcp.server.fastmcp import FastMCP
-from pydantic import BaseModel
-
-from aiq.builder.function import Function
-from aiq.builder.function_base import FunctionBase
-from aiq.builder.workflow import Workflow
-
-logger = logging.getLogger(__name__)
-
-
-def create_function_wrapper(
-    function_name: str,
-    function: FunctionBase,
-    schema: type[BaseModel],
-    is_workflow: bool = False,
-):
-    """Create a wrapper function that exposes the actual parameters of an AIQ Function as an MCP tool.
-
-    Args:
-        function_name: The name of the function/tool
-        function: The AIQ Function object
-        schema: The input schema of the function
-        is_workflow: Whether the function is a Workflow
-
-    Returns:
-        A wrapper function suitable for registration with MCP
-    """
-    # Check if we're dealing with AIQChatRequest - special case
-    is_chat_request = False
-
-    # Check if the schema name is AIQChatRequest
-    if schema.__name__ == "AIQChatRequest" or (hasattr(schema, "__qualname__")
-                                               and "AIQChatRequest" in schema.__qualname__):
-        is_chat_request = True
-        logger.info("Function %s uses AIQChatRequest - creating simplified interface", function_name)
-
-        # For AIQChatRequest, we'll create a simple wrapper with just a query parameter
-        parameters = [Parameter(
-            name="query",
-            kind=Parameter.KEYWORD_ONLY,
-            default=Parameter.empty,
-            annotation=str,
-        )]
-    else:
-        # Regular case - extract parameter information from the input schema
-        # Extract parameter information from the input schema
-        param_fields = schema.model_fields
-
-        parameters = []
-        for name, field in param_fields.items():
-            # Get the field type and convert to appropriate Python type
-            field_type = field.annotation
-
-            # Add the parameter to our list
-            parameters.append(
-                Parameter(
-                    name=name,
-                    kind=Parameter.KEYWORD_ONLY,
-                    default=Parameter.empty if field.is_required else None,
-                    annotation=field_type,
-                ))
-
-    # Create the function signature WITHOUT the ctx parameter
-    # We'll handle this in the wrapper function internally
-    sig = Signature(parameters=parameters, return_annotation=str)
-
-    # Define the actual wrapper function that accepts ctx but doesn't expose it
-    def create_wrapper():
-
-        async def wrapper_with_ctx(**kwargs):
-            """Internal wrapper that will be called by MCP."""
-            # MCP will add a ctx parameter, extract it
-            ctx = kwargs.get("ctx")
-
-            # Remove ctx if present
-            if "ctx" in kwargs:
-                del kwargs["ctx"]
-
-            # Process the function call
-            if ctx:
-                ctx.info("Calling function %s with args: %s", function_name, json.dumps(kwargs, default=str))
-                await ctx.report_progress(0, 100)
-
-            try:
-                # Special handling for AIQChatRequest
-                if is_chat_request:
-                    from aiq.data_models.api_server import AIQChatRequest
-
-                    # Create a chat request from the query string
-                    query = kwargs.get("query", "")
-                    chat_request = AIQChatRequest.from_string(query)
-
-                    # Special handling for Workflow objects
-                    if is_workflow:
-                        # Workflows have a run method that is an async context manager
-                        # that returns an AIQRunner
-                        async with function.run(chat_request) as runner:
-                            # Get the result from the runner
-                            result = await runner.result(to_type=str)
-                    else:
-                        # Regular functions use ainvoke
-                        result = await function.ainvoke(chat_request, to_type=str)
-                else:
-                    # Regular handling
-                    # Handle complex input schema - if we extracted fields from a nested schema,
-                    # we need to reconstruct the input
-                    if len(schema.model_fields) == 1 and len(parameters) > 1:
-                        # Get the field name from the original schema
-                        field_name = next(iter(schema.model_fields.keys()))
-                        field_type = schema.model_fields[field_name].annotation
-
-                        # If it's a pydantic model, we need to create an instance
-                        if hasattr(field_type, "model_validate"):
-                            # Create the nested object
-                            nested_obj = field_type.model_validate(kwargs)
-                            # Call with the nested object
-                            kwargs = {field_name: nested_obj}
-
-                    # Call the AIQ function with the parameters - special handling for Workflow
-                    if is_workflow:
-                        # For workflow with regular input, we'll assume the first parameter is the input
-                        input_value = list(kwargs.values())[0] if kwargs else ""
-
-                        # Workflows have a run method that is an async context manager
-                        # that returns an AIQRunner
-                        async with function.run(input_value) as runner:
-                            # Get the result from the runner
-                            result = await runner.result(to_type=str)
-                    else:
-                        # Regular function call
-                        result = await function.acall_invoke(**kwargs)
-
-                # Report completion
-                if ctx:
-                    await ctx.report_progress(100, 100)
-
-                # Handle different result types for proper formatting
-                if isinstance(result, str):
-                    return result
-                if isinstance(result, (dict, list)):
-                    return json.dumps(result, default=str)
-                return str(result)
-            except Exception as e:
-                if ctx:
-                    ctx.error("Error calling function %s: %s", function_name, str(e))
-                raise
-
-        return wrapper_with_ctx
-
-    # Create the wrapper function
-    wrapper = create_wrapper()
-
-    # Set the signature on the wrapper function (WITHOUT ctx)
-    wrapper.__signature__ = sig
-    wrapper.__name__ = function_name
-
-    # Return the wrapper with proper signature
-    return wrapper
-
-
-def get_function_description(function: FunctionBase) -> str:
-    """
-    Retrieve a human-readable description for an AIQ function or workflow.
-
-    The description is determined using the following precedence:
-       1. If the function is a Workflow and has a 'description' attribute, use it.
-       2. If the Workflow's config has a 'topic', use it.
-       3. If the Workflow's config has a 'description', use it.
-       4. If the function is a regular Function, use its 'description' attribute.
-
-    Args:
-        function: The AIQ FunctionBase instance (Function or Workflow).
-
-    Returns:
-        The best available description string for the function.
-    """
-    function_description = ""
-
-    if isinstance(function, Workflow):
-        config = function.config
-
-        # Workflow doesn't have a description, but probably should
-        if hasattr(function, "description") and function.description:
-            function_description = function.description
-        # Try to get description from config
-        elif hasattr(config, "description") and config.description:
-            function_description = config.description
-        # Try to get anything that might be a description
-        elif hasattr(config, "topic") and config.topic:
-            function_description = config.topic
-
-    elif isinstance(function, Function):
-        function_description = function.description
-
-    return function_description
-
-
-def register_function_with_mcp(mcp: FastMCP, function_name: str, function: FunctionBase) -> None:
-    """Register an AIQ Function as an MCP tool.
-
-    Args:
-        mcp: The FastMCP instance
-        function_name: The name to register the function under
-        function: The AIQ Function to register
-    """
-    logger.info("Registering function %s with MCP", function_name)
-
-    # Get the input schema from the function
-    input_schema = function.input_schema
-    logger.info("Function %s has input schema: %s", function_name, input_schema)
-
-    # Check if we're dealing with a Workflow
-    is_workflow = isinstance(function, Workflow)
-    if is_workflow:
-        logger.info("Function %s is a Workflow", function_name)
-
-    # Get function description
-    function_description = get_function_description(function)
-
-    # Create and register the wrapper function with MCP
-    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow)
-    mcp.tool(name=function_name, description=function_description)(wrapper_func)

aiq/front_ends/register.py

@@ -1,22 +0,0 @@
-# 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.
-
-# pylint: disable=unused-import
-# flake8: noqa
-# isort:skip_file
-
-from .console import register as console_register
-from .fastapi import register as fastapi_register
-from .mcp import register as mcp_register

aiq/front_ends/simple_base/__init__.py

@@ -1,14 +0,0 @@
-# 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.

aiq/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -1,52 +0,0 @@
-# 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.
-
-import logging
-from abc import ABC
-from abc import abstractmethod
-from io import StringIO
-
-import click
-
-from aiq.builder.front_end import FrontEndBase
-from aiq.builder.workflow_builder import WorkflowBuilder
-from aiq.data_models.front_end import FrontEndConfigT
-from aiq.runtime.session import AIQSessionManager
-
-logger = logging.getLogger(__name__)
-
-
-class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
-
-    async def pre_run(self):
-        pass
-
-    async def run(self):
-
-        # Must yield the workflow function otherwise it cleans up
-        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
-
-            if logger.isEnabledFor(logging.INFO):
-                stream = StringIO()
-
-                self.full_config.print_summary(stream=stream)
-
-                click.echo(stream.getvalue())
-
-            await self.run_workflow(builder.build())
-
-    @abstractmethod
-    async def run_workflow(self, session_manager: AIQSessionManager = None):
-        pass

aiq/llm/nim_llm.py

@@ -1,45 +0,0 @@
-# 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 pydantic import PositiveInt
-
-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
-
-
-class NIMModelConfig(LLMBaseConfig, name="nim"):
-    """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
-
-    model_config = ConfigDict(protected_namespaces=())
-
-    api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
-    base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The model name for the hosted NIM.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
-    max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
-
-
-@register_llm_provider(config_type=NIMModelConfig)
-async def nim_model(llm_config: NIMModelConfig, builder: Builder):
-
-    yield LLMProviderInfo(config=llm_config, description="A NIM model for use with an LLM client.")

aiq/llm/openai_llm.py

@@ -1,45 +0,0 @@
-# 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
-
-
-class OpenAIModelConfig(LLMBaseConfig, name="openai"):
-    """An OpenAI LLM provider to be used with an LLM client."""
-
-    model_config = ConfigDict(protected_namespaces=())
-
-    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.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The OpenAI hosted model name.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
-    seed: int | None = Field(default=None, description="Random seed to set for generation.")
-    max_retries: int = Field(default=10, description="The max number of retries for the request.")
-
-
-@register_llm_provider(config_type=OpenAIModelConfig)
-async def openai_llm(config: OpenAIModelConfig, builder: Builder):
-
-    yield LLMProviderInfo(config=config, description="An OpenAI model for use with an LLM client.")

aiq/llm/register.py

@@ -1,22 +0,0 @@
-# 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.
-
-# pylint: disable=unused-import
-# flake8: noqa
-# isort:skip_file
-
-# Import any providers which need to be automatically registered here
-from . import nim_llm
-from . import openai_llm

aiq/llm/utils/__init__.py

@@ -1,14 +0,0 @@
-# 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/llm/utils/env_config_value.py

@@ -1,94 +0,0 @@
-# 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.
-
-import os
-from abc import ABC
-from enum import Enum
-
-
-class EnvConfigValueSource(Enum):
-    ENV_DEFAULT = 1
-    CONSTRUCTOR = 2
-    ENV_OVERRIDE = 3
-
-
-class EnvConfigValue(ABC):
-    """
-    A wrapper for a string used as a configuration value which can be loaded from the system environment or injected via
-    the constructor. This class should be subclassed and the class fields `_ENV_KEY` and `_ENV_KEY_OVERRIDE` can be set
-    to enable environment-loading functionality. Convienience properties are available to check from where the value was
-    loaded.
-    """
-
-    _ENV_KEY: str | None = None
-    _ENV_KEY_OVERRIDE: str | None = None
-    _ALLOW_NONE: bool = False
-
-    def __init__(self, value: str | None = None, use_env: bool = True):
-        """
-        Parameters
-        ----------
-        value : str, optional
-            The value to be contained in the EnvConfigValue. If the value is `None`, an attempt will be made to load it
-            from the environment using `_ENV_KEY`. if the `_ENV_KEY_OVERRIDE` field is not `None`, an attempt will be
-            made to load that environment variable in place of the passed-in value.
-        use_env : bool
-            If False, all environment-loading logic will be bypassed and the passed-in value will be used as-is.
-            defaults to True.
-        """
-
-        self._source = EnvConfigValueSource.CONSTRUCTOR
-
-        if use_env:
-            if value is None and self.__class__._ENV_KEY is not None:
-                value = os.environ.get(self.__class__._ENV_KEY, None)
-                self._source = EnvConfigValueSource.ENV_DEFAULT
-
-            if self.__class__._ENV_KEY_OVERRIDE is not None and self.__class__._ENV_KEY_OVERRIDE in os.environ:
-                value = os.environ[self.__class__._ENV_KEY_OVERRIDE]
-                self._source = EnvConfigValueSource.ENV_OVERRIDE
-
-            if not self.__class__._ALLOW_NONE and value is None:
-
-                message = ("value must not be None, but provided value was None and no environment-based default or "
-                           "override was found.")
-
-                if self.__class__._ENV_KEY is None:
-                    raise ValueError(message)
-
-                raise ValueError(
-                    f"{message} Try passing a value to the constructor, or setting the `{self.__class__._ENV_KEY}` "
-                    "environment variable.")
-
-        else:
-            if not self.__class__._ALLOW_NONE and value is None:
-                raise ValueError("value must not be none")
-
-        assert isinstance(value, str) or value is None
-
-        self._value = value
-        self._use_env = use_env
-
-    @property
-    def source(self) -> EnvConfigValueSource:
-        return self._source
-
-    @property
-    def use_env(self) -> bool:
-        return self._use_env
-
-    @property
-    def value(self) -> str | None:
-        return self._value

aiq/llm/utils/error.py

@@ -1,17 +0,0 @@
-# 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.
-
-IMPORT_ERROR_MESSAGE = (
-    "{package} not found. Install it and other additional dependencies by running the following command:\n")

aiq/memory/__init__.py

@@ -1,20 +0,0 @@
-# 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.
-"""
-AIQ Toolkit Memory Module
-
-This package provides foundational classes and interfaces
-for managing text-based memory in AIQ Toolkit's LLM-based agents.
-"""