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

aiq/agent/react_agent/register.py

@@ -15,6 +15,7 @@
 
 import logging
 
+from pydantic import AliasChoices
 from pydantic import Field
 
 from aiq.agent.base import AGENT_LOG_PREFIX
@@ -42,11 +43,24 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
                                           description="The list of tools to provide to the react agent.")
     llm_name: LLMRef = Field(description="The LLM model to use with the react agent.")
     verbose: bool = Field(default=False, description="Set the verbosity of the react agent's logging.")
-    retry_parsing_errors: bool = Field(default=True, description="Specify retrying when encountering parsing errors.")
-    max_retries: int = Field(default=1, description="Sent the number of retries before raising a parsing error.")
+    retry_agent_response_parsing_errors: bool = Field(
+        default=True,
+        validation_alias=AliasChoices("retry_agent_response_parsing_errors", "retry_parsing_errors"),
+        description="Whether to retry when encountering parsing errors in the agent's response.")
+    parse_agent_response_max_retries: int = Field(
+        default=1,
+        validation_alias=AliasChoices("parse_agent_response_max_retries", "max_retries"),
+        description="Maximum number of times the Agent may retry parsing errors. "
+        "Prevents the Agent from getting into infinite hallucination loops.")
+    tool_call_max_retries: int = Field(default=1, description="The number of retries before raising a tool call error.")
+    max_tool_calls: int = Field(default=15,
+                                validation_alias=AliasChoices("max_tool_calls", "max_iterations"),
+                                description="Maximum number of tool calls before stopping the agent.")
+    pass_tool_call_errors_to_agent: bool = Field(
+        default=True,
+        description="Whether to pass tool call errors to agent. If False, failed tool calls will raise an exception.")
     include_tool_input_schema_in_tool_description: bool = Field(
         default=True, description="Specify inclusion of tool input schemas in the prompt.")
