holmesgpt 0.13.2__py3-none-any.whl → 0.18.4__py3-none-any.whl

holmes/core/transformers/registry.py

@@ -0,0 +1,123 @@
+"""
+Transformer registry for managing available transformers.
+"""
+
+from typing import Any, Dict, List, Optional, Type
+
+from .base import BaseTransformer, TransformerError
+
+
+class TransformerRegistry:
+    """
+    Registry for managing transformer types and creating transformer instances.
+
+    This registry provides a centralized way to register transformer classes
+    and create instances based on configuration.
+    """
+
+    def __init__(self):
+        self._transformers: Dict[str, Type[BaseTransformer]] = {}
+
+    def register(self, transformer_class: Type[BaseTransformer]) -> None:
+        """
+        Register a transformer class, using the transformer's name property.
+
+        Args:
+            transformer_class: The transformer class to register
+
+        Raises:
+            ValueError: If name is already registered or transformer_class is invalid
+        """
+        if not issubclass(transformer_class, BaseTransformer):
+            raise ValueError(
+                f"Transformer class must inherit from BaseTransformer, got {transformer_class}"
+            )
+
+        # Get name from the transformer class
+        try:
+            temp_instance = transformer_class()
+            name = temp_instance.name
+        except Exception:
+            # Fallback to class name if instantiation fails
+            name = transformer_class.__name__
+
+        if name in self._transformers:
+            raise ValueError(f"Transformer '{name}' is already registered")
+
+        self._transformers[name] = transformer_class
+
+    def unregister(self, name: str) -> None:
+        """
+        Unregister a transformer by name.
+
+        Args:
+            name: The name of the transformer to unregister
+
+        Raises:
+            KeyError: If transformer name is not registered
+        """
+        if name not in self._transformers:
+            raise KeyError(f"Transformer '{name}' is not registered")
+
+        del self._transformers[name]
+
+    def create_transformer(
+        self, name: str, config: Optional[Dict[str, Any]] = None
+    ) -> BaseTransformer:
+        """
+        Create a transformer instance by name.
+
+        Args:
+            name: The name of the transformer to create
+            config: Optional configuration for the transformer
+
+        Returns:
+            A new transformer instance
+
+        Raises:
+            KeyError: If transformer name is not registered
+            TransformerError: If transformer creation fails
+        """
+        if name not in self._transformers:
+            raise KeyError(f"Transformer '{name}' is not registered")
+
+        transformer_class = self._transformers[name]
+
+        try:
+            # Handle both old-style dict config and new Pydantic models
+            if config is None:
+                return transformer_class()
+            else:
+                # For Pydantic models, pass config as keyword arguments
+                return transformer_class(**config)
+        except Exception as e:
+            raise TransformerError(f"Failed to create transformer '{name}': {e}") from e
+
+    def is_registered(self, name: str) -> bool:
+        """
+        Check if a transformer is registered.
+
+        Args:
+            name: The name to check
+
+        Returns:
+            True if the transformer is registered, False otherwise
+        """
+        return name in self._transformers
+
+    def list_transformers(self) -> List[str]:
+        """
+        Get a list of all registered transformer names.
+
+        Returns:
+            List of registered transformer names
+        """
+        return list(self._transformers.keys())
+
+    def clear(self) -> None:
+        """Clear all registered transformers."""
+        self._transformers.clear()
+
+
+# Global transformer registry instance
+registry = TransformerRegistry()

holmes/core/transformers/transformer.py

@@ -0,0 +1,32 @@
+"""
+Configuration class for tool transformers.
+"""
+
+import logging
+from typing import Any, Dict
+
+from pydantic import BaseModel, Field, model_validator
+
+from .registry import registry
+
+
+class Transformer(BaseModel):
+    """
+    Configuration for a tool transformer.
+
+    Each transformer config specifies a transformer type and its parameters.
+    This replaces the previous dict-based configuration with proper type safety.
+    """
+
+    name: str = Field(description="Name of the transformer (e.g., 'llm_summarize')")
+    config: Dict[str, Any] = Field(
+        default_factory=dict, description="Configuration parameters for the transformer"
+    )
+
+    @model_validator(mode="after")
+    def validate_transformer(self):
+        """Validate that the transformer name is known to the registry."""
+        if not registry.is_registered(self.name):
+            # Log warning but don't fail validation - allows for graceful degradation
+            logging.warning(f"Transformer '{self.name}' is not registered")
+        return self

