aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/tool/github_tools/get_github_pr.py

@@ -0,0 +1,256 @@
+# 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 typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+
+class GithubListPullsModel(BaseModel):
+    state: Literal["open", "closed", "all"] | None = Field('open', description="Issue state used in issue query filter")
+    head: str | None = Field(None, description="Filters pulls by head user or head organization and branch name")
+    base: str | None = Field(None, description="Filters pull by branch name")
+
+
+class GithubListPullsModelList(BaseModel):
+    filter_params: GithubListPullsModel = Field(description=("A list of query params when fetching pull requests "
+                                                             "each of type GithubListPRModel"))
+
+
+class GithubListPullsToolConfig(FunctionBaseConfig, name="github_list_pulls_tool"):
+    """
+    Tool that lists GitHub Pull Requests based on various filter parameters
+    """
+    repo_name: str = Field(description="The repository name in the format 'owner/repo'")
+    timeout: int = Field(default=300, description="The timeout configuration to use when sending requests.")
+
+
+@register_function(config_type=GithubListPullsToolConfig)
+async def list_github_pulls_async(config: GithubListPullsToolConfig, builder: Builder):
+    """
+    Lists GitHub Pull Requests based on various filter parameters
+
+    """
+    import json
+    import os
+
+    import httpx
+
+    github_pat = os.getenv("GITHUB_PAT")
+    if not github_pat:
+        raise ValueError("GITHUB_PAT environment variable must be set")
+
+    url = f"https://api.github.com/repos/{config.repo_name}/pulls"
+
+    # define the headers for the payload request
+    headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
+
+    async def _github_list_pulls(filter_params) -> dict:
+        async with httpx.AsyncClient(timeout=config.timeout) as client:
+
+            filter_params = filter_params.dict(exclude_unset=True)
+
+            # filter out None values that are explictly set in the request body.
+            filter_params = {k: v for k, v in filter_params.items() if v is not None}
+
+            response = await client.request("GET", url, params=filter_params, headers=headers)
+
+            # Raise an exception for HTTP errors
+            response.raise_for_status()
+
+            # Parse and return the response JSON
+            try:
+                result = response.json()
+
+            except ValueError as e:
+                raise ValueError("The API response is not valid JSON.") from e
+
+        return json.dumps(result)
+
+    yield FunctionInfo.from_fn(_github_list_pulls,
+                               description=(f"Lists GitHub PRs based on filter params "
+                                            f"in the repo named {config.repo_name}"),
+                               input_schema=GithubListPullsModelList)
+
+
+class GithubGetPullModel(BaseModel):
+    pull_number: str = Field(description="The number of the pull request that needs to be fetched")
+
+
+class GithubGetPullToolConfig(FunctionBaseConfig, name="github_get_pull_tool"):
+    """
+    Tool that fetches a particular pull request in a GitHub repository asynchronously.
+    """
+    repo_name: str = "The repository name in the format 'owner/repo'"
+    timeout: int = 300
+
+
+@register_function(config_type=GithubGetPullToolConfig)
+async def get_github_pull_async(config: GithubGetPullToolConfig, builder: Builder):
+    """
+    Fetches a particular pull request in a GitHub repository asynchronously.
+
+    """
+    import json
+    import os
+
+    import httpx
+
+    github_pat = os.getenv("GITHUB_PAT")
+    if not github_pat:
+        raise ValueError("GITHUB_PAT environment variable must be set")
+
+    url = f"https://api.github.com/repos/{config.repo_name}/pulls"
+
+    # define the headers for the payload request
+    headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
+
+    async def _github_get_pull(pull_number) -> list:
+        async with httpx.AsyncClient(timeout=config.timeout) as client:
+            # update the url with the pull number that needs to be updated
+            pull_url = os.path.join(url, pull_number)
+
+            response = await client.request("GET", pull_url, headers=headers)
+
+            # Raise an exception for HTTP errors
+            response.raise_for_status()
+
+            # Parse and return the response JSON
+            try:
+                result = response.json()
+
+            except ValueError as e:
+                raise ValueError("The API response is not valid JSON.") from e
+
+        return json.dumps(result)
+
+    yield FunctionInfo.from_fn(_github_get_pull,
+                               description=(f"Fetches a particular GitHub pull request "
+                                            f"in the repo named {config.repo_name}"),
+                               input_schema=GithubGetPullModel)
+
+
+class GithubGetPullCommitsToolConfig(FunctionBaseConfig, name="github_get_pull_commits_tool"):
+    """
+    Configuration for the GitHub Get Pull Commits Tool.
+    """
+    repo_name: str = "The repository name in the format 'owner/repo'"
+    timeout: int = 300
+
+
+@register_function(config_type=GithubGetPullCommitsToolConfig)
+async def get_github_pull_commits_async(config: GithubGetPullCommitsToolConfig, builder: Builder):
+    """
+    Fetches the commits associated with a particular pull request in a GitHub repository asynchronously.
+
+    """
+    import json
+    import os
+
+    import httpx
+
+    github_pat = os.getenv("GITHUB_PAT")
+    if not github_pat:
+        raise ValueError("GITHUB_PAT environment variable must be set")
+
+    url = f"https://api.github.com/repos/{config.repo_name}/pulls"
+
+    # define the headers for the payload request
+    headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
+
+    async def _github_get_pull(pull_number) -> list:
+        async with httpx.AsyncClient(timeout=config.timeout) as client:
+            # update the url with the pull number that needs to be updated
+            pull_url = os.path.join(url, pull_number)
+            pull_commits_url = os.path.join(pull_url, "commits")
+
+            response = await client.request("GET", pull_commits_url, headers=headers)
+
+            # Raise an exception for HTTP errors
+            response.raise_for_status()
+
+            # Parse and return the response JSON
+            try:
+                result = response.json()
+
+            except ValueError as e:
+                raise ValueError("The API response is not valid JSON.") from e
+
+        return json.dumps(result)
+
+    yield FunctionInfo.from_fn(_github_get_pull,
+                               description=("Fetches the commits for a particular GitHub pull request "
+                                            f" in the repo named {config.repo_name}"),
+                               input_schema=GithubGetPullModel)
+
+
+class GithubGetPullFilesToolConfig(FunctionBaseConfig, name="github_get_pull_files_tool"):
+    """
+    Configuration for the GitHub Get Pull Files Tool.
+    """
+    repo_name: str = "The repository name in the format 'owner/repo'"
+    timeout: int = 300
+
+
+@register_function(config_type=GithubGetPullFilesToolConfig)
+async def get_github_pull_files_async(config: GithubGetPullFilesToolConfig, builder: Builder):
+    """
+    Fetches the files associated with a particular pull request in a GitHub repository asynchronously.
+
+    """
+    import json
+    import os
+
+    import httpx
+
+    github_pat = os.getenv("GITHUB_PAT")
+    if not github_pat:
+        raise ValueError("GITHUB_PAT environment variable must be set")
+
+    url = f"https://api.github.com/repos/{config.repo_name}/pulls"
+
+    # define the headers for the payload request
+    headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
+
+    async def _github_get_pull(pull_number) -> list:
+        async with httpx.AsyncClient(timeout=config.timeout) as client:
+            # update the url with the pull number that needs to be updated
+            pull_url = os.path.join(url, pull_number)
+            pull_files_url = os.path.join(pull_url, "files")
+
+            response = await client.request("GET", pull_files_url, headers=headers)
+
+            # Raise an exception for HTTP errors
+            response.raise_for_status()
+
+            # Parse and return the response JSON
+            try:
+                result = response.json()
+
+            except ValueError as e:
+                raise ValueError("The API response is not valid JSON.") from e
+
+        return json.dumps(result)
+
+    yield FunctionInfo.from_fn(_github_get_pull,
+                               description=("Fetches the files for a particular GitHub pull request "
+                                            f" in the repo named {config.repo_name}"),
+                               input_schema=GithubGetPullModel)