-    max_iterations: int = Field(default=15, description="Number of tool calls before stoping the react agent.")
     description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
     system_prompt: str | None = Field(
         default=None,
@@ -63,29 +77,13 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
     from langchain.schema import BaseMessage
     from langchain_core.messages import trim_messages
-    from langchain_core.prompts import ChatPromptTemplate
-    from langchain_core.prompts import MessagesPlaceholder
     from langgraph.graph.graph import CompiledGraph
 
-    from aiq.agent.react_agent.prompt import USER_PROMPT
-
-    from .agent import ReActAgentGraph
-    from .agent import ReActGraphState
-    from .prompt import react_agent_prompt
-
-    # the ReAct Agent prompt comes from prompt.py, and can be customized there or via config option system_prompt.
-    if config.system_prompt:
-        _prompt_str = config.system_prompt
-        if config.additional_instructions:
-            _prompt_str += f" {config.additional_instructions}"
-        valid_prompt = ReActAgentGraph.validate_system_prompt(config.system_prompt)
-        if not valid_prompt:
-            logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
-            raise ValueError("Invalid system_prompt")
-        prompt = ChatPromptTemplate([("system", config.system_prompt), ("user", USER_PROMPT),
-                                     MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
-    else:
-        prompt = react_agent_prompt
+    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
+
+    prompt = create_react_agent_prompt(config)
 
     # we can choose an LLM for the ReAct agent in the config file
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -96,13 +94,16 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'")
     # configure callbacks, for sending intermediate steps
     # construct the ReAct Agent Graph from the configured llm, prompt, and tools
-    graph: CompiledGraph = await ReActAgentGraph(llm=llm,
-                                                 prompt=prompt,
-                                                 tools=tools,
-                                                 use_tool_schema=config.include_tool_input_schema_in_tool_description,
-                                                 detailed_logs=config.verbose,
-                                                 retry_parsing_errors=config.retry_parsing_errors,
-                                                 max_retries=config.max_retries).build_graph()
+    graph: CompiledGraph = await ReActAgentGraph(
+        llm=llm,
+        prompt=prompt,
+        tools=tools,
+        use_tool_schema=config.include_tool_input_schema_in_tool_description,
+        detailed_logs=config.verbose,
+        retry_agent_response_parsing_errors=config.retry_agent_response_parsing_errors,
+        parse_agent_response_max_retries=config.parse_agent_response_max_retries,
+        tool_call_max_retries=config.tool_call_max_retries,
+        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent).build_graph()
 
     async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
         try:
@@ -117,7 +118,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
             state = ReActGraphState(messages=messages)
 
             # run the ReAct Agent Graph
-            state = await graph.ainvoke(state, config={'recursion_limit': (config.max_iterations + 1) * 2})
+            state = await graph.ainvoke(state, config={'recursion_limit': (config.max_tool_calls + 1) * 2})
             # setting recursion_limit: 4 allows 1 tool call
             #   - allows the ReAct Agent to perform 1 cycle / call 1 single tool,
             #   - but stops the agent when it tries to call a tool a second time
@@ -125,7 +126,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
             # get and return the output from the state
             state = ReActGraphState(**state)
             output_message = state.messages[-1]  # pylint: disable=E1136
-            return AIQChatResponse.from_string(output_message.content)
+            return AIQChatResponse.from_string(str(output_message.content))
 
         except Exception as ex:
             logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
@@ -139,10 +140,10 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
     else:
 
         async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().convert(input_message, to_type=AIQChatRequest)
+            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
 
             oai_output = await _response_fn(oai_input)
 
-            return GlobalTypeConverter.get().convert(oai_output, to_type=str)
+            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
 
         yield FunctionInfo.from_fn(_str_api_fn, description=config.description)

aiq/agent/rewoo_agent/agent.py

@@ -30,12 +30,11 @@ from langgraph.graph import StateGraph
 from pydantic import BaseModel
 from pydantic import Field
 
+from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
 from aiq.agent.base import AGENT_LOG_PREFIX
-from aiq.agent.base import AGENT_RESPONSE_LOG_MESSAGE
 from aiq.agent.base import INPUT_SCHEMA_MESSAGE
 from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
 from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
-from aiq.agent.base import TOOL_RESPONSE_LOG_MESSAGE
 from aiq.agent.base import AgentDecision
 from aiq.agent.base import BaseAgent
 
@@ -65,7 +64,7 @@ class ReWOOAgentGraph(BaseAgent):
                  solver_prompt: ChatPromptTemplate,
                  tools: list[BaseTool],
                  use_tool_schema: bool = True,
-                 callbacks: list[AsyncCallbackHandler] = None,
+                 callbacks: list[AsyncCallbackHandler] | None = None,
                  detailed_logs: bool = False):
         super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
 
@@ -91,7 +90,7 @@ class ReWOOAgentGraph(BaseAgent):
 
         logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
 
-    def _get_tool(self, tool_name):
+    def _get_tool(self, tool_name: str):
         try:
             return self.tools_dict.get(tool_name)
         except Exception as ex:
@@ -180,22 +179,24 @@ class ReWOOAgentGraph(BaseAgent):
             logger.debug("%s Starting the ReWOO Planner Node", AGENT_LOG_PREFIX)
 
             planner = self.planner_prompt | self.llm
-            task = state.task.content
+            task = str(state.task.content)
             if not task:
                 logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
                 return {"result": NO_INPUT_ERROR_MESSAGE}
 
-            plan = ""
-            async for event in planner.astream({"task": task}, config=RunnableConfig(callbacks=self.callbacks)):
-                plan += event.content
+            plan = await self._stream_llm(
+                planner,
+                {"task": task},
+                RunnableConfig(callbacks=self.callbacks)  # type: ignore
+            )
 
-            steps = self._parse_planner_output(plan)
+            steps = self._parse_planner_output(str(plan.content))
 
             if self.detailed_logs:
-                agent_response_log_message = AGENT_RESPONSE_LOG_MESSAGE % (task, plan)
+                agent_response_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(plan.content))
                 logger.info("ReWOO agent planner output: %s", agent_response_log_message)
 
-            return {"plan": AIMessage(content=plan), "steps": steps}
+            return {"plan": plan, "steps": steps}
 
         except Exception as ex:
             logger.exception("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
@@ -213,10 +214,20 @@ class ReWOOAgentGraph(BaseAgent):
                              current_step)
                 raise RuntimeError(f"ReWOO Executor is invoked with an invalid step number: {current_step}")
 
-            step_info = state.steps.content[current_step]["evidence"]
-            placeholder = step_info.get("placeholder", "")
-            tool = step_info.get("tool", "")
-            tool_input = step_info.get("tool_input", "")
+            steps_content = state.steps.content
+            if isinstance(steps_content, list) and current_step < len(steps_content):
+                step = steps_content[current_step]
+                if isinstance(step, dict) and "evidence" in step:
+                    step_info = step["evidence"]
+                    placeholder = step_info.get("placeholder", "")
+                    tool = step_info.get("tool", "")
+                    tool_input = step_info.get("tool_input", "")
+                else:
+                    logger.error("%s Invalid step format at index %s", AGENT_LOG_PREFIX, current_step)
+                    return {"intermediate_results": state.intermediate_results}
+            else:
+                logger.error("%s Invalid steps content or index %s", AGENT_LOG_PREFIX, current_step)
+                return {"intermediate_results": state.intermediate_results}
 
             intermediate_results = state.intermediate_results
 
@@ -250,12 +261,10 @@ class ReWOOAgentGraph(BaseAgent):
 
             # Run the tool. Try to use structured input, if possible
             tool_input_parsed = self._parse_tool_input(tool_input)
-            tool_response = await requested_tool.ainvoke(tool_input_parsed,
-                                                         config=RunnableConfig(callbacks=self.callbacks))
-
-            # some tools, such as Wikipedia, will return an empty response when no search results are found
-            if tool_response is None or tool_response == "":
-                tool_response = "The tool provided an empty response.\n"
+            tool_response = await self._call_tool(requested_tool,
+                                                  tool_input_parsed,
+                                                  RunnableConfig(callbacks=self.callbacks),
+                                                  max_retries=3)
 
             # ToolMessage only accepts str or list[str | dict] as content.
             # Convert into list if the response is a dict.
@@ -264,15 +273,8 @@ class ReWOOAgentGraph(BaseAgent):
 
             tool_response_message = ToolMessage(name=tool, tool_call_id=tool, content=tool_response)
 
-            logger.debug("%s Successfully called the tool", AGENT_LOG_PREFIX)
             if self.detailed_logs:
-                # The tool response can be very large, so we log only the first 1000 characters
-                tool_response_str = tool_response_message.content
-                tool_response_str = tool_response_str[:1000] + "..." if len(
-                    tool_response_str) > 1000 else tool_response_str
-                tool_response_log_message = TOOL_RESPONSE_LOG_MESSAGE % (
-                    requested_tool.name, tool_input_parsed, tool_response_str)
-                logger.info("ReWOO agent executor output: %s", tool_response_log_message)
+                self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
 
             intermediate_results[placeholder] = tool_response_message
             return {"intermediate_results": intermediate_results}
@@ -308,16 +310,15 @@ class ReWOOAgentGraph(BaseAgent):
                 tool = step_info.get("tool")
                 plan += f"Plan: {_plan}\n{placeholder} = {tool}[{tool_input}]"
 
-            task = state.task.content
+            task = str(state.task.content)
             solver_prompt = self.solver_prompt.partial(plan=plan)
             solver = solver_prompt | self.llm
-            output_message = ""
-            async for event in solver.astream({"task": task}, config=RunnableConfig(callbacks=self.callbacks)):
-                output_message += event.content
 
-            output_message = AIMessage(content=output_message)
+            output_message = await self._stream_llm(solver, {"task": task},
+                                                    RunnableConfig(callbacks=self.callbacks))  # type: ignore
+
             if self.detailed_logs:
-                solver_output_log_message = AGENT_RESPONSE_LOG_MESSAGE % (task, output_message.content)
+                solver_output_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(output_message.content))
                 logger.info("ReWOO agent solver output: %s", solver_output_log_message)
 
             return {"result": output_message}

