aiqtoolkit 1.2.0rc3__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/agent/rewoo_agent/prompt.py

@@ -1,108 +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.
-
-# flake8: noqa
-from langchain_core.prompts.chat import ChatPromptTemplate
-
-PLANNER_SYSTEM_PROMPT = """
-For the following task, make plans that can solve the problem step by step. For each plan, indicate \
-which external tool together with tool input to retrieve evidence. You can store the evidence into a \
-placeholder #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)
-
-You may ask the human to the following tools:
-
-{tools}
-
-The tools should be one of the following: [{tool_names}]
-
-You are not required to use all the tools listed. Choose only the ones that best fit the needs of each plan step.
-
-Your output must be a JSON array where each element represents one planning step. Each step must be an object with exactly two keys:
-
-1. "plan": A string that describes in detail the action or reasoning for that step.
-
-2. "evidence": An object representing the external tool call associated with that plan step. This object must have the following keys:
-
-   -"placeholder": A string that identifies the evidence placeholder (e.g., "#E1", "#E2", etc.). The numbering should be sequential based on the order of steps.
-
-   -"tool": A string specifying the name of the external tool used.
-
-   -"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the tool.
-
-Do not include any additional keys or characters in your output, and do not wrap your response with markdown formatting. Your output must be strictly valid JSON.
-
-Important instructions:
-
-Do not output any additional text, comments, or markdown formatting.
-
-Do not include any explanation or reasoning text outside of the JSON array.
-
-The output must be a valid JSON array that can be parsed directly.
-
-Here is an example of how a valid JSON output should look:
-
-[
-  \'{{
-    "plan": "Calculate the result of 2023 minus 25.",
-    "evidence": \'{{
-      "placeholder": "#E1",
-      "tool": "calculator_subtract",
-      "tool_input": [2023, 25]
-    }}\'
-  }}\',
-  \'{{
-    "plan": "Retrieve the year represented by the result stored in #E1.",
-    "evidence": \'{{
-      "placeholder": "#E2",
-      "tool": "haystack_chitchat_agent",
-      "tool_input": "Response with the result number contained in #E1"
-    }}\'
-  }}\',
-  \'{{
-    "plan": "Search for the CEO of Golden State Warriors in the year stored in #E2.",
-    "evidence": \'{{
-      "placeholder": "#E3",
-      "tool": "internet_search",
-      "tool_input": "Who was the CEO of Golden State Warriors in the year #E2?"
-    }}\'
-  }}\'
-]
-
-Begin!
-"""
-
-PLANNER_USER_PROMPT = """
-task: {task}
-"""
-
-rewoo_planner_prompt = ChatPromptTemplate([("system", PLANNER_SYSTEM_PROMPT), ("user", PLANNER_USER_PROMPT)])
-
-SOLVER_SYSTEM_PROMPT = """
-Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
-retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
-contain irrelevant information.
-
-Now solve the question or task according to provided Evidence above. Respond with the answer
-directly with no extra words.
-
-"""
-SOLVER_USER_PROMPT = """
-plan: {plan}
-task: {task}
-
-Response:
-"""
-
-rewoo_solver_prompt = ChatPromptTemplate([("system", SOLVER_SYSTEM_PROMPT), ("user", SOLVER_USER_PROMPT)])

aiq/agent/rewoo_agent/register.py

