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

aiq/agent/tool_calling_agent/agent.py

@@ -1,123 +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_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
-
-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_RESPONSE_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:
-                    # 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)
-                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,105 +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.agent.base import AGENT_LOG_PREFIX
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function_info import FunctionInfo
-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 react 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 .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/builder/builder.py

@@ -1,223 +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 asyncio
-import typing
-from abc import ABC
-from abc import abstractmethod
-from collections.abc import Sequence
-from pathlib import Path
-
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function import Function
-from aiq.data_models.component_ref import EmbedderRef
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.component_ref import MemoryRef
-from aiq.data_models.component_ref import RetrieverRef
-from aiq.data_models.embedder import EmbedderBaseConfig
-from aiq.data_models.evaluator import EvaluatorBaseConfig
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.data_models.function_dependencies import FunctionDependencies
-from aiq.data_models.llm import LLMBaseConfig
-from aiq.data_models.memory import MemoryBaseConfig
-from aiq.data_models.retriever import RetrieverBaseConfig
-from aiq.memory.interfaces import MemoryEditor
-from aiq.retriever.interface import AIQRetriever
-
-
-class UserManagerHolder():
-
-    def __init__(self, context: AIQContext) -> None:
-        self._context = context
-
-    def get_id(self):
-        return self._context.user_manager.get_id()
-
-
-class Builder(ABC):  # pylint: disable=too-many-public-methods
-
-    @abstractmethod
-    async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
-        pass
-
-    @abstractmethod
-    def get_function(self, name: str | FunctionRef) -> Function:
-        pass
-
-    def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
-
-        return [self.get_function(name) for name in function_names]
-
-    @abstractmethod
-    def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
-        pass
-
-    @abstractmethod
-    async def set_workflow(self, config: FunctionBaseConfig) -> Function:
-        pass
-
-    @abstractmethod
-    def get_workflow(self) -> Function:
-        pass
-
-    @abstractmethod
-    def get_workflow_config(self) -> FunctionBaseConfig:
-        pass
-
-    def get_tools(self, tool_names: Sequence[str | FunctionRef],
-                  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
-
-        return [self.get_tool(fn_name=n, wrapper_type=wrapper_type) for n in tool_names]
-
-    @abstractmethod
-    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
-        pass
-
-    @abstractmethod
-    async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
-        pass
-
-    async def get_llms(self, llm_names: Sequence[str | LLMRef],
-                       wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
-
-        coros = [self.get_llm(llm_name=n, wrapper_type=wrapper_type) for n in llm_names]
-
-        llms = await asyncio.gather(*coros, return_exceptions=False)
-
-        return list(llms)
-
-    @abstractmethod
-    async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
-        pass
-
-    @abstractmethod
-    def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
-        pass
-
-    @abstractmethod
-    async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
-        pass
-
-    async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
-                            wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
-
-        coros = [self.get_embedder(embedder_name=n, wrapper_type=wrapper_type) for n in embedder_names]
-
-        embedders = await asyncio.gather(*coros, return_exceptions=False)
-
-        return list(embedders)
-
-    @abstractmethod
-    async def get_embedder(self, embedder_name: str | EmbedderRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
-        pass
-
-    @abstractmethod
-    def get_embedder_config(self, embedder_name: str | EmbedderRef) -> EmbedderBaseConfig:
-        pass
-
-    @abstractmethod
-    async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig):
-        pass
-
-    def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
-        """
-        Return a list of memory clients for the specified names.
-        """
-        return [self.get_memory_client(n) for n in memory_names]
-
-    @abstractmethod
-    def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
-        """
-        Return the instantiated memory client for the given name.
-        """
-        pass
-
-    @abstractmethod
-    def get_memory_client_config(self, memory_name: str | MemoryRef) -> MemoryBaseConfig:
-        pass
-
-    @abstractmethod
-    async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
-        pass
-
-    async def get_retrievers(self,
-                             retriever_names: Sequence[str | RetrieverRef],
-                             wrapper_type: LLMFrameworkEnum | str | None = None):
-
-        tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
-
-        retrievers = await asyncio.gather(*tasks, return_exceptions=False)
-
-        return list(retrievers)
-
-    @typing.overload
-    async def get_retriever(self, retriever_name: str | RetrieverRef,
-                            wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
-        ...
-
-    @typing.overload
-    async def get_retriever(self, retriever_name: str | RetrieverRef, wrapper_type: None) -> AIQRetriever:
-        ...
-
-    @typing.overload
-    async def get_retriever(self, retriever_name: str | RetrieverRef) -> AIQRetriever:
-        ...
-
-    @abstractmethod
-    async def get_retriever(self,
-                            retriever_name: str | RetrieverRef,
-                            wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:
-        pass
-
-    @abstractmethod
-    async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
-        pass
-
-    @abstractmethod
-    def get_user_manager(self) -> UserManagerHolder:
-        pass
-
-    @abstractmethod
-    def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
-        pass
-
-
-class EvalBuilder(Builder):
-
-    @abstractmethod
-    async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
-        pass
-
-    @abstractmethod
-    def get_evaluator(self, evaluator_name: str) -> typing.Any:
-        pass
-
-    @abstractmethod
-    def get_evaluator_config(self, evaluator_name: str) -> EvaluatorBaseConfig:
-        pass
-
-    @abstractmethod
-    def get_max_concurrency(self) -> int:
-        pass
-
-    @abstractmethod
-    def get_output_dir(self) -> Path:
-        pass
-
-    @abstractmethod
-    def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
-        pass