nvidia-nat 1.2.0rc5__py3-none-any.whl → 1.2.0rc6__py3-none-any.whl

{aiq → nat}/agent/react_agent/agent.py

@@ -33,18 +33,18 @@ from langchain_core.tools import BaseTool
 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 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 AgentDecision
-from aiq.agent.dual_node import DualNodeAgent
-from aiq.agent.react_agent.output_parser import ReActOutputParser
-from aiq.agent.react_agent.output_parser import ReActOutputParserException
-from aiq.agent.react_agent.prompt import SYSTEM_PROMPT
-from aiq.agent.react_agent.prompt import USER_PROMPT
-from aiq.agent.react_agent.register import ReActAgentWorkflowConfig
+from nat.agent.base import AGENT_CALL_LOG_MESSAGE
+from nat.agent.base import AGENT_LOG_PREFIX
+from nat.agent.base import INPUT_SCHEMA_MESSAGE
+from nat.agent.base import NO_INPUT_ERROR_MESSAGE
+from nat.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
+from nat.agent.base import AgentDecision
+from nat.agent.dual_node import DualNodeAgent
+from nat.agent.react_agent.output_parser import ReActOutputParser
+from nat.agent.react_agent.output_parser import ReActOutputParserException
+from nat.agent.react_agent.prompt import SYSTEM_PROMPT
+from nat.agent.react_agent.prompt import USER_PROMPT
+from nat.agent.react_agent.register import ReActAgentWorkflowConfig
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/agent/react_agent/register.py

@@ -18,23 +18,23 @@ import logging
 from pydantic import AliasChoices
 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
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
 
 class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
     """
-    Defines an AIQ Toolkit function that uses a ReAct Agent performs reasoning inbetween tool calls, and utilizes the
+    Defines a NAT function that uses a ReAct Agent performs reasoning inbetween tool calls, and utilizes the
     tool names and descriptions to select the optimal tool.
     """
 
@@ -78,10 +78,10 @@ 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
+    from nat.agent.base import AGENT_LOG_PREFIX
+    from nat.agent.react_agent.agent import ReActAgentGraph
+    from nat.agent.react_agent.agent import ReActGraphState
+    from nat.agent.react_agent.agent import create_react_agent_prompt
 
     prompt = create_react_agent_prompt(config)
 
@@ -105,7 +105,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         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:
+    async def _response_fn(input_message: ChatRequest) -> ChatResponse:
         try:
             # initialize the starting state with the user query
             messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
@@ -126,21 +126,21 @@ 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(str(output_message.content))
+            return ChatResponse.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)
             # 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.")
+                return ChatResponse.from_string(str(ex))
+            return ChatResponse.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_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
 
             oai_output = await _response_fn(oai_input)
 

{aiq → nat}/agent/reasoning_agent/reasoning_agent.py

@@ -19,21 +19,21 @@ from collections.abc import AsyncGenerator
 
 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.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
 
 class ReasoningFunctionConfig(FunctionBaseConfig, name="reasoning_agent"):
     """
-    Defines an AIQ Toolkit function that performs reasoning on the input data.
+    Defines a NAT function that performs reasoning on the input data.
     Output is passed to the next function in the workflow.
 
     Designed to be used with an InterceptingFunction.
@@ -85,7 +85,7 @@ 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
+    from nat.agent.base import AGENT_LOG_PREFIX
 
     def remove_r1_think_tags(text: str):
         pattern = r'(<think>)?.*?</think>\s*(.*)'
@@ -135,12 +135,12 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
     if augmented_function.has_streaming_output:
 
         async def streaming_inner(
-                input_message: AIQChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
+                input_message: ChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
             """
             Perform reasoning on the input text.
 
             Args:
-                input_message (AIQChatRequest): The input text to reason on.
+                input_message (ChatRequest): The input text to reason on.
             """
 
             input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
@@ -178,12 +178,12 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
 
     if augmented_function.has_single_output:
 
-        async def single_inner(input_message: AIQChatRequest) -> augmented_function.single_output_type:
+        async def single_inner(input_message: ChatRequest) -> augmented_function.single_output_type:
             """
             Perform reasoning on the input text.
 
             Args:
