diff --git a/python/src/kagent/cli.py b/python/src/kagent/cli.py index 8996bd589c..d4c69d38e5 100644 --- a/python/src/kagent/cli.py +++ b/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() diff --git a/python/src/kagent/tools/argo/_argo_crds.py b/python/src/kagent/tools/argo/_argo_crds.py index effa75d2e3..dea3f205db 100644 --- a/python/src/kagent/tools/argo/_argo_crds.py +++ b/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) diff --git a/python/src/kagent/tools/argo/_argo_rollouts_k8sgw_installation.py b/python/src/kagent/tools/argo/_argo_rollouts_k8sgw_installation.py index e690d94f83..d50540d08f 100644 --- a/python/src/kagent/tools/argo/_argo_rollouts_k8sgw_installation.py +++ b/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,7 +35,7 @@ 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" @@ -41,12 +43,12 @@ def verify_gateway_plugin( 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", ) diff --git a/python/src/kagent/tools/helm/__init__.py b/python/src/kagent/tools/helm/__init__.py new file mode 100644 index 0000000000..b88b32bd84 --- /dev/null +++ b/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"] diff --git a/python/src/kagent/tools/helm/_helm.py b/python/src/kagent/tools/helm/_helm.py new file mode 100644 index 0000000000..1689478ab3 --- /dev/null +++ b/python/src/kagent/tools/helm/_helm.py @@ -0,0 +1,227 @@ +from typing import Annotated, Optional + +from autogen_core.tools import FunctionTool + +from .._utils import create_typed_fn_tool +from ..common.shell import run_command + + +def _helm_list( + namespace: Annotated[Optional[str], "The namespace to list the helm charts from"] = None, + all_namespaces: Annotated[Optional[bool], "If true, list releases from all namespaces"] = False, + all: Annotated[Optional[bool], "If true, show all releases without any filter applied"] = False, + uninstalled: Annotated[Optional[bool], "If true, list uninstalled releases"] = False, + uninstalling: Annotated[Optional[bool], "If true, list uninstalling releases"] = False, + failed: Annotated[Optional[bool], "If true, list failed releases"] = False, + deployed: Annotated[Optional[bool], "If true, list deployed releases"] = False, + filter: Annotated[ + Optional[str], + "a regular expression (Perl compatible). Any releases that match the expression will be included in the results", + ] = None, + pending: Annotated[Optional[bool], "If true, list pending releases"] = False, + output: Annotated[ + Optional[str], + "The output format of the helm list command, one of: table, json, yaml. Prefer table for human readability.", + ] = None, +) -> str: + args: list[str] = [] + if namespace: + args.append(f"-n {namespace}") + + if all_namespaces: + args.append("-A") + + if all: + args.append("-a") + + if uninstalled: + args.append("--uninstalled") + + if uninstalling: + args.append("--uninstalling") + + if failed: + args.append("--failed") + + if deployed: + args.append("--deployed") + + if pending: + args.append("--pending") + + if output: + args.append(f"-o {output}") + + if filter: + args.append(f"-f {filter}") + + return _run_helm_command(f"list {' '.join(args)}") + + +helm_list = FunctionTool( + _helm_list, + description=""" +This command lists all of the releases for a specified namespace (uses current namespace context if namespace not specified). + +By default, it lists only releases that are deployed or failed. Flags like +'--uninstalled' and '--all' will alter this behavior. Such flags can be combined: +'--uninstalled --failed'. + +By default, items are sorted alphabetically. Use the '-d' flag to sort by +release date. + +If the --filter flag is provided, it will be treated as a filter. Filters are +regular expressions (Perl compatible) that are applied to the list of releases. +Only items that match the filter will be returned. + + $ helm list --filter 'ara[a-z]+' + NAME UPDATED CHART + maudlin-arachnid 2020-06-18 14:17:46.125134977 +0000 UTC alpine-0.1.0 + """, + name="helm_list", +) + +HelmList, HelmListConfig = create_typed_fn_tool(helm_list, "kagent.tools.helm.List", "HelmList") + + +def _helm_upgrade( + name: Annotated[str, "The name of the release"], + chart: Annotated[str, "The chart to install"], + namespace: Annotated[str, "The namespace to install the release in"], + create_namespace: Annotated[Optional[bool], "If true, create the namespace if it does not exist"] = False, + set: Annotated[ + Optional[list[str]], + "A list of key-value pairs to set on the release. (can specify multiple or separate values with commas: key1=val1,key2=val2)", + ] = None, + values: Annotated[ + Optional[list[str]], + "A list of files to use as the value source. (can specify multiple or separate values with commas: myvalues.yaml,override.yaml)", + ] = None, + version: Annotated[Optional[str], "The version of the chart to install"] = None, + install: Annotated[Optional[bool], "If true, install the release if it does not exist"] = False, + dry_run: Annotated[ + Optional[bool], "If true, show which releases will be uninstalled without actually uninstalling them" + ] = False, + wait: Annotated[Optional[bool], "If true, wait for the release to be deployed"] = False, +) -> str: + args: list[str] = [name, chart, "-n", namespace] + if create_namespace: + args.append("--create-namespace") + if set: + args.append(f"--set {','.join(set)}") + if values: + args.append(f"--values {','.join(values)}") + if version: + args.append(f"--version {version}") + if install: + args.append("--install") + if dry_run: + args.append("--dry-run") + if wait: + args.append("--wait") + + return _run_helm_command(f"upgrade {' '.join(args)}") + + +helm_upgrade = FunctionTool( + _helm_upgrade, + description=""" +This command upgrades a release to a new version of a chart. + +The upgrade arguments must be a release and chart. The chart +argument can be either: a chart reference('example/mariadb'), a path to a chart directory, +a packaged chart, or a fully qualified URL. For chart references, the latest +version will be specified unless the '--version' flag is set. + +To override values in a chart, use either the '--values' flag and pass in a file +or use the '--set' flag and pass configuration from the command line, to force string +values, use '--set-string'. You can use '--set-file' to set individual +values from a file when the value itself is too long for the command line +or is dynamically generated. You can also use '--set-json' to set json values +(scalars/objects/arrays) from the command line. + +You can specify the '--values'/'-f' flag multiple times. The priority will be given to the +last (right-most) file specified. For example, if both myvalues.yaml and override.yaml +contained a key called 'Test', the value set in override.yaml would take precedence: + + $ helm upgrade -f myvalues.yaml -f override.yaml redis ./redis + +You can specify the '--set' flag multiple times. The priority will be given to the +last (right-most) set specified. For example, if both 'bar' and 'newbar' values are +set for a key called 'foo', the 'newbar' value would take precedence: + + $ helm upgrade --set foo=bar --set foo=newbar redis ./redis + +There are six different ways you can express the chart you want to install: + +1. By chart reference: helm install mymaria example/mariadb +2. By path to a packaged chart: helm install mynginx ./nginx-1.2.3.tgz +3. By path to an unpacked chart directory: helm install mynginx ./nginx +4. By absolute URL: helm install mynginx https://example.com/charts/nginx-1.2.3.tgz +5. By chart reference and repo url: helm install --repo https://example.com/charts/ mynginx nginx +6. By OCI registries: helm install mynginx --version 1.2.3 oci://example.com/charts/nginx + +CHART REFERENCES + +A chart reference is a convenient way of referencing a chart in a chart repository. + +When you use a chart reference with a repo prefix ('example/mariadb'), Helm will look in the local +configuration for a chart repository named 'example', and will then look for a +chart in that repository whose name is 'mariadb'. It will install the latest stable version of that chart +until you specify '--devel' flag to also include development version (alpha, beta, and release candidate releases), or +supply a version number with the '--version' flag. + +To see the list of chart repositories, use 'helm repo list'. To search for +charts in a repository, use 'helm search'. +""", + name="helm_upgrade", +) + +HelmUpgrade, HelmUpgradeConfig = create_typed_fn_tool(helm_upgrade, "kagent.tools.helm.Upgrade", "HelmUpgrade") + + +def _helm_uninstall( + name: Annotated[str, "The name of the release"], + namespace: Annotated[str, "The namespace to uninstall the release from"], + dry_run: Annotated[ + Optional[bool], "If true, show which releases will be uninstalled without actually uninstalling them" + ] = False, + wait: Annotated[Optional[bool], "If true, wait for the release to be uninstalled"] = False, +) -> str: + args: list[str] = [name, "-n", namespace] + if dry_run: + args.append("--dry-run") + if wait: + args.append("--wait") + return _run_helm_command(f"uninstall {' '.join(args)}") + + +helm_uninstall = FunctionTool( + _helm_uninstall, + description=""" +This command takes a release name and uninstalls the release. + +It removes all of the resources associated with the last release of the chart +as well as the release history, freeing it up for future use. + +Use the '--dry-run' flag to see which releases will be uninstalled without actually +uninstalling them. + +Usage: + helm uninstall RELEASE_NAME [...] [flags] + +""", + name="helm_uninstall", +) + + +HelmUninstall, HelmUninstallConfig = create_typed_fn_tool( + helm_uninstall, "kagent.tools.helm.Uninstall", "HelmUninstall" +) + + +def _run_helm_command(command: str) -> str: + # Split the command and remove empty strings + cmd_parts = command.split(" ") + cmd_parts = [part for part in cmd_parts if part] # Remove empty strings from the list + return run_command("helm", cmd_parts) diff --git a/python/src/kagent/tools/istio/_istioctl.py b/python/src/kagent/tools/istio/_istioctl.py index 8c2cc48787..4622ca9d29 100644 --- a/python/src/kagent/tools/istio/_istioctl.py +++ b/python/src/kagent/tools/istio/_istioctl.py @@ -91,7 +91,7 @@ async def _proxy_config( "The type of configuration to get, the allowed values are: all, bootstrap, cluster, ecds, listener, log, route, secret", ] = "all", ) -> str: - return _run_istioctl_command(f"proxy-config {config_type} {'-n ' + ns if ns else ''} {pod_name}") + return _run_istioctl_command(f"proxy-config {config_type} {pod_name}{'.' + ns if ns else ''}") async def _generate_manifest( diff --git a/python/src/kagent/tools/k8s/_kubectl.py b/python/src/kagent/tools/k8s/_kubectl.py index 4516588099..8f8a4e9dfe 100644 --- a/python/src/kagent/tools/k8s/_kubectl.py +++ b/python/src/kagent/tools/k8s/_kubectl.py @@ -140,10 +140,13 @@ def _get_resource_yaml( resource_type: Annotated[ str, "The type of resource to get the YAML definition for (deployment, service, pod, node, ...)" ], - name: Annotated[str, "The name of the resource to get the YAML definition for"], + name: Annotated[ + Optional[str], + "The name of the resource to get the YAML definition for. If not provided, all resources of the given type will be returned.", + ], ns: Annotated[Optional[str], "The namespace of the resource to get the definition for"], ) -> str: - return _run_kubectl_command(f"get {resource_type} {name} {f'-n {ns} ' if ns else ''} -o yaml") + return _run_kubectl_command(f"get {resource_type} {name if name else ''} {f'-n {ns} ' if ns else ''} -o yaml") def _execute_command( @@ -155,7 +158,10 @@ def _execute_command( def _get_resources( - name: Annotated[str, "The name of the resource to get information about"], + name: Annotated[ + Optional[str], + "The name of the resource to get information about. If not provided, all resources of the given type will be returned.", + ], resource_type: Annotated[ str, "The type of resource to get information about (deployment, service, pod, node, ...)" ],