From e936ce27313c866da968bff78c5c211a62b9662d Mon Sep 17 00:00:00 2001 From: Nicolas Herment Date: Thu, 30 Jan 2025 10:02:16 +0100 Subject: [PATCH 1/4] Prompt improvement for custom sections in RCA (#259) --- holmes/core/investigation_structured_output.py | 6 +++--- holmes/plugins/prompts/_general_instructions.jinja2 | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/holmes/core/investigation_structured_output.py b/holmes/core/investigation_structured_output.py index 3d467da..e47d28e 100644 --- a/holmes/core/investigation_structured_output.py +++ b/holmes/core/investigation_structured_output.py @@ -2,9 +2,9 @@ DEFAULT_SECTIONS = { "Alert Explanation": "1-2 sentences explaining the alert itself - note don't say \"The alert indicates a warning event related to a Kubernetes pod doing blah\" rather just say \"The pod XYZ did blah\" because that is what the user actually cares about", - "Investigation": "what you checked and found", - "Conclusions and Possible Root causes": "what conclusions can you reach based on the data you found? what are possible root causes (if you have enough conviction to say) or what uncertainty remains", - "Next Steps": "what you would do next to troubleshoot this issue, any commands that could be run to fix it, or other ways to solve it (prefer giving precise bash commands when possible)", + "Investigation": "What you checked and found", + "Conclusions and Possible Root causes": "What conclusions can you reach based on the data you found? what are possible root causes (if you have enough conviction to say) or what uncertainty remains. Don't say root cause but 'possible root causes'. Be clear to distinguish between what you know for certain and what is a possible explanation", + "Next Steps": "What you would do next to troubleshoot this issue, any commands that could be run to fix it, or other ways to solve it (prefer giving precise bash commands when possible)", "Related logs": "Truncate and share the most relevant logs, especially if these explain the root cause. For example: \nLogs from pod robusta-holmes:\n```\n```\n. Always embed the surroundding +/- 5 log lines to any relevant logs. ", "App or Infra?": "Explain whether the issue is more likely an infrastructure or an application level issue and why you think that.", "External links": "Provide links to external sources. Where to look when investigating this issue. For example provide links to relevant runbooks, etc. Add a short sentence describing each link." diff --git a/holmes/plugins/prompts/_general_instructions.jinja2 b/holmes/plugins/prompts/_general_instructions.jinja2 index 8126075..996cbfe 100644 --- a/holmes/plugins/prompts/_general_instructions.jinja2 +++ b/holmes/plugins/prompts/_general_instructions.jinja2 @@ -8,7 +8,6 @@ In general: * in this case, try to find substrings or search for the correct spellings * always provide detailed information like exact resource names, versions, labels, etc * even if you found the root cause, keep investigating to find other possible root causes and to gather data for the answer like exact names -* when giving an answer don't say root cause but "possible root causes" and be clear to distinguish between what you know for certain and what is a possible explanation * if a runbook url is present as well as tool that can fetch it, you MUST fetch the runbook before beginning your investigation. * if you don't know, say that the analysis was inconclusive. * if there are multiple possible causes list them in a numbered list. From ecf5d33b945af6cc1f7d551d0b367e5a153b8588 Mon Sep 17 00:00:00 2001 From: moshemorad Date: Thu, 30 Jan 2025 13:31:50 +0200 Subject: [PATCH 2/4] Add better example config for tools (#260) --- holmes/core/tools.py | 3 + .../toolsets/grafana/base_grafana_toolset.py | 18 ++++- holmes/plugins/toolsets/grafana/common.py | 13 +-- .../toolsets/grafana/toolset_grafana_loki.py | 20 ++++- holmes/plugins/toolsets/opensearch.py | 81 +++++++++++++++---- .../default_toolset_installation_guide.jinja2 | 4 + holmes/utils/holmes_sync_toolsets.py | 19 ++--- 7 files changed, 113 insertions(+), 45 deletions(-) diff --git a/holmes/core/tools.py b/holmes/core/tools.py index eff1fed..5030d12 100644 --- a/holmes/core/tools.py +++ b/holmes/core/tools.py @@ -377,6 +377,9 @@ def check_prerequisites(self): self._status = ToolsetStatusEnum.ENABLED + def get_example_config(self) -> Dict[str, Any]: + return {} + class YAMLToolset(Toolset): tools: List[YAMLTool] diff --git a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py index 9f29efc..dedb174 100644 --- a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py +++ b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py @@ -1,5 +1,5 @@ import logging -from typing import Any +from typing import Any, ClassVar, Type from holmes.core.tools import ( Tool, Toolset, @@ -11,6 +11,8 @@ class BaseGrafanaToolset(Toolset): + config_class: ClassVar[Type[GrafanaConfig]] = GrafanaConfig + def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool]): super().__init__( name=name, @@ -21,7 +23,8 @@ def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool] tags=[ ToolsetTag.CORE, ], - enabled=False + enabled=False, + is_default=True, ) def prerequisites_callable(self, config: dict[str, Any]) -> bool: @@ -30,10 +33,17 @@ def prerequisites_callable(self, config: dict[str, Any]) -> bool: return False try: - self._grafana_config = GrafanaConfig(**config) - is_healthy = get_health(self._grafana_config.url, self._grafana_config.api_key) + self._grafana_config = BaseGrafanaToolset.config_class(**config) + is_healthy = get_health( + self._grafana_config.url, self._grafana_config.api_key + ) return is_healthy except Exception: logging.exception("Failed to set up grafana toolset") return False + + def get_example_config(self): + example_config = GrafanaConfig(api_key="YOUR API KEY", url="YOUR GRAFANA URL") + return example_config.model_dump() + \ No newline at end of file diff --git a/holmes/plugins/toolsets/grafana/common.py b/holmes/plugins/toolsets/grafana/common.py index e5c5e1a..df1c214 100644 --- a/holmes/plugins/toolsets/grafana/common.py +++ b/holmes/plugins/toolsets/grafana/common.py @@ -1,23 +1,12 @@ from typing import Dict, Optional, Union import uuid import time -import os from pydantic import BaseModel - -GRAFANA_URL_ENV_NAME = "GRAFANA_URL" -GRAFANA_API_KEY_ENV_NAME = "GRAFANA_API_KEY" ONE_HOUR_IN_SECONDS = 3600 -class GrafanaLokiConfig(BaseModel): - pod_name_search_key: str = "pod" - namespace_search_key: str = "namespace" - node_name_search_key: str = "node" - - class GrafanaConfig(BaseModel): - loki: GrafanaLokiConfig = GrafanaLokiConfig() api_key: str url: str @@ -61,5 +50,5 @@ def get_datasource_id(dict: Dict, param: str) -> str: return f"uid/{datasource_id}" except: pass - + return datasource_id diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py index 0023359..7816e85 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py @@ -5,6 +5,7 @@ from holmes.core.tools import Tool, ToolParameter from holmes.plugins.toolsets.grafana.base_grafana_toolset import BaseGrafanaToolset from holmes.plugins.toolsets.grafana.common import ( + GrafanaConfig, get_datasource_id, get_param_or_raise, process_timestamps, @@ -17,6 +18,12 @@ ) +class GrafanaLokiConfig(GrafanaConfig): + pod_name_search_key: str = "pod" + namespace_search_key: str = "namespace" + node_name_search_key: str = "node" + + class ListLokiDatasources(Tool): def __init__(self, toolset: BaseGrafanaToolset): @@ -84,7 +91,7 @@ def invoke(self, params: Dict) -> str: api_key=self._toolset._grafana_config.api_key, loki_datasource_id=get_datasource_id(params, "loki_datasource_id"), node_name=get_param_or_raise(params, "node_name"), - node_name_search_key=self._toolset._grafana_config.loki.node_name_search_key, + node_name_search_key=self._toolset._grafana_config.node_name_search_key, start=start, end=end, limit=int(get_param_or_raise(params, "limit")), @@ -208,8 +215,8 @@ def invoke(self, params: Dict) -> str: loki_datasource_id=get_datasource_id(params, "loki_datasource_id"), pod_regex=get_param_or_raise(params, "pod_regex"), namespace=get_param_or_raise(params, "namespace"), - namespace_search_key=self._toolset._grafana_config.loki.namespace_search_key, - pod_name_search_key=self._toolset._grafana_config.loki.pod_name_search_key, + namespace_search_key=self._toolset._grafana_config.namespace_search_key, + pod_name_search_key=self._toolset._grafana_config.pod_name_search_key, start=start, end=end, limit=int(get_param_or_raise(params, "limit")), @@ -221,6 +228,8 @@ def get_parameterized_one_liner(self, params: Dict) -> str: class GrafanaLokiToolset(BaseGrafanaToolset): + config_class = GrafanaLokiConfig + def __init__(self): super().__init__( name="grafana/loki", @@ -233,3 +242,8 @@ def __init__(self): GetLokiLogsByLabel(self), ], ) + + def get_example_config(self): + example_config = GrafanaLokiConfig(api_key="YOUR API KEY", url="YOUR GRAFANA URL") + return example_config.model_dump() + \ No newline at end of file diff --git a/holmes/plugins/toolsets/opensearch.py b/holmes/plugins/toolsets/opensearch.py index 84d4c73..60e2714 100644 --- a/holmes/plugins/toolsets/opensearch.py +++ b/holmes/plugins/toolsets/opensearch.py @@ -1,7 +1,7 @@ import logging from typing import Any, Dict, List, Optional -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict from holmes.core.tools import ( CallablePrerequisite, Tool, @@ -12,18 +12,46 @@ from opensearchpy import OpenSearch +class OpenSearchHttpAuth(BaseModel): + username: str + password: str + + +class OpenSearchHost(BaseModel): + host: str + port: int = 9200 + + +class OpenSearchCluster(BaseModel): + hosts: list[OpenSearchHost] + headers: Optional[dict[str, Any]] = None + use_ssl: bool = True + ssl_assert_hostname: bool = False + verify_certs: bool = False + ssl_show_warn: bool = False + http_auth: Optional[OpenSearchHttpAuth] = None + + +class OpenSearchConfig(BaseModel): + opensearch_clusters: list[OpenSearchCluster] + + class OpenSearchClient: def __init__(self, **kwargs): - + # Handle http_auth explicitly if "http_auth" in kwargs: http_auth = kwargs.pop("http_auth") if isinstance(http_auth, dict): - kwargs["http_auth"] = (http_auth.get("username"), http_auth.get("password")) + kwargs["http_auth"] = ( + http_auth.get("username"), + http_auth.get("password"), + ) # Initialize OpenSearch client self.client = OpenSearch(**kwargs) -def get_client(clients:List[OpenSearchClient], host:Optional[str]): + +def get_client(clients: List[OpenSearchClient], host: Optional[str]): if len(clients) == 1: return clients[0] @@ -133,7 +161,7 @@ def __init__(self): enabled=False, description="Provide cluster metadata information like health, shards, settings.", docs_url="https://opensearch.org/docs/latest/clients/python-low-level/", - icon_url="https://upload.wikimedia.org/wikipedia/commons/9/91/Opensearch_Logo.svg", + icon_url="https://opensearch.org/assets/brand/PNG/Mark/opensearch_mark_default.png", prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)], tools=[ ListShards(self), @@ -143,21 +171,40 @@ def __init__(self): tags=[ ToolsetTag.CORE, ], - is_default=False, + is_default=True, ) def prerequisites_callable(self, config: dict[str, Any]) -> bool: if not config: return False - clusters_configs: list[dict[str, Any]] = config.get("opensearch_clusters", []) - for cluster in clusters_configs: - try: - logging.info(f"Setting up OpenSearch client") - client = OpenSearchClient(**cluster) - if client.client.cluster.health(params={"timeout": 5}): - self.clients.append(client) - except Exception: - logging.exception("Failed to set up opensearch client") - - return len(self.clients) > 0 + try: + os_config = OpenSearchConfig(**config) + + for cluster in os_config.opensearch_clusters: + try: + logging.info(f"Setting up OpenSearch client") + cluster_kwargs = cluster.model_dump() + client = OpenSearchClient(**cluster_kwargs) + if client.client.cluster.health(params={"timeout": 5}): + self.clients.append(client) + except Exception: + logging.exception("Failed to set up opensearch client") + + return len(self.clients) > 0 + except Exception: + logging.exception("Failed to set up grafana toolset") + return False + + def get_example_config(self) -> Dict[str, Any]: + example_config = OpenSearchConfig( + opensearch_clusters=[ + OpenSearchCluster( + hosts=[OpenSearchHost(host="YOUR OPENSEACH HOST")], + headers={"Authorization": "{{ env.OPENSEARCH_BEARER_TOKEN }}"}, + use_ssl=True, + ssl_assert_hostname=False, + ) + ] + ) + return example_config.model_dump() diff --git a/holmes/utils/default_toolset_installation_guide.jinja2 b/holmes/utils/default_toolset_installation_guide.jinja2 index 1cb32dd..44014ff 100644 --- a/holmes/utils/default_toolset_installation_guide.jinja2 +++ b/holmes/utils/default_toolset_installation_guide.jinja2 @@ -25,6 +25,10 @@ holmes: toolsets: {{toolset_name}}: enabled: true + {% if example_config %} + config: + {{ example_config | indent(8) }} + {% endif %} ``` {% endif %} diff --git a/holmes/utils/holmes_sync_toolsets.py b/holmes/utils/holmes_sync_toolsets.py index a485b03..7ddec7e 100644 --- a/holmes/utils/holmes_sync_toolsets.py +++ b/holmes/utils/holmes_sync_toolsets.py @@ -1,4 +1,7 @@ from datetime import datetime +from typing import Any + +import yaml from holmes.config import Config @@ -47,19 +50,17 @@ def holmes_sync_toolsets_status(dal: SupabaseDal, config: Config) -> None: def render_default_installation_instructions_for_toolset(toolset: Toolset) -> str: env_vars = toolset.get_environment_variables() - context = { + context: dict[str, Any] = { "env_vars": env_vars if env_vars else [], "toolset_name": toolset.name, "enabled": toolset.enabled, - "default_toolset": toolset.is_default, + "example_config": yaml.dump(toolset.get_example_config()), } - if toolset.is_default: - installation_instructions = load_and_render_prompt( - "file://holmes/utils/default_toolset_installation_guide.jinja2", context - ) - return installation_instructions - installation_instructions = load_and_render_prompt( - "file://holmes/utils/installation_guide.jinja2", context + template = ( + "file://holmes/utils/default_toolset_installation_guide.jinja2" + if toolset.is_default + else "file://holmes/utils/installation_guide.jinja2" ) + installation_instructions = load_and_render_prompt(template, context) return installation_instructions From 415f63778447fd28ccdb63c93b822710b20aa4d2 Mon Sep 17 00:00:00 2001 From: moshemorad Date: Thu, 30 Jan 2025 14:29:50 +0200 Subject: [PATCH 3/4] Update docs (#261) --- README.md | 32 +++++++++---------- .../toolsets/grafana/base_grafana_toolset.py | 3 +- .../toolsets/grafana/toolset_grafana_loki.py | 1 + .../toolsets/grafana/toolset_grafana_tempo.py | 1 + 4 files changed, 19 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 0f3618a..f4dffe8 100644 --- a/README.md +++ b/README.md @@ -631,31 +631,31 @@ Using Grafana Loki HolmesGPT can consult logs from [Loki](https://grafana.com/oss/loki/) by proxying through a [Grafana](https://grafana.com/oss/grafana/) instance. -There are 2 parts to configuring access to Grafana Loki: Access/Authentication and search terms. +To configure loki toolset: -For access and authentication, add the following environment variables: - -* `GRAFANA_URL` - e.g. https://my-org.grafana.net -* `GRAFANA_API_KEY` - e.g. glsa_bsm6ZS_sdfs25f +```yaml +toolsets: + grafana/loki: + enabled: true + config: + api_key: "{{ env.GRAFANA_API_KEY }}" + url: "http://loki-url" +``` For search terms, you can optionally tweak the search terms used by the toolset. -This is done by appending the following to your Holmes configuration file: +This is done by appending the following to your Holmes grafana/loki configuration: ```yaml -grafana: - url: https://my-org.grafana.net # - api_key: glsa_bsm6ZS_sdfs25f - loki: - pod_name_search_key: "pod" - namespace_search_key: "namespace" - node_name_search_key: "node" +pod_name_search_key: "pod" +namespace_search_key: "namespace" +node_name_search_key: "node" ``` > You only need to tweak the configuration file if your Loki logs settings for pod, namespace and node differ from the above defaults. -The Loki toolset is configured the using the same Grafana settings as the Grafana Tempo toolset. +
Using Grafana Tempo @@ -664,8 +664,6 @@ HolmesGPT can fetch trace information from Grafana Tempo to debug performance re Tempo is configured the using the same Grafana settings as the Grafana Loki toolset. -grafana: - url: https://my-org.grafana.net #
@@ -875,7 +873,7 @@ Configure Slack to send notifications to specific channels. Provide your Slack t OpenSearch Integration The OpenSearch toolset (`opensearch`) allows Holmes to consult an opensearch cluster for its health, settings and shards information. -The toolset supports multiple opensearch or elasticsearch clusters that are configured by editing Holmes' configuration file (or in cluster to the configuration secret): +The toolset supports multiple opensearch or elasticsearch clusters that are configured by editing Holmes' configuration file: ``` opensearch_clusters: diff --git a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py index dedb174..66599a6 100644 --- a/holmes/plugins/toolsets/grafana/base_grafana_toolset.py +++ b/holmes/plugins/toolsets/grafana/base_grafana_toolset.py @@ -13,11 +13,12 @@ class BaseGrafanaToolset(Toolset): config_class: ClassVar[Type[GrafanaConfig]] = GrafanaConfig - def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool]): + def __init__(self, name: str, description: str, icon_url: str, tools: list[Tool], doc_url: str): super().__init__( name=name, description=description, icon_url=icon_url, + docs_url=doc_url, prerequisites=[CallablePrerequisite(callable=self.prerequisites_callable)], tools=tools, tags=[ diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py index 7816e85..2c1a002 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_loki.py @@ -235,6 +235,7 @@ def __init__(self): name="grafana/loki", description="Fetchs kubernetes pods and node logs from Loki", icon_url="https://grafana.com/media/docs/loki/logo-grafana-loki.png", + doc_url="https://grafana.com/oss/loki/", tools=[ ListLokiDatasources(self), GetLokiLogsByNode(self), diff --git a/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py b/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py index 815ee27..956ff60 100644 --- a/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py +++ b/holmes/plugins/toolsets/grafana/toolset_grafana_tempo.py @@ -137,6 +137,7 @@ def __init__(self): name="grafana/tempo", description="Fetchs kubernetes traces from Tempo", icon_url="https://grafana.com/static/assets/img/blog/tempo.png", + doc_url="https://grafana.com/oss/tempo/", tools=[ ListAllDatasources(self), GetTempoTracesByMinDuration(self), From 99ac32ca6d4e95c03baff9ba38a56e594c6349d7 Mon Sep 17 00:00:00 2001 From: Nicolas Herment Date: Thu, 30 Jan 2025 14:06:49 +0100 Subject: [PATCH 4/4] fix: issue with sections handling messing workload health structured output (#262) The code to parse the sections for the custom RCA output was in the wrong place, causing it meddle with the structured output from the workload health checks. This PR moves this code out of the tool_calling_llm.py into the investigation helper --- holmes/core/investigation.py | 11 ++-- .../core/investigation_structured_output.py | 53 ++++++++++++++----- holmes/core/models.py | 3 +- holmes/core/tool_calling_llm.py | 22 ++------ server.py | 1 - 5 files changed, 52 insertions(+), 38 deletions(-) diff --git a/holmes/core/investigation.py b/holmes/core/investigation.py index c7d7879..3b1407f 100644 --- a/holmes/core/investigation.py +++ b/holmes/core/investigation.py @@ -1,8 +1,7 @@ -from typing import Optional -from rich.console import Console from holmes.common.env_vars import HOLMES_POST_PROCESSING_PROMPT from holmes.config import Config +from holmes.core.investigation_structured_output import process_response_into_sections from holmes.core.issue import Issue from holmes.core.models import InvestigateRequest, InvestigationResult from holmes.core.supabase_dal import SupabaseDal @@ -36,13 +35,15 @@ def investigate_issues(investigate_request: InvestigateRequest, dal: SupabaseDal issue, prompt=investigate_request.prompt_template, post_processing_prompt=HOLMES_POST_PROCESSING_PROMPT, - sections=investigate_request.sections, instructions=resource_instructions, global_instructions=global_instructions ) + + (text_response, sections) = process_response_into_sections(investigation.result) + return InvestigationResult( - analysis=investigation.result, - sections=investigation.sections, + analysis=text_response, + sections=sections, tool_calls=investigation.tool_calls or [], instructions=investigation.instructions, ) diff --git a/holmes/core/investigation_structured_output.py b/holmes/core/investigation_structured_output.py index e47d28e..abdf6e8 100644 --- a/holmes/core/investigation_structured_output.py +++ b/holmes/core/investigation_structured_output.py @@ -1,6 +1,15 @@ -from typing import Any, Dict +from typing import Any, Dict, Optional, Tuple, Union +import json -DEFAULT_SECTIONS = { +from pydantic import RootModel + +InputSectionsDataType = Dict[str, str] + +OutputSectionsDataType = Optional[Dict[str, Union[str, None]]] + +SectionsData = RootModel[OutputSectionsDataType] + +DEFAULT_SECTIONS:InputSectionsDataType = { "Alert Explanation": "1-2 sentences explaining the alert itself - note don't say \"The alert indicates a warning event related to a Kubernetes pod doing blah\" rather just say \"The pod XYZ did blah\" because that is what the user actually cares about", "Investigation": "What you checked and found", "Conclusions and Possible Root causes": "What conclusions can you reach based on the data you found? what are possible root causes (if you have enough conviction to say) or what uncertainty remains. Don't say root cause but 'possible root causes'. Be clear to distinguish between what you know for certain and what is a possible explanation", @@ -10,7 +19,7 @@ "External links": "Provide links to external sources. Where to look when investigating this issue. For example provide links to relevant runbooks, etc. Add a short sentence describing each link." } -def get_output_format_for_investigation(sections: Dict[str, str]) -> Dict[str, Any]: +def get_output_format_for_investigation(sections: InputSectionsDataType) -> Dict[str, Any]: properties = {} required_fields = [] @@ -34,12 +43,32 @@ def get_output_format_for_investigation(sections: Dict[str, str]) -> Dict[str, A return output_format -def combine_sections(sections: Any) -> str: - if isinstance(sections, dict): - content = '' - for section_title, section_content in sections.items(): - if section_content: - # content = content + f'\n# {" ".join(section_title.split("_")).title()}\n{section_content}' - content = content + f'\n# {section_title}\n{section_content}\n' - return content - return f"{sections}" +def combine_sections(sections: Dict) -> str: + content = '' + for section_title, section_content in sections.items(): + if section_content: + content = content + f'\n# {section_title}\n{section_content}\n' + return content + + +def process_response_into_sections(response: Any) -> Tuple[str, OutputSectionsDataType]: + if isinstance(response, dict): + # No matter if the result is already structured, we want to go through the code below to validate the JSON + response = json.dumps(response) + + if not isinstance(response, str): + # if it's not a string, we make it so as it'll be parsed later + response = str(response) + + + try: + parsed_json = json.loads(response) + # TODO: force dict values into a string would make this more resilient as SectionsData only accept none/str as values + sections = SectionsData(root=parsed_json).root + if sections: + combined = combine_sections(sections) + return (combined, sections) + except Exception: + pass + + return (response, None) diff --git a/holmes/core/models.py b/holmes/core/models.py index 71683cf..1a8ed0e 100644 --- a/holmes/core/models.py +++ b/holmes/core/models.py @@ -1,3 +1,4 @@ +from holmes.core.investigation_structured_output import InputSectionsDataType from holmes.core.tool_calling_llm import ToolCallResult from typing import Optional, List, Dict, Any, Union from pydantic import BaseModel, model_validator @@ -21,7 +22,7 @@ class InvestigateRequest(BaseModel): include_tool_calls: bool = False include_tool_call_results: bool = False prompt_template: str = "builtin://generic_investigation.jinja2" - sections: Optional[Dict[str, str]] = None + sections: Optional[InputSectionsDataType] = None # TODO in the future # response_handler: ... diff --git a/holmes/core/tool_calling_llm.py b/holmes/core/tool_calling_llm.py index b0c1a44..2183409 100644 --- a/holmes/core/tool_calling_llm.py +++ b/holmes/core/tool_calling_llm.py @@ -3,7 +3,7 @@ import logging import textwrap from typing import List, Optional, Dict, Type, Union -from holmes.core.investigation_structured_output import DEFAULT_SECTIONS, get_output_format_for_investigation, combine_sections +from holmes.core.investigation_structured_output import DEFAULT_SECTIONS, InputSectionsDataType, get_output_format_for_investigation from holmes.core.performance_timing import PerformanceTiming from holmes.utils.tags import format_tags_in_string, parse_messages_tags from holmes.plugins.prompts import load_and_render_prompt @@ -27,14 +27,11 @@ class ToolCallResult(BaseModel): description: str result: str - class LLMResult(BaseModel): tool_calls: Optional[List[ToolCallResult]] = None - sections: Optional[Dict[str, Union[str, None]]] = None result: Optional[str] = None unprocessed_result: Optional[str] = None instructions: List[str] = [] - # TODO: clean up these two prompt: Optional[str] = None messages: Optional[List[dict]] = None @@ -159,22 +156,12 @@ def call( tools_to_call = getattr(response_message, "tool_calls", None) text_response = response_message.content - sections:Optional[Dict[str, str]] = None - if isinstance(text_response, str): - try: - parsed_json = json.loads(text_response) - text_response = parsed_json - except json.JSONDecodeError: - pass - if not isinstance(text_response, str): - sections = text_response - text_response = combine_sections(sections) if not tools_to_call: # For chatty models post process and summarize the result # this only works for calls where user prompt is explicitly passed through if post_process_prompt and user_prompt: - logging.info(f"Running post processing on investigation.") + logging.info("Running post processing on investigation.") raw_response = text_response post_processed_response = self._post_processing_call( prompt=user_prompt, @@ -185,7 +172,6 @@ def call( perf_timing.end() return LLMResult( result=post_processed_response, - sections=sections, unprocessed_result=raw_response, tool_calls=tool_calls, prompt=json.dumps(messages, indent=2), @@ -195,7 +181,6 @@ def call( perf_timing.end() return LLMResult( result=text_response, - sections=sections, tool_calls=tool_calls, prompt=json.dumps(messages, indent=2), messages=messages, @@ -231,7 +216,6 @@ def _invoke_tool( logging.warning( f"Failed to parse arguments for tool: {tool_name}. args: {tool_to_call.function.arguments}" ) - tool_call_id = tool_to_call.id tool = self.tool_executor.get_tool_by_name(tool_name) @@ -358,7 +342,7 @@ def investigate( console: Optional[Console] = None, global_instructions: Optional[Instructions] = None, post_processing_prompt: Optional[str] = None, - sections: Optional[Dict[str, str]] = None + sections: Optional[InputSectionsDataType] = None ) -> LLMResult: runbooks = self.runbook_manager.get_instructions_for_issue(issue) diff --git a/server.py b/server.py index 8df2261..15e8d46 100644 --- a/server.py +++ b/server.py @@ -148,7 +148,6 @@ def workload_health_check(request: WorkloadHealthRequest): system_prompt = load_and_render_prompt(request.prompt_template, context={'alerts': workload_alerts}) - ai = config.create_toolcalling_llm(dal=dal) structured_output = {"type": "json_object"}