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

aiq/agent/base.py

@@ -148,7 +148,7 @@ class BaseAgent(ABC):
         """
         last_exception = None
 
-        for attempt in range(max_retries + 1):
+        for attempt in range(1, max_retries + 1):
             try:
                 response = await tool.ainvoke(tool_input, config=config)
 
@@ -162,17 +162,18 @@ class BaseAgent(ABC):
 
             except Exception as e:
                 last_exception = e
-                logger.warning("%s Tool call attempt %d/%d failed for tool %s: %s",
-                               AGENT_LOG_PREFIX,
-                               attempt + 1,
-                               max_retries + 1,
-                               tool.name,
-                               str(e))
 
                 # If this was the last attempt, don't sleep
                 if attempt == max_retries:
                     break
 
+                logger.warning("%s Tool call attempt %d/%d failed for tool %s: %s",
+                               AGENT_LOG_PREFIX,
+                               attempt,
+                               max_retries,
+                               tool.name,
+                               str(e))
+
                 # Exponential backoff: 2^attempt seconds
                 sleep_time = 2**attempt
                 logger.debug("%s Retrying tool call for %s in %d seconds...", AGENT_LOG_PREFIX, tool.name, sleep_time)

aiq/agent/react_agent/agent.py

@@ -193,12 +193,11 @@ class ReActAgentGraph(DualNodeAgent):
                                  ex.observation,
                                  output_message.content)
                     if attempt == self.parse_agent_response_max_retries:
-                        logger.error(
+                        logger.warning(
                             "%s Failed to parse agent output after %d attempts, consider enabling or "
                             "increasing parse_agent_response_max_retries",
                             AGENT_LOG_PREFIX,
-                            attempt,
-                            exc_info=True)
+                            attempt)
                         # the final answer goes in the "messages" state channel
                         combined_content = str(ex.observation) + '\n' + str(output_message.content)
                         output_message.content = combined_content

aiq/agent/react_agent/register.py

@@ -18,7 +18,6 @@ import logging
 from pydantic import AliasChoices
 from pydantic import Field
 
-from aiq.agent.base import AGENT_LOG_PREFIX
 from aiq.builder.builder import Builder
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.builder.function_info import FunctionInfo
@@ -79,6 +78,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
     from langchain_core.messages import trim_messages
     from langgraph.graph.graph import CompiledGraph
 
+    from aiq.agent.base import AGENT_LOG_PREFIX
     from aiq.agent.react_agent.agent import ReActAgentGraph
     from aiq.agent.react_agent.agent import ReActGraphState
     from aiq.agent.react_agent.agent import create_react_agent_prompt

aiq/agent/reasoning_agent/reasoning_agent.py

@@ -19,7 +19,6 @@ from collections.abc import AsyncGenerator
 
 from pydantic import Field
 
-from aiq.agent.base import AGENT_LOG_PREFIX
 from aiq.builder.builder import Builder
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.builder.function_info import FunctionInfo
@@ -86,6 +85,8 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
     from langchain_core.language_models import BaseChatModel
     from langchain_core.prompts import PromptTemplate
 
+    from aiq.agent.base import AGENT_LOG_PREFIX
+
     def remove_r1_think_tags(text: str):
         pattern = r'(<think>)?.*?</think>\s*(.*)'
 

aiq/agent/tool_calling_agent/register.py

@@ -17,7 +17,6 @@ import logging
 
 from pydantic import Field
 
-from aiq.agent.base import AGENT_LOG_PREFIX
 from aiq.builder.builder import Builder
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.builder.function_info import FunctionInfo
@@ -49,6 +48,8 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     from langchain_core.messages.human import HumanMessage
     from langgraph.graph.graph import CompiledGraph
 
+    from aiq.agent.base import AGENT_LOG_PREFIX
+
     from .agent import ToolCallAgentGraph
     from .agent import ToolCallAgentGraphState
 

aiq/authentication/api_key/api_key_auth_provider.py

@@ -28,9 +28,13 @@ logger = logging.getLogger(__name__)
 
 class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
-    def __init__(self, config: APIKeyAuthProviderConfig, config_name: str | None = None) -> None:
-        assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyConfig")
+    # fmt: off
+    def __init__(self,
+                 config: APIKeyAuthProviderConfig,
+                 config_name: str | None = None) -> None:  # pylint: disable=unused-argument
+        assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyAuthProviderConfig")
         super().__init__(config)
+    # fmt: on
 
     async def _construct_authentication_header(self) -> BearerTokenCred:
         """