aiq/agent/rewoo_agent/register.py

@@ -150,9 +150,9 @@ async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
     else:
 
         async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().convert(input_message, to_type=AIQChatRequest)
+            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
             oai_output = await _response_fn(oai_input)
 
-            return GlobalTypeConverter.get().convert(oai_output, to_type=str)
+            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
 
         yield FunctionInfo.from_fn(_str_api_fn, description=config.description)

aiq/agent/tool_calling_agent/agent.py

@@ -25,9 +25,8 @@ from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel
 from pydantic import Field
 
+from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
 from aiq.agent.base import AGENT_LOG_PREFIX
-from aiq.agent.base import AGENT_RESPONSE_LOG_MESSAGE
-from aiq.agent.base import TOOL_RESPONSE_LOG_MESSAGE
 from aiq.agent.base import AgentDecision
 from aiq.agent.dual_node import DualNodeAgent
 
@@ -62,7 +61,7 @@ class ToolCallAgentGraph(DualNodeAgent):
             response = await self.llm.ainvoke(state.messages, config=RunnableConfig(callbacks=self.callbacks))
             if self.detailed_logs:
                 agent_input = "\n".join(str(message.content) for message in state.messages)
-                logger.info(AGENT_RESPONSE_LOG_MESSAGE, agent_input, response)
+                logger.info(AGENT_CALL_LOG_MESSAGE, agent_input, response)
 
             state.messages += [response]
             return state
@@ -102,10 +101,7 @@ class ToolCallAgentGraph(DualNodeAgent):
 
             for response in tool_response.get('messages'):
                 if self.detailed_logs:
