docent-python 0.1.41a0__py3-none-any.whl

docent/_log_util/logger.py

@@ -0,0 +1,141 @@
+import logging
+import sys
+from dataclasses import dataclass
+from typing import Any, Dict, Literal, MutableMapping, Optional, Tuple
+
+
+@dataclass
+class ColorCode:
+    fore: str
+    style: str = ""
+
+
+class Colors:
+    # Foreground colors
+    BLACK = ColorCode("\033[30m")
+    RED = ColorCode("\033[31m")
+    GREEN = ColorCode("\033[32m")
+    YELLOW = ColorCode("\033[33m")
+    BLUE = ColorCode("\033[34m")
+    MAGENTA = ColorCode("\033[35m")
+    CYAN = ColorCode("\033[36m")
+    WHITE = ColorCode("\033[37m")
+    BRIGHT_MAGENTA = ColorCode("\033[95m")
+    BRIGHT_CYAN = ColorCode("\033[96m")
+
+    # Styles
+    BOLD = "\033[1m"
+    RESET = "\033[0m"
+
+    @staticmethod
+    def apply(text: str, color: ColorCode) -> str:
+        return f"{color.style}{color.fore}{text}{Colors.RESET}"
+
+
+class ColoredFormatter(logging.Formatter):
+    COLORS: Dict[int, ColorCode] = {
+        logging.DEBUG: Colors.BLUE,
+        logging.INFO: Colors.GREEN,
+        logging.WARNING: Colors.YELLOW,
+        logging.ERROR: Colors.RED,
+        logging.CRITICAL: ColorCode("\033[31m", Colors.BOLD),
+    }
+
+    # Available highlight colors
+    HIGHLIGHT_COLORS: Dict[str, ColorCode] = {
+        "magenta": ColorCode(Colors.BRIGHT_MAGENTA.fore, Colors.BOLD),
+        "cyan": ColorCode(Colors.BRIGHT_CYAN.fore, Colors.BOLD),
+        "yellow": ColorCode(Colors.YELLOW.fore, Colors.BOLD),
+        "red": ColorCode(Colors.RED.fore, Colors.BOLD),
+    }
+
+    def __init__(self, fmt: Optional[str] = None) -> None:
+        super().__init__(
+            fmt or "%(asctime)s [%(levelname)s] %(namespace)s: %(message)s", datefmt="%H:%M:%S"
+        )
+
+    def format(self, record: logging.LogRecord) -> str:
+        # Add namespace to extra fields if not present
+        if not getattr(record, "namespace", None):
+            record.__dict__["namespace"] = record.name
+
+        # Color the level name
+        record.levelname = Colors.apply(record.levelname, self.COLORS[record.levelno])
+
+        # Color the namespace
+        record.__dict__["namespace"] = Colors.apply(record.__dict__["namespace"], Colors.CYAN)
+
+        # Check if highlight flag is set
+        highlight = getattr(record, "highlight", None)
+        if highlight:
+            # Get the highlight color or default to magenta
+            color_name = highlight if isinstance(highlight, str) else "magenta"
+            highlight_color = self.HIGHLIGHT_COLORS.get(
+                color_name, self.HIGHLIGHT_COLORS["magenta"]
+            )
+
+            # Apply highlight to the message
+            original_message = record.getMessage()
+            record.msg = Colors.apply(original_message, highlight_color)
+            if record.args:
+                record.args = ()
+
+        return super().format(record)
+
+
+class LoggerAdapter(logging.LoggerAdapter[logging.Logger]):
+    """
+    Logger adapter that allows highlighting specific log messages.
+    """
+
+    def process(
+        self, msg: Any, kwargs: MutableMapping[str, Any]
+    ) -> Tuple[Any, MutableMapping[str, Any]]:
+        # Pass highlight flag through to the record
+        return msg, kwargs
+
+    def highlight(
+        self,
+        msg: object,
+        *args: Any,
+        color: Literal["magenta", "cyan", "yellow", "red", "green"] = "magenta",
+        **kwargs: Any,
+    ) -> None:
+        """
+        Log a highlighted message.
+
+        Args:
+            msg: The message format string
+            color: The color to highlight with (magenta, cyan, yellow, red)
+            *args: The args for the message format string
+            **kwargs: Additional logging kwargs
+        """
+        kwargs.setdefault("extra", {})
+        if isinstance(kwargs["extra"], dict):
+            kwargs["extra"]["highlight"] = color
+        return self.info(msg, *args, **kwargs)
+
+
+def get_logger(namespace: str) -> LoggerAdapter:
+    """
+    Get a colored logger for the specified namespace.
+
+    Args:
+        namespace: The namespace for the logger
+
+    Returns:
+        A configured logger instance with highlighting support
+    """
+    logger = logging.getLogger(namespace)
+
+    # Only add handler if it doesn't exist
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(ColoredFormatter())
+        logger.addHandler(handler)
+
+        # Set default level to INFO
+        logger.setLevel(logging.INFO)
+
+    # Wrap with adapter to support highlighting
+    return LoggerAdapter(logger, {})

