From ec514f0fa9f684e32802bc69d95426928307d2ee Mon Sep 17 00:00:00 2001 From: droid-dhruv Date: Wed, 20 Aug 2025 14:01:48 +0530 Subject: [PATCH 1/3] added methods in sdk for ai to list and execute tools --- drdroid_debug_toolkit/sdk_v2.py | 227 +++++++++++++++++++++++++++++++- 1 file changed, 226 insertions(+), 1 deletion(-) diff --git a/drdroid_debug_toolkit/sdk_v2.py b/drdroid_debug_toolkit/sdk_v2.py index dbd3a5f..2f39e28 100644 --- a/drdroid_debug_toolkit/sdk_v2.py +++ b/drdroid_debug_toolkit/sdk_v2.py @@ -14,6 +14,7 @@ from .core.sources.posthog_sdk import PostHogSDK from .core.sources.sql_database_connection_sdk import SqlDatabaseConnectionSDK from .core.sources.clickhouse_sdk import ClickHouseSDK +from .core.protos.literal_pb2 import LiteralType logger = logging.getLogger(__name__) @@ -216,4 +217,228 @@ def list_available_tasks(self, source_name: str) -> List[str]: if not method.startswith('_') and callable(getattr(sdk, method))] except Exception as e: logger.error(f"Failed to list tasks for {source_name}: {e}") - return [] \ No newline at end of file + return [] + + # ============================================================================ + # LLM Tooling + # ============================================================================ + + def _convert_literal_type_to_json_type(self, literal_type: Any) -> str: + """Convert protobuf LiteralType to JSON Schema type string. + + Fallbacks to "string" for unknown/unsupported types. + """ + try: + if literal_type == LiteralType.STRING: + return "string" + if literal_type == LiteralType.LONG: + return "integer" + if literal_type == LiteralType.DOUBLE: + return "number" + if literal_type == LiteralType.BOOLEAN: + return "boolean" + if literal_type == LiteralType.STRING_ARRAY: + return "array" + return "string" + except Exception: + return "string" + + def _convert_form_field_to_json_schema(self, field: Any) -> Dict[str, Any]: + """Convert a FormField proto to JSON Schema property definition.""" + try: + description = "" + try: + if getattr(field, "description", None) and getattr(field.description, "value", None): + description = field.description.value + elif getattr(field, "display_name", None) and getattr(field.display_name, "value", None): + description = field.display_name.value + except Exception: + description = "" + + field_schema: Dict[str, Any] = { + "type": self._convert_literal_type_to_json_type(getattr(field, "data_type", None)), + "description": description, + } + + # Handle array item typing for STRING_ARRAY + if getattr(field, "data_type", None) == LiteralType.STRING_ARRAY: + field_schema["items"] = {"type": "string"} + + # Default value (Literal) + try: + if hasattr(field, "HasField") and field.HasField("default_value"): + default_val = field.default_value + if default_val.type == LiteralType.STRING and getattr(default_val, "string", None): + field_schema["default"] = default_val.string.value + elif default_val.type == LiteralType.LONG and getattr(default_val, "long", None): + field_schema["default"] = default_val.long.value + elif default_val.type == LiteralType.BOOLEAN and getattr(default_val, "boolean", None): + field_schema["default"] = default_val.boolean.value + elif default_val.type == LiteralType.DOUBLE and getattr(default_val, "double", None): + field_schema["default"] = default_val.double.value + except Exception: + # best-effort; ignore default extraction failures + pass + + return field_schema + except Exception as e: + logger.error(f"Error converting form field: {e}") + return {"type": "string", "description": ""} + + def get_llm_tools(self, source_name: str) -> List[Dict[str, Any]]: + """ + Generate LLM-friendly tool definitions for a given source. + + The returned format is suitable for passing directly to an LLM that + supports tool/function calling. Each tool corresponds to a task defined + by the underlying SourceManager for the given source. + + Args: + source_name: Name of the source (e.g., "grafana") + + Returns: + A list of dicts with keys: name, description, parameters + """ + try: + sdk = self.get_sdk(source_name) + # Prefer using the BaseSDK helper to retrieve the proper manager instance + source_manager = sdk._get_source_manager(source_name) + + # Build a clean source prefix for tool names + source_prefix = source_name.lower().replace(" ", "_").replace("-", "_") + + tools: List[Dict[str, Any]] = [] + + task_type_callable_map = getattr(source_manager, "task_type_callable_map", {}) + task_proto = getattr(source_manager, "task_proto", None) + + for task_type, task_info in task_type_callable_map.items(): + try: + # Derive a human-readable task type name + task_type_name = str(task_type).lower() + if task_proto is not None and hasattr(task_proto, "TaskType") and hasattr(task_proto.TaskType, "Name"): + try: + task_type_name = task_proto.TaskType.Name(task_type).lower() + if task_type_name.startswith("task_type_"): + task_type_name = task_type_name[10:] + elif task_type_name.startswith("task_"): + task_type_name = task_type_name[5:] + except Exception: + pass + + tool_name = f"{source_prefix}_{task_type_name}" + + # Safety truncate + if len(tool_name) > 60: + tool_name = tool_name[:60] + + # Convert form fields to JSON Schema + properties: Dict[str, Any] = {} + required: List[str] = [] + + for field in task_info.get("form_fields", []): + try: + field_key = ( + field.key_name.value + if getattr(field, "key_name", None) and getattr(field.key_name, "value", None) + else f"field_{len(properties)}" + ) + field_schema = self._convert_form_field_to_json_schema(field) + properties[field_key] = field_schema + + # Required if not is_optional (default to required if flag missing/False) + is_optional = bool(getattr(field, "is_optional", False)) + if not is_optional: + required.append(field_key) + except Exception as field_err: + logger.error(f"Error processing form field for tool '{tool_name}': {field_err}") + continue + + tool_def = { + "name": tool_name, + "description": task_info.get("display_name", f"{source_name.title()} task"), + "parameters": { + "type": "object", + "properties": properties, + "required": required, + }, + } + + tools.append(tool_def) + except Exception as task_err: + logger.error(f"Error creating LLM tool for task_type {task_type}: {task_err}") + continue + + return tools + except Exception as e: + logger.error(f"Failed to generate LLM tools for {source_name}: {e}") + return [] + + def execute_llm_tool(self, source_name: str, tool_name: str, **kwargs) -> Dict[str, Any]: + """ + Execute an LLM tool by name with the provided arguments. + + This method allows an LLM to execute any tool returned by get_llm_tools(). + The tool_name should match the "name" field from the tool definition. + + Args: + source_name: Name of the source (e.g., "grafana") + tool_name: Name of the tool to execute (e.g., "grafana_datasource_query_execution") + **kwargs: Arguments to pass to the tool + + Returns: + Dictionary containing the execution result + + Raises: + ValidationError: If source or tool is not found + TaskExecutionError: If tool execution fails + """ + try: + # Get the SDK for the source + sdk = self.get_sdk(source_name) + source_manager = sdk._get_source_manager(source_name) + + # Extract task type from tool name + # Tool name format: {source}_{task_type} + task_type_name = tool_name.replace(f"{source_name.lower()}_", "", 1) + + # Find the matching task type in the source manager + task_type_callable_map = getattr(source_manager, "task_type_callable_map", {}) + task_proto = getattr(source_manager, "task_proto", None) + + matching_task_type = None + for task_type, task_info in task_type_callable_map.items(): + # Try to match by task type name + try: + if task_proto and hasattr(task_proto, "TaskType") and hasattr(task_proto.TaskType, "Name"): + current_task_name = task_proto.TaskType.Name(task_type).lower() + if current_task_name.startswith("task_type_"): + current_task_name = current_task_name[10:] + elif current_task_name.startswith("task_"): + current_task_name = current_task_name[5:] + + if current_task_name == task_type_name: + matching_task_type = task_type + break + except Exception: + pass + + # Fallback: try string comparison + if str(task_type).lower() == task_type_name: + matching_task_type = task_type + break + + if matching_task_type is None: + available_tools = [f"{source_name}_{str(t).lower()}" for t in task_type_callable_map.keys()] + raise ValidationError( + f"Tool '{tool_name}' not found for source '{source_name}'. " + f"Available tools: {available_tools}" + ) + + # Execute the task using the existing SDK infrastructure + return sdk.execute_task(source_name, task_type_name, **kwargs) + + except ValidationError: + raise + except Exception as e: + raise TaskExecutionError(f"LLM tool execution failed: {e}") \ No newline at end of file From 5c004e4ac01d73b2e0bf9e72abc8195cf29ecb86 Mon Sep 17 00:00:00 2001 From: droid-dhruv Date: Wed, 20 Aug 2025 14:32:17 +0530 Subject: [PATCH 2/3] minor fix --- drdroid_debug_toolkit/sdk_v2.py | 58 +++++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/drdroid_debug_toolkit/sdk_v2.py b/drdroid_debug_toolkit/sdk_v2.py index 2f39e28..d900e99 100644 --- a/drdroid_debug_toolkit/sdk_v2.py +++ b/drdroid_debug_toolkit/sdk_v2.py @@ -4,6 +4,7 @@ import logging from typing import Dict, Any, List, Union +from google.protobuf.struct_pb2 import Struct from .exceptions import ConfigurationError, ConnectionError, ValidationError, TaskExecutionError from .core.sdk_factory import SDKFactory @@ -15,6 +16,9 @@ from .core.sources.sql_database_connection_sdk import SqlDatabaseConnectionSDK from .core.sources.clickhouse_sdk import ClickHouseSDK from .core.protos.literal_pb2 import LiteralType +from .core.protos.playbooks.playbook_pb2 import PlaybookTask +from .core.utils.proto_utils import dict_to_proto, proto_to_dict +from .core.integrations.utils.executor_utils import check_multiple_task_results logger = logging.getLogger(__name__) @@ -435,8 +439,58 @@ def execute_llm_tool(self, source_name: str, tool_name: str, **kwargs) -> Dict[s f"Available tools: {available_tools}" ) - # Execute the task using the existing SDK infrastructure - return sdk.execute_task(source_name, task_type_name, **kwargs) + # Build a PlaybookTask for execution through the SourceManager + # Create minimal task dict structure expected by SourceManager.get_resolved_task + # 1) Determine source enum value + from .core.protos.base_pb2 import Source as SourceEnum, TimeRange + try: + source_enum_value = getattr(SourceEnum, source_name.upper()) + except AttributeError: + raise ValidationError(f"Unsupported source '{source_name}'") + + # 2) Prepare task dictionary + # Map kwargs directly as the task-type specific message payload + source_key = source_name.lower() + task_dict = { + "source": source_enum_value, + source_key: { + "type": matching_task_type, + task_type_name: kwargs or {} + } + } + + # 3) Convert to PlaybookTask proto + task_proto: PlaybookTask = dict_to_proto(task_dict, PlaybookTask) + + # 4) Create a default time range (last 1 hour) + from datetime import datetime, timedelta + end_time = int(datetime.now().timestamp()) + start_time = int((datetime.now() - timedelta(hours=1)).timestamp()) + tr = TimeRange(time_geq=start_time, time_lt=end_time) + + # 5) Resolve task and execute via SourceManager + resolved_task, resolved_source_task, task_local_variable_map = source_manager.get_resolved_task( + Struct(), task_proto + ) + + playbook_task_result = source_manager.task_type_callable_map[matching_task_type]['executor']( + tr, resolved_source_task, sdk._get_connector(source_name) + ) + + # 6) Post-process result and return as dict + if check_multiple_task_results(playbook_task_result): + task_results = [] + for result in playbook_task_result: + processed_result = source_manager.postprocess_task_result( + result, resolved_task, task_local_variable_map + ) + task_results.append(processed_result) + return {"results": [proto_to_dict(r) for r in task_results]} + else: + processed_result = source_manager.postprocess_task_result( + playbook_task_result, resolved_task, task_local_variable_map + ) + return proto_to_dict(processed_result) except ValidationError: raise From 30e82f729f4b44527aea980ddaffa7f61df8fe20 Mon Sep 17 00:00:00 2001 From: Karanveer Singh Sirohi Date: Thu, 21 Aug 2025 13:06:36 +0530 Subject: [PATCH 3/3] adds a get_openai_functions method --- drdroid_debug_toolkit/sdk_v2.py | 34 +++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/drdroid_debug_toolkit/sdk_v2.py b/drdroid_debug_toolkit/sdk_v2.py index d900e99..d941ef0 100644 --- a/drdroid_debug_toolkit/sdk_v2.py +++ b/drdroid_debug_toolkit/sdk_v2.py @@ -378,6 +378,40 @@ def get_llm_tools(self, source_name: str) -> List[Dict[str, Any]]: logger.error(f"Failed to generate LLM tools for {source_name}: {e}") return [] + def get_openai_functions(self, source_name: str) -> List[Dict[str, Any]]: + """ + Generate OpenAI function format tools for a given source. + + Returns tools in the OpenAI function calling format, eliminating the need + for manual conversion loops. + + Args: + source_name: Name of the source (e.g., "grafana") + + Returns: + A list of OpenAI function format dicts with structure: + { + "type": "function", + "function": { + "name": "tool_name", + "description": "tool_description", + "parameters": {...} + } + } + """ + tools = self.get_llm_tools(source_name) + return [ + { + "type": "function", + "function": { + "name": tool["name"], + "description": tool["description"], + "parameters": tool["parameters"] + } + } + for tool in tools + ] + def execute_llm_tool(self, source_name: str, tool_name: str, **kwargs) -> Dict[str, Any]: """ Execute an LLM tool by name with the provided arguments.