aiq/tool/github_tools/update_github_issue.py

@@ -0,0 +1,100 @@
+# 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 typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+
+class GithubUpdateIssueModel(BaseModel):
+    issue_number: str = Field(description="The issue number that will be updated")
+    title: str | None = Field(None, description="The title of the GitHub Issue")
+    body: str | None = Field(None, description="The body of the GitHub Issue")
+    state: Literal["open", "closed"] | None = Field(None, description="The new state of the issue")
+
+    state_reason: Literal["completed", "not_planned", "reopened", None] | None = Field(
+        None, description="The reason for changing the state of the issue")
+
+    labels: list[str] | None = Field(None, description="A list of labels to assign to the issue")
+    assignees: list[str] | None = Field(None, description="A list of assignees to assign to the issue")
+
+
+class GithubUpdateIssueModelList(BaseModel):
+    issues: list[GithubUpdateIssueModel] = Field(description=("A list of GitHub issues each "
+                                                              "of type GithubUpdateIssueModel"))
+
+
+class GithubUpdateIssueToolConfig(FunctionBaseConfig, name="github_update_issue_tool"):
+    """
+    Tool that updates an issue in a GitHub repository asynchronously.
+    """
+    repo_name: str = "The repository name in the format 'owner/repo'"
+    timeout: int = 300
+
+
+@register_function(config_type=GithubUpdateIssueToolConfig)
+async def update_github_issue_async(config: GithubUpdateIssueToolConfig, builder: Builder):
+    """
+    Updates an issue in a GitHub repository asynchronously.
+    """
+    import json
+    import os
+
+    import httpx
+
+    github_pat = os.getenv("GITHUB_PAT")
+    if not github_pat:
+        raise ValueError("GITHUB_PAT environment variable must be set")
+
+    url = f"https://api.github.com/repos/{config.repo_name}/issues"
+
+    # define the headers for the payload request
+    headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
+
+    async def _github_update_issue(issues) -> list:
+        results = []
+        async with httpx.AsyncClient(timeout=config.timeout) as client:
+            for issue in issues:
+                payload = issue.dict(exclude_unset=True)
+
+                # update the url with the issue number that needs to be updated
+                issue_number = payload.pop("issue_number")
+                issue_url = os.path.join(url, issue_number)
+
+                response = await client.request("PATCH", issue_url, json=payload, headers=headers)
+
+                # Raise an exception for HTTP errors
+                response.raise_for_status()
+
+                # Parse and return the response JSON
+                try:
+                    result = response.json()
+                    results.append(result)
+
+                except ValueError as e:
+                    raise ValueError("The API response is not valid JSON.") from e
+
+        return json.dumps(results)
+
+    yield FunctionInfo.from_fn(_github_update_issue,
+                               description=(f"Updates a GitHub issue in the "
+                                            f"repo named {config.repo_name}"),
+                               input_schema=GithubUpdateIssueModelList)