holmes/core/truncation/compaction.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Optional
+
+import litellm
+from litellm.types.utils import ModelResponse
+
+from holmes.core.llm import LLM
+from holmes.plugins.prompts import load_and_render_prompt
+
+
+def strip_system_prompt(
+    conversation_history: list[dict],
+) -> tuple[list[dict], Optional[dict]]:
+    if not conversation_history:
+        return conversation_history, None
+    first_message = conversation_history[0]
+    if first_message and first_message.get("role") == "system":
+        return conversation_history[1:], first_message
+    return conversation_history[:], None
+
+
+def find_last_user_prompt(conversation_history: list[dict]) -> Optional[dict]:
+    if not conversation_history:
+        return None
+    last_user_prompt: Optional[dict] = None
+    for message in conversation_history:
+        if message.get("role") == "user":
+            last_user_prompt = message
+    return last_user_prompt
+
+
+def compact_conversation_history(
+    original_conversation_history: list[dict], llm: LLM
+) -> list[dict]:
+    """
+    The compacted conversation history contains:
+      1. Original system prompt, uncompacted (if present)
+      2. Last user prompt, uncompacted (if present)
+      3. Compacted conversation history (role=assistant)
+      4. Compaction message (role=system)
+    """
+    conversation_history, system_prompt_message = strip_system_prompt(
+        original_conversation_history
+    )
+    compaction_instructions = load_and_render_prompt(
+        prompt="builtin://conversation_history_compaction.jinja2", context={}
+    )
+    conversation_history.append({"role": "user", "content": compaction_instructions})
+
+    # Set modify_params to handle providers like Anthropic that require tools
+    # when conversation history contains tool calls
+    original_modify_params = litellm.modify_params
+    try:
+        litellm.modify_params = True  # necessary when using anthropic
+        response: ModelResponse = llm.completion(
+            messages=conversation_history, drop_params=True
+        )  # type: ignore
+    finally:
+        litellm.modify_params = original_modify_params
+    response_message = None
+    if (
+        response
+        and response.choices
+        and response.choices[0]
+        and response.choices[0].message  # type:ignore
+    ):
+        response_message = response.choices[0].message  # type:ignore
+    else:
+        logging.error(
+            "Failed to compact conversation history. Unexpected LLM's response for compaction"
+        )
+        return original_conversation_history
+
+    compacted_conversation_history: list[dict] = []
+    if system_prompt_message:
+        compacted_conversation_history.append(system_prompt_message)
+
+    last_user_prompt = find_last_user_prompt(original_conversation_history)
+    if last_user_prompt:
+        compacted_conversation_history.append(last_user_prompt)
+
+    compacted_conversation_history.append(
+        response_message.model_dump(
+            exclude_defaults=True, exclude_unset=True, exclude_none=True
+        )
+    )
+
+    compacted_conversation_history.append(
+        {
+            "role": "system",
+            "content": "The conversation history has been compacted to preserve available space in the context window. Continue.",
+        }
+    )
+    return compacted_conversation_history

holmes/core/truncation/dal_truncation_utils.py

@@ -0,0 +1,23 @@
+from holmes.common.env_vars import MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION
+
+
+def truncate_string(data_str: str) -> str:
+    if data_str and len(data_str) > MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION:
+        return (
+            data_str[:MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION]
+            + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS"
+        )
+    return data_str
+
+
+def truncate_evidences_entities_if_necessary(evidence_list: list[dict]):
+    if (
+        not MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION
+        or MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION <= 0
+    ):
+        return
+
+    for evidence in evidence_list:
+        data = evidence.get("data")
+        if data:
+            evidence["data"] = truncate_string(str(data))

holmes/core/truncation/input_context_window_limiter.py