aiq/builder/function.py

@@ -76,11 +76,16 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
         -------
         _T
             The converted value.
+
+        Raises
+        ------
+        ValueError
+            If the value cannot be converted to the specified type (when `to_type` is specified).
         """
 
         return self._converter.convert(value, to_type=to_type)
 
-    def try_convert(self, value: typing.Any, to_type: type[_T]) -> _T:
+    def try_convert(self, value: typing.Any, to_type: type[_T]) -> _T | typing.Any:
         """
         Converts the given value to the specified type using graceful error handling.
         If conversion fails, returns the original value and continues processing.
@@ -94,7 +99,7 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
         Returns
         -------
-        _T
+        _T | typing.Any
             The converted value, or original value if conversion fails.
         """
         return self._converter.try_convert(value, to_type=to_type)
@@ -129,17 +134,22 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
         -------
         typing.Any
             The output of the function optionally converted to the specified type.
+
+        Raises
+        ------
+        ValueError
+            If the output of the function cannot be converted to the specified type.
         """
 
         with self._context.push_active_function(self.instance_name,
                                                 input_data=value) as manager:  # Set the current invocation context
             try:
-                converted_input: InputT = self._convert_input(value)  # type: ignore
+                converted_input: InputT = self._convert_input(value)
 
                 result = await self._ainvoke(converted_input)
 
                 if to_type is not None and not isinstance(result, to_type):
-                    result = self._converter.try_convert(result, to_type=to_type)
+                    result = self.convert(result, to_type)
 
                 manager.set_output(result)
 
@@ -215,18 +225,23 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
         ------
         typing.Any
             The output of the function optionally converted to the specified type.
+
+        Raises
+        ------
+        ValueError
+            If the output of the function cannot be converted to the specified type (when `to_type` is specified).
         """
 
         with self._context.push_active_function(self.instance_name, input_data=value) as manager:
             try:
-                converted_input: InputT = self._convert_input(value)  # type: ignore
+                converted_input: InputT = self._convert_input(value)
 
                 # Collect streaming outputs to capture the final result
                 final_output: list[typing.Any] = []
 
                 async for data in self._astream(converted_input):
                     if to_type is not None and not isinstance(data, to_type):
-                        converted_data = self._converter.try_convert(data, to_type=to_type)
+                        converted_data = self.convert(data, to_type=to_type)
                         final_output.append(converted_data)
                         yield converted_data
                     else:

aiq/builder/function_base.py

@@ -350,7 +350,7 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
         # output because the ABC has it.
         return True
 
-    def _convert_input(self, value: typing.Any):
+    def _convert_input(self, value: typing.Any) -> InputT:
         if (isinstance(value, self.input_class)):
             return value
 
@@ -373,4 +373,8 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
             return value
 
         # Fallback to the converter
-        return self._converter.try_convert(value, to_type=self.input_class)
+        try:
+            return self._converter.convert(value, to_type=self.input_class)
+        except ValueError as e:
+            # Input parsing should yield a TypeError instead of a ValueError
+            raise TypeError from e

aiq/cli/commands/sizing/calc.py

@@ -274,9 +274,12 @@ def calc_command(ctx,
 
         click.echo(tabulate(table, headers=headers, tablefmt="github"))
 
-        # Display slope-based GPU estimates at the end
-        click.echo("")  # Add blank line for separation
-        click.echo(click.style("=== GPU ESTIMATES ===", fg="bright_blue", bold=True))
+        # Display slope-based GPU estimates if they are available
+        if results.gpu_estimates.gpu_estimate_by_llm_latency is not None or \
+                results.gpu_estimates.gpu_estimate_by_wf_runtime is not None:
+            click.echo("")
+            click.echo(click.style("=== GPU ESTIMATES ===", fg="bright_blue", bold=True))
+
         if results.gpu_estimates.gpu_estimate_by_wf_runtime is not None:
             click.echo(
                 click.style(

aiq/cli/commands/start.py

@@ -190,11 +190,6 @@ class StartCommandGroup(click.Group):
         # Override default front end config with values from the config file for serverless execution modes.
         # Check that we have the right kind of front end
         if (not isinstance(config.general.front_end, front_end.config_type)):
-            logger.warning(
-                "The front end type in the config file (%s) does not match the command name (%s). "
-                "Overwriting the config file front end.",
-                config.general.front_end.type,
-                cmd_name)
 
             # Set the front end config
             config.general.front_end = front_end.config_type()

aiq/cli/commands/uninstall.py

@@ -53,13 +53,11 @@ async def uninstall_packages(packages: list[dict[str, str]]) -> None:
         await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
 
 
-@click.group(name=__name__,
-             invoke_without_command=True,
-             help=("Uninstall an AIQ Toolkit plugin packages from the local environment."))
+@click.group(name=__name__, invoke_without_command=True, help=("Uninstall plugin packages from the local environment."))
 @click.argument("packages", type=str)
 def uninstall_command(packages: str) -> None:
     """