@@ -1,158 +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 pydantic import Field
-
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.api_server import AIQChatRequest
-from aiq.data_models.api_server import AIQChatResponse
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.utils.type_converter import GlobalTypeConverter
-
-logger = logging.getLogger(__name__)
-
-
-class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
-    """
-    Defines an AIQ Toolkit function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
-    tool names and descriptions to select the optimal tool.
-    """
-
-    tool_names: list[FunctionRef] = Field(default_factory=list,
-                                          description="The list of tools to provide to the rewoo agent.")
-    llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
-    include_tool_input_schema_in_tool_description: bool = Field(
-        default=True, description="Specify inclusion of tool input schemas in the prompt.")
-    description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
-    planner_prompt: str | None = Field(
-        default=None,
-        description="Provides the PLANNER_PROMPT to use with the agent")  # defaults to PLANNER_PROMPT in prompt.py
-    solver_prompt: str | None = Field(
-        default=None,
-        description="Provides the SOLVER_PROMPT to use with the agent")  # defaults to SOLVER_PROMPT in prompt.py
-    max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
-    use_openai_api: bool = Field(default=False,
-                                 description=("Use OpenAI API for the input/output types to the function. "
-                                              "If False, strings will be used."))
-    additional_instructions: str | None = Field(
-        default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
-
-
-@register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
-async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
-    from langchain.schema import BaseMessage
-    from langchain_core.messages import trim_messages
-    from langchain_core.messages.human import HumanMessage
-    from langchain_core.prompts import ChatPromptTemplate
-    from langgraph.graph.graph import CompiledGraph
-
-    from aiq.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
-    from aiq.agent.rewoo_agent.prompt import SOLVER_USER_PROMPT
-
-    from .agent import ReWOOAgentGraph
-    from .agent import ReWOOGraphState
-    from .prompt import rewoo_planner_prompt
-    from .prompt import rewoo_solver_prompt
-
-    # the ReWOO Agent prompt comes from prompt.py, and can be customized there or via config option planner_prompt and
-    # solver_prompt.
-    if config.planner_prompt:
-        planner_prompt = config.planner_prompt
-        if config.additional_instructions:
-            planner_prompt += f"{config.additional_instructions}"
-        valid = ReWOOAgentGraph.validate_planner_prompt(config.planner_prompt)
-        if not valid:
-            logger.exception("Invalid planner_prompt")
-            raise ValueError("Invalid planner_prompt")
-        planner_prompt = ChatPromptTemplate([("system", config.planner_prompt), ("user", PLANNER_USER_PROMPT)])
-    else:
-        planner_prompt = rewoo_planner_prompt
-
-    if config.solver_prompt:
-        solver_prompt = config.solver_prompt
-        if config.additional_instructions:
-            solver_prompt += f"{config.additional_instructions}"
-        valid = ReWOOAgentGraph.validate_solver_prompt(config.solver_prompt)
-        if not valid:
-            logger.exception("Invalid solver_prompt")
-            raise ValueError("Invalid solver_prompt")
-        solver_prompt = ChatPromptTemplate([("system", config.solver_prompt), ("user", SOLVER_USER_PROMPT)])
-    else:
-        solver_prompt = rewoo_solver_prompt
-
-    # we can choose an LLM for the ReWOO agent in the config file
-    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-
-    # the agent can run any installed tool, simply install the tool and add it to the config file
-    # the sample tool provided can easily be copied or changed
-    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-    if not tools:
-        raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
-
-    # construct the ReWOO Agent Graph from the configured llm, prompt, and tools
-    graph: CompiledGraph = await ReWOOAgentGraph(llm=llm,
-                                                 planner_prompt=planner_prompt,
-                                                 solver_prompt=solver_prompt,
-                                                 tools=tools,
-                                                 use_tool_schema=config.include_tool_input_schema_in_tool_description,
-                                                 detailed_logs=config.verbose).build_graph()
-
-    async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
-        try:
-            # initialize the starting state with the user query
-            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
-                                                        max_tokens=config.max_history,
-                                                        strategy="last",
-                                                        token_counter=len,
-                                                        start_on="human",
-                                                        include_system=True)
-            task = HumanMessage(content=messages[0].content)
-            state = ReWOOGraphState(task=task)
-
-            # run the ReWOO Agent Graph
-            state = await graph.ainvoke(state)
-
-            # get and return the output from the state
-            state = ReWOOGraphState(**state)
-            output_message = state.result.content  # pylint: disable=E1101
-            return AIQChatResponse.from_string(output_message)
-
-        except Exception as ex:
-            logger.exception("ReWOO Agent failed with exception: %s", ex, exc_info=ex)
-            # here, we can implement custom error messages
-            if config.verbose:
-                return AIQChatResponse.from_string(str(ex))
-            return AIQChatResponse.from_string("I seem to be having a problem.")
-
-    if (config.use_openai_api):
-        yield FunctionInfo.from_fn(_response_fn, description=config.description)
-
-    else:
-
-        async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
-            oai_output = await _response_fn(oai_input)
-
-            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