-                input_message (AIQChatRequest): The input text to reason on.
+                input_message (ChatRequest): The input text to reason on.
             """
 
             input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])

{aiq → nat}/agent/rewoo_agent/agent.py

@@ -30,13 +30,13 @@ 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 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 AgentDecision
-from aiq.agent.base import BaseAgent
+from nat.agent.base import AGENT_CALL_LOG_MESSAGE
+from nat.agent.base import AGENT_LOG_PREFIX
+from nat.agent.base import INPUT_SCHEMA_MESSAGE
+from nat.agent.base import NO_INPUT_ERROR_MESSAGE
+from nat.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
+from nat.agent.base import AgentDecision
+from nat.agent.base import BaseAgent
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/agent/rewoo_agent/prompt.py

@@ -13,9 +13,6 @@
 # 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 \
@@ -29,19 +26,25 @@ 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:
+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:
+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.
+   -"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.
+   -"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.
+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:
 
@@ -87,8 +90,6 @@ 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 \
@@ -104,5 +105,3 @@ task: {task}
 
 Response:
 """
-
-rewoo_solver_prompt = ChatPromptTemplate([("system", SOLVER_SYSTEM_PROMPT), ("user", SOLVER_USER_PROMPT)])

{aiq → nat}/agent/rewoo_agent/register.py

@@ -15,25 +15,26 @@
 
 import logging
 
+from pydantic import AliasChoices
 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
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.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
+    Defines a NAT function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
     tool names and descriptions to select the optimal tool.
     """
 
@@ -54,51 +55,48 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
     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.")
+    additional_planner_instructions: str | None = Field(
+        default=None,
+        validation_alias=AliasChoices("additional_planner_instructions", "additional_instructions"),
+        description="Additional instructions to provide to the agent in addition to the base planner prompt.")
+    additional_solver_instructions: str | None = Field(
+        default=None,
+        description="Additional instructions to provide to the agent in addition to the base solver prompt.")
 
 
 @register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
-async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
+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 nat.agent.rewoo_agent.prompt import PLANNER_SYSTEM_PROMPT
+    from nat.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
+    from nat.agent.rewoo_agent.prompt import SOLVER_SYSTEM_PROMPT
+    from nat.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
+
+    # the ReWOO Agent prompts are defined in prompt.py, and can be customized there or by modifying the config option
+    # planner_prompt and solver_prompt.
+    planner_system_prompt = PLANNER_SYSTEM_PROMPT if config.planner_prompt is None else config.planner_prompt
+    if config.additional_planner_instructions:
+        planner_system_prompt += f"{config.additional_planner_instructions}"
+    if not ReWOOAgentGraph.validate_planner_prompt(planner_system_prompt):
+        logger.exception("Invalid planner prompt")
+        raise ValueError("Invalid planner prompt")
+    planner_prompt = ChatPromptTemplate([("system", planner_system_prompt), ("user", PLANNER_USER_PROMPT)])
+
+    solver_system_prompt = SOLVER_SYSTEM_PROMPT if config.solver_prompt is None else config.solver_prompt
+    if config.additional_solver_instructions:
+        solver_system_prompt += f"{config.additional_solver_instructions}"
+    if not ReWOOAgentGraph.validate_solver_prompt(solver_system_prompt):
+        logger.exception("Invalid solver prompt")
+        raise ValueError("Invalid solver prompt")
+    solver_prompt = ChatPromptTemplate([("system", solver_system_prompt), ("user", SOLVER_USER_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)
@@ -117,7 +115,7 @@ async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
                                                  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:
+    async def _response_fn(input_message: ChatRequest) -> ChatResponse:
         try:
             # initialize the starting state with the user query
             messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
@@ -135,14 +133,14 @@ async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
             # 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)
+            return ChatResponse.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.")
+                return ChatResponse.from_string(str(ex))
+            return ChatResponse.from_string("I seem to be having a problem.")
 
     if (config.use_openai_api):
         yield FunctionInfo.from_fn(_response_fn, description=config.description)
@@ -150,7 +148,7 @@ async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
     else:
 
         async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
+            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
             oai_output = await _response_fn(oai_input)
 
             return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)

{aiq → nat}/agent/tool_calling_agent/agent.py

@@ -25,10 +25,10 @@ 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
+from nat.agent.base import AGENT_CALL_LOG_MESSAGE
+from nat.agent.base import AGENT_LOG_PREFIX
+from nat.agent.base import AgentDecision
+from nat.agent.dual_node import DualNodeAgent
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/agent/tool_calling_agent/register.py

@@ -17,13 +17,13 @@ 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
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
@@ -48,7 +48,7 @@ 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 nat.agent.base import AGENT_LOG_PREFIX
 
     from .agent import ToolCallAgentGraph
     from .agent import ToolCallAgentGraphState

{aiq → nat}/authentication/api_key/api_key_auth_provider.py

@@ -17,11 +17,11 @@ 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
+from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from nat.authentication.interfaces import AuthProviderBase
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import BearerTokenCred
+from nat.data_models.authentication import HeaderAuthScheme
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +51,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
         """
 