-    Uninstall AIQ Toolkit plugin packages from the local environment.
+    Uninstall plugin packages from the local environment.
     """
 
     packages = packages.split()

aiq/data_models/api_server.py

@@ -121,28 +121,22 @@ class AIQChatRequest(BaseModel):
     # Optional fields (OpenAI Chat Completions API compatible)
     model: str | None = Field(default=None, description="name of the model to use")
     frequency_penalty: float | None = Field(default=0.0,
-                                            ge=-2.0,
-                                            le=2.0,
                                             description="Penalty for new tokens based on frequency in text")
     logit_bias: dict[str, float] | None = Field(default=None,
                                                 description="Modify likelihood of specified tokens appearing")
     logprobs: bool | None = Field(default=None, description="Whether to return log probabilities")
-    top_logprobs: int | None = Field(default=None, ge=0, le=20, description="Number of most likely tokens to return")
-    max_tokens: int | None = Field(default=None, ge=1, description="Maximum number of tokens to generate")
-    n: int | None = Field(default=1, ge=1, le=128, description="Number of chat completion choices to generate")
-    presence_penalty: float | None = Field(default=0.0,
-                                           ge=-2.0,
-                                           le=2.0,
-                                           description="Penalty for new tokens based on presence in text")
+    top_logprobs: int | None = Field(default=None, description="Number of most likely tokens to return")
+    max_tokens: int | None = Field(default=None, description="Maximum number of tokens to generate")
+    n: int | None = Field(default=1, description="Number of chat completion choices to generate")
+    presence_penalty: float | None = Field(default=0.0, description="Penalty for new tokens based on presence in text")
     response_format: dict[str, typing.Any] | None = Field(default=None, description="Response format specification")
     seed: int | None = Field(default=None, description="Random seed for deterministic sampling")
     service_tier: typing.Literal["auto", "default"] | None = Field(default=None,
                                                                    description="Service tier for the request")
-    stop: str | list[str] | None = Field(default=None, description="Up to 4 sequences where API will stop generating")
     stream: bool | None = Field(default=False, description="Whether to stream partial message deltas")
     stream_options: dict[str, typing.Any] | None = Field(default=None, description="Options for streaming")
-    temperature: float | None = Field(default=1.0, ge=0.0, le=2.0, description="Sampling temperature between 0 and 2")
-    top_p: float | None = Field(default=None, ge=0.0, le=1.0, description="Nucleus sampling parameter")
+    temperature: float | None = Field(default=1.0, description="Sampling temperature between 0 and 2")
+    top_p: float | None = Field(default=None, description="Nucleus sampling parameter")
     tools: list[dict[str, typing.Any]] | None = Field(default=None, description="List of tools the model may call")
     tool_choice: str | dict[str, typing.Any] | None = Field(default=None, description="Controls which tool is called")
     parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling")

aiq/data_models/component_ref.py

@@ -130,7 +130,7 @@ class ObjectStoreRef(ComponentRef):
     """
 
     @property
