Skip to content

Commit 9efedb4

Browse files
nhermentnherment
committed
chore: address PR comments
1 parent af0d9ec commit 9efedb4

File tree

4 files changed

+51
-59
lines changed

4 files changed

+51
-59
lines changed

‎holmes/plugins/toolsets/newrelic/newrelic.py‎

Lines changed: 35 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111
from pydantic import BaseModel
1212
from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
1313
from holmes.plugins.toolsets.prometheus.data_compression import (
14-
raw_metric_to_compressed_metric,
15-
summarize_metrics,
14+
simplify_prometheus_metric_object,
15+
compact_metrics,
1616
)
1717
from holmes.plugins.toolsets.prometheus.model import PromResponse
1818
from holmes.plugins.toolsets.utils import toolset_name_for_one_liner
@@ -80,42 +80,39 @@ def __init__(self, toolset: "NewRelicToolset"):
8080
)
8181
self._toolset = toolset
8282

83-
def compress_metrics_data(self, response: PromResponse) -> Optional[str]:
83+
def compact_metrics_data(self, response: PromResponse) -> Optional[str]:
8484
llm_data: Optional[str] = None
8585
try:
86-
if self._toolset.config and self._toolset.compress_metrics:
87-
metrics = [
88-
raw_metric_to_compressed_metric(metric, remove_labels=set())
89-
for metric in response.data.result
90-
]
91-
92-
compressed_data = summarize_metrics(metrics)
93-
original_size = len(json.dumps(response.to_json()))
94-
compressed_size = len(json.dumps(compressed_data))
95-
compression_ratio = (
96-
(1 - compressed_size / original_size) * 100
97-
if original_size > 0
98-
else 0
99-
)
86+
metrics = [
87+
simplify_prometheus_metric_object(metric, remove_labels=set())
88+
for metric in response.data.result
89+
]
90+
91+
compacted_data = compact_metrics(metrics)
92+
original_size = len(json.dumps(response.to_json()))
93+
compacted_size = len(json.dumps(compacted_data))
94+
compaction_ratio = (
95+
(1 - compacted_size / original_size) * 100 if original_size > 0 else 0
96+
)
10097

101-
if compression_ratio > self._toolset.compress_metrics_minimum_ratio:
102-
# below this amount it's likely not worth mutating the response
103-
llm_data = compressed_data
104-
logging.info(
105-
f"Compressed Newrelic metrics: {original_size:,} → {compressed_size:,} chars "
106-
f"({compression_ratio:.1f}% reduction)"
107-
)
108-
else:
109-
logging.info(
110-
f"Compressed Newrelic metrics: {original_size:,} → {compressed_size:,} chars "
111-
f"({compression_ratio:.1f}% reduction). Original data will be used instead."
112-
)
98+
if compaction_ratio > self._toolset.compact_metrics_minimum_ratio:
99+
# below this amount it's likely not worth mutating the response
100+
llm_data = compacted_data
101+
logging.debug(
102+
f"Compressed Newrelic metrics: {original_size:,} → {compacted_size:,} chars "
103+
f"({compaction_ratio:.1f}% reduction)"
104+
)
105+
else:
106+
logging.debug(
107+
f"Compressed Newrelic metrics: {original_size:,} → {compacted_size:,} chars "
108+
f"({compaction_ratio:.1f}% reduction). Original data will be used instead."
109+
)
113110
except Exception:
114111
logging.warning("Failed to compress newrelic data", exc_info=True)
115112

116113
return llm_data
117114