-        from aiq.authentication.interfaces import AUTHORIZATION_HEADER
+        from nat.authentication.interfaces import AUTHORIZATION_HEADER
 
         config: APIKeyAuthProviderConfig = self.config
 

{aiq → nat}/authentication/api_key/api_key_auth_provider_config.py

@@ -20,11 +20,11 @@ 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
+from nat.authentication.exceptions.api_key_exceptions import APIKeyFieldError
+from nat.authentication.exceptions.api_key_exceptions import HeaderNameFieldError
+from nat.authentication.exceptions.api_key_exceptions import HeaderPrefixFieldError
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import HeaderAuthScheme
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/authentication/api_key/register.py

@@ -13,14 +13,14 @@
 # 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
+from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from nat.builder.builder import Builder
+from nat.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
+    from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
 
     yield APIKeyAuthProvider(config=config)

{aiq → nat}/authentication/http_basic_auth/http_basic_auth_provider.py

@@ -15,14 +15,14 @@
 
 from pydantic import SecretStr
 
-from aiq.authentication.interfaces import AuthProviderBase
-from aiq.builder.context import AIQContext
-from aiq.data_models.authentication import AuthenticatedContext
-from aiq.data_models.authentication import AuthFlowType
-from aiq.data_models.authentication import AuthProviderBaseConfig
-from aiq.data_models.authentication import AuthResult
-from aiq.data_models.authentication import BasicAuthCred
-from aiq.data_models.authentication import BearerTokenCred
+from nat.authentication.interfaces import AuthProviderBase
+from nat.builder.context import Context
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import BasicAuthCred
+from nat.data_models.authentication import BearerTokenCred
 
 
 class HTTPBasicAuthProvider(AuthProviderBase):
@@ -43,11 +43,11 @@ class HTTPBasicAuthProvider(AuthProviderBase):
         Performs simple HTTP Authentication using the provided user ID.
         """
 
-        context = AIQContext.get()
+        context = Context.get()
 
         if user_id is None and hasattr(context, "metadata") and hasattr(
                 context.metadata, "cookies") and context.metadata.cookies is not None:
-            session_id = context.metadata.cookies.get("aiqtoolkit-session", None)
+            session_id = context.metadata.cookies.get("nat-session", None)
             if not session_id:
                 raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
 

{aiq → nat}/authentication/http_basic_auth/register.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from aiq.builder.builder import Builder
-from aiq.cli.register_workflow import register_auth_provider
-from aiq.data_models.authentication import AuthProviderBaseConfig
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+from nat.data_models.authentication import AuthProviderBaseConfig
 
 
 class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
@@ -25,6 +25,6 @@ class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
 @register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
 async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
 
-    from aiq.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
+    from nat.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
 
     yield HTTPBasicAuthProvider(config)

{aiq → nat}/authentication/interfaces.py

@@ -17,11 +17,11 @@ import typing
 from abc import ABC
 from abc import abstractmethod
 
-from aiq.data_models.authentication import AuthenticatedContext
-from aiq.data_models.authentication import AuthFlowType
-from aiq.data_models.authentication import AuthProviderBaseConfig
-from aiq.data_models.authentication import AuthProviderBaseConfigT
-from aiq.data_models.authentication import AuthResult
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthProviderBaseConfigT
+from nat.data_models.authentication import AuthResult
 
 AUTHORIZATION_HEADER = "Authorization"
 
@@ -75,7 +75,7 @@ class FlowHandlerBase(ABC):
 
     Each front end will define a FlowHandler that will implement the authenticate method.
 
-    The `authenticate` method will be stored as the callback in the AIQContextState.user_auth_callback
+    The `authenticate` method will be stored as the callback in the ContextState.user_auth_callback
     """
 
     @abstractmethod