aiq/tool/mcp/__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/tool/mcp/mcp_client.py

@@ -0,0 +1,220 @@
+# 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 __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from enum import Enum
+from typing import Any
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+from mcp.types import TextContent
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import create_model
+
+logger = logging.getLogger(__name__)
+
+
+def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
+    """
+    Create a pydantic model from the input schema of the MCP tool
+    """
+    _type_map = {
+        "string": str,
+        "number": float,
+        "integer": int,
+        "boolean": bool,
+        "array": list,
+        "null": None,
+        "object": dict,
+    }
+
+    properties = mcp_input_schema.get("properties", {})
+    schema_dict = {}
+
+    def _generate_valid_classname(class_name: str):
+        return class_name.replace('_', ' ').replace('-', ' ').title().replace(' ', '')
+
+    def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
+        json_type = field_properties.get("type", "string")
+        enum_vals = field_properties.get("enum")
+
+        if enum_vals:
+            enum_name = f"{field_name.capitalize()}Enum"
+            field_type = Enum(enum_name, {item: item for item in enum_vals})
+
+        elif json_type == "object" and "properties" in field_properties:
+            field_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
+        elif json_type == "array" and "items" in field_properties:
+            item_properties = field_properties.get("items", {})
+            if item_properties.get("type") == "object":
+                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
+            else:
+                item_type = _type_map.get(json_type, Any)
+            field_type = list[item_type]
+        else:
+            field_type = _type_map.get(json_type, Any)
+
+        default_value = field_properties.get("default", ...)
+        nullable = field_properties.get("nullable", False)
+        description = field_properties.get("description", "")
+
+        field_type = field_type | None if nullable else field_type
+
+        return field_type, Field(default=default_value, description=description)
+
+    for field_name, field_props in properties.items():
+        schema_dict[field_name] = _generate_field(field_name=field_name, field_properties=field_props)
+    return create_model(f"{_generate_valid_classname(name)}InputSchema", **schema_dict)
+
+
+class MCPSSEClient:
+    """
+    Client for creating a session and connecting to an MCP server using SSE
+
+    Args:
+      url (str): The url of the MCP server
+    """
+
+    def __init__(self, url: str):
+        self.url = url
+
+    @asynccontextmanager
+    async def connect_to_sse_server(self):
+        """
+        Establish a session with an MCP SSE server within an aync context
+        """
+        async with sse_client(url=self.url) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                yield session
+
+
+class MCPBuilder(MCPSSEClient):
+    """
+    Builder class used to connect to an MCP Server and generate ToolClients
+
+    Args:
+        url (str): The url of the MCP server
+    """
+
+    def __init__(self, url):
+        super().__init__(url)
+        self._tools = None
+
+    async def get_tools(self):
+        """
+        Retrieve a dictionary of all tools served by the MCP server.
+        """
+        async with self.connect_to_sse_server() as session:
+            response = await session.list_tools()
+
+        return {
+            tool.name: MCPToolClient(self.url, tool.name, tool.description, tool_input_schema=tool.inputSchema)
+            for tool in response.tools
+        }
+
+    async def get_tool(self, tool_name: str) -> MCPToolClient:
+        """
+        Get an MCP Tool by name.
+
+        Args:
+            tool_name (str): Name of the tool to load.
+
+        Returns:
+            MCPToolClient for the configured tool.
+
+        Raise:
+            ValueError if no tool is available with that name.
+        """
+        if not self._tools:
+            self._tools = await self.get_tools()
+
+        tool = self._tools.get(tool_name)
+        if not tool:
+            raise ValueError(f"Tool {tool_name} not available at {self.url}")
+        return tool
+
+    async def call_tool(self, tool_name: str, tool_args: dict | None):
+        async with self.connect_to_sse_server() as session:
+            result = await session.call_tool(tool_name, tool_args)
+            return result
+
+
+class MCPToolClient(MCPSSEClient):
+    """
+    Client wrapper used to call an MCP tool.
+
+    Args:
+        url (str): The url of the MCP server
+        tool_name (str): The name of the tool to wrap
+        tool_description (str): The description of the tool provided by the MCP server.
+        tool_input_schema (dict): The input schema for the tool.
+    """
+
+    def __init__(self, url: str, tool_name: str, tool_description: str | None, tool_input_schema: dict | None = None):
+        super().__init__(url)
+        self._tool_name = tool_name
+        self._tool_description = tool_description
+        self._input_schema = model_from_mcp_schema(self._tool_name, tool_input_schema) if tool_input_schema else None
+
+    @property
+    def name(self):
+        """Returns the name of the tool."""
+        return self._tool_name
+
+    @property
+    def description(self):
+        """
+        Returns the tool's description. If none was provided. Provides a simple description using the tool's name
+        """
+        if not self._tool_description:
+            return f"MCP Tool {self._tool_name}"
+        return self._tool_description
+
+    @property
+    def input_schema(self):
+        """
+        Returns the tool's input_schema.
+        """
+        return self._input_schema
+
+    def set_description(self, description: str):
+        """
+        Manually define the tool's description using the provided string.
+        """
+        self._tool_description = description
+
+    async def acall(self, tool_args: dict) -> str:
+        """
+        Call the MCP tool with the provided arguments.
+
+        Args:
+            tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
+        """
+        async with self.connect_to_sse_server() as session:
+            result = await session.call_tool(self._tool_name, tool_args)
+
+        output = []
+        for res in result.content:
+            if isinstance(res, TextContent):
+                output.append(res.text)
+            else:
+                # Log non-text content for now
+                logger.warning("Got not-text output from %s of type %s", self.name, type(res))
+        return "\n".join(output)

