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

holmes/plugins/toolsets/kafka.py

@@ -1,7 +1,10 @@
 import logging
+from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import yaml  # type: ignore
+from confluent_kafka import Consumer
+from confluent_kafka._model import Node
 from confluent_kafka.admin import (
     AdminClient,
     BrokerMetadata,
@@ -17,18 +20,16 @@ from confluent_kafka.admin import (
     PartitionMetadata,
     TopicMetadata,
 )
-from confluent_kafka import Consumer
-from confluent_kafka._model import Node
-from enum import Enum
 from confluent_kafka.admin import _TopicPartition as TopicPartition
 from pydantic import BaseModel, ConfigDict
 
 from holmes.core.tools import (
     CallablePrerequisite,
     StructuredToolResult,
+    StructuredToolResultStatus,
     Tool,
+    ToolInvokeContext,
     ToolParameter,
-    ToolResultStatus,
     Toolset,
     ToolsetTag,
 )
@@ -153,15 +154,13 @@ class ListKafkaConsumers(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             kafka_cluster_name = get_param_or_raise(params, "kafka_cluster_name")
             client = self.get_kafka_client(kafka_cluster_name)
             if client is None:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="No admin_client on toolset. This toolset is misconfigured.",
                     params=params,
                 )
@@ -190,7 +189,7 @@ class ListKafkaConsumers(BaseKafkaTool):
             if errors_text:
                 result_text = result_text + "\n\n" + errors_text
             return StructuredToolResult(
-                status=ToolResultStatus.SUCCESS,
+                status=StructuredToolResultStatus.SUCCESS,
                 data=result_text,
                 params=params,
             )
@@ -198,7 +197,7 @@ class ListKafkaConsumers(BaseKafkaTool):
             error_msg = f"Failed to list consumer groups: {str(e)}"
             logging.error(error_msg)
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=error_msg,
                 params=params,
             )
@@ -228,16 +227,14 @@ class DescribeConsumerGroup(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         group_id = params["group_id"]
         try:
             kafka_cluster_name = get_param_or_raise(params, "kafka_cluster_name")
             client = self.get_kafka_client(kafka_cluster_name)
             if client is None:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="No admin_client on toolset. This toolset is misconfigured.",
                     params=params,
                 )
@@ -247,13 +244,13 @@ class DescribeConsumerGroup(BaseKafkaTool):
             if futures.get(group_id):
                 group_metadata = futures.get(group_id).result()
                 return StructuredToolResult(
-                    status=ToolResultStatus.SUCCESS,
-                    data=yaml.dump(convert_to_dict(group_metadata)),
+                    status=StructuredToolResultStatus.SUCCESS,
+                    data=convert_to_dict(group_metadata),
                     params=params,
                 )
             else:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="Group not found",
                     params=params,
                 )
@@ -261,7 +258,7 @@ class DescribeConsumerGroup(BaseKafkaTool):
             error_msg = f"Failed to describe consumer group {group_id}: {str(e)}"
             logging.error(error_msg)
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=error_msg,
                 params=params,
             )
@@ -286,30 +283,28 @@ class ListTopics(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         try:
             kafka_cluster_name = get_param_or_raise(params, "kafka_cluster_name")
             client = self.get_kafka_client(kafka_cluster_name)
             if client is None:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="No admin_client on toolset. This toolset is misconfigured.",
                     params=params,
                 )
 
             topics = client.list_topics()
             return StructuredToolResult(
-                status=ToolResultStatus.SUCCESS,
-                data=yaml.dump(convert_to_dict(topics)),
+                status=StructuredToolResultStatus.SUCCESS,
+                data=convert_to_dict(topics),
                 params=params,
             )
         except Exception as e:
             error_msg = f"Failed to list topics: {str(e)}"
             logging.error(error_msg)
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=error_msg,
                 params=params,
             )
@@ -344,16 +339,14 @@ class DescribeTopic(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         topic_name = params["topic_name"]
         try:
             kafka_cluster_name = get_param_or_raise(params, "kafka_cluster_name")
             client = self.get_kafka_client(kafka_cluster_name)
             if client is None:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="No admin_client on toolset. This toolset is misconfigured.",
                     params=params,
                 )
@@ -373,15 +366,15 @@ class DescribeTopic(BaseKafkaTool):
                 result["configuration"] = convert_to_dict(config)
 
             return StructuredToolResult(
-                status=ToolResultStatus.SUCCESS,
-                data=yaml.dump(result),
+                status=StructuredToolResultStatus.SUCCESS,
+                data=result,
                 params=params,
             )
         except Exception as e:
             error_msg = f"Failed to describe topic {topic_name}: {str(e)}"
             logging.error(error_msg, exc_info=True)
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=error_msg,
                 params=params,
             )
@@ -469,16 +462,14 @@ class FindConsumerGroupsByTopic(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         topic_name = params["topic_name"]
         try:
             kafka_cluster_name = get_param_or_raise(params, "kafka_cluster_name")
             client = self.get_kafka_client(kafka_cluster_name)
             if client is None:
                 return StructuredToolResult(
-                    status=ToolResultStatus.ERROR,
+                    status=StructuredToolResultStatus.ERROR,
                     error="No admin_client on toolset. This toolset is misconfigured.",
                     params=params,
                 )
@@ -530,7 +521,7 @@ class FindConsumerGroupsByTopic(BaseKafkaTool):
                 result_text = result_text + "\n\n" + errors_text
 
             return StructuredToolResult(
-                status=ToolResultStatus.SUCCESS,
+                status=StructuredToolResultStatus.SUCCESS,
                 data=result_text,
                 params=params,
             )
@@ -540,7 +531,7 @@ class FindConsumerGroupsByTopic(BaseKafkaTool):
             )
             logging.error(error_msg)
             return StructuredToolResult(
-                status=ToolResultStatus.ERROR,
+                status=StructuredToolResultStatus.ERROR,
                 error=error_msg,
                 params=params,
             )
@@ -559,12 +550,10 @@ class ListKafkaClusters(BaseKafkaTool):
             toolset=toolset,
         )
 
-    def _invoke(
-        self, params: dict, user_approved: bool = False
-    ) -> StructuredToolResult:
+    def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolResult:
         cluster_names = list(self.toolset.clients.keys())
         return StructuredToolResult(
-            status=ToolResultStatus.SUCCESS,
+            status=StructuredToolResultStatus.SUCCESS,
             data="Available Kafka Clusters:\n" + "\n".join(cluster_names),
             params=params,
         )
@@ -610,6 +599,8 @@ class KafkaToolset(Toolset):
                     admin_config = {
                         "bootstrap.servers": cluster.kafka_broker,
                         "client.id": cluster.kafka_client_id,
+                        "socket.timeout.ms": 15000,  # 15 second timeout
+                        "api.version.request.timeout.ms": 15000,  # 15 second API version timeout
                     }
 
                     if cluster.kafka_security_protocol:
@@ -623,6 +614,9 @@ class KafkaToolset(Toolset):
                         admin_config["sasl.password"] = cluster.kafka_password
 
                     client = AdminClient(admin_config)
+                    # Test the connection by trying to list topics with a timeout
+                    # This will fail fast if the broker is not reachable
+                    _ = client.list_topics(timeout=10)  # 10 second timeout
                     self.clients[cluster.name] = client  # Store in dictionary
                 except Exception as e:
                     message = (