holmesgpt 0.11.5__py3-none-any.whl

holmes/core/llm.py

@@ -0,0 +1,250 @@
+import json
+import logging
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional, Type, Union
+
+from litellm.types.utils import ModelResponse
+import sentry_sdk
+
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from pydantic import BaseModel
+import litellm
+import os
+from holmes.common.env_vars import (
+    THINKING,
+    TEMPERATURE,
+)
+
+
+def environ_get_safe_int(env_var, default="0"):
+    try:
+        return max(int(os.environ.get(env_var, default)), 0)
+    except ValueError:
+        return int(default)
+
+
+OVERRIDE_MAX_OUTPUT_TOKEN = environ_get_safe_int("OVERRIDE_MAX_OUTPUT_TOKEN")
+OVERRIDE_MAX_CONTENT_SIZE = environ_get_safe_int("OVERRIDE_MAX_CONTENT_SIZE")
+
+
+class LLM:
+    @abstractmethod
+    def __init__(self):
+        self.model: str  # type: ignore
+
+    @abstractmethod
+    def get_context_window_size(self) -> int:
+        pass
+
+    @abstractmethod
+    def get_maximum_output_token(self) -> int:
+        pass
+
+    @abstractmethod
+    def count_tokens_for_message(self, messages: list[dict]) -> int:
+        pass
+
+    @abstractmethod
+    def completion(
+        self,
+        messages: List[Dict[str, Any]],
+        tools: Optional[List[Dict[str, Any]]] = [],
+        tool_choice: Optional[Union[str, dict]] = None,
+        response_format: Optional[Union[dict, Type[BaseModel]]] = None,
+        temperature: Optional[float] = None,
+        drop_params: Optional[bool] = None,
+        stream: Optional[bool] = None,
+    ) -> Union[ModelResponse, CustomStreamWrapper]:
+        pass
+
+
+class DefaultLLM(LLM):
+    model: str
+    api_key: Optional[str]
+    base_url: Optional[str]
+    args: Dict
+
+    def __init__(self, model: str, api_key: Optional[str] = None, args: Dict = {}):
+        self.model = model
+        self.api_key = api_key
+        self.args = args
+
+        if not args:
+            self.check_llm(self.model, self.api_key)
+
+    def check_llm(self, model: str, api_key: Optional[str]):
+        logging.debug(f"Checking LiteLLM model {model}")
+        # TODO: this WAS a hack to get around the fact that we can't pass in an api key to litellm.validate_environment
+        # so without this hack it always complains that the environment variable for the api key is missing
+        # to fix that, we always set an api key in the standard format that litellm expects (which is ${PROVIDER}_API_KEY)
+        # TODO: we can now handle this better - see https://github.com/BerriAI/litellm/issues/4375#issuecomment-2223684750
+        lookup = litellm.get_llm_provider(self.model)
+        if not lookup:
+            raise Exception(f"Unknown provider for model {model}")
+        provider = lookup[1]
+        if provider == "watsonx":
+            # NOTE: LiteLLM's validate_environment does not currently include checks for IBM WatsonX.
+            # The following WatsonX-specific variables are set based on documentation from:
+            # https://docs.litellm.ai/docs/providers/watsonx
+            # Required variables for WatsonX:
+            # - WATSONX_URL: Base URL of your WatsonX instance (required)
+            # - WATSONX_APIKEY or WATSONX_TOKEN: IBM Cloud API key or IAM auth token (one is required)
+            model_requirements = {"missing_keys": [], "keys_in_environment": True}
+            if api_key:
+                os.environ["WATSONX_APIKEY"] = api_key
+            if "WATSONX_URL" not in os.environ:
+                model_requirements["missing_keys"].append("WATSONX_URL")  # type: ignore
+                model_requirements["keys_in_environment"] = False
+            if "WATSONX_APIKEY" not in os.environ and "WATSONX_TOKEN" not in os.environ:
+                model_requirements["missing_keys"].extend(  # type: ignore
+                    ["WATSONX_APIKEY", "WATSONX_TOKEN"]
+                )
+                model_requirements["keys_in_environment"] = False
+            # WATSONX_PROJECT_ID is required because we don't let user pass it to completion call directly
+            if "WATSONX_PROJECT_ID" not in os.environ:
+                model_requirements["missing_keys"].append("WATSONX_PROJECT_ID")  # type: ignore
+                model_requirements["keys_in_environment"] = False
+            # https://docs.litellm.ai/docs/providers/watsonx#usage---models-in-deployment-spaces
+            # using custom watsonx deployments might require to set WATSONX_DEPLOYMENT_SPACE_ID env
+            if "watsonx/deployment/" in self.model:
+                logging.warning(
+                    "Custom WatsonX deployment detected. You may need to set the WATSONX_DEPLOYMENT_SPACE_ID "
+                    "environment variable for proper functionality. For more information, refer to the documentation: "
+                    "https://docs.litellm.ai/docs/providers/watsonx#usage---models-in-deployment-spaces"
+                )
+        elif provider == "bedrock" and os.environ.get("AWS_PROFILE"):
+            model_requirements = {"keys_in_environment": True, "missing_keys": []}
+        else:
+            #
+            api_key_env_var = f"{provider.upper()}_API_KEY"
+            if api_key:
+                os.environ[api_key_env_var] = api_key
+            model_requirements = litellm.validate_environment(model=model)
+
+        if not model_requirements["keys_in_environment"]:
+            raise Exception(
+                f"model {model} requires the following environment variables: {model_requirements['missing_keys']}"
+            )
+
+    def _strip_model_prefix(self) -> str:
+        """
+        Helper function to strip 'openai/' prefix from model name if it exists.
+        model cost is taken from here which does not have the openai prefix
+        https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json
+        """
+        model_name = self.model
+        if model_name.startswith("openai/"):
+            model_name = model_name[len("openai/") :]  # Strip the 'openai/' prefix
+        elif model_name.startswith("bedrock/"):
+            model_name = model_name[len("bedrock/") :]  # Strip the 'bedrock/' prefix
+        elif model_name.startswith("vertex_ai/"):
+            model_name = model_name[
+                len("vertex_ai/") :
+            ]  # Strip the 'vertex_ai/' prefix
+
+        return model_name
+
+        # this unfortunately does not seem to work for azure if the deployment name is not a well-known model name
+        # if not litellm.supports_function_calling(model=model):
+        #    raise Exception(f"model {model} does not support function calling. You must use HolmesGPT with a model that supports function calling.")
+
+    def get_context_window_size(self) -> int:
+        if OVERRIDE_MAX_CONTENT_SIZE:
+            logging.debug(
+                f"Using override OVERRIDE_MAX_CONTENT_SIZE {OVERRIDE_MAX_CONTENT_SIZE}"
+            )
+            return OVERRIDE_MAX_CONTENT_SIZE
+
+        model_name = os.environ.get("MODEL_TYPE", self._strip_model_prefix())
+        try:
+            return litellm.model_cost[model_name]["max_input_tokens"]
+        except Exception:
+            logging.warning(
+                f"Couldn't find model's name {model_name} in litellm's model list, fallback to 128k tokens for max_input_tokens"
+            )
+            return 128000
+
+    @sentry_sdk.trace
+    def count_tokens_for_message(self, messages: list[dict]) -> int:
+        total_token_count = 0
+        for message in messages:
+            if "token_count" in message and message["token_count"]:
+                total_token_count += message["token_count"]
+            else:
+                # message can be counted by this method only if message contains a "content" key
+                if "content" in message:
+                    if isinstance(message["content"], str):
+                        message_to_count = [
+                            {"type": "text", "text": message["content"]}
+                        ]
+                    elif isinstance(message["content"], list):
+                        message_to_count = [
+                            {"type": "text", "text": json.dumps(message["content"])}
+                        ]
+                    elif isinstance(message["content"], dict):
+                        if "type" not in message["content"]:
+                            message_to_count = [
+                                {"type": "text", "text": json.dumps(message["content"])}
+                            ]
+                    token_count = litellm.token_counter(
+                        model=self.model, messages=message_to_count
+                    )
+                    message["token_count"] = token_count
+                    total_token_count += token_count
+        return total_token_count
+
+    def completion(
+        self,
+        messages: List[Dict[str, Any]],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, dict]] = None,
+        response_format: Optional[Union[dict, Type[BaseModel]]] = None,
+        temperature: Optional[float] = None,
+        drop_params: Optional[bool] = None,
+        stream: Optional[bool] = None,
+    ) -> Union[ModelResponse, CustomStreamWrapper]:
+        tools_args = {}
+        if tools and len(tools) > 0 and tool_choice == "auto":
+            tools_args["tools"] = tools
+            tools_args["tool_choice"] = tool_choice  # type: ignore
+
+        if THINKING:
+            self.args.setdefault("thinking", json.loads(THINKING))
+
+        if self.args.get("thinking", None):
+            litellm.modify_params = True
+
+        result = litellm.completion(
+            model=self.model,
+            api_key=self.api_key,
+            messages=messages,
+            temperature=temperature or self.args.pop("temperature", TEMPERATURE),
+            response_format=response_format,
+            drop_params=drop_params,
+            stream=stream,
+            **tools_args,
+            **self.args,
+        )
+
+        if isinstance(result, ModelResponse):
+            return result
+        elif isinstance(result, CustomStreamWrapper):
+            return result
+        else:
+            raise Exception(f"Unexpected type returned by the LLM {type(result)}")
+
+    def get_maximum_output_token(self) -> int:
+        if OVERRIDE_MAX_OUTPUT_TOKEN:
+            logging.debug(
+                f"Using OVERRIDE_MAX_OUTPUT_TOKEN {OVERRIDE_MAX_OUTPUT_TOKEN}"
+            )
+            return OVERRIDE_MAX_OUTPUT_TOKEN
+
+        model_name = os.environ.get("MODEL_TYPE", self._strip_model_prefix())
+        try:
+            return litellm.model_cost[model_name]["max_output_tokens"]
+        except Exception:
+            logging.warning(
+                f"Couldn't find model's name {model_name} in litellm's model list, fallback to 4096 tokens for max_output_tokens"
+            )
+            return 4096