-                    # The tool response can be very large, so we log only the first 1000 characters
-                    response.content = response.content[:1000] + "..." if len(
-                        response.content) > 1000 else response.content
-                    logger.info(TOOL_RESPONSE_LOG_MESSAGE, tools, tool_input, response.content)
+                    self._log_tool_response(str(tools), str(tool_input), response.content)
                 state.messages += [response]
 
             return state

aiq/agent/tool_calling_agent/register.py

@@ -38,7 +38,7 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
     tool_names: list[FunctionRef] = Field(default_factory=list,
                                           description="The list of tools to provide to the tool calling agent.")
     llm_name: LLMRef = Field(description="The LLM model to use with the tool calling agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the react agent's logging.")
+    verbose: bool = Field(default=False, description="Set the verbosity of the tool calling agent's logging.")
     handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     description: str = Field(default="Tool Calling Agent Workflow", description="Description of this functions use.")
     max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.")

aiq/authentication/__init__.py

@@ -0,0 +1,14 @@
+# 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/authentication/api_key/__init__.py

@@ -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/authentication/api_key/api_key_auth_provider.py

@@ -0,0 +1,92 @@
+# 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 logging
+
+from pydantic import SecretStr
+
+from aiq.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from aiq.authentication.interfaces import AuthProviderBase
+from aiq.data_models.authentication import AuthResult
+from aiq.data_models.authentication import BearerTokenCred
+from aiq.data_models.authentication import HeaderAuthScheme
+
+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")
+        super().__init__(config)
+
+    async def _construct_authentication_header(self) -> BearerTokenCred:
+        """
+        Constructs the authenticated HTTP header based on the authentication scheme.
+        Basic Authentication follows the OpenAPI 3.0 Basic Authentication standard as well as RFC 7617.
+
+        Args:
+            header_auth_scheme (HeaderAuthScheme): The HTTP authentication scheme to use.
+                                             Supported schemes: BEARER, X_API_KEY, BASIC, CUSTOM.
+
+        Returns:
+            BearerTokenCred: The HTTP headers containing the authentication credentials.
+                             Returns None if the scheme is not supported or configuration is invalid.
+
+        """
+
+        from aiq.authentication.interfaces import AUTHORIZATION_HEADER
+
+        config: APIKeyAuthProviderConfig = self.config
+
+        header_auth_scheme = config.auth_scheme
+
+        if header_auth_scheme == HeaderAuthScheme.BEARER:
+            return BearerTokenCred(token=SecretStr(f"{config.raw_key}"),
+                                   scheme=HeaderAuthScheme.BEARER.value,
+                                   header_name=AUTHORIZATION_HEADER)
+
+        if header_auth_scheme == HeaderAuthScheme.X_API_KEY:
+            return BearerTokenCred(token=SecretStr(f"{config.raw_key}"),
+                                   scheme=HeaderAuthScheme.X_API_KEY.value,
+                                   header_name='')
+
+        if header_auth_scheme == HeaderAuthScheme.CUSTOM:
+            if not config.custom_header_name:
+                raise ValueError('custom_header_name required when using header_auth_scheme=CUSTOM')
+
+            if not config.custom_header_prefix:
+                raise ValueError('custom_header_prefix required when using header_auth_scheme=CUSTOM')
+
+            return BearerTokenCred(token=SecretStr(f"{config.raw_key}"),
+                                   scheme=config.custom_header_prefix,
+                                   header_name=config.custom_header_name)
+
+        raise ValueError(f"Unsupported header auth scheme: {header_auth_scheme}")
+
+    async def authenticate(self, user_id: str | None = None) -> AuthResult | None:
+        """
+        Authenticate the user using the API key credentials.
+
+        Args:
+            user_id (str): The user ID to authenticate.
+
+        Returns:
+            AuthenticatedContext: The authenticated context containing headers, query params, cookies, etc.
+        """
+
+        headers = await self._construct_authentication_header()
+
+        return AuthResult(credentials=[headers])

aiq/authentication/api_key/api_key_auth_provider_config.py

