chore: move utils functions out of main.py

holmes/main.py

@@ -16,15 +16,11 @@
 import logging
 import socket
 import uuid
-import warnings
 from datetime import datetime
-from enum import Enum
 from pathlib import Path
 from typing import List, Optional
 
 import typer
-from rich.console import Console
-from rich.logging import RichHandler
 from rich.markdown import Markdown
 from rich.rule import Rule
 
@@ -37,14 +33,16 @@
 )
 from holmes.core.prompt import build_initial_ask_messages
 from holmes.core.resource_instruction import ResourceInstructionDocument
-from holmes.core.tool_calling_llm import LLMResult
 from holmes.core.tools import pretty_print_toolset_status
+from holmes.core.tracing import SpanType, TracingFactory
 from holmes.interactive import run_interactive_loop
 from holmes.plugins.destinations import DestinationType
 from holmes.plugins.interfaces import Issue
-from holmes.core.tracing import TracingFactory, SpanType
 from holmes.plugins.prompts import load_and_render_prompt
 from holmes.plugins.sources.opsgenie import OPSGENIE_TEAM_INTEGRATION_KEY_HELP
+from holmes.utils.console.consts import system_prompt_help
+from holmes.utils.console.logging import init_logging
+from holmes.utils.console.result import handle_result
 from holmes.utils.file_utils import write_json_file
 
 app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False)
@@ -71,94 +69,6 @@
 app.add_typer(toolset_app, name="toolset")
 
 
-class Verbosity(Enum):
-    NORMAL = 0
-    LOG_QUERIES = 1  # TODO: currently unused
-    VERBOSE = 2
-    VERY_VERBOSE = 3
-
-
-def cli_flags_to_verbosity(verbose_flags: List[bool]) -> Verbosity:
-    if verbose_flags is None or len(verbose_flags) == 0:
-        return Verbosity.NORMAL
-    elif len(verbose_flags) == 1:
-        return Verbosity.LOG_QUERIES
-    elif len(verbose_flags) == 2:
-        return Verbosity.VERBOSE
-    else:
-        return Verbosity.VERY_VERBOSE
-
-
-def suppress_noisy_logs():
-    # disable INFO logs from OpenAI
-    logging.getLogger("httpx").setLevel(logging.WARNING)
-    # disable INFO logs from LiteLLM
-    logging.getLogger("LiteLLM").setLevel(logging.WARNING)
-    # disable INFO logs from AWS (relevant when using bedrock)
-    logging.getLogger("boto3").setLevel(logging.WARNING)
-    logging.getLogger("botocore").setLevel(logging.WARNING)
-    # when running in --verbose mode we don't want to see DEBUG logs from these libraries
-    logging.getLogger("openai._base_client").setLevel(logging.INFO)
-    logging.getLogger("httpcore").setLevel(logging.INFO)
-    logging.getLogger("markdown_it").setLevel(logging.INFO)
-    # suppress UserWarnings from the slack_sdk module
-    warnings.filterwarnings("ignore", category=UserWarning, module="slack_sdk.*")
-
-
-def init_logging(verbose_flags: Optional[List[bool]] = None):
-    verbosity = cli_flags_to_verbosity(verbose_flags)  # type: ignore
-
-    if verbosity == Verbosity.VERY_VERBOSE:
-        logging.basicConfig(
-            level=logging.DEBUG,
-            format="%(message)s",
-            handlers=[
-                RichHandler(
-                    show_level=False,
-                    markup=True,
-                    show_time=False,
-                    show_path=False,
-                    console=Console(width=None),
-                )
-            ],
-        )
-    elif verbosity == Verbosity.VERBOSE:
-        logging.basicConfig(
-            level=logging.INFO,
-            format="%(message)s",
-            handlers=[
-                RichHandler(
-                    show_level=False,
-                    markup=True,
-                    show_time=False,
-                    show_path=False,
-                    console=Console(width=None),
-                )
-            ],
-        )
-        logging.getLogger().setLevel(logging.DEBUG)
-        suppress_noisy_logs()
-    else:
-        logging.basicConfig(
-            level=logging.INFO,
-            format="%(message)s",
-            handlers=[
-                RichHandler(
-                    show_level=False,
-                    markup=True,
-                    show_time=False,
-                    show_path=False,
-                    console=Console(width=None),
-                )
-            ],
-        )
-        suppress_noisy_logs()
-
-    logging.debug(f"verbosity is {verbosity}")
-
-    return Console()
-
-
 # Common cli options
 # The defaults for options that are also in the config file MUST be None or else the cli defaults will override settings in the config file
 opt_api_key: Optional[str] = typer.Option(
@@ -234,9 +144,6 @@ def init_logging(verbose_flags: Optional[List[bool]] = None):
     help="Additional documents to provide the LLM (typically URLs to runbooks)",
 )
 
-# Common help texts
-system_prompt_help = "Advanced. System prompt for LLM. Values starting with builtin:// are loaded from holmes/plugins/prompts, values starting with file:// are loaded from the given path, other values are interpreted as a prompt string"
-
 
 def parse_documents(documents: Optional[str]) -> List[ResourceInstructionDocument]:
     resource_documents = []
