param updates + safer code

Author: Avi-Robusta

src/robusta/core/model/base_params.py

@@ -182,23 +182,14 @@ class ConversationType(str, Enum):
     ISSUE = "issue"
 
 
-class HolmesChatSource(str, Enum):
-    """
-    Conversation types for holmes_conversation action
-    """
-
-    UI = "ui"
-    RELAY = "relay"
-
-
 class HolmesChatParams(HolmesParams):
     """
     :var ask: User's prompt for holmes
     """
 
     ask: str
     conversation_history: Optional[list[dict]] = None
-    source: HolmesChatSource = HolmesChatSource.UI
+    render_graph_images: bool = False
 
 
 class HolmesIssueChatParams(HolmesChatParams):

src/robusta/core/playbooks/internal/ai_integration.py

@@ -8,7 +8,6 @@
     AIInvestigateParams,
     ChartValuesFormat,
     HolmesChatParams,
-    HolmesChatSource,
     HolmesConversationParams,
     HolmesIssueChatParams,
     HolmesWorkloadHealthChatParams,
@@ -357,29 +356,34 @@ def holmes_chat(event: ExecutionBaseEvent, params: HolmesChatParams):
         result.raise_for_status()
         holmes_result = HolmesChatResult(**json.loads(result.text))
         holmes_result.files = []
-        if params.source == HolmesChatSource.RELAY:
-            for tool in holmes_result.tool_calls:
-                if tool.tool_name != "execute_prometheus_range_query":
-                    continue
-                try:
+        if params.render_graph_images:
+            try:
+                for tool in holmes_result.tool_calls:
+                    if tool.tool_name != "execute_prometheus_range_query":
+                        continue
+
                     json_content = json.loads(tool.result)
                     query_result = PrometheusQueryResult(data=json_content.get("data", {}))
                     try:
                         output_type_str = json_content.get("output_type", "Plain")
                         output_type = ChartValuesFormat[output_type_str]
                     except KeyError:
                         output_type = ChartValuesFormat.Plain  # fallback in case of an invalid string
+
                     chart = build_chart_from_prometheus_result(
                         query_result, json_content.get("description", "graph"), values_format=output_type
                     )
                     contents = convert_svg_to_png(chart.render())
                     name = json_content.get("description", "graph").replace(" ", "_")
                     holmes_result.files.append(FileBlock(f"{name}.png", contents))
-                except Exception as e:
-                    logging.exception(f"Failed to parse JSON: {e}\nRaw content:\n{tool.result}")
-            holmes_result.tool_calls = [
-                tool for tool in holmes_result.tool_calls if tool.tool_name != "execute_prometheus_range_query"
-            ]
+
+                holmes_result.tool_calls = [
+                    tool for tool in holmes_result.tool_calls if tool.tool_name != "execute_prometheus_range_query"
+                ]
+
+            except Exception:
+                logging.exception(f"Failed to convert tools to images")
+
         finding = Finding(
             title="AI Ask Chat",
             aggregation_key="HolmesChatResult",