Add helm tools

python/src/kagent/cli.py

@@ -3,6 +3,23 @@
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel
 
+from kagent.tools.argo._argo_crds import (
+    ArgoCRDTool,
+    ArgoCRDToolConfig,
+    ArgoCRDToolInput,
+)
+from kagent.tools.argo._argo_rollouts_k8sgw_installation import (
+    check_plugin_logs,
+    verify_gateway_plugin,
+)
+from kagent.tools.argo._kubectl_argo_rollouts import (
+    pause_rollout,
+    promote_rollout,
+    set_rollout_image,
+    verify_argo_rollouts_controller_install,
+    verify_kubectl_plugin_install,
+)
+from kagent.tools.helm._helm import helm_list, helm_uninstall, helm_upgrade
 from kagent.tools.istio._istio_crds import (
     IstioCRDTool,
     IstioCRDToolConfig,
@@ -28,6 +45,7 @@
     create_resource,
     delete_resource,
     describe_resource,
+    get_available_api_resources,
     get_cluster_configuration,
     get_events,
     get_pod_logs,
@@ -39,7 +57,6 @@
     remove_label,
     rollout,
     scale,
-    get_available_api_resources,
 )
 from kagent.tools.prometheus._prometheus import (
     AlertmanagersInput,
@@ -77,20 +94,6 @@
     TSDBStatusInput,
     TSDBStatusTool,
 )
-from kagent.tools.argo import (
-    PauseRollout,
-    PromoteRollout,
-    SetRolloutImage,
-    VerifyKubectlPluginInstall,
-    VerifyArgoRolloutsControllerInstall,
-    CheckPluginLogsTool,
-    VerifyGatewayPluginTool,
-)
-from kagent.tools.argo._argo_crds import (
-    ArgoCRDTool,
-    ArgoCRDToolConfig,
-    ArgoCRDToolInput,
-)
 
 app = typer.Typer()
 
@@ -199,19 +202,21 @@ def targets_tool(input: TargetsInput):
 
 @app.command()
 def argo():
+    mcp.add_tool(verify_gateway_plugin._func, verify_gateway_plugin.name, verify_gateway_plugin.description)
+    mcp.add_tool(check_plugin_logs._func, check_plugin_logs.name, check_plugin_logs.description)
     mcp.add_tool(
-        VerifyKubectlPluginInstall._func, VerifyKubectlPluginInstall.name, VerifyKubectlPluginInstall.description
+        verify_kubectl_plugin_install._func,
+        verify_kubectl_plugin_install.name,
+        verify_kubectl_plugin_install.description,
     )
     mcp.add_tool(
-        VerifyArgoRolloutsControllerInstall._func,
-        VerifyArgoRolloutsControllerInstall.name,
-        VerifyArgoRolloutsControllerInstall.description,
+        verify_argo_rollouts_controller_install._func,
+        verify_argo_rollouts_controller_install.name,
+        verify_argo_rollouts_controller_install.description,
     )
-    mcp.add_tool(CheckPluginLogsTool._func, CheckPluginLogsTool.name, CheckPluginLogsTool.description)
-    mcp.add_tool(VerifyGatewayPluginTool._func, VerifyGatewayPluginTool.name, VerifyGatewayPluginTool.description)
-    mcp.add_tool(PauseRollout._func, PauseRollout.name, PauseRollout.description)
-    mcp.add_tool(PromoteRollout._func, PromoteRollout.name, PromoteRollout.description)
-    mcp.add_tool(SetRolloutImage._func, SetRolloutImage.name, SetRolloutImage.description)
+    mcp.add_tool(pause_rollout._func, pause_rollout.name, pause_rollout.description)
+    mcp.add_tool(promote_rollout._func, promote_rollout.name, promote_rollout.description)
+    mcp.add_tool(set_rollout_image._func, set_rollout_image.name, set_rollout_image.description)
 
     cfg = ArgoCRDToolConfig(model="gpt-4o-mini", openai_api_key=None)
 
