Main 1907 workload health endpoint

.gitignore

@@ -155,3 +155,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+.vscode

holmes/core/runbooks.py

@@ -1,9 +1,7 @@
-import logging
-import textwrap
 from typing import List
 from holmes.core.issue import Issue
 from holmes.plugins.runbooks import Runbook
-import jinja2
+
 
 # TODO: our default prompt has a lot of kubernetes specific stuff - see if we can get that into the runbook
 class RunbookManager:

holmes/core/supabase_dal.py

@@ -12,12 +12,13 @@
 from holmes.common.env_vars import (ROBUSTA_CONFIG_PATH, ROBUSTA_ACCOUNT_ID, STORE_URL, STORE_API_KEY, STORE_EMAIL,
                                     STORE_PASSWORD)
 
+from datetime import datetime, timedelta
 
 SUPABASE_TIMEOUT_SECONDS = int(os.getenv("SUPABASE_TIMEOUT_SECONDS", 3600))
 
 ISSUES_TABLE = "Issues"
 EVIDENCE_TABLE = "Evidence"
-
+RUNBOOKS_TABLE = "HolmesRunbooks"
 
 class RobustaConfig(BaseModel):
     sinks_config: List[Dict[str, Dict]]
@@ -98,7 +99,7 @@ def get_issue_data(self, issue_id: str) -> Optional[Dict]:
             issue_response = (
                 self.client
                     .table(ISSUES_TABLE)
-                    .select(f"*")
+                    .select("*")
                     .filter("id", "eq", issue_id)
                     .execute()
             )
@@ -113,9 +114,81 @@ def get_issue_data(self, issue_id: str) -> Optional[Dict]:
         evidence = (
             self.client
             .table(EVIDENCE_TABLE)
-            .select(f"*")
+            .select("*")
             .filter("issue_id", "eq", issue_id)
             .execute()
         )
-        issue_data["evidence"] = evidence.data
-        return issue_data
+        enrichment_blacklist = {"text_file", "graph", "ai_analysis", "holmes"}
+        data = [enrich for enrich in evidence.data if enrich.get("enrichment_type") not in enrichment_blacklist]
+
+        issue_data["evidence"] = data
+        return issue_data
+
+    def get_resource_instructions(self, type: str, name: str) -> List[str]:
+        if not self.enabled or not name:
+            return []
+
+        res = (
+            self.client
+            .table(RUNBOOKS_TABLE)
+            .select("runbook")
+            .eq("account_id", self.account_id)
+            .eq("subject_type", type)
+            .eq("subject_name", name)
+            .execute()
+        )
+        if res.data:
+            return res.data[0].get("runbook").get("instructions")
+
+        return []
+
+    def get_workload_issues(self, resource: dict, since_hours: float) -> List[str]:
+        if not self.enabled or not resource:
+            return []
+
+        cluster = resource.get("cluster")
+        if not cluster:
+            logging.debug("Missing workload cluster for issues.")
+            return []
+
+        since: str = (datetime.now() - timedelta(hours=since_hours)).isoformat()
+
+        svc_key = f"{resource.get('namespace', '')}/{resource.get('kind', '')}/{resource.get('name', '')}"
+        logging.debug(f"getting issues for workload {svc_key}")
+        try:
+            res = (
+                self.client
+                .table(ISSUES_TABLE)
+                .select("id, creation_date, aggregation_key")
+                .eq("account_id", self.account_id)
+                .eq("cluster", cluster)
+                .eq("service_key", svc_key)
+                .gte("creation_date", since)
+                .order("creation_date")
+                .execute()
+            )
+
+            if not res.data:
+                return []
+
+            issue_dict = dict()
+            for issue in res.data:
+                issue_dict[issue.get("aggregation_key")] = issue.get("id")
+
+            unique_issues: list[str] = list(issue_dict.values())
+
+            res = (
+                self.client
+                .table(EVIDENCE_TABLE)
+                .select("data, enrichment_type")
+                .in_("issue_id", unique_issues)
+                .execute()
+            )
+
+            enrichment_blacklist = {"text_file", "graph", "ai_analysis", "holmes"}
+            data = [evidence.get("data") for evidence in res.data if evidence.get("enrichment_type") not in enrichment_blacklist]
+            return data
+
+        except:
+            logging.exception("failed to fetch workload issues data")
+            return []