-    @typing.override
+    @override
     def component_group(self):
         return ComponentGroup.OBJECT_STORES
 

aiq/data_models/discovery_metadata.py

@@ -21,6 +21,7 @@ import typing
 from enum import Enum
 from functools import lru_cache
 from pathlib import Path
+from types import ModuleType
 from typing import TYPE_CHECKING
 
 from pydantic import BaseModel
@@ -115,6 +116,55 @@ class DiscoveryMetadata(BaseModel):
                     return data.get(root_package, None)
         return None
 
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_module(module: ModuleType | None) -> str:
+        """Get the distribution name from the config type using the mapping of module names to distro names.
+
+        Args:
+            module (ModuleType): A registered component's module.
+
+        Returns:
+            str: The distribution name of the AIQ Toolkit component.
+        """
+        from aiq.runtime.loader import get_all_aiq_entrypoints_distro_mapping
+
+        if module is None:
+            return "aiqtoolkit"
+
+        # Get the mapping of module names to distro names
+        mapping = get_all_aiq_entrypoints_distro_mapping()
+        module_package = module.__package__
+
+        if module_package is None:
+            return "aiqtoolkit"
+
+        # Traverse the module package parts in reverse order to find the distro name
+        # This is because the module package is the root package for the AIQ Toolkit component
+        # and the distro name is the name of the package that contains the component
+        module_package_parts = module_package.split(".")
+        for part_idx in range(len(module_package_parts), 0, -1):
+            candidate_module_name = ".".join(module_package_parts[0:part_idx])
+            candidate_distro_name = mapping.get(candidate_module_name, None)
+            if candidate_distro_name is not None:
+                return candidate_distro_name
+
+        return "aiqtoolkit"
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_config_type(config_type: type["TypedBaseModelT"]) -> str:
+        """Get the distribution name from the config type using the mapping of module names to distro names.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+
+        Returns:
+            str: The distribution name of the AIQ Toolkit component.
+        """
+        module = inspect.getmodule(config_type)
+        return DiscoveryMetadata.get_distribution_name_from_module(module)
+
     @staticmethod
     @lru_cache
     def get_distribution_name(root_package: str) -> str:
@@ -123,6 +173,7 @@ class DiscoveryMetadata(BaseModel):
         root package name 'aiq'. They provide mapping in a metadata file
         for optimized installation.
         """
+
         distro_name = DiscoveryMetadata.get_distribution_name_from_private_data(root_package)
         return distro_name if distro_name else root_package
 
@@ -142,8 +193,7 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(config_type)
-            root_package: str = module.__package__.split(".")[0]
-            distro_name = DiscoveryMetadata.get_distribution_name(root_package)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_config_type(config_type)
 
             if not distro_name:
                 # raise an exception
@@ -187,12 +237,13 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(fn)
-            root_package: str = module.__package__.split(".")[0]
-            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
+
             try:
-                version = importlib.metadata.version(root_package) if root_package != "" else ""
+                # version = importlib.metadata.version(root_package) if root_package != "" else ""
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
             except importlib.metadata.PackageNotFoundError:
-                logger.warning("Package metadata not found for %s", root_package)
+                logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
             logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
@@ -201,7 +252,7 @@ class DiscoveryMetadata(BaseModel):
         if isinstance(wrapper_type, LLMFrameworkEnum):
             wrapper_type = wrapper_type.value
 
-        return DiscoveryMetadata(package=root_package,
+        return DiscoveryMetadata(package=distro_name,
                                  version=version,
                                  component_type=component_type,
                                  component_name=wrapper_type,
@@ -220,7 +271,6 @@ class DiscoveryMetadata(BaseModel):
         """
 
         try:
-            package_name = DiscoveryMetadata.get_distribution_name(package_name)
             try:
                 metadata = importlib.metadata.metadata(package_name)
                 description = metadata.get("Summary", "")
@@ -263,12 +313,11 @@ class DiscoveryMetadata(BaseModel):
 
         try:
             module = inspect.getmodule(config_type)