@@ -250,35 +157,6 @@ def parse_documents(documents: Optional[str]) -> List[ResourceInstructionDocumen
     return resource_documents
 
 
-def handle_result(
-    result: LLMResult,
-    console: Console,
-    destination: DestinationType,
-    config: Config,
-    issue: Issue,
-    show_tool_output: bool,
-    add_separator: bool,
-):
-    if destination == DestinationType.CLI:
-        if show_tool_output and result.tool_calls:
-            for tool_call in result.tool_calls:
-                console.print("[bold magenta]Used Tool:[/bold magenta]", end="")
-                # we need to print this separately with markup=False because it contains arbitrary text and we don't want console.print to interpret it
-                console.print(
-                    f"{tool_call.description}. Output=\n{tool_call.result}",
-                    markup=False,
-                )
-
-        console.print("[bold green]AI:[/bold green]", end=" ")
-        console.print(Markdown(result.result))  # type: ignore
-        if add_separator:
-            console.print(Rule())
-
-    elif destination == DestinationType.SLACK:
-        slack = config.create_slack_destination()
-        slack.send_issue(issue, result)
-
-
 # TODO: add streaming output
 @app.command()
 def ask(

holmes/utils/console/consts.py

@@ -0,0 +1,2 @@
+# Common help texts
+system_prompt_help = "Advanced. System prompt for LLM. Values starting with builtin:// are loaded from holmes/plugins/prompts, values starting with file:// are loaded from the given path, other values are interpreted as a prompt string"

holmes/utils/console/logging.py

@@ -0,0 +1,95 @@
+import logging
+import warnings
+from enum import Enum
+from typing import List, Optional
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+
+class Verbosity(Enum):
+    NORMAL = 0
+    LOG_QUERIES = 1  # TODO: currently unused
+    VERBOSE = 2
+    VERY_VERBOSE = 3
+
+
+def cli_flags_to_verbosity(verbose_flags: List[bool]) -> Verbosity:
+    if verbose_flags is None or len(verbose_flags) == 0:
+        return Verbosity.NORMAL
+    elif len(verbose_flags) == 1:
+        return Verbosity.LOG_QUERIES
+    elif len(verbose_flags) == 2:
+        return Verbosity.VERBOSE
+    else:
+        return Verbosity.VERY_VERBOSE
+
+
+def suppress_noisy_logs():
+    # disable INFO logs from OpenAI
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    # disable INFO logs from LiteLLM
+    logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+    # disable INFO logs from AWS (relevant when using bedrock)
+    logging.getLogger("boto3").setLevel(logging.WARNING)
+    logging.getLogger("botocore").setLevel(logging.WARNING)
+    # when running in --verbose mode we don't want to see DEBUG logs from these libraries
+    logging.getLogger("openai._base_client").setLevel(logging.INFO)
+    logging.getLogger("httpcore").setLevel(logging.INFO)
+    logging.getLogger("markdown_it").setLevel(logging.INFO)
+    # suppress UserWarnings from the slack_sdk module
+    warnings.filterwarnings("ignore", category=UserWarning, module="slack_sdk.*")
+
+
+def init_logging(verbose_flags: Optional[List[bool]] = None):
+    verbosity = cli_flags_to_verbosity(verbose_flags)  # type: ignore
+
+    if verbosity == Verbosity.VERY_VERBOSE:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format="%(message)s",
+            handlers=[
+                RichHandler(
+                    show_level=False,
+                    markup=True,
+                    show_time=False,
+                    show_path=False,
+                    console=Console(width=None),
+                )
+            ],
+        )
+    elif verbosity == Verbosity.VERBOSE:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(message)s",
+            handlers=[
+                RichHandler(
+                    show_level=False,
+                    markup=True,
+                    show_time=False,
+                    show_path=False,
+                    console=Console(width=None),
+                )
+            ],
+        )
+        logging.getLogger().setLevel(logging.DEBUG)
+        suppress_noisy_logs()
+    else:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(message)s",
+            handlers=[
+                RichHandler(
+                    show_level=False,
+                    markup=True,
+                    show_time=False,
+                    show_path=False,
+                    console=Console(width=None),
+                )
+            ],
+        )
+        suppress_noisy_logs()
+
+    logging.debug(f"verbosity is {verbosity}")
+
+    return Console()

holmes/utils/console/result.py

@@ -0,0 +1,37 @@
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.rule import Rule
+
+from holmes.config import Config
+from holmes.core.tool_calling_llm import LLMResult
+from holmes.plugins.destinations import DestinationType
+from holmes.plugins.interfaces import Issue
+
+
+def handle_result(
+    result: LLMResult,
+    console: Console,
+    destination: DestinationType,
+    config: Config,
+    issue: Issue,
+    show_tool_output: bool,
+    add_separator: bool,
+):
+    if destination == DestinationType.CLI:
+        if show_tool_output and result.tool_calls:
+            for tool_call in result.tool_calls:
+                console.print("[bold magenta]Used Tool:[/bold magenta]", end="")
+                # we need to print this separately with markup=False because it contains arbitrary text and we don't want console.print to interpret it
+                console.print(
+                    f"{tool_call.description}. Output=\n{tool_call.result}",
+                    markup=False,
+                )
+
+        console.print("[bold green]AI:[/bold green]", end=" ")
+        console.print(Markdown(result.result))  # type: ignore
+        if add_separator:
+            console.print(Rule())
+
+    elif destination == DestinationType.SLACK:
+        slack = config.create_slack_destination()
+        slack.send_issue(issue, result)