holmes/core/tool_calling_llm.py

@@ -31,6 +31,7 @@ class LLMResult(BaseModel):
     tool_calls: Optional[List[ToolCallResult]] = None
     result: Optional[str] = None
     unprocessed_result: Optional[str] = None
+    instructions: List[str] = []
 
     # TODO: clean up these two
     prompt: Optional[str] = None
@@ -239,11 +240,13 @@ def __init__(
         self.runbook_manager = runbook_manager
 
     def investigate(
-        self, issue: Issue, prompt: str, console: Console, post_processing_prompt: Optional[str] = None
+        self, issue: Issue, prompt: str, console: Console, instructions: List[str] = [], post_processing_prompt: Optional[str] = None
     ) -> LLMResult:
         environment = jinja2.Environment()
         system_prompt_template = environment.from_string(prompt)
         runbooks = self.runbook_manager.get_instructions_for_issue(issue)
+        runbooks.extend(instructions)
+
         if runbooks:
             console.print(
                 f"[bold]Analyzing with {len(runbooks)} runbooks: {runbooks}[/bold]"
@@ -252,12 +255,23 @@ def investigate(
             console.print(
                 f"[bold]No runbooks found for this issue. Using default behaviour. (Add runbooks to guide the investigation.)[/bold]"
             )
-        system_prompt = system_prompt_template.render(issue=issue, runbooks=runbooks)
-        user_prompt = f"{issue.raw}"
+        system_prompt = system_prompt_template.render(issue=issue)
+
+        user_prompt = ""
+        if runbooks:
+            for i in runbooks:
+                user_prompt += f"* {i}\n"
+
+            user_prompt = f'My instructions to check \n"""{user_prompt}"""'
+
+        user_prompt = f"{user_prompt}\n This is context from the issue {issue.raw}"
         logging.debug(
             "Rendered system prompt:\n%s", textwrap.indent(system_prompt, "    ")
         )
         logging.debug(
             "Rendered user prompt:\n%s", textwrap.indent(user_prompt, "    ")
         )
-        return self.call(system_prompt, user_prompt, post_processing_prompt)
+
+        res = self.call(system_prompt, user_prompt, post_processing_prompt)
+        res.instructions = runbooks
+        return res

holmes/main.py

@@ -338,7 +338,7 @@ def alertmanager(
         console.print(
             f"[bold yellow]Analyzing issue {i+1}/{len(issues)}: {issue.name}...[/bold yellow]"
         )
-        result = ai.investigate(issue, system_prompt, console, post_processing_prompt)
+        result = ai.investigate(issue, system_prompt, console, [], post_processing_prompt)
         results.append({"issue": issue.model_dump(), "result": result.model_dump()})
         handle_result(result, console, destination, config, issue, False, True)
 
@@ -450,7 +450,7 @@ def jira(
         console.print(
             f"[bold yellow]Analyzing Jira ticket {i+1}/{len(issues)}: {issue.name}...[/bold yellow]"
         )
-        result = ai.investigate(issue, system_prompt, console, post_processing_prompt)
+        result = ai.investigate(issue, system_prompt, console, [], post_processing_prompt)
 
         console.print(Rule())
         console.print(f"[bold green]AI analysis of {issue.url}[/bold green]")
@@ -539,7 +539,7 @@ def github(
     )
     for i, issue in enumerate(issues):
         console.print(f"[bold yellow]Analyzing GitHub issue {i+1}/{len(issues)}: {issue.name}...[/bold yellow]")
-        result = ai.investigate(issue, system_prompt, console, post_processing_prompt)
+        result = ai.investigate(issue, system_prompt, console, [], post_processing_prompt)
 
         console.print(Rule())
         console.print(f"[bold green]AI analysis of {issue.url}[/bold green]")
@@ -616,7 +616,7 @@ def pagerduty(
     results = []
     for i, issue in enumerate(issues):
         console.print(f"[bold yellow]Analyzing PagerDuty incident {i+1}/{len(issues)}: {issue.name}...[/bold yellow]")
-        result = ai.investigate(issue, system_prompt, console, post_processing_prompt)
+        result = ai.investigate(issue, system_prompt, console, [], post_processing_prompt)
 
         console.print(Rule())
         console.print(f"[bold green]AI analysis of {issue.url}[/bold green]")
@@ -694,7 +694,7 @@ def opsgenie(
     )
     for i, issue in enumerate(issues):
         console.print(f"[bold yellow]Analyzing OpsGenie alert {i+1}/{len(issues)}: {issue.name}...[/bold yellow]")
-        result = ai.investigate(issue, system_prompt, console, post_processing_prompt)
+        result = ai.investigate(issue, system_prompt, console, [], post_processing_prompt)
 
         console.print(Rule())
         console.print(f"[bold green]AI analysis of {issue.url}[/bold green]")

holmes/plugins/prompts/generic_ask.jinja2

@@ -51,4 +51,11 @@ Relevant logs:
 2021-01-01T00:00:00.000Z [ERROR] Missing required field 'email' in request body
 ```
 
-Validation error led to unhandled Java exception causing a crash.
+Validation error led to unhandled Java exception causing a crash.
+
+{% if alerts %}
+Here are issues and configuration changes that happend to this kubernetes workload in recent time. Check if these can help you understand the issue.
+{% for a in alerts %}
+{{ a }}
+{% endfor %}
+{% endif %}

holmes/plugins/prompts/generic_investigation.jinja2

@@ -4,6 +4,8 @@ Do not say 'based on the tool output' or explicitly refer to tools at all.
 
 Provide an excruciatingly terse analysis of the following {{ issue.source_type }} alert/issue and why it is firing.
 
+If the user provides you with extra instructions in a triple quotes section, ALWAYS perform their instructions and then perform your investigation.
+
 In general:
 * when it can provide extra information, first run as many tools as you need to gather more information, then respond. 
 * if possible, do so repeatedly with different tool calls each time to gather more information.
@@ -77,10 +79,3 @@ User: Why did the webserver-example app crash?
 *Details:* Validation error led to unhandled Java exception causing a crash: `2021-01-01T00:00:00.000Z [ERROR] Missing required field 'email' in request body`
 
 End of Examples
-
-{% if runbooks %}
-Here are runbooks for this specific investigation. Please follow them if relevant.
-{% for r in runbooks %}
-* {{ r }}
-{% endfor %}
-{% endif %}

server.py

@@ -6,16 +6,16 @@
     print("added custom certificate")
 
 # DO NOT ADD ANY IMPORTS OR CODE ABOVE THIS LINE
-# IMPORTING ABOVE MIGHT INITIALIZE AN HTTPS CLIENT THAT DOESN'T TRUST THE CUSTOM CERTIFICATE
-
+# IMPORTING ABOVE MIGHT INITIALIZE AN HTTPS CLIENT THAT DOESN'T TRUST THE CUSTOM CERTIFICATEE
+import jinja2
 import logging
 import uvicorn
 import colorlog
 
 from holmes.core.tool_calling_llm import ToolCallResult
 from typing import List, Union, Dict, Any, Optional
-
-from fastapi import FastAPI
+from litellm.exceptions import AuthenticationError
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from rich.console import Console
 
@@ -40,6 +40,18 @@ class InvestigateRequest(BaseModel):
     # response_handler: ...
 
 
+class WorkloadHealthRequest(BaseModel):
+    ask: str
+    resource: dict
+    alert_history_since_hours: float = 24
+    alert_history: bool = True
+    stored_instrucitons: bool = True
+    instructions: Optional[List[str]] = []
+    include_tool_calls: bool = False
+    include_tool_call_results: bool = False
+    prompt_template: str = "builtin://generic_ask.jinja2"
+
+
 def init_logging():
     logging_level = os.environ.get("LOG_LEVEL", "INFO")
     logging_format = "%(log_color)s%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s"
@@ -67,41 +79,67 @@ def init_logging():
 class InvestigationResult(BaseModel):
     analysis: Optional[str] = None
     tool_calls: List[ToolCallResult] = []
+    instructions: List[str] = []
 
 
 @app.post("/api/investigate")
 def investigate_issues(investigate_request: InvestigateRequest):
-    context = fetch_context_data(investigate_request.context)
-    raw_data = investigate_request.model_dump()
-    if context:
-        raw_data["extra_context"] = context
-
-    ai = config.create_issue_investigator(console, allowed_toolsets=ALLOWED_TOOLSETS)
-    issue = Issue(
-        id=context['id'] if context else "",
-        name=investigate_request.title,
-        source_type=investigate_request.source,
-        source_instance_id=investigate_request.source_instance_id,
-        raw=raw_data,
-    )
-    investigation = ai.investigate(
-        issue,
-        prompt=load_prompt(investigate_request.prompt_template),
-        console=console,
-        post_processing_prompt=HOLMES_POST_PROCESSING_PROMPT
-    )
-    return InvestigationResult(
-        analysis=investigation.result,
-        tool_calls=investigation.tool_calls,
-    )
-
-
-def fetch_context_data(context: Dict[str, Any]) -> dict:
-    for context_item in context.keys():
-        if context_item == "robusta_issue_id":
-            # Note we only accept a single robusta_issue_id. I don't think it
-            # makes sense to have several of them in the context structure.
-            return dal.get_issue_data(context[context_item])
+    try:
+        context = dal.get_issue_data(investigate_request.context.get("robusta_issue_id"))
+
+        instructions = dal.get_resource_instructions("alert", investigate_request.context.get("issue_type"))
+        raw_data = investigate_request.model_dump()
+        if context:
+            raw_data["extra_context"] = context
+
+        ai = config.create_issue_investigator(console, allowed_toolsets=ALLOWED_TOOLSETS)
+        issue = Issue(
+            id=context['id'] if context else "",
+            name=investigate_request.title,
+            source_type=investigate_request.source,
+            source_instance_id=investigate_request.source_instance_id,
+            raw=raw_data,
+        )
+        investigation = ai.investigate(
+            issue,
+            prompt=load_prompt(investigate_request.prompt_template),
+            console=console,
+            post_processing_prompt=HOLMES_POST_PROCESSING_PROMPT,
+            instructions=instructions,
+        )
+
+        return InvestigationResult(
+            analysis=investigation.result,
+            tool_calls=investigation.tool_calls,
+            instructions=investigation.instructions
+        )
+    except AuthenticationError as e:
+        raise HTTPException(status_code=401, detail=e.message)
+
+
+@app.post("/api/workload_health_check")
+def workload_health_check(request: WorkloadHealthRequest):
+
+    try:
+        workload_alerts: list[str] = []
+        if request.alert_history:
+            workload_alerts = dal.get_workload_issues(request.resource, request.alert_history_since_hours)
+
+        system_prompt = load_prompt(request.prompt_template)
+        system_prompt = jinja2.Environment().from_string(system_prompt)
+        system_prompt = system_prompt.render(alerts=workload_alerts)
+
+        ai = config.create_toolcalling_llm(console, allowed_toolsets=ALLOWED_TOOLSETS)
+
+        ai_call = ai.call(system_prompt, request.ask, HOLMES_POST_PROCESSING_PROMPT)
+
+        return InvestigationResult(
+            analysis=ai_call.result,
+            tool_calls=ai_call.tool_calls,
+            instructions=ai_call.instructions
+        )
+    except AuthenticationError as e:
+        raise HTTPException(status_code=401, detail=e.message)          
 
 
 if __name__ == "__main__":