add interactive AI-powered deubgging tool az aks agent

Author: mainred

src/aks-preview/azext_aks_preview/_params.py

@@ -23,6 +23,7 @@
     validate_nat_gateway_idle_timeout,
     validate_nat_gateway_managed_outbound_ip_count,
 )
+from azure.cli.core.api import get_config_dir
 from azure.cli.core.commands.parameters import (
     edge_zone_type,
     file_type,
@@ -2744,6 +2745,55 @@ def load_arguments(self, _):
                 help="Name of the load balancer configuration. Required.",
             )
 
+    with self.argument_context("aks agent") as c:
+        c.positional("prompt", help="Ask any question and answer using available tools.")
+        c.argument(
+            "api_key",
+            default=None,
+            required=False,
+            help="API key to use for the LLM (if not given, uses environment variables AZURE_API_KEY, OPENAI_API_KEY).",
+        )
+        c.argument(
+            "model",
+            default=None,
+            required=False,
+            help="Model to use for the LLM. For example, azure/<deployment_name> for Azure OpenAI, or <model_name> for OpenAI.",
+        )
+        c.argument(
+            "interactive",
+            type=bool,
+            default=True,
+            required=False,
+            help="Enabled interactive mode. If set to false, the agent will not prompt for input and will run in batch mode.",
+        )
+        c.argument(
+            "max_steps",
+            type=int,
+            default=10,
+            required=False,
+            help="Maximum number of steps the LLM can take to investigate the issue.",
+        )
+        c.argument(
+            "echo",
+            type=bool,
+            default=True,
+            required=False,
+            help="Echo back the question provided to HolmesGPT in the output.",
+        )
+        c.argument(
+            "show_tool_output",
+            type=bool,
+            default=True,
+            required=False,
+            help="Show the output of each tool that was called.",
+        )
+        c.argument(
+            "config_file",
+            type=bool,
+            default=os.path.join(get_config_dir(), "agent.config"),
+            required=False,
+            help="Show the output of each tool that was called.",
+        )
 
 def _get_default_install_location(exe_name):
     system = platform.system()

src/aks-preview/azext_aks_preview/commands.py

@@ -186,6 +186,7 @@ def load_command_table(self, _):
         g.custom_command(
             "operation-abort", "aks_operation_abort", supports_no_wait=True
         )
+        g.custom_command("agent", "aks_agent")
 
     # AKS maintenance configuration commands
     with self.command_group(

src/aks-preview/azext_aks_preview/custom.py

@@ -8,13 +8,26 @@
 import json
 import os
 import os.path
+from pathlib import Path
 import platform
-import ssl
+import socket
 import sys
 import threading
 import time
+import typer
+import uuid
 import webbrowser
 
+from holmes.config import Config
+from holmes.core.prompt import build_initial_ask_messages
+from holmes.interactive import run_interactive_loop
+from holmes.logging import init_logging
+from holmes.plugins.interfaces import Issue
+from holmes.plugins.prompts import load_and_render_prompt
+from holmes.utils.console.logging import init_logging
+from holmes.utils.console.result import handle_result
+from holmes.plugins.destinations import DestinationType
+
 from azext_aks_preview._client_factory import (
     CUSTOM_MGMT_AKS_PREVIEW,
     cf_agent_pools,
@@ -4333,3 +4346,102 @@ def aks_loadbalancer_rebalance_nodes(
     }
 
     return aks_loadbalancer_rebalance_internal(managed_clusters_client, parameters)
+
+
+def aks_agent(
+        prompt,
+        api_key,
+        model,
+        interactive,
+        max_steps,
+        echo,
+        show_tool_output,
+        config_file,
+        ):
+
+    # TODO: make log verbose configurable, currently disbled by [].
+    console = init_logging([])
+
+    # Detect and read piped input
+    piped_data = None
+    if not sys.stdin.isatty():
+        piped_data = sys.stdin.read().strip()
+        if interactive:
+            console.print(
+                "[bold yellow]Interactive mode disabled when reading piped input[/bold yellow]"
+            )
+            interactive = False
+
+    config_file = Path(config_file)
+    config = Config.load_from_file(
+        config_file,
+        api_key=api_key,
+        model=model,
+        max_steps=max_steps,
+    )
+
+    ai = config.create_console_toolcalling_llm(
+        dal=None
+    )
+    template_context = {
+        "toolsets": ai.tool_executor.toolsets,
+        "runbooks": config.get_runbook_catalog(),
+    }
+    # TODO: extend the system prompt with AKS context
+    system_prompt= "builtin://generic_ask.jinja2"
+    system_prompt_rendered = load_and_render_prompt(system_prompt, template_context)  # type: ignore
+
+    if not prompt and not interactive and not piped_data:
+        raise typer.BadParameter(
+            "Either the 'prompt' argument must be provided (unless using --interactive mode)."
+        )
+
+    # Handle piped data
+    if piped_data:
+        if prompt:
+            # User provided both piped data and a prompt
+            prompt = f"Here's some piped output:\n\n{piped_data}\n\n{prompt}"
+        else:
+            # Only piped data, no prompt - ask what to do with it
+            prompt = f"Here's some piped output:\n\n{piped_data}\n\nWhat can you tell me about this output?"
+
+    if echo and not interactive and prompt:
+        console.print("[bold yellow]User:[/bold yellow] " + prompt)
+
+    # TODO: add refresh-toolset to refresh the toolset if it has changed
+    if interactive:
+        run_interactive_loop(
+            ai,
+            console,
+            system_prompt_rendered,
+            prompt,
+        )
+        return
+
+    messages = build_initial_ask_messages(
+        console,
+        system_prompt_rendered,
+        prompt,  # type: ignore
+    )
+
+    response = ai.call(messages)
+
+
+    messages = response.messages  # type: ignore # Update messages with the full history
+
+    issue = Issue(
+        id=str(uuid.uuid4()),
+        name=prompt, 
+        source_type="holmes-ask",
+        raw={"prompt": prompt, "full_conversation": messages},
+        source_instance_id=socket.gethostname(),
+    )
+    handle_result(
+        response,
+        console,
+        DestinationType.CLI, 
+        config,
+        issue,
+        show_tool_output,
+        False, 
+    )

src/aks-preview/setup.py

@@ -7,7 +7,7 @@
 
 from codecs import open as open1
 
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
 VERSION = "18.0.0b18"
 
@@ -23,7 +23,9 @@
     "License :: OSI Approved :: MIT License",
 ]
 
-DEPENDENCIES = []
+DEPENDENCIES = [
+    "holmesgpt @ git+ssh://[email protected]/robusta-dev/holmesgpt@create-console-package", # will use a official pypi package once available
+]
 
 with open1("README.rst", "r", encoding="utf-8") as f:
     README = f.read()