@@ -1,119 +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=R0917
-import logging
-
-from langchain_core.callbacks.base import AsyncCallbackHandler
-from langchain_core.language_models import BaseChatModel
-from langchain_core.messages.base import BaseMessage
-from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import BaseTool
-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 AgentDecision
-from aiq.agent.dual_node import DualNodeAgent
-
-logger = logging.getLogger(__name__)
-
-
-class ToolCallAgentGraphState(BaseModel):
-    """State schema for the Tool Calling Agent Graph"""
-    messages: list[BaseMessage] = Field(default_factory=list)  # input and output of the Agent
-
-
-class ToolCallAgentGraph(DualNodeAgent):
-    """Configurable LangGraph Tool Calling Agent. A Tool Calling Agent requires an LLM which supports tool calling.
-    A tool Calling Agent utilizes the tool input parameters to select the optimal tool.  Supports handling tool errors.
-    Argument "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
-
-    def __init__(self,
-                 llm: BaseChatModel,
-                 tools: list[BaseTool],
-                 callbacks: list[AsyncCallbackHandler] = None,
-                 detailed_logs: bool = False,
-                 handle_tool_errors: bool = True):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
-        self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
-        logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)
-
-    async def agent_node(self, state: ToolCallAgentGraphState):
-        try:
-            logger.debug('%s Starting the Tool Calling Agent Node', AGENT_LOG_PREFIX)
-            if len(state.messages) == 0:
-                raise RuntimeError('No input received in state: "messages"')
-            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_CALL_LOG_MESSAGE, agent_input, response)
-
-            state.messages += [response]
-            return state
-        except Exception as ex:
-            logger.exception("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
-
-    async def conditional_edge(self, state: ToolCallAgentGraphState):
-        try:
-            logger.debug("%s Starting the Tool Calling Conditional Edge", AGENT_LOG_PREFIX)
-            last_message = state.messages[-1]
-            if last_message.tool_calls:
-                # the agent wants to call a tool
-                logger.debug('%s Agent is calling a tool', AGENT_LOG_PREFIX)
-                return AgentDecision.TOOL
-            if self.detailed_logs:
-                logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.messages[-1].content)
-            return AgentDecision.END
-        except Exception as ex:
-            logger.exception("%s Failed to determine whether agent is calling a tool: %s",
-                             AGENT_LOG_PREFIX,
-                             ex,
-                             exc_info=True)
-            logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
-            return AgentDecision.END
-
-    async def tool_node(self, state: ToolCallAgentGraphState):
-        try:
-            logger.debug("%s Starting Tool Node", AGENT_LOG_PREFIX)
-            tool_calls = state.messages[-1].tool_calls
-            tools = [tool.get('name') for tool in tool_calls]
-            tool_input = state.messages[-1]
-            tool_response = await self.tool_caller.ainvoke(input={"messages": [tool_input]},
-                                                           config=RunnableConfig(callbacks=self.callbacks,
-                                                                                 configurable={}))
-            # this configurable = {} argument is needed due to a bug in LangGraph PreBuilt ToolNode ^
-
-            for response in tool_response.get('messages'):
-                if self.detailed_logs:
-                    self._log_tool_response(str(tools), str(tool_input), response.content)
-                state.messages += [response]
-
-            return state
-        except Exception as ex:
-            logger.exception("%s Failed to call tool_node: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex
-
-    async def build_graph(self):
-        try:
-            await super()._build_graph(state_schema=ToolCallAgentGraphState)
-            logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
-            return self.graph
-        except Exception as ex:
-            logger.exception("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex

aiq/agent/tool_calling_agent/register.py

@@ -1,106 +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 pydantic import Field
-
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
-
-logger = logging.getLogger(__name__)
-
-
-class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent"):
-    """
-    A Tool Calling Agent requires an LLM which supports tool calling. A tool Calling Agent utilizes the tool
-    input parameters to select the optimal tool.  Supports handling tool errors.
-    """
-
-    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 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.")
-
-
-@register_function(config_type=ToolCallAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
-async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, builder: Builder):
-    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
-
-    # 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)
-    # the agent can run any installed tool, simply install the tool and add it to the config file
-    # the sample tools provided can easily be copied or changed
-    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-    if not tools:
-        raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'")
-
-    # some LLMs support tool calling
-    # these models accept the tool's input schema and decide when to use a tool based on the input's relevance
-    try:
-        # in tool calling agents, we bind the tools to the LLM, to pass the tools' input schemas at runtime
-        llm = llm.bind_tools(tools)
-    except NotImplementedError as ex:
-        logger.error("%s Failed to bind tools: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-        raise ex
-
-    # construct the Tool Calling Agent Graph from the configured llm, and tools
-    graph: CompiledGraph = await ToolCallAgentGraph(llm=llm,
-                                                    tools=tools,
-                                                    detailed_logs=config.verbose,
-                                                    handle_tool_errors=config.handle_tool_errors).build_graph()
-
-    async def _response_fn(input_message: str) -> str:
-        try:
-            # initialize the starting state with the user query
-            input_message = HumanMessage(content=input_message)
-            state = ToolCallAgentGraphState(messages=[input_message])
-
-            # run the Tool Calling Agent Graph
-            state = await graph.ainvoke(state, config={'recursion_limit': (config.max_iterations + 1) * 2})
-            # setting recursion_limit: 4 allows 1 tool call
-            #   - allows the Tool Calling Agent to perform 1 cycle / call 1 single tool,
-            #   - but stops the agent when it tries to call a tool a second time
-
-            # get and return the output from the state
-            state = ToolCallAgentGraphState(**state)
-            output_message = state.messages[-1]  # pylint: disable=E1136
-            return output_message.content
-        except Exception as ex:
-            logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            if config.verbose:
-                return str(ex)
-            return "I seem to be having a problem."
-
-    try:
-        yield FunctionInfo.from_fn(_response_fn, description=config.description)
-    except GeneratorExit:
-        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX, exc_info=True)
-    finally:
-        logger.debug("%s Cleaning up react_agent workflow.", AGENT_LOG_PREFIX)

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

@@ -1,92 +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 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])