From edb93f6247d3515ee274048eaaa2a4d771578312 Mon Sep 17 00:00:00 2001
From: Avi-Robusta <97387909+Avi-Robusta@users.noreply.github.com>
Date: Sun, 1 Dec 2024 15:17:13 +0200
Subject: [PATCH] [MAIN-2448] - Slack labels to ai (#1640)

* added slack labels

* updated holmes convo string

* rebase fix

* fixing slack labels

* log fix

* checking label size before all the params
---
 src/robusta/core/model/base_params.py         |  2 ++
 .../core/playbooks/internal/ai_integration.py | 10 ++++++-
 src/robusta/integrations/slack/sender.py      | 27 +++++++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/robusta/core/model/base_params.py b/src/robusta/core/model/base_params.py
index ab76c82a2..f0ebd1688 100644
--- a/src/robusta/core/model/base_params.py
+++ b/src/robusta/core/model/base_params.py
@@ -152,11 +152,13 @@ class HolmesIssueChatParamsContext(BaseModel):
     :var investigation_result: HolmesInvestigationResult object that contains investigation saved to Evidence table by frontend for the issue.
     :var issue_type: aggregation key of the issue
     :var robusta_issue_id: id of the issue
+    :var labels: labels from the issue
     """
 
     investigation_result: HolmesInvestigationResult
     issue_type: str
     robusta_issue_id: Optional[str] = None
+    labels: Optional[Dict[str, str]] = None
 
 
 # will be deprecated later alongside with holmes_conversation action
diff --git a/src/robusta/core/playbooks/internal/ai_integration.py b/src/robusta/core/playbooks/internal/ai_integration.py
index 1b05504a6..4c580b266 100644
--- a/src/robusta/core/playbooks/internal/ai_integration.py
+++ b/src/robusta/core/playbooks/internal/ai_integration.py
@@ -50,6 +50,7 @@ def ask_holmes(event: ExecutionBaseEvent, params: AIInvestigateParams):
     subject = params.resource.dict() if params.resource else {}
 
     try:
+        params.ask = add_labels_to_ask(params)
         holmes_req = HolmesRequest(
             source=params.context.get("source", "unknown source") if params.context else "unknown source",
             title=investigation__title,
@@ -162,7 +163,14 @@ def holmes_workload_health(event: ExecutionBaseEvent, params: HolmesWorkloadHeal
 
 
 def build_conversation_title(params: HolmesConversationParams) -> str:
-    return f"{params.resource}, {params.ask} for issue {params.context.robusta_issue_id}"
+    return f"{params.resource}, {params.ask} for issue '{params.context.robusta_issue_id}'"
+
+
+def add_labels_to_ask(params: HolmesConversationParams) -> str:
+    label_string = f"the alert has the following labels: {params.context.get('labels')}" if params.context.get("labels") else ""
+    ask = f"{params.ask}, {label_string}" if label_string else params.ask
+    logging.debug(f"holmes ask query: {ask}")
+    return ask
 
 
 # old version of holmes conversation API
diff --git a/src/robusta/integrations/slack/sender.py b/src/robusta/integrations/slack/sender.py
index 482ff1afd..94cde1bf3 100644
--- a/src/robusta/integrations/slack/sender.py
+++ b/src/robusta/integrations/slack/sender.py
@@ -1,3 +1,4 @@
+import copy
 import logging
 import ssl
 import tempfile
@@ -308,6 +309,31 @@ def __send_blocks_to_slack(
                 f"error sending message to slack\ne={e}\ntext={message}\nchannel={channel}\nblocks={*output_blocks,}\nattachment_blocks={*attachment_blocks,}"
             )
 
+
+    def __limit_labels_size(self, labels: dict, max_size: int = 1000) -> dict:
+        # slack can only send 2k tokens in a callback so the labels are limited in size
+
+        low_priority_labels = ["job", "prometheus", "severity", "service"]
+        current_length = len(str(labels))
+        if current_length <= max_size:
+            return labels
+        
+        limited_labels = copy.deepcopy(labels)
+
+        # first remove the low priority labels if needed
+        for key in low_priority_labels:
+            if current_length <= max_size:
+                break
+            if key in limited_labels:
+                del limited_labels[key]
+                current_length = len(str(limited_labels))
+
+        while current_length > max_size and limited_labels:
+            limited_labels.pop(next(iter(limited_labels)))
+            current_length = len(str(limited_labels))
+
+        return limited_labels
+
     def __create_holmes_callback(self, finding: Finding) -> CallbackBlock:
         resource = ResourceInfo(
             name=finding.subject.name if finding.subject.name else "",
@@ -321,6 +347,7 @@ def __create_holmes_callback(self, finding: Finding) -> CallbackBlock:
             "robusta_issue_id": str(finding.id),
             "issue_type": finding.aggregation_key,
             "source": finding.source.name,
+            "labels": self.__limit_labels_size(labels=finding.subject.labels)
         }
 
         return CallbackBlock(