-            root_package: str = module.__package__.split(".")[0]
-            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
             try:
-                version = importlib.metadata.version(root_package) if root_package != "" else ""
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
             except importlib.metadata.PackageNotFoundError:
-                logger.warning("Package metadata not found for %s", root_package)
+                logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
             logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
@@ -279,7 +328,7 @@ class DiscoveryMetadata(BaseModel):
 
         description = generate_config_type_docs(config_type=config_type)
 
-        return DiscoveryMetadata(package=root_package,
+        return DiscoveryMetadata(package=distro_name,
                                  version=version,
                                  component_type=component_type,
                                  component_name=component_name,

aiq/front_ends/console/console_front_end_plugin.py

@@ -15,12 +15,10 @@
 
 import asyncio
 import logging
-from io import StringIO
 
 import click
 from colorama import Fore
 
-from aiq.builder.workflow_builder import WorkflowBuilder
 from aiq.data_models.interactive import HumanPromptModelType
 from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import HumanResponseText
@@ -61,27 +59,9 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
         if (not self.front_end_config.input_query and not self.front_end_config.input_file):
             raise click.UsageError("Must specify either --input_query or --input_file")
 
-    async def run(self):
-
-        # Must yield the workflow function otherwise it cleans up
-        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
-
-            session_manager: AIQSessionManager = None
-
-            if logger.isEnabledFor(logging.INFO):
-                stream = StringIO()
-
-                self.full_config.print_summary(stream=stream)
-
-                click.echo(stream.getvalue())
-
-                workflow = builder.build()
-                session_manager = AIQSessionManager(workflow)
-
-            await self.run_workflow(session_manager)
-
-    async def run_workflow(self, session_manager: AIQSessionManager = None):
+    async def run_workflow(self, session_manager: AIQSessionManager):
 
+        assert session_manager is not None, "Session manager must be provided"
         runner_outputs = None
 
         if (self.front_end_config.input_query):

aiq/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -45,8 +45,10 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
                 click.echo(stream.getvalue())
 
-            await self.run_workflow(builder.build())
+            workflow = builder.build()
+            session_manager = AIQSessionManager(workflow)
+            await self.run_workflow(session_manager)
 
     @abstractmethod
-    async def run_workflow(self, session_manager: AIQSessionManager = None):
+    async def run_workflow(self, session_manager: AIQSessionManager):
         pass

aiq/object_store/in_memory_object_store.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
+
 from aiq.builder.builder import Builder
 from aiq.cli.register_workflow import register_object_store
 from aiq.data_models.object_store import KeyAlreadyExistsError
@@ -37,37 +39,37 @@ class InMemoryObjectStore(ObjectStore):
     """
 
     def __init__(self) -> None:
+        self._lock = asyncio.Lock()
         self._store: dict[str, ObjectStoreItem] = {}
 
     @override
     async def put_object(self, key: str, item: ObjectStoreItem) -> None:
-        if key in self._store:
-            raise KeyAlreadyExistsError(key)
-
-        self._store[key] = item
-        return
+        async with self._lock:
+            if key in self._store:
+                raise KeyAlreadyExistsError(key)
+            self._store[key] = item
 
     @override
     async def upsert_object(self, key: str, item: ObjectStoreItem) -> None:
-        self._store[key] = item
-        return
+        async with self._lock:
+            self._store[key] = item
 
     @override
     async def get_object(self, key: str) -> ObjectStoreItem:
-
-        if key not in self._store:
-            raise NoSuchKeyError(key)
-
-        return self._store[key]
+        async with self._lock:
+            value = self._store.get(key)
+            if value is None:
+                raise NoSuchKeyError(key)
+            return value
 
     @override
     async def delete_object(self, key: str) -> None:
-        if key not in self._store:
+        try:
+            async with self._lock:
+                self._store.pop(key)
+        except KeyError:
             raise NoSuchKeyError(key)
 
-        self._store.pop(key)
-        return
-
 
 @register_object_store(config_type=InMemoryObjectStoreConfig)
 async def in_memory_object_store(config: InMemoryObjectStoreConfig, builder: Builder):