holmesgpt 0.11.5__py3-none-any.whl

holmes/plugins/toolsets/opensearch/opensearch_utils.py

@@ -0,0 +1,166 @@
+import json
+import logging
+from typing import Optional, Any, cast
+from urllib.parse import urljoin
+
+from pydantic import BaseModel
+import requests  # type: ignore
+
+from holmes.core.tools import Toolset
+
+
+class OpenSearchLoggingLabelsConfig(BaseModel):
+    pod: str = "kubernetes.pod_name"
+    namespace: str = "kubernetes.namespace_name"
+    timestamp: str = "@timestamp"
+    message: str = "message"
+
+
+class BaseOpenSearchConfig(BaseModel):
+    opensearch_url: str
+    index_pattern: str
+    opensearch_auth_header: Optional[str] = None
+    # Setting to None will disable the cache
+    fields_ttl_seconds: Optional[int] = 14400  # 4 hours
+
+
+class OpenSearchLoggingConfig(BaseOpenSearchConfig):
+    labels: OpenSearchLoggingLabelsConfig = OpenSearchLoggingLabelsConfig()
+
+
+def add_auth_header(auth_header: Optional[str]) -> dict[str, Any]:
+    results = {}
+    if auth_header:
+        results["Authorization"] = auth_header
+    return results
+
+
+def get_search_url(config: BaseOpenSearchConfig) -> str:
+    return urljoin(config.opensearch_url, f"/{config.index_pattern}/_search")
+
+
+def opensearch_health_check(config: BaseOpenSearchConfig) -> tuple[bool, str]:
+    url = get_search_url(config)
+    try:
+        headers = {"Content-Type": "application/json"}
+        headers.update(add_auth_header(config.opensearch_auth_header))
+        health_response = requests.get(
+            url=url,
+            verify=True,
+            data=json.dumps({"size": 1}),
+            headers=headers,
+        )
+        health_response.raise_for_status()
+        return True, ""
+    except Exception as e:
+        logging.info("Failed to initialize opensearch toolset", exc_info=True)
+        return False, f"Failed to initialize opensearch toolset. url={url}. {str(e)}"
+
+
+def format_log_to_json(log_line: Any) -> str:
+    try:
+        return json.dumps(log_line)
+    except Exception:
+        # Handle potential serialization errors (e.g., non-serializable objects)
+        return str(log_line)
+
+
+def format_logs(
+    logs: list[dict[str, Any]],
+    config: OpenSearchLoggingConfig,
+) -> str:
+    if not logs or not isinstance(logs, list):
+        return ""
+
+    # Get field names from config or use defaults
+    message_field = config.labels.message
+
+    formatted_lines = []
+
+    for hit in logs:
+        # Ensure hit is a dictionary and has _source
+        if not isinstance(hit, dict):
+            formatted_lines.append(format_log_to_json(hit))
+            continue
+        source = hit.get("_source")
+        if not isinstance(source, dict):
+            formatted_lines.append(format_log_to_json(hit))
+            continue
+
+        message = source.get(message_field, None)
+
+        if message and not isinstance(message, str):
+            message = str(message)  # Convert non-strings
+
+        if message:
+            formatted_lines.append(message)
+        else:
+            # fallback displaying the logs line as-is
+            formatted_lines.append(format_log_to_json(hit))
+
+    return "\n".join(formatted_lines)
+
+
+def build_query(
+    config: OpenSearchLoggingConfig,
+    namespace: str,
+    pod_name: str,
+    start_time: Optional[str] = None,
+    end_time: Optional[str] = None,
+    match: Optional[str] = None,
+    limit: Optional[int] = None,
+) -> dict[str, Any]:
+    size = limit if limit else 5000
+
+    pod_field = config.labels.pod
+    namespace_field = config.labels.namespace
+    timestamp_field = config.labels.timestamp
+    message_field = config.labels.message
+
+    must_constraints: list[dict] = [
+        {"term": {f"{pod_field}.keyword": pod_name}},
+        {"term": {f"{namespace_field}.keyword": namespace}},
+    ]
+
+    query = {
+        "size": size,
+        "sort": [{timestamp_field: {"order": "asc"}}],
+        "query": {"bool": {"must": must_constraints}},
+    }
+
+    # Add timestamp range if provided
+    if start_time or end_time:
+        range_query: dict = {"range": {timestamp_field: {}}}
+        if start_time:
+            range_query["range"][timestamp_field]["gte"] = start_time
+        if end_time:
+            range_query["range"][timestamp_field]["lte"] = end_time
+
+        must_constraints.append(range_query)
+
+    # Add message filter if provided
+    if match:
+        must_constraints.append({"match_phrase": {message_field: match}})
+
+    return query
+
+
+class BaseOpenSearchToolset(Toolset):
+    def get_example_config(self) -> dict[str, Any]:
+        example_config = BaseOpenSearchConfig(
+            opensearch_url="YOUR OPENSEARCH URL",
+            index_pattern="YOUR OPENSEARCH INDEX NAME",
+            opensearch_auth_header="YOUR OPENSEARCH 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 traces URL. Check your config"
+        else:
+            self.config = BaseOpenSearchConfig(**config)
+            return opensearch_health_check(self.config)
+
+    @property
+    def opensearch_config(self) -> BaseOpenSearchConfig:
+        return cast(BaseOpenSearchConfig, self.config)