@@ -0,0 +1,124 @@
+# 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 logging
+import re
+import string
+
+from pydantic import Field
+from pydantic import field_validator
+
+from aiq.authentication.exceptions.api_key_exceptions import APIKeyFieldError
+from aiq.authentication.exceptions.api_key_exceptions import HeaderNameFieldError
+from aiq.authentication.exceptions.api_key_exceptions import HeaderPrefixFieldError
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.data_models.authentication import HeaderAuthScheme
+
+logger = logging.getLogger(__name__)
+
+# Strict RFC 7230 compliant header name regex
+HEADER_NAME_REGEX = re.compile(r"^[!#$%&'*+\-.^_`|~0-9a-zA-Z]+$")
+
+
+class APIKeyAuthProviderConfig(AuthProviderBaseConfig, name="api_key"):
+    """
+    API Key authentication configuration model.
+    """
+
+    raw_key: str = Field(description=("Raw API token or credential to be injected into the request parameter. "
+                                      "Used for 'bearer','x-api-key','custom', and other schemes. "))
+
+    auth_scheme: HeaderAuthScheme = Field(default=HeaderAuthScheme.BEARER,
+                                          description=("The HTTP authentication scheme to use. "
+                                                       "Supported schemes: BEARER, X_API_KEY, BASIC, CUSTOM."))
+
+    custom_header_name: str | None = Field(description="The HTTP header name that MUST be used in conjunction "
+                                           "with the custom_header_prefix when HeaderAuthScheme is CUSTOM.",
+                                           default=None)
+    custom_header_prefix: str | None = Field(description="The HTTP header prefix that MUST be used in conjunction "
+                                             "with the custom_header_name when HeaderAuthScheme is CUSTOM.",
+                                             default=None)
+
+    @field_validator('raw_key')
+    @classmethod
+    def validate_raw_key(cls, value: str) -> str:
+        if not value:
+            raise APIKeyFieldError('value_missing', 'raw_key field value is required.')
+
+        if len(value) < 8:
+            raise APIKeyFieldError(
+                'value_too_short',
+                'raw_key field value must be at least 8 characters long for security. '
+                f'Got: {len(value)} characters.')
+
+        if len(value.strip()) != len(value):
+            raise APIKeyFieldError('whitespace_found',
+                                   'raw_key field value cannot have leading or trailing whitespace.')
+
+        if any(c in string.whitespace for c in value):
+            raise APIKeyFieldError('contains_whitespace', 'raw_key must not contain any '
+                                   'whitespace characters.')
+
+        return value
+
+    @field_validator('custom_header_name')
+    @classmethod
+    def validate_custom_header_name(cls, value: str) -> str:
+        if not value:
+            raise HeaderNameFieldError('value_missing', 'custom_header_name is required.')
+
+        if value != value.strip():
+            raise HeaderNameFieldError('whitespace_found',
+                                       'custom_header_name field value cannot have leading or trailing whitespace.')
+
+        if any(c in string.whitespace for c in value):
+            raise HeaderNameFieldError('contains_whitespace',
+                                       'custom_header_name must not contain any whitespace characters.')
+
+        if not HEADER_NAME_REGEX.fullmatch(value):
+            raise HeaderNameFieldError(
+                'invalid_format',
+                'custom_header_name must match the HTTP token syntax: ASCII letters, digits, or allowed symbols.')
+
+        return value
+
+    @field_validator('custom_header_prefix')
+    @classmethod
+    def validate_custom_header_prefix(cls, value: str) -> str:
+        if not value:
+            raise HeaderPrefixFieldError('value_missing', 'custom_header_prefix is required.')
+
+        if value != value.strip():
+            raise HeaderPrefixFieldError(
+                'whitespace_found', 'custom_header_prefix field value cannot have '
+                'leading or trailing whitespace.')
+
+        if any(c in string.whitespace for c in value):
+            raise HeaderPrefixFieldError('contains_whitespace',
+                                         'custom_header_prefix must not contain any whitespace characters.')
+
+        if not value.isascii():
+            raise HeaderPrefixFieldError('invalid_format', 'custom_header_prefix must be ASCII.')
+
+        return value
+
+    @field_validator('raw_key', mode='after')
+    @classmethod
+    def validate_raw_key_after(cls, value: str) -> str:
+        if not value:
+            raise APIKeyFieldError('value_missing', 'raw_key field value is '
+                                   'required after construction.')
+
+        return value

aiq/authentication/api_key/register.py

@@ -0,0 +1,26 @@
+# 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.
+
+from aiq.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_auth_provider
+
+
+@register_auth_provider(config_type=APIKeyAuthProviderConfig)
+async def api_key_client(config: APIKeyAuthProviderConfig, builder: Builder):
+
+    from aiq.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
+
+    yield APIKeyAuthProvider(config=config)

aiq/authentication/exceptions/__init__.py

@@ -0,0 +1,14 @@
+# 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/authentication/exceptions/api_key_exceptions.py

@@ -0,0 +1,38 @@
+# 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.
+
+
+class APIKeyFieldError(Exception):
+    """Raised when API Key Config api_key field validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class HeaderNameFieldError(Exception):
+    """Raised when API Key Config header_name field validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class HeaderPrefixFieldError(Exception):
+    """Raised when API Key Config header_prefix field validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)