From fc06e8e929d54e2ca282479e46e8ad7fa2886de7 Mon Sep 17 00:00:00 2001 From: Mohse Morad Date: Thu, 30 Jan 2025 12:59:38 +0200 Subject: [PATCH] Add better example config for tools --- holmes/core/tools.py | 3 + .../toolsets/grafana/base_grafana_toolset.py | 18 ++++- holmes/plugins/toolsets/grafana/common.py | 13 +-- .../toolsets/grafana/toolset_grafana_loki.py | 20 ++++- holmes/plugins/toolsets/opensearch.py | 81 +++++++++++++++---- .../default_toolset_installation_guide.jinja2 | 4 + holmes/utils/holmes_sync_toolsets.py | 19 ++--- 7 files changed, 113 insertions(+), 45 deletions(-) diff --git a/holmes/core/tools.py b/holmes/core/tools.py index eff1fed2..5030d120 100644 --- a/holmes/core/tools.py +++ b/holmes/core/tools.py @@ -377,6 +377,9 @@ def check_prerequisites(self): self._status = ToolsetStatusEnum.ENABLED + def get_example_config(self) -> Dict[str, Any]: + return {} + class YAMLToolset(Toolset): tools: List[YAMLTool] diff --git a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py index 9f29efc3..dedb1743 100644 --- a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py +++ b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py @@ -1,5 +1,5 @@ import logging -from typing import Any +from typing import Any, ClassVar, Type from holmes.core.tools import ( Tool, Toolset, @@ -11,6 +11,8 @@ class BaseGrafanaToolset(Toolset): + config_class: ClassVar[Type[GrafanaConfig]] = GrafanaConfig + def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool]): super().__init__( name=name, @@ -21,7 +23,8 @@ def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool] tags=[ ToolsetTag.CORE, ], - enabled=False + enabled=False, + is_default=True, ) def prerequisites_callable(self, config: dict[str, Any]) -> bool: @@ -30,10 +33,17 @@ def prerequisites_callable(self, config: dict[str, Any]) -> bool: return False try: - self._grafana_config = GrafanaConfig(**config) - is_healthy = get_health(self._grafana_config.url, self._grafana_config.api_key) + self._grafana_config = BaseGrafanaToolset.config_class(**config) + is_healthy = get_health( + self._grafana_config.url, self._grafana_config.api_key + ) return is_healthy except Exception: logging.exception("Failed to set up grafana toolset") return False + + def get_example_config(self): + example_config = GrafanaConfig(api_key="YOUR API KEY", url="YOUR GRAFANA URL") + return example_config.model_dump() + \ No newline at end of file diff --git a/holmes/plugins/toolsets/grafana/common.py b/holmes/plugins/toolsets/grafana/common.py index e5c5e1a9..df1c2144 100644 --- a/holmes/plugins/toolsets/grafana/common.py +++ b/holmes/plugins/toolsets/grafana/common.py @@ -1,23 +1,12 @@ from typing import Dict, Optional, Union import uuid import time -import os from pydantic import BaseModel - -GRAFANA_URL_ENV_NAME = "GRAFANA_URL" -GRAFANA_API_KEY_ENV_NAME = "GRAFANA_API_KEY" ONE_HOUR_IN_SECONDS = 3600 -class GrafanaLokiConfig(BaseModel): - pod_name_search_key: str = "pod" - namespace_search_key: str = "namespace" - node_name_search_key: str = "node" - - class GrafanaConfig(BaseModel): - loki: GrafanaLokiConfig = GrafanaLokiConfig() api_key: str url: str @@ -61,5 +50,5 @@ def get_datasource_id(dict: Dict, param: str) -> str: return f"uid/{datasource_id}" except: pass - + return datasource_id diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py index 0023359d..7816e85b 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py @@ -5,6 +5,7 @@ from holmes.core.tools import Tool, ToolParameter from holmes.plugins.toolsets.grafana.base_grafana_toolset import BaseGrafanaToolset from holmes.plugins.toolsets.grafana.common import ( + GrafanaConfig, get_datasource_id, get_param_or_raise, process_timestamps, @@ -17,6 +18,12 @@ ) +class GrafanaLokiConfig(GrafanaConfig): + pod_name_search_key: str = "pod" + namespace_search_key: str = "namespace" + node_name_search_key: str = "node" + + class ListLokiDatasources(Tool): def __init__(self, toolset: BaseGrafanaToolset): @@ -84,7 +91,7 @@ def invoke(self, params: Dict) -> str: api_key=self._toolset._grafana_config.api_key, loki_datasource_id=get_datasource_id(params, "loki_datasource_id"), node_name=get_param_or_raise(params, "node_name"), - node_name_search_key=self._toolset._grafana_config.loki.node_name_search_key, + node_name_search_key=self._toolset._grafana_config.node_name_search_key, start=start, end=end, limit=int(get_param_or_raise(params, "limit")), @@ -208,8 +215,8 @@ def invoke(self, params: Dict) -> str: loki_datasource_id=get_datasource_id(params, "loki_datasource_id"), pod_regex=get_param_or_raise(params, "pod_regex"), namespace=get_param_or_raise(params, "namespace"), - namespace_search_key=self._toolset._grafana_config.loki.namespace_search_key, - pod_name_search_key=self._toolset._grafana_config.loki.pod_name_search_key, + namespace_search_key=self._toolset._grafana_config.namespace_search_key, + pod_name_search_key=self._toolset._grafana_config.pod_name_search_key, start=start, end=end, limit=int(get_param_or_raise(params, "limit")), @@ -221,6 +228,8 @@ def get_parameterized_one_liner(self, params: Dict) -> str: class GrafanaLokiToolset(BaseGrafanaToolset): + config_class = GrafanaLokiConfig + def __init__(self): super().__init__( name="grafana/loki", @@ -233,3 +242,8 @@ def __init__(self): GetLokiLogsByLabel(self), ], ) + + def get_example_config(self): + example_config = GrafanaLokiConfig(api_key="YOUR API KEY", url="YOUR GRAFANA URL") + return example_config.model_dump() + \ No newline at end of file diff --git a/holmes/plugins/toolsets/opensearch.py b/holmes/plugins/toolsets/opensearch.py index 84d4c73e..60e27147 100644 --- a/holmes/plugins/toolsets/opensearch.py +++ b/holmes/plugins/toolsets/opensearch.py @@ -1,7 +1,7 @@ import logging from typing import Any, Dict, List, Optional -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from holmes.core.tools import ( CallablePrerequisite, Tool, @@ -12,18 +12,46 @@ from opensearchpy import OpenSearch +class OpenSearchHttpAuth(BaseModel): + username: str + password: str + + +class OpenSearchHost(BaseModel): + host: str + port: int = 9200 + + +class OpenSearchCluster(BaseModel): + hosts: list[OpenSearchHost] + headers: Optional[dict[str, Any]] = None + use_ssl: bool = True + ssl_assert_hostname: bool = False + verify_certs: bool = False + ssl_show_warn: bool = False + http_auth: Optional[OpenSearchHttpAuth] = None + + +class OpenSearchConfig(BaseModel): + opensearch_clusters: list[OpenSearchCluster] + + class OpenSearchClient: def __init__(self, **kwargs): - + # Handle http_auth explicitly if "http_auth" in kwargs: http_auth = kwargs.pop("http_auth") if isinstance(http_auth, dict): - kwargs["http_auth"] = (http_auth.get("username"), http_auth.get("password")) + kwargs["http_auth"] = ( + http_auth.get("username"), + http_auth.get("password"), + ) # Initialize OpenSearch client self.client = OpenSearch(**kwargs) -def get_client(clients:List[OpenSearchClient], host:Optional[str]): + +def get_client(clients: List[OpenSearchClient], host: Optional[str]): if len(clients) == 1: return clients[0] @@ -133,7 +161,7 @@ def __init__(self): enabled=False, description="Provide cluster metadata information like health, shards, settings.", docs_url="https://opensearch.org/docs/latest/clients/python-low-level/", - icon_url="https://upload.wikimedia.org/wikipedia/commons/9/91/Opensearch_Logo.svg", + icon_url="https://opensearch.org/assets/brand/PNG/Mark/opensearch_mark_default.png", prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)], tools=[ ListShards(self), @@ -143,21 +171,40 @@ def __init__(self): tags=[ ToolsetTag.CORE, ], - is_default=False, + is_default=True, ) def prerequisites_callable(self, config: dict[str, Any]) -> bool: if not config: return False - clusters_configs: list[dict[str, Any]] = config.get("opensearch_clusters", []) - for cluster in clusters_configs: - try: - logging.info(f"Setting up OpenSearch client") - client = OpenSearchClient(**cluster) - if client.client.cluster.health(params={"timeout": 5}): - self.clients.append(client) - except Exception: - logging.exception("Failed to set up opensearch client") - - return len(self.clients) > 0 + try: + os_config = OpenSearchConfig(**config) + + for cluster in os_config.opensearch_clusters: + try: + logging.info(f"Setting up OpenSearch client") + cluster_kwargs = cluster.model_dump() + client = OpenSearchClient(**cluster_kwargs) + if client.client.cluster.health(params={"timeout": 5}): + self.clients.append(client) + except Exception: + logging.exception("Failed to set up opensearch client") + + return len(self.clients) > 0 + except Exception: + logging.exception("Failed to set up grafana toolset") + return False + + def get_example_config(self) -> Dict[str, Any]: + example_config = OpenSearchConfig( + opensearch_clusters=[ + OpenSearchCluster( + hosts=[OpenSearchHost(host="YOUR OPENSEACH HOST")], + headers={"Authorization": "{{ env.OPENSEARCH_BEARER_TOKEN }}"}, + use_ssl=True, + ssl_assert_hostname=False, + ) + ] + ) + return example_config.model_dump() diff --git a/holmes/utils/default_toolset_installation_guide.jinja2 b/holmes/utils/default_toolset_installation_guide.jinja2 index 1cb32dd3..44014ffe 100644 --- a/holmes/utils/default_toolset_installation_guide.jinja2 +++ b/holmes/utils/default_toolset_installation_guide.jinja2 @@ -25,6 +25,10 @@ holmes: toolsets: {{toolset_name}}: enabled: true + {% if example_config %} + config: + {{ example_config | indent(8) }} + {% endif %} ``` {% endif %} diff --git a/holmes/utils/holmes_sync_toolsets.py b/holmes/utils/holmes_sync_toolsets.py index a485b038..7ddec7e2 100644 --- a/holmes/utils/holmes_sync_toolsets.py +++ b/holmes/utils/holmes_sync_toolsets.py @@ -1,4 +1,7 @@ from datetime import datetime +from typing import Any + +import yaml from holmes.config import Config @@ -47,19 +50,17 @@ def holmes_sync_toolsets_status(dal: SupabaseDal, config: Config) -> None: def render_default_installation_instructions_for_toolset(toolset: Toolset) -> str: env_vars = toolset.get_environment_variables() - context = { + context: dict[str, Any] = { "env_vars": env_vars if env_vars else [], "toolset_name": toolset.name, "enabled": toolset.enabled, - "default_toolset": toolset.is_default, + "example_config": yaml.dump(toolset.get_example_config()), } - if toolset.is_default: - installation_instructions = load_and_render_prompt( - "file://holmes/utils/default_toolset_installation_guide.jinja2", context - ) - return installation_instructions - installation_instructions = load_and_render_prompt( - "file://holmes/utils/installation_guide.jinja2", context + template = ( + "file://holmes/utils/default_toolset_installation_guide.jinja2" + if toolset.is_default + else "file://holmes/utils/installation_guide.jinja2" ) + installation_instructions = load_and_render_prompt(template, context) return installation_instructions