holmes/core/models.py

@@ -0,0 +1,157 @@
+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
+from enum import Enum
+
+
+class InvestigationResult(BaseModel):
+    analysis: Optional[str] = None
+    sections: Optional[Dict[str, Union[str, None]]] = None
+    tool_calls: List[ToolCallResult] = []
+    instructions: List[str] = []
+
+
+class InvestigateRequest(BaseModel):
+    source: str  # "prometheus" etc
+    title: str
+    description: str
+    subject: dict
+    context: Dict[str, Any]
+    source_instance_id: str = "ApiRequest"
+    include_tool_calls: bool = False
+    include_tool_call_results: bool = False
+    prompt_template: str = "builtin://generic_investigation.jinja2"
+    sections: Optional[InputSectionsDataType] = None
+    model: Optional[str] = None
+    # TODO in the future
+    # response_handler: ...
+
+
+class ToolCallConversationResult(BaseModel):
+    name: str
+    description: str
+    output: str
+
+
+class ConversationInvestigationResponse(BaseModel):
+    analysis: Optional[str] = None
+    tool_calls: List[ToolCallResult] = []
+
+
+class ConversationInvestigationResult(BaseModel):
+    analysis: Optional[str] = None
+    tools: Optional[List[ToolCallConversationResult]] = []
+
+
+class IssueInvestigationResult(BaseModel):
+    """
+    :var result: A dictionary containing the summary of the issue investigation.
+    :var tools: A list of dictionaries where each dictionary contains information
+                about the tool, its name, description and output.
+
+    It is based on the holmes investigation saved to Evidence table.
+    """
+
+    result: str
+    tools: Optional[List[ToolCallConversationResult]] = []
+
+
+class HolmesConversationHistory(BaseModel):
+    ask: str
+    answer: ConversationInvestigationResult
+
+
+# HolmesConversationIssueContext, ConversationType and ConversationRequest classes will be deprecated later
+class HolmesConversationIssueContext(BaseModel):
+    investigation_result: IssueInvestigationResult
+    conversation_history: Optional[List[HolmesConversationHistory]] = []
+    issue_type: str
+    robusta_issue_id: Optional[str] = None
+    source: Optional[str] = None
+
+
+class ConversationType(str, Enum):
+    ISSUE = "issue"
+
+
+class ConversationRequest(BaseModel):
+    user_prompt: str
+    source: Optional[str] = None
+    resource: Optional[dict] = None
+    # ConversationType.ISSUE is default as we gonna deprecate this class and won't add new conversation types
+    conversation_type: Optional[ConversationType] = ConversationType.ISSUE
+    context: HolmesConversationIssueContext
+    include_tool_calls: bool = False
+    include_tool_call_results: bool = False
+
+
+class ChatRequestBaseModel(BaseModel):
+    conversation_history: Optional[list[dict]] = None
+    model: Optional[str] = None
+
+    # In our setup with litellm, the first message in conversation_history
+    # should follow the structure [{"role": "system", "content": ...}],
+    # where the "role" field is expected to be "system".
+    @model_validator(mode="before")
+    def check_first_item_role(cls, values):
+        conversation_history = values.get("conversation_history")
+        if (
+            conversation_history
+            and isinstance(conversation_history, list)
+            and len(conversation_history) > 0
+        ):
+            first_item = conversation_history[0]
+            if not first_item.get("role") == "system":
+                raise ValueError(
+                    "The first item in conversation_history must contain 'role': 'system'"
+                )
+        return values
+
+
+class IssueChatRequest(ChatRequestBaseModel):
+    ask: str
+    investigation_result: IssueInvestigationResult
+    issue_type: str
+
+
+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://kubernetes_workload_ask.jinja2"
+    model: Optional[str] = None
+
+
+class ChatRequest(ChatRequestBaseModel):
+    ask: str
+
+
+class FollowUpAction(BaseModel):
+    id: str
+    action_label: str
+    pre_action_notification_text: str
+    prompt: str
+
+
+class ChatResponse(BaseModel):
+    analysis: str
+    conversation_history: list[dict]
+    tool_calls: Optional[List[ToolCallResult]] = []
+    follow_up_actions: Optional[List[FollowUpAction]] = []
+
+
+class WorkloadHealthInvestigationResult(BaseModel):
+    analysis: Optional[str] = None
+    tools: Optional[List[ToolCallConversationResult]] = []
+
+
+class WorkloadHealthChatRequest(ChatRequestBaseModel):
+    ask: str
+    workload_health_result: WorkloadHealthInvestigationResult
+    resource: dict