@@ -0,0 +1,219 @@
+import logging
+from typing import Any, Optional
+
+import sentry_sdk
+from pydantic import BaseModel
+
+from holmes.common.env_vars import (
+    ENABLE_CONVERSATION_HISTORY_COMPACTION,
+    MAX_OUTPUT_TOKEN_RESERVATION,
+)
+from holmes.core.llm import (
+    LLM,
+    TokenCountMetadata,
+    get_context_window_compaction_threshold_pct,
+)
+from holmes.core.models import TruncationMetadata, TruncationResult
+from holmes.core.truncation.compaction import compact_conversation_history
+from holmes.utils import sentry_helper
+from holmes.utils.stream import StreamEvents, StreamMessage
+
+TRUNCATION_NOTICE = "\n\n[TRUNCATED]"
+
+
+def _truncate_tool_message(
+    msg: dict, allocated_space: int, needed_space: int
+) -> TruncationMetadata:
+    msg_content = msg["content"]
+    tool_call_id = msg["tool_call_id"]
+    tool_name = msg["name"]
+
+    # Ensure the indicator fits in the allocated space
+    if allocated_space > len(TRUNCATION_NOTICE):
+        original = msg_content if isinstance(msg_content, str) else str(msg_content)
+        msg["content"] = (
+            original[: allocated_space - len(TRUNCATION_NOTICE)] + TRUNCATION_NOTICE
+        )
+        end_index = allocated_space - len(TRUNCATION_NOTICE)
+    else:
+        msg["content"] = TRUNCATION_NOTICE[:allocated_space]
+        end_index = allocated_space
+
+    msg.pop("token_count", None)  # Remove token_count if present
+    logging.info(
+        f"Truncating tool message '{tool_name}' from {needed_space} to {allocated_space} tokens"
+    )
+    truncation_metadata = TruncationMetadata(
+        tool_call_id=tool_call_id,
+        start_index=0,
+        end_index=end_index,
+        tool_name=tool_name,
+        original_token_count=needed_space,
+    )
+    return truncation_metadata
+
+
+# TODO: I think there's a bug here because we don't account for the 'role' or json structure like '{...}' when counting tokens
+# However, in practice it works because we reserve enough space for the output tokens that the minor inconsistency does not matter
+# We should fix this in the future
+# TODO: we truncate using character counts not token counts - this means we're overly agressive with truncation - improve it by considering
+# token truncation and not character truncation
+def truncate_messages_to_fit_context(
+    messages: list, max_context_size: int, maximum_output_token: int, count_tokens_fn
+) -> TruncationResult:
+    """
+    Helper function to truncate tool messages to fit within context limits.
+
+    Args:
+        messages: List of message dictionaries with roles and content
+        max_context_size: Maximum context window size for the model
+        maximum_output_token: Maximum tokens reserved for model output
+        count_tokens_fn: Function to count tokens for a list of messages
+
+    Returns:
+        Modified list of messages with truncated tool responses
+
+    Raises:
+        Exception: If non-tool messages exceed available context space
+    """
+    messages_except_tools = [
+        message for message in messages if message["role"] != "tool"
+    ]
+    tokens = count_tokens_fn(messages_except_tools)
+    message_size_without_tools = tokens.total_tokens
+
+    tool_call_messages = [message for message in messages if message["role"] == "tool"]
+
+    reserved_for_output_tokens = min(maximum_output_token, MAX_OUTPUT_TOKEN_RESERVATION)
+    if message_size_without_tools >= (max_context_size - reserved_for_output_tokens):
+        logging.error(
+            f"The combined size of system_prompt and user_prompt ({message_size_without_tools} tokens) exceeds the model's context window for input."
+        )
+        raise Exception(
+            f"The combined size of system_prompt and user_prompt ({message_size_without_tools} tokens) exceeds the maximum context size of {max_context_size - reserved_for_output_tokens} tokens available for input."
+        )
+
+    if len(tool_call_messages) == 0:
+        return TruncationResult(truncated_messages=messages, truncations=[])
+
+    available_space = (
+        max_context_size - message_size_without_tools - reserved_for_output_tokens
+    )
+    remaining_space = available_space
+    tool_call_messages.sort(
+        key=lambda x: count_tokens_fn(
+            [{"role": "tool", "content": x["content"]}]
+        ).total_tokens
+    )
+
+    truncations = []
+
+    # Allocate space starting with small tools and going to larger tools, while maintaining fairness
+    # Small tools can often get exactly what they need, while larger tools may need to be truncated
+    # We ensure fairness (no tool gets more than others that need it) and also maximize utilization (we don't leave space unused)
+    for i, msg in enumerate(tool_call_messages):
+        remaining_tools = len(tool_call_messages) - i
+        max_allocation = remaining_space // remaining_tools
+        needed_space = count_tokens_fn(
+            [{"role": "tool", "content": msg["content"]}]
+        ).total_tokens
+        allocated_space = min(needed_space, max_allocation)
+
+        if needed_space > allocated_space:
+            truncation_metadata = _truncate_tool_message(
+                msg, allocated_space, needed_space
+            )
+            truncations.append(truncation_metadata)
+
+        remaining_space -= allocated_space
+
+    if truncations:
+        sentry_helper.capture_tool_truncations(truncations)
+
+    return TruncationResult(truncated_messages=messages, truncations=truncations)
+
+
+class ContextWindowLimiterOutput(BaseModel):
+    metadata: dict
+    messages: list[dict]
+    events: list[StreamMessage]
+    max_context_size: int
+    maximum_output_token: int
+    tokens: TokenCountMetadata
+    conversation_history_compacted: bool
+
+
+@sentry_sdk.trace
+def limit_input_context_window(
+    llm: LLM, messages: list[dict], tools: Optional[list[dict[str, Any]]]
+) -> ContextWindowLimiterOutput:
+    events = []
+    metadata = {}
+    initial_tokens = llm.count_tokens(messages=messages, tools=tools)  # type: ignore
+    max_context_size = llm.get_context_window_size()
+    maximum_output_token = llm.get_maximum_output_token()
+    conversation_history_compacted = False
+    if ENABLE_CONVERSATION_HISTORY_COMPACTION and (
+        initial_tokens.total_tokens + maximum_output_token
+    ) > (max_context_size * get_context_window_compaction_threshold_pct() / 100):
+        compacted_messages = compact_conversation_history(
+            original_conversation_history=messages, llm=llm
+        )
+        compacted_tokens = llm.count_tokens(compacted_messages, tools=tools)
+        compacted_total_tokens = compacted_tokens.total_tokens
+
+        if compacted_total_tokens < initial_tokens.total_tokens:
+            messages = compacted_messages
+            compaction_message = f"The conversation history has been compacted from {initial_tokens.total_tokens} to {compacted_total_tokens} tokens"
+            logging.info(compaction_message)
+            conversation_history_compacted = True
+            events.append(
+                StreamMessage(
+                    event=StreamEvents.CONVERSATION_HISTORY_COMPACTED,
+                    data={
+                        "content": compaction_message,
+                        "messages": compacted_messages,
+                        "metadata": {
+                            "initial_tokens": initial_tokens.total_tokens,
+                            "compacted_tokens": compacted_total_tokens,
+                        },
+                    },
+                )
+            )
+            events.append(
+                StreamMessage(
+                    event=StreamEvents.AI_MESSAGE,
+                    data={"content": compaction_message},
+                )
+            )
+        else:
+            logging.debug(
+                f"Failed to reduce token count when compacting conversation history. Original tokens:{initial_tokens.total_tokens}. Compacted tokens:{compacted_total_tokens}"
+            )
+
+    tokens = llm.count_tokens(messages=messages, tools=tools)  # type: ignore
+    if (tokens.total_tokens + maximum_output_token) > max_context_size:
+        # Compaction was not sufficient. Truncating messages.
+        truncated_res = truncate_messages_to_fit_context(
+            messages=messages,
+            max_context_size=max_context_size,
+            maximum_output_token=maximum_output_token,
+            count_tokens_fn=llm.count_tokens,
+        )
+        metadata["truncations"] = [t.model_dump() for t in truncated_res.truncations]
+        messages = truncated_res.truncated_messages
+
+        # recount after truncation
+        tokens = llm.count_tokens(messages=messages, tools=tools)  # type: ignore
+    else:
+        metadata["truncations"] = []
+
+    return ContextWindowLimiterOutput(
+        events=events,
+        messages=messages,
+        metadata=metadata,
+        max_context_size=max_context_size,
+        maximum_output_token=maximum_output_token,
+        tokens=tokens,
+        conversation_history_compacted=conversation_history_compacted,
+    )