Printing AI messages

The AIModelPrintingMixin class is a mixin that is used to print messages during the AI model call.

This class is not meant to display the final response to the user, but just to serve as an indicator of what is going on. After the AI model call, the terminal is cleaned up.

The mixin offers three modes:

• normal: The default value. If the response is streamed, it will indicate to the user that we're getting chunks without displaying the full response.
• preview: If the response is streamed, it will display the response from the AI model as it being received.
• silent: We don't display anything to the standard output. Useful in scripts.
• non_tty: If the standard output is not a terminal, we use this mode.

Screen casts

These video examples should show you how the mixin is used by higher-level classes.

Warning

The screen casts are slightly out-of-date: AIPrompt.from_str has been removed, and you should use the AIPrompt constructor directly instead.

Normal mode

Preview mode

Silent mode

conatus.models.printing.AIModelPrintingMixin

AIModelPrintingMixin(
    config: ModelConfig, *, debug_mode: bool | None = None
)

Mixin for printing messages during the AI model call.

PARAMETER	DESCRIPTION
config	The configuration for the model. TYPE: ModelConfig
debug_mode	Whether to run in debug mode. TYPE: bool | None DEFAULT: None

Source code in conatus/models/printing.py

def __init__(
    self,
    config: ModelConfig,
    *,
    debug_mode: bool | None = None,
) -> None:
    """Initialize the printing mixin.

    Args:
        config: The configuration for the model.
        debug_mode: Whether to run in debug mode.
    """
    self.stdout_mode = config.stdout_mode or self.stdout_mode
    # If the output is not a terminal, we don't use ANSI codes
    if not sys.stdout.isatty() and config.stdout_mode != "silent":
        self.stdout_mode = (
            "non_tty"
            if not config.use_mock
            else (config.stdout_mode or "normal")
        )
    if isinstance(config.model_name, str):
        self.model_name = config.model_name
    else:
        self.model_name = "model"
    self.terminal_width = shutil.get_terminal_size().columns
    self.terminal_height = shutil.get_terminal_size().lines
    self.lines_to_clear = 1
    self.i = 0
    self._no_chunks_received: bool = True
    self.last_print_ts = datetime.now(timezone.utc)
    self.delta_between_prints = timedelta(milliseconds=50)
    self.debug_mode = (
        debug_mode
        if debug_mode is not None
        else getattr(self, "debug_mode", False)
    )
    if not self.debug_mode:  # pragma: no branch
        return

    # Initialize the debug buffer to write to debug.txt
    with Path("debug.txt").open(
        "w", encoding="utf-8"
    ) as f:  # pragma: no cover
        _ = f.write(
            "=== Debug Log Started ===\n"
            f"Is stdout a tty? {sys.stdout.isatty()}\n"
            f"Terminal size: {self.terminal_width}x{self.terminal_height}\n"
        )

model_name instance-attribute

model_name: str

The name of the model.

stdout_mode class-attribute instance-attribute

stdout_mode: Literal[
    "silent", "preview", "normal", "non_tty"
] = (stdout_mode or stdout_mode)

The mode for printing the messages.

• 'normal': Notify the user that we're waiting for a response, and then that we're receiving the response, displaying the number of chunks received so far.

• 'preview': Preview the response with a fancy output that updates as the response chunks are received. Only works if the response is a stream. If preview is set and the response is not a stream, it will default to 'normal'.

• 'silent': Do not print anything to the standard output.

Note that if we detect that we are running in a non TTY environment, we will use a special mode called 'non_tty', unless the user asked for 'silent'.

terminal_width instance-attribute

terminal_width: int = columns

The width of the terminal.

terminal_height instance-attribute

terminal_height: int = lines

The height of the terminal.

lines_to_clear instance-attribute

lines_to_clear: int = 1

The number of lines to clear after printing the message.

This number might increase over time if the message gets long.

i instance-attribute

i: int = 0

The index of the response.

last_print_ts instance-attribute

last_print_ts: datetime = now(utc)

The timestamp of the last print.

delta_between_prints instance-attribute

delta_between_prints: timedelta = timedelta(milliseconds=50)

The time delta between prints.

debug_mode class-attribute instance-attribute

debug_mode: bool = (
    debug_mode
    if debug_mode is not None
    else getattr(self, "debug_mode", False)
)

