diff --git a/linter_exclusions.yml b/linter_exclusions.yml index 65581d5385e..035369f37e5 100644 --- a/linter_exclusions.yml +++ b/linter_exclusions.yml @@ -198,6 +198,11 @@ aks update: cluster_service_load_balancer_health_probe_mode: rule_exclusions: - option_length_too_long +aks agent: + parameters: + prompt: + rule_exclusions: + - no_positional_parameters arcdata dc config init: parameters: path: diff --git a/src/aks-preview/HISTORY.rst b/src/aks-preview/HISTORY.rst index 85a80440c87..d21cbe90ed7 100644 --- a/src/aks-preview/HISTORY.rst +++ b/src/aks-preview/HISTORY.rst @@ -11,10 +11,11 @@ To release a new version, please select a new version number (usually plus 1 to Pending +++++++ +* Add framework for interactive AI-powered debugging tool. 18.0.0b26 +++++++ -* Add `az aks identity-binding` command group for identity binding feataure. +* Add `az aks identity-binding` command group for identity binding feature. 18.0.0b25 +++++++ diff --git a/src/aks-preview/azext_aks_preview/_consts.py b/src/aks-preview/azext_aks_preview/_consts.py index 03261a1bdd1..5cec8e35be9 100644 --- a/src/aks-preview/azext_aks_preview/_consts.py +++ b/src/aks-preview/azext_aks_preview/_consts.py @@ -373,3 +373,9 @@ CONST_K8S_EXTENSION_NAME = "k8s-extension" CONST_K8S_EXTENSION_ACTION_MOD_NAME = "azext_k8s_extension.action" CONST_K8S_EXTENSION_FORMAT_MOD_NAME = "azext_k8s_extension._format" + +# aks agent constants +CONST_AGENT_CONFIG_PATH_DIR_ENV_KEY = "HOLMES_CONFIGPATH_DIR" +CONST_AGENT_NAME = "AKS AGENT" +CONST_AGENT_NAME_ENV_KEY = "AGENT_NAME" +CONST_AGENT_CONFIG_FILE_NAME = "aksAgent.config" diff --git a/src/aks-preview/azext_aks_preview/_help.py b/src/aks-preview/azext_aks_preview/_help.py index 086d27e5034..b8287a1df2d 100644 --- a/src/aks-preview/azext_aks_preview/_help.py +++ b/src/aks-preview/azext_aks_preview/_help.py @@ -3943,3 +3943,101 @@ type: string short-summary: Name of the identity binding to show. """ + +# pylint: disable=line-too-long +# helps[ +# "aks agent" +# ] = """ +# type: command +# short-summary: Run AI assistant to analyze and troubleshoot Kubernetes clusters. +# long-summary: |- +# This command allows you to ask questions about your Azure Kubernetes cluster and get answers using AI models. +# Environment variables must be set to use the AI model, please refer to https://docs.litellm.ai/docs/providers to learn more about supported AI providers and models and required environment variables. +# parameters: +# - name: --name -n +# type: string +# short-summary: Name of the managed cluster. +# - name: --resource-group -g +# type: string +# short-summary: Name of the resource group. +# - name: --model +# type: string +# short-summary: Model to use for the LLM. +# - name: --api-key +# type: string +# short-summary: API key to use for the LLM (if not given, uses environment variables AZURE_API_KEY, OPENAI_API_KEY). +# - name: --config-file +# type: string +# short-summary: Path to configuration file. +# - name: --max-steps +# type: int +# short-summary: Maximum number of steps the LLM can take to investigate the issue. +# - name: --no-interactive +# type: bool +# short-summary: Disable interactive mode. When set, the agent will not prompt for input and will run in batch mode. +# - name: --no-echo-request +# type: bool +# short-summary: Disable echoing back the question provided to AKS Agent in the output. +# - name: --show-tool-output +# type: bool +# short-summary: Show the output of each tool that was called during the analysis. +# - name: --refresh-toolsets +# type: bool +# short-summary: Refresh the toolsets status. +# +# examples: +# - name: Ask about pod issues in the cluster with Azure OpenAI +# text: |- +# export AZURE_API_BASE="https://my-azureopenai-service.openai.azure.com/" +# export AZURE_API_VERSION="2025-01-01-preview" +# export AZURE_API_KEY="sk-xxx" +# az aks agent "Why are my pods not starting?" --name MyManagedCluster --resource-group MyResourceGroup --model azure/my-gpt4.1-deployment +# - name: Ask about pod issues in the cluster with OpenAI +# text: |- +# export OPENAI_API_KEY="sk-xxx" +# az aks agent "Why are my pods not starting?" --name MyManagedCluster --resource-group MyResourceGroup --model gpt-4o +# text: az aks agent "Why are my pods not starting?" +# - name: Run in interactive mode without a question +# text: az aks agent "Check the pod status in my cluster" --name MyManagedCluster --resource-group MyResourceGroup --model azure/my-gpt4.1-deployment --api-key "sk-xxx" +# - name: Run in non-interactive batch mode +# text: az aks agent "Diagnose networking issues" --no-interactive --max-steps 15 --model azure/my-gpt4.1-deployment +# - name: Show detailed tool output during analysis +# text: az aks agent "Why is my service workload unavailable in namespace workload-ns?" --show-tool-output --model azure/my-gpt4.1-deployment +# - name: Use custom configuration file +# text: az aks agent "Check kubernetes pod resource usage" --config-file /path/to/custom.config --model azure/my-gpt4.1-deployment +# - name: Run agent with no echo of the original question +# text: az aks agent "What is the status of my cluster?" --no-echo-request --model azure/my-gpt4.1-deployment +# - name: Refresh toolsets to get the latest available tools +# text: az aks agent "What is the status of my cluster?" --refresh-toolsets --model azure/my-gpt4.1-deploymen +# - name: Run agent with config file +# text: | +# az aks agent "Check kubernetes pod resource usage" --config-file /path/to/custom.config +# Here is an example of config file: +# ```json +# model: "gpt-4o" +# api_key: "..." +# # define a list of mcp servers, mcp server can be defined +# mcp_servers: +# aks_mcp: +# description: "The AKS-MCP is a Model Context Protocol (MCP) server that enables AI assistants to interact with Azure Kubernetes Service (AKS) clusters" +# url: "http://localhost:8003/sse" +# +# # try adding your own tools or toggle the built-in toolsets here +# # e.g. query company-specific data, fetch logs from your existing observability tools, etc +# # To check how to add a customized toolset, please refer to https://docs.robusta.dev/master/configuration/holmesgpt/custom_toolsets.html#custom-toolsets +# # To find all built-in toolsets, please refer to https://docs.robusta.dev/master/configuration/holmesgpt/builtin_toolsets.html +# toolsets: +# # add a new json processor toolset +# json_processor: +# description: "A toolset for processing JSON data using jq" +# prerequisites: +# - command: "jq --version" # Ensure jq is installed +# tools: +# - name: "process_json" +# description: "A tool that uses jq to process JSON input" +# command: "echo '{{ json_input }}' | jq '.'" # Example jq command to format JSON +# # disable a built-in toolsets +# aks/core: +# enabled: false +# ``` +# """ diff --git a/src/aks-preview/azext_aks_preview/_params.py b/src/aks-preview/azext_aks_preview/_params.py index 5d11488d6ee..b5b5deb99fe 100644 --- a/src/aks-preview/azext_aks_preview/_params.py +++ b/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, @@ -223,6 +224,7 @@ validate_max_blocked_nodes, validate_resource_group_parameter, validate_location_resource_group_cluster_parameters, + # validate_agent_config_file, ) from azext_aks_preview.azurecontainerstorage._consts import ( CONST_ACSTOR_ALL, @@ -2775,6 +2777,71 @@ def load_arguments(self, _): action="store_true", ) +# pylint: disable=line-too-long +# with self.argument_context("aks agent") as c: +# c.positional( +# "prompt", +# help="Ask any question and answer using available tools.", +# ) +# c.argument( +# "resource_group_name", +# options_list=["--resource-group", "-g"], +# help="Name of resource group.", +# required=False, +# ) +# c.argument( +# "name", +# options_list=["--name", "-n"], +# help="Name of the managed cluster.", +# required=False, +# ) +# 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( +# "config_file", +# default=os.path.join(get_config_dir(), "aksAgent.config"), +# validator=validate_agent_config_file, +# required=False, +# help="Path to the config file.", +# ) +# c.argument( +# "model", +# help="The model to use for the LLM.", +# required=False, +# type=str, +# ) +# c.argument( +# "api-key", +# help="API key to use for the LLM (if not given, uses environment variables AZURE_API_KEY, OPENAI_API_KEY)", +# required=False, +# type=str, +# ) +# c.argument( +# "no_interactive", +# help="Disable interactive mode. When set, the agent will not prompt for input and will run in batch mode.", +# action="store_true", +# ) +# c.argument( +# "no_echo_request", +# help="Disable echoing back the question provided to AKS Agent in the output.", +# action="store_true", +# ) +# c.argument( +# "show_tool_output", +# help="Show the output of each tool that was called.", +# action="store_true", +# ) +# c.argument( +# "refresh_toolsets", +# help="Refresh the toolsets status.", +# action="store_true", +# ) + def _get_default_install_location(exe_name): system = platform.system() diff --git a/src/aks-preview/azext_aks_preview/_validators.py b/src/aks-preview/azext_aks_preview/_validators.py index aad753e004a..1cef9949c38 100644 --- a/src/aks-preview/azext_aks_preview/_validators.py +++ b/src/aks-preview/azext_aks_preview/_validators.py @@ -8,10 +8,12 @@ import os import os.path import re +import yaml from ipaddress import ip_network from math import isclose, isnan from azure.cli.core import keys +from azure.cli.core.api import get_config_dir from azure.cli.core.azclierror import ( ArgumentUsageError, InvalidArgumentValueError, @@ -35,6 +37,7 @@ CONST_NETWORK_POD_IP_ALLOCATION_MODE_STATIC_BLOCK, CONST_NODEPOOL_MODE_GATEWAY, CONST_AZURE_SERVICE_MESH_MAX_EGRESS_NAME_LENGTH, + CONST_AGENT_CONFIG_FILE_NAME, ) from azext_aks_preview._helpers import _fuzzy_match from knack.log import get_logger @@ -977,3 +980,38 @@ def validate_location_resource_group_cluster_parameters(namespace): raise MutuallyExclusiveArgumentError( "Cannot specify --location and --resource-group and --cluster at the same time." ) + + +def _validate_param_yaml_file(yaml_path, param_name): + if not yaml_path: + return + if not os.path.exists(yaml_path): + raise InvalidArgumentValueError( + f"--{param_name}={yaml_path}: file is not found." + ) + if not os.access(yaml_path, os.R_OK): + raise InvalidArgumentValueError( + f"--{param_name}={yaml_path}: file is not readable." + ) + try: + with open(yaml_path, "r") as file: + yaml.safe_load(file) + except yaml.YAMLError as e: + raise InvalidArgumentValueError( + f"--{param_name}={yaml_path}: file is not a valid YAML file: {e}" + ) + except Exception as e: + raise InvalidArgumentValueError( + f"--{param_name}={yaml_path}: An error occurred while reading the config file: {e}" + ) + + +def validate_agent_config_file(namespace): + config_file = namespace.config_file + if not config_file: + return + default_config_path = os.path.join(get_config_dir(), CONST_AGENT_CONFIG_FILE_NAME) + if config_file == default_config_path and not os.path.exists(config_file): + return + + _validate_param_yaml_file(config_file, "config-file") diff --git a/src/aks-preview/azext_aks_preview/agent/__init__.py b/src/aks-preview/azext_aks_preview/agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/aks-preview/azext_aks_preview/agent/agent.py b/src/aks-preview/azext_aks_preview/agent/agent.py new file mode 100644 index 00000000000..b49547b2b29 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/agent/agent.py @@ -0,0 +1,210 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import logging +import os +import socket +import sys +import uuid +from pathlib import Path + +from azext_aks_preview._consts import ( + CONST_AGENT_CONFIG_PATH_DIR_ENV_KEY, + CONST_AGENT_NAME, + CONST_AGENT_NAME_ENV_KEY, +) +from azure.cli.core.api import get_config_dir +from azure.cli.core.commands.client_factory import get_subscription_id +from knack.util import CLIError + +from .prompt import AKS_CONTEXT_PROMPT +from .telemetry import CLITelemetryClient + + +# NOTE(mainred): holmes leverage the log handler RichHandler to provide colorful, readable and well-formatted logs +# making the interactive mode more user-friendly. +# And we removed exising log handlers to avoid duplicate logs. +# Also make the console log consistent, we remove the telemetry and data logger to skip redundant logs. +def init_log(): + # NOTE(mainred): we need to disable INFO logs from LiteLLM before LiteLLM library is loaded, to avoid logging the + # debug logs from heading of LiteLLM. + logging.getLogger("LiteLLM").setLevel(logging.WARNING) + logging.getLogger("telemetry.main").setLevel(logging.WARNING) + logging.getLogger("telemetry.process").setLevel(logging.WARNING) + logging.getLogger("telemetry.save").setLevel(logging.WARNING) + logging.getLogger("telemetry.client").setLevel(logging.WARNING) + logging.getLogger("az_command_data_logger").setLevel(logging.WARNING) + + from holmes.utils.console.logging import init_logging + + # TODO: make log verbose configurable, currently disabled by []. + return init_logging([]) + + +# pylint: disable=too-many-locals +def aks_agent( + cmd, + resource_group_name, + name, + prompt, + model, + api_key, + max_steps, + config_file, + no_interactive, + no_echo_request, + show_tool_output, + refresh_toolsets, +): + """ + Interact with the AKS agent using a prompt or piped input. + + :param prompt: The prompt to send to the agent. + :type prompt: str + :param model: The model to use for the LLM. + :type model: str + :param max_steps: Maximum number of steps to take. + :type max_steps: int + :param config_file: Path to the config file. + :type config_file: str + :param no_interactive: Disable interactive mode. + :type no_interactive: bool + :param no_echo_request: Disable echoing back the question provided to AKS Agent in the output. + :type no_echo_request: bool + :param show_tool_output: Whether to show tool output. + :type show_tool_output: bool + :param refresh_toolsets: Refresh the toolsets status. + :type refresh_toolsets: bool + """ + with CLITelemetryClient(): + + if sys.version_info < (3, 10): + raise CLIError( + "Please upgrade the python version to 3.10 or above to use aks agent." + ) + + # reverse the value of the variables so that + interactive = not no_interactive + echo = not no_echo_request + + console = init_log() + + os.environ[CONST_AGENT_CONFIG_PATH_DIR_ENV_KEY] = get_config_dir() + # Holmes library allows the user to specify the agent name through environment variable + # before loading the library. + + os.environ[CONST_AGENT_NAME_ENV_KEY] = CONST_AGENT_NAME + + from holmes.config import Config + from holmes.core.prompt import build_initial_ask_messages + from holmes.interactive import run_interactive_loop + from holmes.plugins.destinations import DestinationType + from holmes.plugins.interfaces import Issue + from holmes.plugins.prompts import load_and_render_prompt + from holmes.utils.console.result import handle_result + + # 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 + + expanded_config_file = Path(os.path.expanduser(config_file)) + + config = Config.load_from_file( + expanded_config_file, + model=model, + api_key=api_key, + max_steps=max_steps, + ) + + ai = config.create_console_toolcalling_llm( + dal=None, + refresh_toolsets=refresh_toolsets, + ) + console.print( + "[bold yellow]This tool uses AI to generate responses and may not always be accurate.[bold yellow]" + ) + + if not prompt and not interactive and not piped_data: + raise CLIError( + "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) + + subscription_id = get_subscription_id(cmd.cli_ctx) + + aks_template_context = { + "cluster_name": name, + "resource_group": resource_group_name, + "subscription_id": subscription_id, + } + + aks_context_prompt = load_and_render_prompt( + AKS_CONTEXT_PROMPT, aks_template_context + ) + + # Variables not exposed to the user. + # Adds a prompt for post processing. + post_processing_prompt = None + # File to append to prompt + + include_file = None + if interactive: + run_interactive_loop( + ai, + console, + prompt, + include_file, + post_processing_prompt, + show_tool_output=show_tool_output, + system_prompt_additions=aks_context_prompt, + ) + return + + messages = build_initial_ask_messages( + console, + prompt, + include_file, + ai.tool_executor, + config.get_runbook_catalog(), + system_prompt_additions=aks_context_prompt, + ) + + response = ai.call(messages) + + messages = response.messages + + 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, + ) diff --git a/src/aks-preview/azext_aks_preview/agent/prompt.py b/src/aks-preview/azext_aks_preview/agent/prompt.py new file mode 100644 index 00000000000..4f40ac24049 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/agent/prompt.py @@ -0,0 +1,94 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +AKS_CONTEXT_PROMPT = """ +# AKS-Specific Context and Workflow + +You are now operating in Azure Kubernetes Service (AKS) mode. All investigations must consider both Azure control plane and Kubernetes data plane components. + +## AKS Context Requirements + +### MANDATORY: Establish AKS Cluster Context +Before any troubleshooting, you MUST establish and validate the AKS cluster context: + +{% if cluster_name and resource_group %} +**User-provided context:** +- Cluster: `{{cluster_name}}` +- Resource Group: `{{resource_group}}` +- Subscription: `{{subscription_id}}` + +⚠️ **MANDATORY Validation** - You MUST perform ALL Context Validation Steps below before proceeding with any investigation. Do not skip validation even when context is provided by the user. +{% else %} +**Auto-discovery required** - Detect AKS context using this priority: + +1. **Primary method**: Check if `aks/core` toolset is available in your toolsets + - If available, use the `aks/core` tools to get cluster context directly + - This is the preferred method as it provides the most reliable context discovery +2. **Fallback method**: If `aks/core` toolset is not available: + - Get current Azure subscription ID + - Extract AKS cluster name from current kubeconfig context + - Find resource group by listing AKS clusters with matching name in the subscription + +**Critical**: You MUST first check toolset availability before choosing the discovery method. + +**Error handling:** If discovery fails (empty response, errors, or toolset unavailable), you MUST: +1. **IMMEDIATELY STOP ALL OPERATIONS** - Do not proceed with any investigation +2. **DO NOT ATTEMPT ANY TROUBLESHOOTING** - No kubectl commands, no Azure commands, nothing +3. **DO NOT INFER THE RESOURCE NAME** - Do not assume any resource name, resource group, or subscription ID +4. **ONLY display the context failure message** on separate lines: +``` +Cluster name: +Resource group: +Subscription ID: + +Please provide the correct cluster context. You can either: +1. Specify the context in this session: "Please use cluster 'my-cluster' in resource group 'my-rg' under subscription 'my-subscription'" +2. Or restart with context: `az aks agent --name --resource-group --subscription ` +``` +**IMPORTANT**: When displaying the CLI command example above, use it EXACTLY as written with the placeholder format ``, ``, ``. + +{% endif %} + +### Context Validation Steps - MANDATORY FOR ALL SCENARIOS +**These steps MUST be performed whether context is user-provided or auto-discovered:** + +1. **Verify cluster exists** in specified resource group/subscription: + - Confirm the AKS cluster can be found under the resource group and subscription + - If cluster is not found, STOP and report the validation failure +2. **Check kubeconfig context** - ensure the current kubectl context matches the target AKS cluster: + - **MANDATORY**: This step MUST be performed even if you're only checking Azure resources + - Get current kubectl context: `kubectl config current-context` + - **ONLY if context doesn't match the target AKS cluster name**: + a. **Attempt to download credentials**: Use `az aks get-credentials` to download cluster credentials + b. **If credential download fails or no tool is available**, you MUST instruct the user to manually download credentials: + ``` + Please manually download AKS credentials: + az aks get-credentials --resource-group {{resource_group}} --name {{cluster_name}} --subscription {{subscription_id}} + ``` + c. **Attempt to switch the kubernetes context**: Use `kubectl config use-context` command (NEVER use `run_bash_command` tool to switch context) + d. **If context switch fails or no tool is available**, you MUST instruct the user to manually switch context: + ``` + Please manually switch to the correct kubectl context: + kubectl config use-context {{cluster_name}} + ``` + - **Verify the current context is now set to the cluster name**: Run `kubectl config current-context` and confirm it matches the target AKS cluster name + - **If context already matches**: Skip credential download and proceed + - **This ensures the kubectl context is actively switched to the target cluster for any future Kubernetes operations in the session** + +**CRITICAL**: Before performing ANY Kubernetes operations (kubectl commands, checking pods, services, deployments, etc.), you MUST ALWAYS verify that the current kubectl context matches the target AKS cluster name. If it doesn't match, you MUST download the correct credentials and switch context before proceeding. This validation is required EVERY TIME you need to interact with Kubernetes resources, even if you've already validated Azure resources in the same session. + +**Only proceed with investigation after ALL validation steps pass successfully.** + +### AKS Investigation Approach +- **Start with cluster health** (nodes, system pods, control plane) +- **Check Azure-specific components** (load balancers, NSGs, managed identity) +- **Check Kubernetes-specific components** (deployments, services, ingress, namespaces, RBAC) +- **Analyze both Azure and Kubernetes logs** +- **Use AKS-aware tools** from available toolsets +- **Consider AKS limitations and best practices** + +**Note**: "Cluster" in this context refers to both the Azure-managed AKS cluster AND the Kubernetes resources running within it. Both layers must be validated before proceeding. + +""" diff --git a/src/aks-preview/azext_aks_preview/agent/telemetry.py b/src/aks-preview/azext_aks_preview/agent/telemetry.py new file mode 100644 index 00000000000..67025fd4903 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/agent/telemetry.py @@ -0,0 +1,75 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import datetime +import logging +import os +import platform + +from applicationinsights import TelemetryClient +from azure.cli.core.telemetry import _get_hash_mac_address, _get_user_agent + +DEFAULT_INSTRUMENTATION_KEY = "c301e561-daea-42d9-b9d1-65fca4166704" +APPLICATIONINSIGHTS_INSTRUMENTATION_KEY_ENV = "APPLICATIONINSIGHTS_INSTRUMENTATION_KEY" + + +class CLITelemetryClient: + def __init__(self): + instrumentation_key = self._get_application_insights_instrumentation_key() + self._telemetry_client = TelemetryClient( + instrumentation_key=instrumentation_key + ) + self.start_time = datetime.datetime.utcnow() + self.end_time = "" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.end_time = datetime.datetime.utcnow() + self.track_agent_started() + self.flush() + + def track(self, event_name, properties=None): + if properties is None: + properties = {} + properties.update(self._generate_payload()) + self._telemetry_client.track_trace(event_name, properties, logging.INFO) + + def track_agent_started(self): + timestamp_properties = { + "time.start": str(self.start_time), + "time.end": str(self.end_time), + } + self.track("AgentCLIStartup", properties=timestamp_properties) + + def flush(self): + self._telemetry_client.flush() + + def _generate_payload(self): + extension_name = "aks-preview" + try: + from azure.cli.core.extension import get_extension + + ext_name = "aks-preview" + ext = get_extension(ext_name) + extension_name = f"aks-preview@{ext.version}" + except: # pylint: disable=W0702 + pass + + return { + "device.id": _get_hash_mac_address(), + "service.name": "aks agent", + "OS.Type": platform.system().lower(), # eg. darwin, windows + "OS.Version": platform.version().lower(), # eg. 10.0.14942 + "OS.Platform": platform.platform().lower(), # eg. windows-10-10.0.19041-sp0 + "UserAgent": _get_user_agent(), + "extensionname": extension_name, # extension and version + } + + def _get_application_insights_instrumentation_key(self) -> str: + return os.getenv( + APPLICATIONINSIGHTS_INSTRUMENTATION_KEY_ENV, DEFAULT_INSTRUMENTATION_KEY + ) diff --git a/src/aks-preview/azext_aks_preview/commands.py b/src/aks-preview/azext_aks_preview/commands.py index 2c92399a949..f40ccf6cb15 100644 --- a/src/aks-preview/azext_aks_preview/commands.py +++ b/src/aks-preview/azext_aks_preview/commands.py @@ -188,6 +188,7 @@ def load_command_table(self, _): "operation-abort", "aks_operation_abort", supports_no_wait=True ) g.custom_command("bastion", "aks_bastion") + # g.custom_command("agent", "aks_agent") # AKS maintenance configuration commands with self.command_group( diff --git a/src/aks-preview/azext_aks_preview/custom.py b/src/aks-preview/azext_aks_preview/custom.py index a367a4207c7..d9a17887244 100644 --- a/src/aks-preview/azext_aks_preview/custom.py +++ b/src/aks-preview/azext_aks_preview/custom.py @@ -85,6 +85,8 @@ add_virtual_node_role_assignment, enable_addons, ) +from azext_aks_preview.agent.agent import aks_agent as aks_agent_internal + from azext_aks_preview.aks_diagnostics import aks_kanalyze_cmd, aks_kollect_cmd from azext_aks_preview.aks_draft.commands import ( aks_draft_cmd_create, @@ -4398,3 +4400,36 @@ def aks_bastion(cmd, client, resource_group_name, name, bastion=None, port=None, aks_identity_binding_delete = aks_ib_cmd_delete aks_identity_binding_show = aks_ib_cmd_show aks_identity_binding_list = aks_ib_cmd_list + + +# pylint: disable=unused-argument +def aks_agent( + cmd, + client, + prompt, + model, + max_steps, + config_file, + resource_group_name=None, + name=None, + api_key=None, + no_interactive=False, + no_echo_request=False, + show_tool_output=False, + refresh_toolsets=False, +): + + aks_agent_internal( + cmd, + resource_group_name, + name, + prompt, + model, + api_key, + max_steps, + config_file, + no_interactive, + no_echo_request, + show_tool_output, + refresh_toolsets, + ) diff --git a/src/aks-preview/setup.py b/src/aks-preview/setup.py index f5dd594b79e..9dc62a95f23 100644 --- a/src/aks-preview/setup.py +++ b/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.0b26" @@ -23,7 +23,9 @@ "License :: OSI Approved :: MIT License", ] -DEPENDENCIES = [] +DEPENDENCIES = [ + "holmesgpt==0.12.4; python_version >= '3.10'", +] with open1("README.rst", "r", encoding="utf-8") as f: README = f.read()