holmesgpt 0.14.1a0__py3-none-any.whl → 0.14.3a0__py3-none-any.whl

holmes/__init__.py

@@ -1,5 +1,5 @@
 # This is patched by github actions during release
-__version__ = "0.14.1-alpha"
+__version__ = "0.14.3-alpha"
 
 # Re-export version functions from version module for backward compatibility
 from .version import (

holmes/clients/robusta_client.py

@@ -1,8 +1,8 @@
 import logging
-from typing import List, Optional
+from typing import List, Optional, Dict, Any
 import requests  # type: ignore
 from functools import cache
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 from holmes.common.env_vars import ROBUSTA_API_ENDPOINT
 
 HOLMES_GET_INFO_URL = f"{ROBUSTA_API_ENDPOINT}/api/holmes/get_info"
@@ -17,6 +17,9 @@ class HolmesInfo(BaseModel):
 class RobustaModelsResponse(BaseModel):
     model_config = ConfigDict(extra="ignore")
     models: List[str]
+    models_args: Dict[str, Any] = Field(
+        default_factory=dict, alias="models_holmes_args"
+    )
     default_model: Optional[str] = None
 
 

holmes/common/env_vars.py

@@ -73,11 +73,17 @@ LOG_LLM_USAGE_RESPONSE = load_bool("LOG_LLM_USAGE_RESPONSE", False)
 # For CLI only, enable user approval for potentially sensitive commands that would otherwise be rejected
 ENABLE_CLI_TOOL_APPROVAL = load_bool("ENABLE_CLI_TOOL_APPROVAL", True)
 
-MAX_GRAPH_POINTS = float(os.environ.get("MAX_GRAPH_POINTS", 200))
+MAX_GRAPH_POINTS = float(os.environ.get("MAX_GRAPH_POINTS", 100))
 
 # Limit each tool response to N% of the total context window.
 # Number between 0 and 100
 # Setting to either 0 or any number above 100 disables the logic that limits tool response size
 TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT = float(
-    os.environ.get("TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", 10)
+    os.environ.get("TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", 15)
 )
+
+MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION = int(
+    os.environ.get("MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", 3000)
+)
+
+DISABLE_PROMETHEUS_TOOLSET = load_bool("DISABLE_PROMETHEUS_TOOLSET", False)

holmes/config.py

@@ -9,11 +9,8 @@ import sentry_sdk
 import yaml  # type: ignore
 from pydantic import BaseModel, ConfigDict, FilePath, PrivateAttr, SecretStr
 
-
+from holmes.common.env_vars import ROBUSTA_CONFIG_PATH
 from holmes.core.llm import DefaultLLM, LLMModelRegistry
-from holmes.common.env_vars import (
-    ROBUSTA_CONFIG_PATH,
-)
 from holmes.core.tools_utils.tool_executor import ToolExecutor
 from holmes.core.toolset_manager import ToolsetManager
 from holmes.plugins.runbooks import (
@@ -33,8 +30,8 @@ if TYPE_CHECKING:
     from holmes.plugins.sources.pagerduty import PagerDutySource
     from holmes.plugins.sources.prometheus.plugin import AlertManagerSource
 
-from holmes.core.supabase_dal import SupabaseDal
 from holmes.core.config import config_path_dir
+from holmes.core.supabase_dal import SupabaseDal
 from holmes.utils.definitions import RobustaConfig
 from holmes.utils.pydantic_utils import RobustaBaseConfig, load_model_from_file
 
@@ -129,9 +126,9 @@ class Config(RobustaBaseConfig):
         return self._llm_model_registry
 
     def log_useful_info(self):
-        if self.llm_model_registry and self.llm_model_registry.models:
+        if self.llm_model_registry.models:
             logging.info(
-                f"loaded models: {list(self.llm_model_registry.models.keys())}"
+                f"Loaded models: {list(self.llm_model_registry.models.keys())}"
             )
         else:
             logging.warning("No llm models were loaded")

holmes/core/conversations.py

@@ -262,7 +262,10 @@ def build_issue_chat_messages(
 
 
 def add_or_update_system_prompt(
-    conversation_history: List[Dict[str, str]], ai: ToolCallingLLM, config: Config
+    conversation_history: List[Dict[str, str]],
+    ai: ToolCallingLLM,
+    config: Config,
+    additional_system_prompt: Optional[str] = None,
 ):
     """Either add the system prompt or replace an existing system prompt.
     As a 'defensive' measure, this code will only replace an existing system prompt if it is the
@@ -278,6 +281,9 @@ def add_or_update_system_prompt(
 
     system_prompt = load_and_render_prompt(template_path, context)
 
+    if additional_system_prompt:
+        system_prompt = system_prompt + "\n" + additional_system_prompt
+
     if not conversation_history or len(conversation_history) == 0:
         conversation_history.append({"role": "system", "content": system_prompt})
     elif conversation_history[0]["role"] == "system":
@@ -303,6 +309,7 @@ def build_chat_messages(
     ai: ToolCallingLLM,
     config: Config,
     global_instructions: Optional[Instructions] = None,
+    additional_system_prompt: Optional[str] = None,
 ) -> List[dict]:
     """
     This function generates a list of messages for general chat conversation and ensures that the message sequence adheres to the model's context window limitations
@@ -358,7 +365,10 @@ def build_chat_messages(
         conversation_history = conversation_history.copy()
 
     conversation_history = add_or_update_system_prompt(
-        conversation_history=conversation_history, ai=ai, config=config
+        conversation_history=conversation_history,
+        ai=ai,
+        config=config,
+        additional_system_prompt=additional_system_prompt,
     )
 
     ask = add_global_instructions_to_user_prompt(ask, global_instructions)

holmes/core/feedback.py

@@ -0,0 +1,191 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Callable, Optional
+
+from .llm import LLM
+
+DEFAULT_PRIVACY_NOTICE_BANNER = "Your feedback will be used to improve Holmesgpt's performance. Please avoid sharing sensitive personal information. By continuing, you consent to this data usage."
+PRIVACY_NOTICE_BANNER = os.environ.get(
+    "PRIVACY_NOTICE_BANNER", DEFAULT_PRIVACY_NOTICE_BANNER
+)
+
+
+class FeedbackInfoBase(ABC):
+    """Abstract base class for all feedback-related classes that must implement to_dict()."""
+
+    @abstractmethod
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation. Must be implemented by all subclasses."""
+        pass
+
+
+class FeedbackLLM(FeedbackInfoBase):
+    """Class to represent a LLM in the feedback."""
+
+    def __init__(self, model: str, max_context_size: int):
+        self.model = model
+        self.max_context_size = max_context_size
+
+    def update_from_llm(self, llm: LLM):
+        self.model = llm.model
+        self.max_context_size = llm.get_context_window_size()
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation."""
+        return self.__dict__
+
+
+# TODO: extend the FeedbackLLMResponse to include each tool call results details used for evaluate the overall response.
+# Currenlty tool call details in plan:
+# - toolcall parameter and success/failure, toolcall truncation size
+# - Holmes plan (todo list)
+# - Holmes intermediate output
+class FeedbackLLMResponse(FeedbackInfoBase):
+    """Class to represent a LLM response in the feedback"""
+
+    def __init__(self, user_ask: str, response: str):
+        self.user_ask = user_ask
+        self.response = response
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation."""
+        return self.__dict__
+
+
+class FeedbackMetadata(FeedbackInfoBase):
+    """Class to store feedback metadata."""
+
+    def __init__(self):
+        # In iteration mode, there can be multiple ask and response pairs.
+        self.llm_responses = []
+        self.llm = FeedbackLLM("", 0)
+
+    def add_llm_response(self, user_ask: str, response: str) -> None:
+        """Add a LLM response to the metadata."""
+        llm_response = FeedbackLLMResponse(user_ask, response)
+        self.llm_responses.append(llm_response)
+
+    def update_llm(self, llm: LLM) -> None:
+        """Update the LLM information in the metadata."""
+        self.llm.update_from_llm(llm)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation."""
+        return {
+            "llm_responses": [resp.to_dict() for resp in self.llm_responses],
+            "llm": self.llm.to_dict(),
+        }
+
+
+class UserFeedback(FeedbackInfoBase):
+    """Class to store user rate and comment to the AI response."""
+
+    def __init__(self, is_positive: bool, comment: Optional[str]):
+        self.is_positive = is_positive
+        self.comment = comment
+
+    @property
+    def rating_text(self) -> str:
+        """Return human-readable rating text."""
+        return "useful" if self.is_positive else "not useful"
+
+    @property
+    def rating_emoji(self) -> str:
+        """Return emoji representation of the rating."""
+        return "👍" if self.is_positive else "👎"
+
+    def __str__(self) -> str:
+        """Return string representation of the feedback."""
+        if self.comment:
+            return f"Rating: {self.rating_text}. Comment: {self.comment}"
+        else:
+            return f"Rating: {self.rating_text}. No additional comment."
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation."""
+        return {
+            "is_positive": self.is_positive,
+            "comment": self.comment,
+        }
+
+
+class Feedback(FeedbackInfoBase):
+    """Class to store overall feedback data used to evaluate the AI response."""
+
+    def __init__(self):
+        self.metadata = FeedbackMetadata()
+        self.user_feedback: Optional[UserFeedback] = None
+
+    def set_user_feedback(self, user_feedback: UserFeedback) -> None:
+        """Set the user feedback."""
+        self.user_feedback = user_feedback
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary representation."""
+        return {
+            "metadata": self.metadata.to_dict(),
+            "user_feedback": self.user_feedback.to_dict()
+            if self.user_feedback
+            else None,
+        }
+
+
+FeedbackCallback = Callable[[Feedback], None]
+
+
+def feedback_callback_example(feedback: Feedback) -> None:
+    """
+    Example implementation of a feedback callback function.
+
+    This function demonstrates how to process feedback data using to_dict() methods
+    and could be used for:
+    - Logging feedback to files or databases
+    - Sending feedback to analytics services
+    - Training data collection
+    - User satisfaction monitoring
+
+    Args:
+        feedback: Feedback object containing user feedback and metadata
+    """
+    print("\n=== Feedback Received ===")
+
+    # Convert entire feedback to dict first - this is the main data structure
+    feedback_dict = feedback.to_dict()
+    print(f"Complete feedback dictionary keys: {list(feedback_dict.keys())}")
+
+    # How to check user feedback using to_dict()
+    print("\n1. Checking User Feedback:")
+    user_feedback_dict = (
+        feedback.user_feedback.to_dict() if feedback.user_feedback else None
+    )
+    if user_feedback_dict:
+        print(f"   User feedback dict: {user_feedback_dict}")
+        print(f"   Is positive: {user_feedback_dict['is_positive']}")
+        print(f"   Comment: {user_feedback_dict['comment'] or 'None'}")
+        # You can also access properties through the object:
+        print(f"   Rating emoji: {feedback.user_feedback.rating_emoji}")  # type: ignore
+        print(f"   Rating text: {feedback.user_feedback.rating_text}")  # type: ignore
+    else:
+        print("   No user feedback provided (user_feedback is None)")
+
+    # How to check LLM information using to_dict()
+    print("\n2. Checking LLM Information:")
+    metadata_dict = feedback.metadata.to_dict()
+    llm_dict = metadata_dict["llm"]
+    print(f"   LLM dict: {llm_dict}")
+    print(f"   Model: {llm_dict['model']}")
+    print(f"   Max context size: {llm_dict['max_context_size']}")
+
+    # How to check ask and response pairs using to_dict()
+    print("\n3. Checking Ask and Response History:")
+    llm_responses_dict = metadata_dict["llm_responses"]
+    print(f"   Number of exchanges: {len(llm_responses_dict)}")
+
+    for i, response_dict in enumerate(llm_responses_dict, 1):
+        print(f"   Exchange {i} dict: {list(response_dict.keys())}")
+        user_ask = response_dict["user_ask"]
+        ai_response = response_dict["response"]
+        print(f"     User ask: {user_ask}")
+        print(f"     AI response: {ai_response}")
+
+    print("=== End Feedback ===\n")

holmes/core/llm.py

@@ -1,9 +1,10 @@
 import json
 import logging
 from abc import abstractmethod
+from math import floor
 from typing import Any, Dict, List, Optional, Type, Union, TYPE_CHECKING
 
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import ModelResponse, TextCompletionResponse
 import sentry_sdk
 
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@@ -90,9 +91,13 @@ class DefaultLLM(LLM):
         self.args = args or {}
         self.tracer = tracer
         self.name = name
-
+        self.update_custom_args()
         self.check_llm(self.model, self.api_key, self.api_base, self.api_version)
 
+    def update_custom_args(self):
+        self.max_context_size = self.args.get("custom_args", {}).get("max_context_size")
+        self.args.pop("custom_args", None)
+
     def check_llm(
         self,
         model: str,
@@ -178,6 +183,9 @@ class DefaultLLM(LLM):
         return list(dict.fromkeys(names_to_try))
 
     def get_context_window_size(self) -> int:
+        if self.max_context_size:
+            return self.max_context_size
+
         if OVERRIDE_MAX_CONTENT_SIZE:
             logging.debug(
                 f"Using override OVERRIDE_MAX_CONTENT_SIZE {OVERRIDE_MAX_CONTENT_SIZE}"
@@ -285,6 +293,8 @@ class DefaultLLM(LLM):
             raise Exception(f"Unexpected type returned by the LLM {type(result)}")
 
     def get_maximum_output_token(self) -> int:
+        max_output_tokens = floor(min(64000, self.get_context_window_size() / 5))
+
         if OVERRIDE_MAX_OUTPUT_TOKEN:
             logging.debug(
                 f"Using OVERRIDE_MAX_OUTPUT_TOKEN {OVERRIDE_MAX_OUTPUT_TOKEN}"
@@ -294,17 +304,22 @@ class DefaultLLM(LLM):
         # Try each name variant
         for name in self._get_model_name_variants_for_lookup():
             try:
-                return litellm.model_cost[name]["max_output_tokens"]
+                litellm_max_output_tokens = litellm.model_cost[name][
+                    "max_output_tokens"
+                ]
+                if litellm_max_output_tokens < max_output_tokens:
+                    max_output_tokens = litellm_max_output_tokens
+                return max_output_tokens
             except Exception:
                 continue
 
         # Log which lookups we tried
         logging.warning(
             f"Couldn't find model {self.model} in litellm's model list (tried: {', '.join(self._get_model_name_variants_for_lookup())}), "
-            f"using default 4096 tokens for max_output_tokens. "
+            f"using {max_output_tokens} tokens for max_output_tokens. "
             f"To override, set OVERRIDE_MAX_OUTPUT_TOKEN environment variable to the correct value for your model."
         )
-        return 4096
+        return max_output_tokens
 
     def _add_cache_control_to_last_message(
         self, messages: List[Dict[str, Any]]
@@ -342,7 +357,7 @@ class DefaultLLM(LLM):
         if content is None:
             return
 
-        if isinstance(content, str):
+        if isinstance(content, str) and content:
             # Convert string to structured format with cache_control
             target_msg["content"] = [
                 {
@@ -424,7 +439,8 @@ class LLMModelRegistry:
 
             for model in robusta_models.models:
                 logging.info(f"Loading Robusta AI model: {model}")
-                self._llms[model] = self._create_robusta_model_entry(model)
+                args = robusta_models.models_args.get(model)
+                self._llms[model] = self._create_robusta_model_entry(model, args)
 
             if robusta_models.default_model:
                 logging.info(
@@ -492,7 +508,7 @@ class LLMModelRegistry:
             )
 
         model_key, first_model_params = next(iter(self._llms.items()))
-        logging.info(f"Using first available model: {model_key}")
+        logging.debug(f"Using first available model: {model_key}")
         return first_model_params.copy()
 
     def get_llm(self, name: str) -> LLM:  # TODO: fix logic
@@ -509,13 +525,17 @@ class LLMModelRegistry:
 
         return models
 
-    def _create_robusta_model_entry(self, model_name: str) -> dict[str, Any]:
-        return self._create_model_entry(
+    def _create_robusta_model_entry(
+        self, model_name: str, args: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        entry = self._create_model_entry(
             model="gpt-4o",  # Robusta AI model is using openai like API.
             model_name=model_name,
             base_url=f"{ROBUSTA_API_ENDPOINT}/llm/{model_name}",
             is_robusta_model=True,
         )
+        entry["custom_args"] = args or {}  # type: ignore[assignment]
+        return entry
 
     def _create_model_entry(
         self,
@@ -530,3 +550,25 @@ class LLMModelRegistry:
             "is_robusta_model": is_robusta_model,
             "model": model,
         }
+
+
+def get_llm_usage(
+    llm_response: Union[ModelResponse, CustomStreamWrapper, TextCompletionResponse],
+) -> dict:
+    usage: dict = {}
+    if (
+        (
+            isinstance(llm_response, ModelResponse)
+            or isinstance(llm_response, TextCompletionResponse)
+        )
+        and hasattr(llm_response, "usage")
+        and llm_response.usage
+    ):  # type: ignore
+        usage["prompt_tokens"] = llm_response.usage.prompt_tokens  # type: ignore
+        usage["completion_tokens"] = llm_response.usage.completion_tokens  # type: ignore
+        usage["total_tokens"] = llm_response.usage.total_tokens  # type: ignore
+    elif isinstance(llm_response, CustomStreamWrapper):
+        complete_response = litellm.stream_chunk_builder(chunks=llm_response)  # type: ignore
+        if complete_response:
+            return get_llm_usage(complete_response)
+    return usage

holmes/core/models.py

@@ -1,9 +1,87 @@
+import json
 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, Field
 from enum import Enum
 
+from holmes.core.tools import StructuredToolResult, StructuredToolResultStatus
+
+
+class TruncationMetadata(BaseModel):
+    tool_call_id: str
+    start_index: int
+    end_index: int
+    tool_name: str
+    original_token_count: int
+
+
+class TruncationResult(BaseModel):
+    truncated_messages: list[dict]
+    truncations: list[TruncationMetadata]
+
+
+class ToolCallResult(BaseModel):
+    tool_call_id: str
+    tool_name: str
+    description: str
+    result: StructuredToolResult
+    size: Optional[int] = None
+
+    def as_tool_call_message(self):
+        return {
+            "tool_call_id": self.tool_call_id,
+            "role": "tool",
+            "name": self.tool_name,
+            "content": format_tool_result_data(self.result),
+        }
+
+    def as_tool_result_response(self):
+        result_dump = self.result.model_dump()
+        result_dump["data"] = self.result.get_stringified_data()
+
+        return {
+            "tool_call_id": self.tool_call_id,
+            "tool_name": self.tool_name,
+            "description": self.description,
+            "role": "tool",
+            "result": result_dump,
+        }
+
+    def as_streaming_tool_result_response(self):
+        result_dump = self.result.model_dump()
+        result_dump["data"] = self.result.get_stringified_data()
+
+        return {
+            "tool_call_id": self.tool_call_id,
+            "role": "tool",
+            "description": self.description,
+            "name": self.tool_name,
+            "result": result_dump,
+        }
+
+
+def format_tool_result_data(tool_result: StructuredToolResult) -> str:
+    tool_response = tool_result.data
+    if isinstance(tool_result.data, str):
+        tool_response = tool_result.data
+    else:
+        try:
+            if isinstance(tool_result.data, BaseModel):
+                tool_response = tool_result.data.model_dump_json(indent=2)
+            else:
+                tool_response = json.dumps(tool_result.data, indent=2)
+        except Exception:
+            tool_response = str(tool_result.data)
+    if tool_result.status == StructuredToolResultStatus.ERROR:
+        tool_response = f"{tool_result.error or 'Tool execution failed'}:\n\n{tool_result.data or ''}".strip()
+
+    if tool_result.params:
+        tool_response = (
+            f"Params used for the tool call: {json.dumps(tool_result.params)}. The tool call output follows on the next line.\n"
+            + tool_response
+        )
+    return tool_response
+
 
 class InvestigationResult(BaseModel):
     analysis: Optional[str] = None
@@ -87,10 +165,31 @@ class ConversationRequest(BaseModel):
     include_tool_call_results: bool = False
 
 
+class PendingToolApproval(BaseModel):
+    """Represents a tool call that requires user approval."""
+
+    tool_call_id: str
+    tool_name: str
+    description: str
+    params: Dict[str, Any]
+
+
+class ToolApprovalDecision(BaseModel):
+    """Represents a user's decision on a tool approval."""
+
+    tool_call_id: str
+    approved: bool
+
+
 class ChatRequestBaseModel(BaseModel):
     conversation_history: Optional[list[dict]] = None
     model: Optional[str] = None
     stream: bool = Field(default=False)
+    enable_tool_approval: Optional[bool] = (
+        False  # Optional boolean for backwards compatibility
+    )
+    tool_decisions: Optional[List[ToolApprovalDecision]] = None
+    additional_system_prompt: Optional[str] = None
 
     # In our setup with litellm, the first message in conversation_history
     # should follow the structure [{"role": "system", "content": ...}],
@@ -146,6 +245,7 @@ class ChatResponse(BaseModel):
     conversation_history: list[dict]
     tool_calls: Optional[List[ToolCallResult]] = []
     follow_up_actions: Optional[List[FollowUpAction]] = []
+    pending_approvals: Optional[List[PendingToolApproval]] = None
     metadata: Optional[Dict[Any, Any]] = None