holmes/core/openai_formatting.py

@@ -0,0 +1,51 @@
+import re
+
+# parses both simple types: "int", "array", "string"
+# but also arrays of those simpler types: "array[int]", "array[string]", etc.
+pattern = r"^(array\[(?P<inner_type>\w+)\])|(?P<simple_type>\w+)$"
+
+
+def type_to_open_ai_schema(type_value):
+    match = re.match(pattern, type_value.strip())
+
+    if not match:
+        raise ValueError(f"Invalid type format: {type_value}")
+
+    if match.group("inner_type"):
+        return {"type": "array", "items": {"type": match.group("inner_type")}}
+
+    else:
+        return {"type": match.group("simple_type")}
+
+
+def format_tool_to_open_ai_standard(
+    tool_name: str, tool_description: str, tool_parameters: dict
+):
+    tool_properties = {}
+    for param_name, param_attributes in tool_parameters.items():
+        tool_properties[param_name] = type_to_open_ai_schema(param_attributes.type)
+        if param_attributes.description is not None:
+            tool_properties[param_name]["description"] = param_attributes.description
+
+    result = {
+        "type": "function",
+        "function": {
+            "name": tool_name,
+            "description": tool_description,
+            "parameters": {
+                "properties": tool_properties,
+                "required": [
+                    param_name
+                    for param_name, param_attributes in tool_parameters.items()
+                    if param_attributes.required
+                ],
+                "type": "object",
+            },
+        },
+    }
+
+    # gemini doesnt have parameters object if it is without params
+    if tool_properties is None or tool_properties == {}:
+        result["function"].pop("parameters")  # type: ignore
+
+    return result