Whether to run in debug mode.

print_before_sending

print_before_sending(message: str | None = None) -> None

Print a user-facing message before the AI model call.

PARAMETER	DESCRIPTION
message	The message to print. If not provided, the default message will be used. TYPE: str | None DEFAULT: None

Source code in conatus/models/printing.py

def print_before_sending(
    self,
    message: str | None = None,
) -> None:
    """Print a user-facing message before the AI model call.

    Args:
        message: The message to print. If not provided, the default
            message will be used.
    """
    message = (
        f"Waiting for response from {self.model_name}..."
        if message is None
        else message
    )
    optional_newline = "\n" if self.stdout_mode == "preview" else ""
    if self.stdout_mode != "silent":
        self._stdout_write(
            grey(message) + optional_newline,
            flush=True,
        )
    self.i = 0
    self._no_chunks_received = False

write_preview_response

write_preview_response(
    response: IncompleteAIResponse[Any] | AIResponse[Any],
) -> None

Write the preview response.

This method handles real-time updating of the AI model's response, showing:

- The current message content
- Any tool/function calls
- Token usage statistics

The display is updated in-place using ANSI escape sequences for smooth updates. We also calculate the number of lines to clear after writing the new content, as well as the current terminal width.

Example Output:

Message: I think the best approach would be...
Tool calls: search_database({"query": "user preferences"})
Total tokens: 147

Pro tip: To debug, use time.sleep(1) to see how the preview response is updated.

PARAMETER	DESCRIPTION
response	The current incomplete response from the AI model. TYPE: IncompleteAIResponse[Any] | AIResponse[Any]

Source code in conatus/models/printing.py

def write_preview_response(
    self,
    response: IncompleteAIResponse[Any] | AIResponse[Any],  # pyright: ignore[reportExplicitAny]
) -> None:
    """Write the preview response.

    This method handles real-time updating of the AI model's response,
    showing:

        - The current message content
        - Any tool/function calls
        - Token usage statistics

    The display is updated in-place using ANSI escape sequences for smooth
    updates. We also calculate the number of lines to clear after writing
    the new content, as well as the current terminal width.

    Example Output:

    ```
    Message: I think the best approach would be...
    Tool calls: search_database({"query": "user preferences"})
    Total tokens: 147
    ```

    Pro tip: To debug, use `time.sleep(1)` to see how the preview response
    is updated.

    Args:
        response: The current incomplete response from the AI model.
    """
    self.i += 1

    response_text = response.message_received.all_text or " "
    # If not enough time has passed since last print, skip this update
    if not self._should_update_preview(response_text):
        return None

    # If we're not in preview mode, we will write something simpler
    if self.stdout_mode != "preview":
        return self._write_non_preview_response()

    # Prepare the message content for display
    max_len = self.terminal_width - 15
    message_content_lines = self._prepare_message_content_lines(
        response_text, max_len
    )

    # Calculate display parameters
    display_params = self._calculate_display_parameters(
        message_content_lines, response, max_len
    )

    # Update the display - explicitly cast to int to avoid type errors
    new_lines_to_clear: int = display_params["new_lines_to_clear"]
    self._update_terminal_cursor(new_lines_to_clear)
    self._write_message_content(message_content_lines, display_params)

    # Update lines to clear for next time
    self.lines_to_clear = new_lines_to_clear - 1

    return None

clean_after_receiving

clean_after_receiving() -> None

Clean the terminal after the AI model call.

Source code in conatus/models/printing.py

def clean_after_receiving(self) -> None:
    """Clean the terminal after the AI model call."""
    self._no_chunks_received = True
    self.i = 0
    if self.stdout_mode not in {"silent", "non_tty"}:
        self._clean_preview_response(debug_mode=self.debug_mode)

write_debug

write_debug(message: str) -> None

Write debug information to the debug buffer.

PARAMETER	DESCRIPTION
message	The debug message to write. TYPE: str

Source code in conatus/models/printing.py

def write_debug(self, message: str) -> None:
    """Write debug information to the debug buffer.

    Args:
        message: The debug message to write.
    """
    with Path("debug.txt").open(
        "a", encoding="utf-8"
    ) as f:  # pragma: no cover
        _ = f.write(f"[{self.i}] {message}\n")
        f.flush()