holmesgpt 0.11.5__py3-none-any.whl

holmes/plugins/toolsets/opensearch/opensearch.py

@@ -0,0 +1,245 @@
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+from opensearchpy import OpenSearch
+from pydantic import BaseModel, ConfigDict
+
+from holmes.core.tools import (
+    CallablePrerequisite,
+    StructuredToolResult,
+    Tool,
+    ToolParameter,
+    ToolResultStatus,
+    Toolset,
+    ToolsetTag,
+)
+from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR
+
+
+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"),
+                )
+        # Initialize OpenSearch client
+        self.hosts = [host.get("host") for host in kwargs.get("hosts", [])]
+        self.client = OpenSearch(**kwargs)
+
+
+def get_client(clients: List[OpenSearchClient], host: Optional[str]):
+    if len(clients) == 1:
+        return clients[0]
+
+    if not host:
+        raise Exception("Missing host to resolve opensearch client")
+
+    for client in clients:
+        found = any(host in client.hosts for client in clients)
+        if found:
+            return client
+
+    raise Exception(
+        f"Failed to resolve opensearch client. Could not find a matching host: {host}"
+    )
+
+
+class BaseOpenSearchTool(Tool):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    toolset: "OpenSearchToolset"
+
+
+class ListShards(BaseOpenSearchTool):
+    def __init__(self, toolset: "OpenSearchToolset"):
+        super().__init__(
+            name="opensearch_list_shards",
+            description="List the shards within an opensearch cluster",
+            parameters={
+                "host": ToolParameter(
+                    description="The cluster host",
+                    type="string",
+                    required=True,
+                )
+            },
+            toolset=toolset,
+        )
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        client = get_client(self.toolset.clients, host=params.get("host", ""))
+        shards = client.client.cat.shards()
+        return StructuredToolResult(
+            status=ToolResultStatus.SUCCESS,
+            data=str(shards),
+            params=params,
+        )
+
+    def get_parameterized_one_liner(self, params: Dict) -> str:
+        return f"opensearch ListShards({params.get('host')})"
+
+
+class GetClusterSettings(BaseOpenSearchTool):
+    def __init__(self, toolset: "OpenSearchToolset"):
+        super().__init__(
+            name="opensearch_get_cluster_settings",
+            description="Retrieve the cluster's settings",
+            parameters={
+                "host": ToolParameter(
+                    description="The cluster host",
+                    type="string",
+                    required=True,
+                )
+            },
+            toolset=toolset,
+        )
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        client = get_client(self.toolset.clients, host=params.get("host"))
+        response = client.client.cluster.get_settings(
+            include_defaults=True, flat_settings=True
+        )
+        return StructuredToolResult(
+            status=ToolResultStatus.SUCCESS,
+            data=str(response),
+            params=params,
+        )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return f"opensearch GetClusterSettings({params.get('host')})"
+
+
+class GetClusterHealth(BaseOpenSearchTool):
+    def __init__(self, toolset: "OpenSearchToolset"):
+        super().__init__(
+            name="opensearch_get_cluster_health",
+            description="Retrieve the cluster's health",
+            parameters={
+                "host": ToolParameter(
+                    description="The cluster host",
+                    type="string",
+                    required=True,
+                )
+            },
+            toolset=toolset,
+        )
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        client = get_client(self.toolset.clients, host=params.get("host", ""))
+        health = client.client.cluster.health()
+        return StructuredToolResult(
+            status=ToolResultStatus.SUCCESS,
+            data=str(health),
+            params=params,
+        )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return f"opensearch GetClusterHealth({params.get('host')})"
+
+
+class ListOpenSearchHosts(BaseOpenSearchTool):
+    def __init__(self, toolset: "OpenSearchToolset"):
+        super().__init__(
+            name="opensearch_list_hosts",
+            description="List all OpenSearch hosts in the cluster",
+            parameters={},
+            toolset=toolset,
+        )
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        hosts = [host for client in self.toolset.clients for host in client.hosts]
+        return StructuredToolResult(
+            status=ToolResultStatus.SUCCESS,
+            data=str(hosts),
+            params=params,
+        )
+
+    def get_parameterized_one_liner(self, params: Dict) -> str:
+        return "opensearch ListOpenSearchHosts()"
+
+
+class OpenSearchToolset(Toolset):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    clients: List[OpenSearchClient] = []
+
+    def __init__(self):
+        super().__init__(
+            name="opensearch/status",
+            enabled=False,
+            description="Provide cluster metadata information like health, shards, settings.",
+            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/opensearch.html",
+            icon_url="https://opensearch.org/assets/brand/PNG/Mark/opensearch_mark_default.png",
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
+            tools=[
+                ListShards(self),
+                GetClusterSettings(self),
+                GetClusterHealth(self),
+                ListOpenSearchHosts(self),
+            ],
+            tags=[
+                ToolsetTag.CORE,
+            ],
+        )
+
+    def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
+        if not config:
+            return False, TOOLSET_CONFIG_MISSING_ERROR
+
+        try:
+            os_config = OpenSearchConfig(**config)
+            errors = []
+            for cluster in os_config.opensearch_clusters:
+                try:
+                    logging.info("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 as e:
+                    message = f"Failed to set up opensearch client {str(cluster.hosts)}. {str(e)}"
+                    logging.exception(message)
+                    errors.append(message)
+
+            return len(self.clients) > 0, "\n".join(errors)
+        except Exception as e:
+            logging.exception("Failed to set up OpenSearch toolset")
+            return False, str(e)
+
+    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()

holmes/plugins/toolsets/opensearch/opensearch_logs.py

@@ -0,0 +1,151 @@
+import logging
+from typing import Any, Dict, Optional, Tuple
+
+import requests  # type: ignore
+from requests import RequestException  # type: ignore
+from urllib.parse import urljoin
+
+from holmes.core.tools import (
+    CallablePrerequisite,
+    StructuredToolResult,
+    ToolResultStatus,
+    ToolsetTag,
+)
+from holmes.plugins.toolsets.logging_utils.logging_api import (
+    BasePodLoggingToolset,
+    FetchPodLogsParams,
+    PodLoggingTool,
+    process_time_parameters,
+)
+from holmes.plugins.toolsets.opensearch.opensearch_utils import (
+    OpenSearchLoggingConfig,
+    add_auth_header,
+    build_query,
+    format_logs,
+    opensearch_health_check,
+)
+
+LOGS_FIELDS_CACHE_KEY = "cached_logs_fields"
+
+
+class OpenSearchLogsToolset(BasePodLoggingToolset):
+    """Implementation of the unified logging API for OpenSearch logs"""
+
+    def __init__(self):
+        super().__init__(
+            name="opensearch/logs",
+            description="OpenSearch integration to fetch logs",
+            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/opensearch_logs.html",
+            icon_url="https://opensearch.org/wp-content/uploads/2025/01/opensearch_mark_default.png",
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
+            tools=[
+                PodLoggingTool(self),
+            ],
+            tags=[
+                ToolsetTag.CORE,
+            ],
+        )
+
+    def get_example_config(self) -> Dict[str, Any]:
+        example_config = OpenSearchLoggingConfig(
+            opensearch_url="YOUR OPENSEARCH LOGS URL",
+            index_pattern="YOUR OPENSEARCH LOGS INDEX NAME",
+            opensearch_auth_header="YOUR OPENSEARCH LOGS AUTH HEADER (Optional)",
+        )
+        return example_config.model_dump()
+
+    def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
+        if not config:
+            return False, "Missing OpenSearch configuration. Check your config."
+
+        self.config = OpenSearchLoggingConfig(**config)
+
+        return opensearch_health_check(self.config)
+
+    @property
+    def opensearch_config(self) -> Optional[OpenSearchLoggingConfig]:
+        return self.config
+
+    def fetch_pod_logs(self, params: FetchPodLogsParams) -> StructuredToolResult:
+        if not self.opensearch_config:
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error="Missing OpenSearch configuration",
+                params=params.model_dump(),
+            )
+
+        try:
+            start_time = None
+            end_time = None
+            if params.start_time or params.end_time:
+                start_time, end_time = process_time_parameters(
+                    params.start_time, params.end_time
+                )
+
+            query = build_query(
+                config=self.opensearch_config,
+                namespace=params.namespace,
+                pod_name=params.pod_name,
+                start_time=start_time,
+                end_time=end_time,
+                match=params.filter,
+                limit=params.limit,
+            )
+
+            headers = {"Content-Type": "application/json"}
+            headers.update(
+                add_auth_header(self.opensearch_config.opensearch_auth_header)
+            )
+
+            url = urljoin(
+                self.opensearch_config.opensearch_url,
+                f"/{self.opensearch_config.index_pattern}/_search",
+            )
+            logs_response = requests.post(
+                url=url,
+                timeout=180,
+                verify=True,
+                json=query,
+                headers=headers,
+            )
+
+            if logs_response.status_code == 200:
+                response = logs_response.json()
+                logs = format_logs(
+                    logs=response.get("hits", {}).get("hits", []),
+                    config=self.opensearch_config,
+                )
+                return StructuredToolResult(
+                    status=ToolResultStatus.SUCCESS,
+                    data=logs,
+                    params=params.model_dump(),
+                )
+            else:
+                return StructuredToolResult(
+                    status=ToolResultStatus.ERROR,
+                    return_code=logs_response.status_code,
+                    error=logs_response.text,
+                    params=params.model_dump(),
+                )
+
+        except requests.Timeout:
+            logging.warning("Timeout while fetching OpenSearch logs", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error="Request timed out while fetching OpenSearch logs",
+                params=params.model_dump(),
+            )
+        except RequestException as e:
+            logging.warning("Failed to fetch OpenSearch logs", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Network error while fetching OpenSearch logs: {str(e)}",
+                params=params.model_dump(),
+            )
+        except Exception as e:
+            logging.warning("Failed to process OpenSearch logs", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Unexpected error: {str(e)}",
+                params=params.model_dump(),
+            )