holmes/core/performance_timing.py

@@ -0,0 +1,72 @@
+import time
+import logging
+
+from functools import wraps
+
+from holmes.common.env_vars import (
+    LOG_PERFORMANCE,
+)
+
+
+class PerformanceTiming:
+    def __init__(self, name):
+        self.ended = False
+
+        self.name = name
+        self.start_time = time.time()
+        self.last_measure_time = self.start_time
+        self.last_measure_label = "Start"
+        self.timings = []
+
+    def measure(self, label):
+        if not LOG_PERFORMANCE:
+            return
+        if self.ended:
+            raise Exception("cannot measure a perf timing that is already ended")
+        current_time = time.time()
+
+        time_since_start = int((current_time - self.start_time) * 1000)
+        time_since_last = int((current_time - self.last_measure_time) * 1000)
+
+        self.timings.append((label, time_since_last, time_since_start))
+
+        self.last_measure_time = current_time
+        self.last_measure_label = label
+
+    def end(self, custom_message: str = ""):
+        if not LOG_PERFORMANCE:
+            return
+        self.ended = True
+        current_time = time.time()
+        time_since_start = int((current_time - self.start_time) * 1000)
+        message = f"{self.name} {custom_message} {time_since_start}ms"
+        logging.info(message)
+        if LOG_PERFORMANCE:
+            for label, time_since_last, time_since_start in self.timings:
+                logging.info(
+                    f"\t{self.name}({label}) +{time_since_last}ms  {time_since_start}ms"
+                )
+
+
+def log_function_timing(label=None):
+    def decorator(func):
+        @wraps(func)
+        def function_timing_wrapper(*args, **kwargs):
+            start_time = time.perf_counter()
+            result = func(*args, **kwargs)
+            end_time = time.perf_counter()
+            total_time = int((end_time - start_time) * 1000)
+
+            function_identifier = (
+                f'"{label}: {func.__name__}()"' if label else f'"{func.__name__}()"'
+            )
+            logging.info(f"Function {function_identifier} took {total_time}ms")
+            return result
+
+        return function_timing_wrapper
+
+    if callable(label):
+        func = label
+        label = None
+        return decorator(func)
+    return decorator