118-
def to_new_relic_records(
115+
def to_prometheus_records(
119116
self,
120117
records: List[Dict[str, Any]],
121118
params: Optional[Dict[str, Any]] = None,
@@ -175,15 +172,13 @@ def _invoke(
175172
if qtype == "metrics" or "timeseries" in query.lower():
176173
enriched_params = dict(params)
177174
enriched_params["query"] = query
178-
prom_data = self.to_new_relic_records(result, params=enriched_params)
175+
prom_data = self.to_prometheus_records(result, params=enriched_params)
179176

180177
return_result = prom_data.to_json()
181-
if len(return_result.get("data", {}).get("results", [])):
182-
return_result = result # type: ignore[assignment]
183178
return StructuredToolResult(
184179
status=StructuredToolResultStatus.SUCCESS,
185180
data=json.dumps(return_result, indent=2),
186-
llm_data=self.compress_metrics_data(prom_data),
181+
llm_data=self.compact_metrics_data(prom_data),
187182
params=params,
188183
)
189184

@@ -246,16 +241,16 @@ class NewrelicConfig(BaseModel):
246241
nr_api_key: Optional[str] = None
247242
nr_account_id: Optional[str] = None
248243
is_eu_datacenter: Optional[bool] = False
249-
compress_metrics: bool = True
250-
compress_metrics_minimum_ratio: int = 30 # 20 means 20% size reduction
244+
compact_metrics: bool = True
245+
compact_metrics_minimum_ratio: int = 30 # 20 means 20% size reduction
251246

252247

253248
class NewRelicToolset(Toolset):
254249
nr_api_key: Optional[str] = None
255250
nr_account_id: Optional[str] = None
256251
is_eu_datacenter: bool = False
257-
compress_metrics: bool = True
258-
compress_metrics_minimum_ratio: int = 30
252+
compact_metrics: bool = True
253+
compact_metrics_minimum_ratio: int = 30
259254

260255
def __init__(self):
261256
super().__init__(
@@ -286,10 +281,8 @@ def prerequisites_callable(
286281
self.nr_account_id = nr_config.nr_account_id
287282
self.nr_api_key = nr_config.nr_api_key
288283
self.is_eu_datacenter = nr_config.is_eu_datacenter or False
289-
self.compress_metrics = nr_config.compress_metrics
290-
self.compress_metrics_minimum_ratio = (
291-
nr_config.compress_metrics_minimum_ratio
292-
)
284+
self.compact_metrics = nr_config.compact_metrics
285+
self.compact_metrics_minimum_ratio = nr_config.compact_metrics_minimum_ratio
293286

294287
if not self.nr_account_id or not self.nr_api_key:
295288
return False, "New Relic account ID or API key is missing"

‎holmes/plugins/toolsets/prometheus/data_compression.py‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ class Group(BaseModel):
1616
metrics: Sequence[Union["Group", CompressedMetric]]
1717

1818

19+
INDENT_SPACES = " "
20+
21+
1922
def format_labels(
2023
labels: set[tuple[str, Any]], section_name: str, line_prefix: str = ""
2124
) -> list[str]:
@@ -38,9 +41,6 @@ def format_labels(
3841
return lines
3942

4043

41-
INDENT_SPACES = " "
42-
43-
4444
def format_data(data: Union[Group, CompressedMetric], line_prefix: str = "") -> str:
4545
lines = []
4646
if isinstance(data, CompressedMetric):
@@ -84,7 +84,8 @@ def format_zero_values_data(
8484
) -> str:
8585
lines = []
8686
if isinstance(data, CompressedMetric):
87-
lines.extend(format_labels(labels=data.labels, section_name=""))
87+
if data.labels:
88+
lines.extend(format_labels(labels=data.labels, section_name=""))
8889
txt = "\n".join([line_prefix + line for line in lines])
8990
return txt
9091

@@ -133,9 +134,7 @@ def format_zero_values_data(
133134

134135

135136
def format_zero_values_metrics(metrics: list[Union[Group, CompressedMetric]]) -> str:
136-
formatted_string = (
137-
"# Metrics with the following hierarchised labels have NO DATA:\n"
138-
)
137+
formatted_string = "# Metrics with the following hierarchised labels have all values set to ZERO:\n"
139138
for metric in metrics:
140139
formatted_string += (
141140
format_zero_values_data(metric, line_prefix=INDENT_SPACES) + "\n"
@@ -155,7 +154,7 @@ def format_compressed_metrics(metrics: list[Union[Group, CompressedMetric]]) ->
155154
return formatted_string
156155

157156

158-
def raw_metric_to_compressed_metric(
157+
def simplify_prometheus_metric_object(
159158
raw_metric: PromSeries, remove_labels: set[tuple[str, Any]]
160159
) -> CompressedMetric:
161160
labels: set[tuple[str, Any]] = set()
@@ -287,7 +286,7 @@ def set_default(obj):
287286
raise TypeError
288287

289288

290-
def summarize_metrics(metrics_to_process: list[CompressedMetric]) -> str:
289+
def compact_metrics(metrics_to_process: list[CompressedMetric]) -> str:
291290
summarized_text = ""
292291
filtered_metrics = pre_filter_metrics(metrics=metrics_to_process)
293292
if (

‎holmes/plugins/toolsets/prometheus/prometheus.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@
2222
)
2323
from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION
2424
from holmes.plugins.toolsets.prometheus.data_compression import (
25-
raw_metric_to_compressed_metric,
26-
summarize_metrics,
25+
simplify_prometheus_metric_object,
26+
compact_metrics,
2727
)
2828
from holmes.plugins.toolsets.prometheus.model import PromSeries
2929
from holmes.plugins.toolsets.prometheus.utils import parse_duration_to_seconds
@@ -1386,13 +1386,13 @@ def _invoke(
13861386
PromSeries(**metric) for metric in metrics_list_dict
13871387
]
13881388
metrics = [
1389-
raw_metric_to_compressed_metric(
1389+
simplify_prometheus_metric_object(
13901390
metric, remove_labels=set()
13911391
)
13921392
for metric in raw_metrics
13931393
]
13941394

1395-
compressed_data = summarize_metrics(metrics)
1395+
compressed_data = compact_metrics(metrics)
13961396
response_data["raw_data"] = result_data
13971397
original_size = len(json.dumps(result_data))
13981398
compressed_size = len(json.dumps(compressed_data))

‎tests/plugins/toolsets/prometheus/test_data_compression.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
format_compressed_metrics,
88
format_data,
99
group_metrics,
10-
raw_metric_to_compressed_metric,
11-
summarize_metrics,
10+
simplify_prometheus_metric_object,
11+
compact_metrics,
1212
)
1313
from holmes.plugins.toolsets.prometheus.model import PromSeries
1414

@@ -413,11 +413,11 @@ def test_format_data_realistic(self):
413413

414414
raw_metrics = [PromSeries(**metric) for metric in metrics_list_dict]
415415
metrics = [
416-
raw_metric_to_compressed_metric(metric, remove_labels=set())
416+
simplify_prometheus_metric_object(metric, remove_labels=set())
417417
for metric in raw_metrics
418418
]
419419

420-
formatted_data = summarize_metrics(metrics)
420+
formatted_data = compact_metrics(metrics)
421421

422422
ratio = len(formatted_data) / len(json.dumps(data, indent=2))
423423
# print(formatted_data)

0 commit comments

Comments
 (0)