holmesgpt 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

holmes/plugins/toolsets/prometheus/prometheus.py

@@ -3,14 +3,15 @@ import logging
 import os
 import re
 import time
+import dateutil.parser
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 from urllib.parse import urljoin
 
 import requests  # type: ignore
 from pydantic import BaseModel, field_validator, Field, model_validator
 from requests import RequestException
-from requests_aws4auth import AWS4Auth
-
+from prometrix.connect.aws_connect import AWSPrometheusConnect
+from prometrix.models.prometheus_config import PrometheusConfig as BasePrometheusConfig
 from holmes.core.tools import (
     CallablePrerequisite,
     StructuredToolResult,
@@ -29,10 +30,10 @@ from holmes.plugins.toolsets.utils import (
     toolset_name_for_one_liner,
 )
 from holmes.utils.cache import TTLCache
-from holmes.common.env_vars import IS_OPENSHIFT
+from holmes.common.env_vars import IS_OPENSHIFT, MAX_GRAPH_POINTS
 from holmes.common.openshift import load_openshift_token
 from holmes.plugins.toolsets.logging_utils.logging_api import (
-    DEFAULT_TIME_SPAN_SECONDS,
+    DEFAULT_GRAPH_TIME_SPAN_SECONDS,
 )
 from holmes.utils.keygen_utils import generate_random_key
 
@@ -81,34 +82,105 @@ class PrometheusConfig(BaseModel):
     def is_amp(self) -> bool:
         return False
 
-    def get_auth(self) -> Any:
-        return None
-
 
 class AMPConfig(PrometheusConfig):
-    aws_access_key: str
-    aws_secret_access_key: str
+    aws_access_key: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
     aws_region: str
     aws_service_name: str = "aps"
-    healthcheck: str = "api/v1/query?query=up"  # Override for AMP
+    healthcheck: str = "api/v1/query?query=up"
     prometheus_ssl_enabled: bool = False
+    assume_role_arn: Optional[str] = None
+
+    # Refresh the AWS client (and its STS creds) every N seconds (default: 15 minutes)
+    refresh_interval_seconds: int = 900
+
+    _aws_client: Optional[AWSPrometheusConnect] = None
+    _aws_client_created_at: float = 0.0
 
     def is_amp(self) -> bool:
         return True
 
-    def get_auth(self):
-        return AWS4Auth(
-            self.aws_access_key,  # type: ignore
-            self.aws_secret_access_key,  # type: ignore
-            self.aws_region,  # type: ignore
-            self.aws_service_name,  # type: ignore
-        )
+    def _should_refresh_client(self) -> bool:
+        if not self._aws_client:
+            return True
+        return (
+            time.time() - self._aws_client_created_at
+        ) >= self.refresh_interval_seconds
+
+    def get_aws_client(self) -> Optional[AWSPrometheusConnect]:
+        if not self._aws_client or self._should_refresh_client():
+            try:
+                base_config = BasePrometheusConfig(
+                    url=self.prometheus_url,
+                    disable_ssl=not self.prometheus_ssl_enabled,
+                    additional_labels=self.additional_labels,
+                )
+                self._aws_client = AWSPrometheusConnect(
+                    access_key=self.aws_access_key,
+                    secret_key=self.aws_secret_access_key,
+                    token=None,
+                    region=self.aws_region,
+                    service_name=self.aws_service_name,
+                    assume_role_arn=self.assume_role_arn,
+                    config=base_config,
+                )
+                self._aws_client_created_at = time.time()
+            except Exception:
+                logging.exception("Failed to create/refresh AWS client")
+                return self._aws_client
+        return self._aws_client
 
 
 class BasePrometheusTool(Tool):
     toolset: "PrometheusToolset"
 
 
+def do_request(
+    config,  # PrometheusConfig | AMPConfig
+    url: str,
+    params: Optional[Dict] = None,
+    data: Optional[Dict] = None,
+    timeout: int = 60,
+    verify: Optional[bool] = None,
+    headers: Optional[Dict] = None,
+    method: str = "GET",
+) -> requests.Response:
+    """
+    Route a request through either:
+      - AWSPrometheusConnect (SigV4) when config is AMPConfig
+      - plain requests otherwise
+
+    method defaults to GET so callers can omit it for reads.
+    """
+    if verify is None:
+        verify = config.prometheus_ssl_enabled
+    if headers is None:
+        headers = config.headers or {}
+
+    if isinstance(config, AMPConfig):
+        client = config.get_aws_client()  # cached AWSPrometheusConnect
+        return client.signed_request(  # type: ignore
+            method=method,
+            url=url,
+            data=data,
+            params=params,
+            verify=verify,
+            headers=headers,
+        )
+
+    # Non-AMP: plain HTTP
+    return requests.request(
+        method=method,
+        url=url,
+        headers=headers,
+        params=params,
+        data=data,
+        timeout=timeout,
+        verify=verify,
+    )
+
+
 def filter_metrics_by_type(metrics: Dict, expected_type: str):
     return {
         metric_name: metric_data
@@ -133,14 +205,18 @@ METRICS_SUFFIXES_TO_STRIP = ["_bucket", "_count", "_sum"]
 def fetch_metadata(
     prometheus_url: str,
     headers: Optional[Dict],
-    auth=None,
+    config,
     verify_ssl: bool = True,
 ) -> Dict:
     metadata_url = urljoin(prometheus_url, "api/v1/metadata")
-    metadata_response = requests.get(
-        metadata_url, headers=headers, timeout=60, verify=verify_ssl, auth=auth
+    metadata_response = do_request(
+        config=config,
+        url=metadata_url,
+        headers=headers,
+        timeout=60,
+        verify=verify_ssl,
+        method="GET",
     )
-
     metadata_response.raise_for_status()
 
     metadata = metadata_response.json()["data"]
@@ -163,14 +239,20 @@ def fetch_metadata_with_series_api(
     prometheus_url: str,
     metric_name: str,
     headers: Dict,
-    auth=None,
+    config,
     verify_ssl: bool = True,
 ) -> Dict:
     url = urljoin(prometheus_url, "api/v1/series")
     params: Dict = {"match[]": f'{{__name__=~".*{metric_name}.*"}}', "limit": "10000"}
 
-    response = requests.get(
-        url, headers=headers, timeout=60, params=params, auth=auth, verify=verify_ssl
+    response = do_request(
+        config=config,
+        url=url,
+        headers=headers,
+        params=params,
+        timeout=60,
+        verify=verify_ssl,
+        method="GET",
     )
     response.raise_for_status()
     metrics = response.json()["data"]
@@ -199,6 +281,42 @@ def result_has_data(result: Dict) -> bool:
     return False
 
 
+def adjust_step_for_max_points(
+    start_timestamp: str,
+    end_timestamp: str,
+    step: float,
+) -> float:
+    """
+    Adjusts the step parameter to ensure the number of data points doesn't exceed max_points.
+    Max points is controlled by the PROMETHEUS_MAX_GRAPH_POINTS environment variable (default: 300).
+
+    Args:
+        start_timestamp: RFC3339 formatted start time
+        end_timestamp: RFC3339 formatted end time
+        step: The requested step duration in seconds
+
+    Returns:
+        Adjusted step value in seconds that ensures points <= max_points
+    """
+
+    start_dt = dateutil.parser.parse(start_timestamp)
+    end_dt = dateutil.parser.parse(end_timestamp)
+
+    time_range_seconds = (end_dt - start_dt).total_seconds()
+
+    current_points = time_range_seconds / step
+
+    # If current points exceed max, adjust the step
+    if current_points > MAX_GRAPH_POINTS:
+        adjusted_step = time_range_seconds / MAX_GRAPH_POINTS
+        logging.info(
+            f"Adjusting step from {step}s to {adjusted_step}s to limit points from {current_points:.0f} to {MAX_GRAPH_POINTS}"
+        )
+        return adjusted_step
+
+    return step
+
+
 def add_prometheus_auth(prometheus_auth_header: Optional[str]) -> Dict[str, Any]:
     results = {}
     if prometheus_auth_header:
@@ -212,7 +330,7 @@ def fetch_metrics_labels_with_series_api(
     cache: Optional[TTLCache],
     metrics_labels_time_window_hrs: Union[int, None],
     metric_name: str,
-    auth=None,
+    config=None,
     verify_ssl: bool = True,
 ) -> dict:
     """This is a slow query. Takes 5+ seconds to run"""
@@ -229,13 +347,14 @@ def fetch_metrics_labels_with_series_api(
         params["end"] = int(time.time())
         params["start"] = params["end"] - (metrics_labels_time_window_hrs * 60 * 60)
 
-    series_response = requests.get(
+    series_response = do_request(
+        config=config,
         url=series_url,
         headers=headers,
         params=params,
-        auth=auth,
         timeout=60,
         verify=verify_ssl,
+        method="GET",
     )
     series_response.raise_for_status()
     series = series_response.json()["data"]
@@ -261,7 +380,7 @@ def fetch_metrics_labels_with_labels_api(
     metrics_labels_time_window_hrs: Union[int, None],
     metric_names: List[str],
     headers: Dict,
-    auth=None,
+    config=None,
     verify_ssl: bool = True,
 ) -> dict:
     metrics_labels = {}
@@ -281,13 +400,14 @@ def fetch_metrics_labels_with_labels_api(
             params["end"] = int(time.time())
             params["start"] = params["end"] - (metrics_labels_time_window_hrs * 60 * 60)
 
-        response = requests.get(
+        response = do_request(
+            config=config,
             url=url,
             headers=headers,
             params=params,
-            auth=auth,
             timeout=60,
             verify=verify_ssl,
+            method="GET",
         )
         response.raise_for_status()
         labels = response.json()["data"]
@@ -308,7 +428,7 @@ def fetch_metrics(
     should_fetch_labels_with_labels_api: bool,
     should_fetch_metadata_with_series_api: bool,
     headers: Dict,
-    auth=None,
+    config=None,
     verify_ssl: bool = True,
 ) -> dict:
     metrics = None
@@ -318,7 +438,7 @@ def fetch_metrics(
             prometheus_url=prometheus_url,
             metric_name=metric_name,
             headers=headers,
-            auth=auth,
+            config=config,
             verify_ssl=verify_ssl,
         )
         should_fetch_labels = False  # series API returns the labels
@@ -326,7 +446,7 @@ def fetch_metrics(
         metrics = fetch_metadata(
             prometheus_url=prometheus_url,
             headers=headers,
-            auth=auth,
+            config=config,
             verify_ssl=verify_ssl,
         )
         metrics = filter_metrics_by_name(metrics, metric_name)
@@ -340,7 +460,7 @@ def fetch_metrics(
                 metrics_labels_time_window_hrs=metrics_labels_time_window_hrs,
                 metric_names=list(metrics.keys()),
                 headers=headers,
-                auth=auth,
+                config=config,
                 verify_ssl=verify_ssl,
             )
         else:
@@ -350,7 +470,7 @@ def fetch_metrics(
                 metrics_labels_time_window_hrs=metrics_labels_time_window_hrs,
                 metric_name=metric_name,
                 headers=headers,
-                auth=auth,
+                config=config,
                 verify_ssl=verify_ssl,
             )
 
@@ -371,7 +491,9 @@ class ListPrometheusRules(BasePrometheusTool):
         )
         self._cache = None
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if not self.toolset.config or not self.toolset.config.prometheus_url:
             return StructuredToolResult(
                 status=ToolResultStatus.ERROR,
@@ -402,13 +524,14 @@ class ListPrometheusRules(BasePrometheusTool):
 
             rules_url = urljoin(prometheus_url, "api/v1/rules")
 
-            rules_response = requests.get(
+            rules_response = do_request(
+                config=self.toolset.config,
                 url=rules_url,
                 params=params,
-                auth=self.toolset.config.get_auth(),
                 timeout=180,
                 verify=self.toolset.config.prometheus_ssl_enabled,
                 headers=self.toolset.config.headers,
+                method="GET",
             )
             rules_response.raise_for_status()
             data = rules_response.json()["data"]
@@ -467,7 +590,9 @@ class ListAvailableMetrics(BasePrometheusTool):
         )
         self._cache = None
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if not self.toolset.config or not self.toolset.config.prometheus_url:
             return StructuredToolResult(
                 status=ToolResultStatus.ERROR,
@@ -500,12 +625,13 @@ class ListAvailableMetrics(BasePrometheusTool):
                 should_fetch_labels_with_labels_api=self.toolset.config.fetch_labels_with_labels_api,
                 should_fetch_metadata_with_series_api=self.toolset.config.fetch_metadata_with_series_api,
                 headers=self.toolset.config.headers,
-                auth=self.toolset.config.get_auth(),
+                config=self.toolset.config,
                 verify_ssl=self.toolset.config.prometheus_ssl_enabled,
             )
 
-            if params.get("type_filter"):
-                metrics = filter_metrics_by_type(metrics, params.get("type_filter"))
+            type_filter = params.get("type_filter")
+            if type_filter:
+                metrics = filter_metrics_by_type(metrics, type_filter)
 
             output = ["Metric | Description | Type | Labels"]
             output.append("-" * 100)
@@ -572,7 +698,9 @@ class ExecuteInstantQuery(BasePrometheusTool):
             toolset=toolset,
         )
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if not self.toolset.config or not self.toolset.config.prometheus_url:
             return StructuredToolResult(
                 status=ToolResultStatus.ERROR,
@@ -587,12 +715,14 @@ class ExecuteInstantQuery(BasePrometheusTool):
 
             payload = {"query": query}
 
-            response = requests.post(
+            response = do_request(
+                config=self.toolset.config,
                 url=url,
                 headers=self.toolset.config.headers,
-                auth=self.toolset.config.get_auth(),
                 data=payload,
                 timeout=60,
+                verify=self.toolset.config.prometheus_ssl_enabled,
+                method="POST",
             )
 
             if response.status_code == 200:
@@ -684,7 +814,7 @@ class ExecuteRangeQuery(BasePrometheusTool):
                 ),
                 "start": ToolParameter(
                     description=standard_start_datetime_tool_param_description(
-                        DEFAULT_TIME_SPAN_SECONDS
+                        DEFAULT_GRAPH_TIME_SPAN_SECONDS
                     ),
                     type="string",
                     required=False,
@@ -708,7 +838,9 @@ class ExecuteRangeQuery(BasePrometheusTool):
             toolset=toolset,
         )
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if not self.toolset.config or not self.toolset.config.prometheus_url:
             return StructuredToolResult(
                 status=ToolResultStatus.ERROR,
@@ -723,9 +855,16 @@ class ExecuteRangeQuery(BasePrometheusTool):
             (start, end) = process_timestamps_to_rfc3339(
                 start_timestamp=params.get("start"),
                 end_timestamp=params.get("end"),
-                default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS,
+                default_time_span_seconds=DEFAULT_GRAPH_TIME_SPAN_SECONDS,
             )
             step = params.get("step", "")
+
+            step = adjust_step_for_max_points(
+                start_timestamp=start,
+                end_timestamp=end,
+                step=float(step) if step else MAX_GRAPH_POINTS,
+            )
+
             description = params.get("description", "")
             output_type = params.get("output_type", "Plain")
             payload = {
@@ -735,12 +874,14 @@ class ExecuteRangeQuery(BasePrometheusTool):
                 "step": step,
             }
 
-            response = requests.post(
+            response = do_request(
+                config=self.toolset.config,
                 url=url,
                 headers=self.toolset.config.headers,
-                auth=self.toolset.config.get_auth(),
                 data=payload,
                 timeout=120,
+                verify=self.toolset.config.prometheus_ssl_enabled,
+                method="POST",
             )
 
             if response.status_code == 200:
@@ -823,7 +964,7 @@ class PrometheusToolset(Toolset):
         super().__init__(
             name="prometheus/metrics",
             description="Prometheus integration to fetch metadata and execute PromQL queries",
-            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/prometheus.html",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/prometheus/",
             icon_url="https://upload.wikimedia.org/wikipedia/commons/3/38/Prometheus_software_logo.svg",
             prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
             tools=[
@@ -847,10 +988,8 @@ class PrometheusToolset(Toolset):
     def determine_prometheus_class(
         self, config: dict[str, Any]
     ) -> Type[Union[PrometheusConfig, AMPConfig]]:
-        has_aws_credentials = (
-            "aws_access_key" in config or "aws_secret_access_key" in config
-        )
-        return AMPConfig if has_aws_credentials else PrometheusConfig
+        has_aws_fields = "aws_region" in config
+        return AMPConfig if has_aws_fields else PrometheusConfig
 
     def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]:
         try:
@@ -904,12 +1043,13 @@ class PrometheusToolset(Toolset):
 
         url = urljoin(self.config.prometheus_url, self.config.healthcheck)
         try:
-            response = requests.get(
+            response = do_request(
+                config=self.config,
                 url=url,
                 headers=self.config.headers,
-                auth=self.config.get_auth(),
                 timeout=10,
                 verify=self.config.prometheus_ssl_enabled,
+                method="GET",
             )
 
             if response.status_code == 200:
@@ -926,6 +1066,7 @@ class PrometheusToolset(Toolset):
                 f"Failed to initialize using url={url}",
             )
         except Exception as e:
+            logging.exception("Failed to initialize Prometheus")
             return (
                 False,
                 f"Failed to initialize using url={url}. Unexpected error: {str(e)}",

holmes/plugins/toolsets/rabbitmq/toolset_rabbitmq.py

@@ -63,7 +63,9 @@ class ListConfiguredClusters(BaseRabbitMQTool):
             toolset=toolset,
         )
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         if not self.toolset.config:
             raise ValueError("RabbitMQ is not configured.")
 
@@ -101,7 +103,9 @@ class GetRabbitMQClusterStatus(BaseRabbitMQTool):
             toolset=toolset,
         )
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         try:
             # Fetch node details which include partition info
             cluster_config = self._get_cluster_config(
@@ -130,7 +134,7 @@ class RabbitMQToolset(Toolset):
         super().__init__(
             name="rabbitmq/core",
             description="Provides tools to interact with RabbitMQ to diagnose cluster health, node status, and specifically network partitions (split-brain).",
-            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/rabbitmq.html",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/rabbitmq/",
             icon_url="https://cdn.worldvectorlogo.com/logos/rabbitmq.svg",
             prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)],
             tools=[

holmes/plugins/toolsets/robusta/robusta.py

@@ -45,7 +45,9 @@ class FetchRobustaFinding(Tool):
             logging.error(error)
             return {"error": error}
 
-    def _invoke(self, params: Dict) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         finding_id = params[PARAM_FINDING_ID]
         try:
             finding = self._fetch_finding(finding_id)
@@ -113,7 +115,9 @@ class FetchResourceRecommendation(Tool):
             )
         return None
 
-    def _invoke(self, params: Dict) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         try:
             recommendations = self._resource_recommendation(params)
             if recommendations:
@@ -171,7 +175,9 @@ class FetchConfigurationChanges(Tool):
             )
         return None
 
-    def _invoke(self, params: Dict) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         try:
             changes = self._fetch_change_history(params)
             if changes:
@@ -213,7 +219,7 @@ class RobustaToolset(Toolset):
         super().__init__(
             icon_url="https://cdn.prod.website-files.com/633e9bac8f71dfb7a8e4c9a6/646be7710db810b14133bdb5_logo.svg",
             description="Fetches alerts metadata and change history",
-            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/robusta.html",
+            docs_url="https://holmesgpt.dev/data-sources/builtin-toolsets/robusta/",
             name="robusta",
             prerequisites=[dal_prereq],
             tools=[

holmes/plugins/toolsets/runbook/runbook_fetcher.py

@@ -35,7 +35,9 @@ class RunbookFetcher(Tool):
             toolset=toolset,  # type: ignore
         )
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         link: str = params["link"]
 
         search_paths = [DEFAULT_RUNBOOK_SEARCH_PATH]
@@ -126,7 +128,7 @@ class RunbookToolset(Toolset):
             tools=[
                 RunbookFetcher(self),
             ],
-            docs_url="https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/runbook.html",
+            docs_url="https://holmesgpt.dev/data-sources/",
             tags=[
                 ToolsetTag.CORE,
             ],

holmes/plugins/toolsets/servicenow/servicenow.py

@@ -115,7 +115,9 @@ class ReturnChangesInTimerange(ServiceNowBaseTool):
         start = params.get("start", "last hour")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Change Requests ({start})"
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         parsed_params = {}
         try:
             (start, _) = process_timestamps_to_rfc3339(
@@ -158,7 +160,9 @@ class ReturnChange(ServiceNowBaseTool):
         sys_id = params.get("sys_id", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Get Change Details ({sys_id})"
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         try:
             url = "https://{instance}.service-now.com/api/now/v2/table/change_request/{sys_id}".format(
                 instance=self.toolset.config.get("instance"),
@@ -190,7 +194,9 @@ class ReturnChangesWithKeyword(ServiceNowBaseTool):
         keyword = params.get("keyword", "")
         return f"{toolset_name_for_one_liner(self.toolset.name)}: Search Changes ({keyword})"
 
-    def _invoke(self, params: Any) -> StructuredToolResult:
+    def _invoke(
+        self, params: dict, user_approved: bool = False
+    ) -> StructuredToolResult:
         parsed_params = {}
         try:
             url = f"https://{self.toolset.config.get('instance')}.service-now.com/api/now/v2/table/change_request"

holmes/plugins/toolsets/slab.yaml

@@ -1,7 +1,7 @@
 toolsets:
   slab:
     description: "Fetches slab pages"
-    docs_url: "https://docs.robusta.dev/master/configuration/holmesgpt/toolsets/slab.html"
+    docs_url: "https://holmesgpt.dev/data-sources/builtin-toolsets/slab/"
     icon_url: "https://platform.robusta.dev/demos/slab-mark.svg"
     tags:
       - core

holmes/utils/console/logging.py

@@ -41,9 +41,14 @@ def suppress_noisy_logs():
     warnings.filterwarnings("ignore", category=UserWarning, module="slack_sdk.*")
 
 
-def init_logging(verbose_flags: Optional[List[bool]] = None):
+def init_logging(verbose_flags: Optional[List[bool]] = None, log_costs: bool = False):
     verbosity = cli_flags_to_verbosity(verbose_flags)  # type: ignore
 
+    # Setup cost logger if requested
+    if log_costs:
+        cost_logger = logging.getLogger("holmes.costs")
+        cost_logger.setLevel(logging.DEBUG)
+
     if verbosity == Verbosity.VERY_VERBOSE:
         logging.basicConfig(
             force=True,