docent/data_models/__init__.py

@@ -0,0 +1,14 @@
+from docent.data_models.agent_run import AgentRun
+from docent.data_models.citation import InlineCitation
+from docent.data_models.judge import Label
+from docent.data_models.regex import RegexSnippet
+from docent.data_models.transcript import Transcript, TranscriptGroup
+
+__all__ = [
+    "AgentRun",
+    "InlineCitation",
+    "Label",
+    "RegexSnippet",
+    "Transcript",
+    "TranscriptGroup",
+]

docent/data_models/_tiktoken_util.py

@@ -0,0 +1,91 @@
+import tiktoken
+
+MAX_TOKENS = 100_000
+
+
+def get_token_count(text: str, model: str = "gpt-4") -> int:
+    """Get the number of tokens in a text under the GPT-4 tokenization scheme."""
+    encoding = tiktoken.encoding_for_model(model)
+    return len(encoding.encode(text))
+
+
+def truncate_to_token_limit(text: str, max_tokens: int, model: str = "gpt-4") -> str:
+    """Truncate text to stay within the specified token limit."""
+    encoding = tiktoken.encoding_for_model(model)
+    tokens = encoding.encode(text)
+
+    if len(tokens) <= max_tokens:
+        return text
+
+    return encoding.decode(tokens[:max_tokens])
+
+
+class MessageRange:
+    """A range of messages in a transcript. start is inclusive, end is exclusive."""
+
+    start: int
+    end: int
+    include_metadata: bool
+    num_tokens: int
+
+    def __init__(self, start: int, end: int, include_metadata: bool, num_tokens: int):
+        self.start = start
+        self.end = end
+        self.include_metadata = include_metadata
+        self.num_tokens = num_tokens
+
+
+def group_messages_into_ranges(
+    token_counts: list[int], metadata_tokens: int, max_tokens: int, margin: int = 50
+) -> list[MessageRange]:
+    """Split a list of messages + metadata into ranges that stay within the specified token limit.
+
+    Always tries to create ranges with metadata included, unless a single message + metadata is too long,
+    in which case you get a lone message with no metadata
+    """
+    ranges: list[MessageRange] = []
+    start_index = 0
+    running_token_count = 0
+
+    i = 0
+    while i < len(token_counts):
+        new_token_count = token_counts[i]
+        if running_token_count + new_token_count + metadata_tokens > max_tokens - margin:
+            if start_index == i:  # a single message + metadata is already too long
+                ranges.append(
+                    MessageRange(
+                        start=i, end=i + 1, include_metadata=False, num_tokens=new_token_count
+                    )
+                )
+                i += 1
+            else:
+                # add all messages from start_index to i-1, with metadata included
+                ranges.append(
+                    MessageRange(
+                        start=start_index,
+                        end=i,
+                        include_metadata=True,
+                        num_tokens=running_token_count + metadata_tokens,
+                    )
+                )
+            running_token_count = 0
+            start_index = i
+        else:
+            running_token_count += new_token_count
+            i += 1
+
+    if running_token_count > 0:
+        include_metadata = running_token_count + metadata_tokens < max_tokens - margin
+        num_tokens = (
+            running_token_count + metadata_tokens if include_metadata else running_token_count
+        )
+        ranges.append(
+            MessageRange(
+                start=start_index,
+                end=len(token_counts),
+                include_metadata=include_metadata,
+                num_tokens=num_tokens,
+            )
+        )
+
+    return ranges