holmes/plugins/toolsets/opensearch/opensearch_traces.py

@@ -0,0 +1,211 @@
+import os
+import logging
+
+from typing import Any, Dict
+
+import requests  # type: ignore
+from cachetools import TTLCache  # type: ignore
+from holmes.core.tools import (
+    CallablePrerequisite,
+    Tool,
+    ToolParameter,
+    ToolsetTag,
+)
+import json
+from requests import RequestException
+
+from holmes.plugins.toolsets.opensearch.opensearch_utils import (
+    BaseOpenSearchToolset,
+    add_auth_header,
+    get_search_url,
+)
+from holmes.core.tools import StructuredToolResult, ToolResultStatus
+
+TRACES_FIELDS_CACHE_KEY = "cached_traces_fields"
+
+
+class GetTracesFields(Tool):
+    def __init__(self, toolset: BaseOpenSearchToolset):
+        super().__init__(
+            name="get_traces_fields",
+            description="Get all the fields in the traces documents",
+            parameters={},
+        )
+        self._toolset = toolset
+        self._cache = None
+
+    def _invoke(self, params: Dict) -> StructuredToolResult:
+        try:
+            if not self._cache and self._toolset.opensearch_config.fields_ttl_seconds:
+                self._cache = TTLCache(
+                    maxsize=5, ttl=self._toolset.opensearch_config.fields_ttl_seconds
+                )
+
+            if self._cache:
+                cached_response = self._cache.get(TRACES_FIELDS_CACHE_KEY, None)
+                if cached_response:
+                    logging.debug("traces fields returned from cache")
+                    return StructuredToolResult(
+                        status=ToolResultStatus.SUCCESS,
+                        data=cached_response,
+                        params=params,
+                    )
+
+            body = {
+                "size": 1,
+                "_source": False,
+                "script_fields": {
+                    "all_fields": {
+                        "script": {
+                            "lang": "painless",
+                            "source": "Map fields = new HashMap(); List stack = new ArrayList(); stack.add(['prefix': '', 'obj': params['_source']]); while (!stack.isEmpty()) { Map item = stack.remove(stack.size() - 1); String prefix = item.get('prefix'); Map obj = item.get('obj'); for (entry in obj.entrySet()) { String fullPath = prefix.isEmpty() ? entry.getKey() : prefix + '.' + entry.getKey(); fields.put(fullPath, true); if (entry.getValue() instanceof Map) { stack.add(['prefix': fullPath, 'obj': entry.getValue()]); } } } return fields.keySet();",
+                        }
+                    }
+                },
+            }
+            headers = {"Content-Type": "application/json"}
+            headers.update(
+                add_auth_header(self._toolset.opensearch_config.opensearch_auth_header)
+            )
+            logs_response = requests.get(
+                url=get_search_url(self._toolset.opensearch_config),
+                timeout=180,
+                verify=True,
+                data=json.dumps(body),
+                headers=headers,
+            )
+            logs_response.raise_for_status()
+            response = json.dumps(logs_response.json())
+            if self._cache:
+                self._cache[TRACES_FIELDS_CACHE_KEY] = response
+            return StructuredToolResult(
+                status=ToolResultStatus.SUCCESS,
+                data=response,
+                params=params,
+            )
+        except requests.Timeout:
+            logging.warning(
+                "Timeout while fetching opensearch traces fields", exc_info=True
+            )
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error="Request timed out while fetching opensearch traces fields",
+                params=params,
+            )
+        except RequestException as e:
+            logging.warning("Failed to fetch opensearch traces fields", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Network error while opensearch traces fields: {str(e)}",
+                params=params,
+            )
+        except Exception as e:
+            logging.warning("Failed to process opensearch traces fields", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Unexpected error: {str(e)}",
+                params=params,
+            )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return "list traces documents fields"
+
+
+class TracesSearchQuery(Tool):
+    def __init__(self, toolset: "OpenSearchTracesToolset"):
+        super().__init__(
+            name="traces_in_range_search",
+            description="Get traces in a specified time range for an opensearch query",
+            parameters={
+                "query": ToolParameter(
+                    description="An OpenSearch search query. It should be a stringified json, matching opensearch search syntax. "
+                    "The query must contain a 'range' on the startTimeMillis field. From a given startTimeMillis, until a given startTimeMillis. This is time in milliseconds",
+                    type="string",
+                ),
+            },
+        )
+        self._toolset = toolset
+        self._cache = None
+
+    def _invoke(self, params: Any) -> StructuredToolResult:
+        err_msg = ""
+        try:
+            body = json.loads(params.get("query"))
+            full_query = body
+            full_query["size"] = int(
+                os.environ.get("OPENSEARCH_TRACES_SEARCH_SIZE", "5000")
+            )
+            logging.debug(f"opensearch traces search query: {full_query}")
+            headers = {"Content-Type": "application/json"}
+            headers.update(
+                add_auth_header(self._toolset.opensearch_config.opensearch_auth_header)
+            )
+
+            logs_response = requests.get(
+                url=get_search_url(self._toolset.opensearch_config),
+                timeout=180,
+                verify=True,
+                data=json.dumps(full_query),
+                headers=headers,
+            )
+            if logs_response.status_code > 300:
+                err_msg = logs_response.text
+
+            logs_response.raise_for_status()
+            return StructuredToolResult(
+                status=ToolResultStatus.SUCCESS,
+                data=json.dumps(logs_response.json()),
+                params=params,
+            )
+        except requests.Timeout:
+            logging.warning(
+                "Timeout while fetching opensearch traces search", exc_info=True
+            )
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Request timed out while fetching opensearch traces search {err_msg}",
+                params=params,
+            )
+        except RequestException as e:
+            logging.warning("Failed to fetch opensearch traces search", exc_info=True)
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Network error while opensearch traces search {err_msg} : {str(e)}",
+                params=params,
+            )
+        except Exception as e:
+            logging.warning(
+                "Failed to process opensearch traces search ", exc_info=True
+            )
+            return StructuredToolResult(
+                status=ToolResultStatus.ERROR,
+                error=f"Unexpected error {err_msg}: {str(e)}",
+                params=params,
+            )
+
+    def get_parameterized_one_liner(self, params) -> str:
+        return f'search traces: query="{params.get("query")}"'
+
+
+class OpenSearchTracesToolset(BaseOpenSearchToolset):
+    def __init__(self):
+        super().__init__(
+            name="opensearch/traces",
+            description="OpenSearch integration to fetch traces",
+            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/opensearch-traces.html",
+            icon_url="https://opensearch.org/assets/brand/PNG/Mark/opensearch_mark_default.png",
+            prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
+            tools=[
+                GetTracesFields(toolset=self),
+                TracesSearchQuery(toolset=self),
+            ],
+            tags=[
+                ToolsetTag.CORE,
+            ],
+        )
+        template_file_path = os.path.abspath(
+            os.path.join(
+                os.path.dirname(__file__), "opensearch_traces_instructions.jinja2"
+            )
+        )
+        self._load_llm_instructions(jinja_template=f"file://{template_file_path}")

holmes/plugins/toolsets/opensearch/opensearch_traces_instructions.jinja2

@@ -0,0 +1,12 @@
+# Traces and Spans
+* All applications are instrumented and has spans and traces.
+* ALWAYS, first, use get_traces_fields to understand the documents fields, then ALWAYS, search for trace documents based on the fields recieved in the previous call. You MUST filter by the fields names that appears here
+* When investigating network errors, latency, or performance issues, ALWAYS check traces and spans
+* The search query should always contain the root 'query' element in the search expression
+* Search traces using the `fetch_opensearch_logs` tool
+* try to search traces based on what you're looking for, and the trace fields available
+* Always build the query with '.keyword' when doing exact match
+* if time range is not specified, get logs for the last hour
+* When investigating numeric errors, these are error codes
+* Always search for traces BEFORE checking logs. If the logs contain information that can lead to another trace search, do another trace search
+* ALWAYS check for traces to learn about the traffic, latency and errors