holmes/core/prompt.py

@@ -0,0 +1,42 @@
+from rich.console import Console
+from typing import Optional, List, Dict
+from pathlib import Path
+
+
+def append_file_to_user_prompt(user_prompt: str, file_path: Path) -> str:
+    with file_path.open("r") as f:
+        user_prompt += f"\n\n<attached-file path='{file_path.absolute()}>'\n{f.read()}\n</attached-file>"
+
+    return user_prompt
+
+
+def append_all_files_to_user_prompt(
+    console: Console, user_prompt: str, file_paths: Optional[List[Path]]
+) -> str:
+    if not file_paths:
+        return user_prompt
+
+    for file_path in file_paths:
+        console.print(f"[bold yellow]Adding file {file_path} to context[/bold yellow]")
+        user_prompt = append_file_to_user_prompt(user_prompt, file_path)
+
+    return user_prompt
+
+
+def build_initial_ask_messages(
+    console: Console,
+    system_prompt_rendered: str,
+    initial_user_prompt: str,
+    file_paths: Optional[List[Path]],
+) -> List[Dict]:
+    """Build the initial messages for the AI call."""
+    user_prompt_with_files = append_all_files_to_user_prompt(
+        console, initial_user_prompt, file_paths
+    )
+
+    messages = [
+        {"role": "system", "content": system_prompt_rendered},
+        {"role": "user", "content": user_prompt_with_files},
+    ]
+
+    return messages

holmes/core/resource_instruction.py

@@ -0,0 +1,17 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class ResourceInstructionDocument(BaseModel):
+    """Represents context necessary for an investigation in the form of a URL
+    It is expected that Holmes will use that URL to fetch additional context about an error.
+    This URL can for example be the location of a runbook
+    """
+
+    url: str
+
+
+class ResourceInstructions(BaseModel):
+    instructions: List[str] = []
+    documents: List[ResourceInstructionDocument] = []

holmes/core/runbooks.py

@@ -0,0 +1,26 @@
+from typing import List
+from holmes.core.issue import Issue
+from holmes.plugins.runbooks import Runbook
+
+
+# TODO: our default prompt has a lot of kubernetes specific stuff - see if we can get that into the runbook
+class RunbookManager:
+    def __init__(self, runbooks: List[Runbook]):
+        self.runbooks = runbooks
+
+    def get_instructions_for_issue(self, issue: Issue) -> List[str]:
+        instructions = []
+        for runbook in self.runbooks:
+            if runbook.match.issue_id and not runbook.match.issue_id.match(issue.id):
+                continue
+            if runbook.match.issue_name and not runbook.match.issue_name.match(
+                issue.name
+            ):
+                continue
+            if runbook.match.source and not runbook.match.source.match(
+                issue.source_type
+            ):
+                continue
+            instructions.append(runbook.instructions)
+
+        return instructions