aiq/tool/mcp/mcp_tool.py

@@ -0,0 +1,75 @@
+# 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 BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+
+from aiq.builder.builder import Builder
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
+    """
+    Function which connects to a Model Context Protocol (MCP) server and wraps the selected tool as an AgentIQ function.
+    """
+    # Add your custom configuration parameters here
+    url: HttpUrl = Field(description="The URL of the MCP server")
+    mcp_tool_name: str = Field(description="The name of the tool served by the MCP Server that you want to use")
+    description: str | None = Field(default=None,
+                                    description="""
+        Description for the tool that will override the description provided by the MCP server. Should only be used if
+        the description provided by the server is poor or nonexistent
+        """)
+
+
+@register_function(config_type=MCPToolConfig)
+async def mcp_tool(config: MCPToolConfig, builder: Builder):
+    """
+    Generate an AgentIQ Function that wraps a tool provided by the MCP server.
+    """
+
+    from aiq.tool.mcp.mcp_client import MCPBuilder
+    from aiq.tool.mcp.mcp_client import MCPToolClient
+
+    client = MCPBuilder(url=str(config.url))
+
+    tool: MCPToolClient = await client.get_tool(config.mcp_tool_name)
+    if config.description:
+        tool.set_description(description=config.description)
+
+    logger.info("Configured to use tool: %s from MCP server at %s", tool.name, str(config.url))
+
+    def _convert_from_str(input_str: str) -> tool.input_schema:
+        return tool.input_schema.model_validate_json(input_str)
+
+    async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
+        if tool_input:
+            args = tool_input.model_dump()
+            return await tool.acall(args)
+
+        _ = tool.input_schema.model_validate(kwargs)
+        return await tool.acall(kwargs)
+
+    yield FunctionInfo.create(single_fn=_response_fn,
+                              description=tool.description,
+                              input_schema=tool.input_schema,
+                              converters=[_convert_from_str])