@@ -277,7 +282,14 @@ def k8s():
     mcp.add_tool(
         get_available_api_resources._func, get_available_api_resources.name, get_available_api_resources.description
     )
+    mcp.run()
 
+
+@app.command()
+def helm():
+    mcp.add_tool(helm_list._func, helm_list.name, helm_list.description)
+    mcp.add_tool(helm_upgrade._func, helm_upgrade.name, helm_upgrade.description)
+    mcp.add_tool(helm_uninstall._func, helm_uninstall.name, helm_uninstall.description)
     mcp.run()
 
 

python/src/kagent/tools/argo/_argo_crds.py

@@ -106,42 +106,16 @@ async def _generate_crd(
         except Exception as e:
             return f"Error generating policy: {str(e)}"
 
-    async def _generate_crd(
-        self, system_prompt: str, policy_description: str, cancellation_token: CancellationToken
-    ) -> str:
-        """
-        Asynchronously generates a Custom Resource Definition (CRD) based on the provided system prompt and policy description.
-
-        Args:
-        system_prompt (str): The system prompt to guide the CRD generation.
-        policy_description (str): The description of the policy to be included in the CRD.
-
-        Returns:
-        str: The generated CRD content or an error message if the generation fails.
-
-        Raises:
-        Exception: If there is an error during the CRD generation process.
-        """
-        try:
-            if cancellation_token.is_cancelled():
-                raise Exception("Operation cancelled")
-
-            result = await self._model_client.create(
-                messages=[SystemMessage(content=system_prompt), UserMessage(content=policy_description, source="user")],
-                json_output=True,
-            )
-            return result.content
-        except Exception as e:
-            return f"Error generating policy: {str(e)}"
-
     async def _generate_rollout_crd(
         self,
         resource_description: Annotated[str, "Detailed description of the Rollout to generate YAML for"],
+        cancellation_token: CancellationToken,
     ) -> str:
-        return await self._generate_crd(ROLLOUT_PROMPT, resource_description)
+        return await self._generate_crd(ROLLOUT_PROMPT, resource_description, cancellation_token)
 
     async def _generate_analysis_template_crd(
         self,
         resource_description: Annotated[str, "Detailed description of the AnalysisTemplate to generate YAML for"],
+        cancellation_token: CancellationToken,
     ) -> str:
-        return await self._generate_crd(ANALYSIS_TEMPLATE_PROMPT, resource_description)
+        return await self._generate_crd(ANALYSIS_TEMPLATE_PROMPT, resource_description, cancellation_token)

python/src/kagent/tools/argo/_argo_rollouts_k8sgw_installation.py

@@ -1,7 +1,9 @@
-from typing import Annotated, Optional
-import re
+import os
 import platform
+import re
+import tempfile
 from dataclasses import dataclass
+from typing import Annotated, Optional
 
 from autogen_core.tools import FunctionTool
 
@@ -18,7 +20,7 @@ class GatewayPluginStatus:
     error_message: Optional[str] = None
 
 
-def verify_gateway_plugin(
+def _verify_gateway_plugin(
     version: Annotated[Optional[str], "Version of the Gateway API plugin to verify. If None, checks latest"] = None,
     namespace: Annotated[str, "Namespace where Argo Rollouts is installed"] = "argo-rollouts",
     should_install: Annotated[bool, "Flag to determine if the plugin should be installed if not present"] = True,
@@ -33,20 +35,20 @@ def verify_gateway_plugin(
         config_map = run_command(command="kubectl", args=cmd)
         if "argoproj-labs/gatewayAPI" not in config_map:
             if should_install:
-                return configure_gateway_plugin(version, namespace)
+                return _configure_gateway_plugin(version, namespace)
             else:
                 return GatewayPluginStatus(
                     installed=False, error_message="Gateway API plugin is not configured and installation is disabled"
                 )
         return GatewayPluginStatus(installed=True, error_message="Gateway API plugin is already configured")
     except Exception as e:
         if should_install:
-            return configure_gateway_plugin(version, namespace)
+            return _configure_gateway_plugin(version, namespace)
         else:
             return GatewayPluginStatus(installed=False, error_message=f"Error verifying plugin: {str(e)}")
 
 
-def configure_gateway_plugin(
+def _configure_gateway_plugin(
     version: Optional[str],
     namespace: str,
 ) -> GatewayPluginStatus:
@@ -55,11 +57,11 @@ def configure_gateway_plugin(
     """
     try:
         # Determine system architecture
-        arch = get_system_architecture()
+        arch = _get_system_architecture()
 
         # If version not specified, get latest from GitHub
         if not version:
-            version = get_latest_version()
+            version = _get_latest_version()
 
         # Create ConfigMap manifest
         config_map = f"""
@@ -75,8 +77,6 @@ def configure_gateway_plugin(
 """
 
         # Create a temp file with the ConfigMap that adds the trafficRouterPlugins for the Kubernetes Gateway API
-        import tempfile
-        import os
 
         with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as temp_file:
             temp_file.write(config_map)
@@ -101,7 +101,7 @@ def configure_gateway_plugin(
         return GatewayPluginStatus(installed=False, error_message=f"Error during plugin configuration: {str(e)}")
 
 
-def get_system_architecture() -> str:
+def _get_system_architecture() -> str:
     """
     Determine the system architecture for plugin download.
     """
@@ -130,7 +130,7 @@ def get_system_architecture() -> str:
         raise ValueError(f"Unsupported system: {system}")
 
 
-def get_latest_version() -> str:
+def _get_latest_version() -> str:
     """
     Get the latest version of the Gateway API plugin from GitHub.
     """
@@ -149,10 +149,10 @@ def get_latest_version() -> str:
         return "0.5.0"  # Default to known stable version
 
 
-def check_plugin_logs(
+def _check_plugin_logs(
     namespace: Annotated[str, "Namespace where Argo Rollouts is installed"] = "argo-rollouts",
     timeout: Annotated[Optional[int], "Timeout in seconds for log checking"] = 60,
-) -> str:
+) -> GatewayPluginStatus:
     """
     Check Argo Rollouts controller logs for plugin installation status.
     """
@@ -181,26 +181,26 @@ def check_plugin_logs(
 
 
 # Create the function tools
-verify_gateway_plugin_tool = FunctionTool(
-    verify_gateway_plugin,
+verify_gateway_plugin = FunctionTool(
+    _verify_gateway_plugin,
     description="Verify and configure Gateway API plugin for Argo Rollouts",
     name="verify_gateway_plugin",
 )
 
-check_plugin_logs_tool = FunctionTool(
-    check_plugin_logs,
+check_plugin_logs = FunctionTool(
+    _check_plugin_logs,
     description="Check Argo Rollouts controller logs for Gateway API plugin installation status",
     name="check_plugin_logs",
 )
 
 VerifyGatewayPluginTool, VerifyGatewayPluginToolConfig = create_typed_fn_tool(
-    verify_gateway_plugin_tool,
+    verify_gateway_plugin,
     "kagent.tools.argo.VerifyGatewayPluginTool",
     "VerifyGatewayPluginTool",
 )
 
 CheckPluginLogsTool, CheckPluginLogsToolConfig = create_typed_fn_tool(
-    check_plugin_logs_tool,
+    check_plugin_logs,
     "kagent.tools.argo.CheckPluginLogsTool",
     "CheckPluginLogsTool",
 )

python/src/kagent/tools/helm/__init__.py

@@ -0,0 +1,3 @@
+from ._helm import HelmList, HelmListConfig, HelmUninstall, HelmUninstallConfig, HelmUpgrade, HelmUpgradeConfig
+
+__all__ = ["HelmUpgrade", "HelmUpgradeConfig", "HelmList", "HelmListConfig", "HelmUninstall", "HelmUninstallConfig"]