unique_toolkit 1.42.9__py3-none-any.whl → 1.43.1__py3-none-any.whl

unique_toolkit/_common/experimental/write_up_agent/services/dataframe_handler/utils.py

@@ -0,0 +1,130 @@
+"""Utility functions for DataFrame operations."""
+
+import re
+
+import pandas as pd
+
+
+def to_snake_case(text: str) -> str:
+    """
+    Convert a string to snake_case.
+
+    This ensures column names are compatible with Jinja template syntax.
+
+    Examples:
+        >>> to_snake_case("MyColumn")
+        'my_column'
+        >>> to_snake_case("my_column")
+        'my_column'
+        >>> to_snake_case("My Column Name")
+        'my_column_name'
+        >>> to_snake_case("column-name")
+        'column_name'
+        >>> to_snake_case("Column_123")
+        'column_123'
+
+    Args:
+        text: String to convert
+
+    Returns:
+        snake_case version of the string
+    """
+    # Replace spaces and hyphens with underscores
+    text = text.replace(" ", "_").replace("-", "_")
+
+    # Insert underscore before uppercase letters (for camelCase/PascalCase)
+    text = re.sub(r"(?<!^)(?=[A-Z])", "_", text)
+
+    # Convert to lowercase
+    text = text.lower()
+
+    # Remove duplicate underscores
+    text = re.sub(r"_+", "_", text)
+
+    # Remove leading/trailing underscores
+    text = text.strip("_")
+
+    return text
+
+
+def from_snake_case_to_display_name(text: str) -> str:
+    """
+    Convert snake_case text back to Title Case for display.
+
+    Args:
+        text: snake_case text to convert
+
+    Returns:
+        Title Case version of the text
+
+    Example:
+        >>> from_snake_case("executive_summary")
+        'Executive Summary'
+        >>> from_snake_case("my_column_name")
+        'My Column Name'
+        >>> from_snake_case("api_design")
+        'Api Design'
+    """
+    # Split on underscores and capitalize each word
+    words = text.split("_")
+    return " ".join(word.capitalize() for word in words)
+
+
+def normalize_column_names(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Convert all DataFrame column names to snake_case.
+
+    This normalization ensures column names are compatible with Jinja template
+    syntax (e.g., {{ row.my_column }} works, but {{ row.My Column }} doesn't).
+
+    Examples:
+        >>> df = pd.DataFrame({"My Column": [1], "AnotherColumn": [2]})
+        >>> normalized = normalize_column_names(df)
+        >>> list(normalized.columns)
+        ['my_column', 'another_column']
+
+    Args:
+        df: Input DataFrame
+
+    Returns:
+        New DataFrame with normalized column names
+    """
+    ## TODO [UN-16142]: Normalization may lead to duplicate column names, we should handle this case
+    normalized_columns = {col: to_snake_case(col) for col in df.columns}
+    return df.rename(columns=normalized_columns)
+
+
+def limit_dataframe_rows(df: pd.DataFrame, max_rows: int) -> pd.DataFrame:
+    """
+    Limit DataFrame to first N rows.
+
+    Args:
+        df: DataFrame to limit
+        max_rows: Maximum number of rows
+
+    Returns:
+        DataFrame with at most max_rows rows
+    """
+    if len(df) <= max_rows:
+        return df.copy()
+    return df.head(max_rows).copy()
+
+
+def dataframe_to_dict_records(
+    df: pd.DataFrame, columns: list[str] | None = None
+) -> list[dict]:
+    """
+    Convert DataFrame to list of dict records.
+
+    Args:
+        df: DataFrame to convert
+        columns: Optional list of columns to include
+
+    Returns:
+        List of dict records
+    """
+    if columns:
+        df = df.loc[:, columns]
+
+    # Replace NaN with None for better serialization
+    return df.where(pd.notna(df), None).to_dict(orient="records")

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/__init__.py

@@ -0,0 +1,27 @@
+"""Generation handler module."""
+
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.config import (
+    GenerationHandlerConfig,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.exceptions import (
+    AggregationError,
+    BatchCreationError,
+    GenerationHandlerError,
+    LLMCallError,
+    PromptBuildError,
+    TokenLimitError,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.service import (
+    GenerationHandler,
+)
+
+__all__ = [
+    "GenerationHandler",
+    "GenerationHandlerConfig",
+    "GenerationHandlerError",
+    "BatchCreationError",
+    "PromptBuildError",
+    "LLMCallError",
+    "AggregationError",
+    "TokenLimitError",
+]

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/config.py

@@ -0,0 +1,56 @@
+"""Configuration for generation handler."""
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.prompts.config import (
+    GenerationHandlerPromptsConfig,
+)
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit._common.validators import LMI, get_LMI_default_field
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
+
+
+class GenerationHandlerConfig(BaseModel):
+    """Configuration for generation handler.
+
+    This configuration controls how groups are batched, how prompts are built,
+    and how the LLM is called for generating summaries.
+    """
+
+    model_config = get_configuration_dict()
+
+    language_model: LMI = get_LMI_default_field(
+        DEFAULT_GPT_4o,
+        description="The language model to use for generating summaries.",
+    )
+
+    common_instruction: str = Field(
+        default="You are a technical writer. Summarize the provided content concisely and clearly.",
+        description="Common instruction applied to all groups",
+    )
+
+    # TODO [UN-16142]: Add default instructions for each group
+    group_specific_instructions: dict[str, str] = Field(
+        default_factory=dict,
+        description=(
+            "Custom instructions per group. "
+            "Keys should be formatted as 'column:value' (e.g., 'section:Introduction')"
+        ),
+    )
+
+    max_tokens_per_batch: int = Field(
+        default=4000,
+        ge=100,
+        description="Maximum tokens per batch for LLM input (affects batching strategy)",
+    )
+
+    max_rows_per_batch: int = Field(
+        default=20,
+        ge=1,
+        description="Maximum rows per batch (secondary limit to tokens)",
+    )
+
+    prompts_config: GenerationHandlerPromptsConfig = Field(
+        default_factory=GenerationHandlerPromptsConfig,
+        description="Configuration for the prompts.",
+    )

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/exceptions.py

@@ -0,0 +1,79 @@
+"""Exceptions for generation handler operations."""
+
+
+class GenerationHandlerError(Exception):
+    """Base exception for all generation handler errors."""
+
+    pass
+
+
+class BatchCreationError(GenerationHandlerError):
+    """Raised when batch creation fails."""
+
+    def __init__(
+        self,
+        message: str,
+        group_key: str | None = None,
+        row_count: int | None = None,
+    ):
+        super().__init__(message)
+        self.group_key = group_key
+        self.row_count = row_count
+
+
+class PromptBuildError(GenerationHandlerError):
+    """Raised when prompt building fails."""
+
+    def __init__(
+        self,
+        message: str,
+        prompt_type: str | None = None,
+        context: dict | None = None,
+    ):
+        super().__init__(message)
+        self.prompt_type = prompt_type
+        self.context = context or {}
+
+
+class LLMCallError(GenerationHandlerError):
+    """Raised when LLM call fails."""
+
+    def __init__(
+        self,
+        message: str,
+        group_key: str | None = None,
+        batch_index: int | None = None,
+        error_details: str | None = None,
+    ):
+        super().__init__(message)
+        self.group_key = group_key
+        self.batch_index = batch_index
+        self.error_details = error_details
+
+
+class AggregationError(GenerationHandlerError):
+    """Raised when aggregating batch results fails."""
+
+    def __init__(
+        self,
+        message: str,
+        group_key: str | None = None,
+        batch_count: int | None = None,
+    ):
+        super().__init__(message)
+        self.group_key = group_key
+        self.batch_count = batch_count
+
+
+class TokenLimitError(GenerationHandlerError):
+    """Raised when token counting or limit validation fails."""
+
+    def __init__(
+        self,
+        message: str,
+        estimated_tokens: int | None = None,
+        max_tokens: int | None = None,
+    ):
+        super().__init__(message)
+        self.estimated_tokens = estimated_tokens
+        self.max_tokens = max_tokens

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/prompts/config.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.experimental.write_up_agent.utils import template_loader
+from unique_toolkit._common.pydantic.rjsf_tags import RJSFMetaTag
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+
+PARENT_DIR = Path(__file__).parent
+
+_SYSTEM_PROMPT_TEMPLATE = template_loader(PARENT_DIR, "system_prompt.j2")
+_USER_PROMPT_TEMPLATE = template_loader(PARENT_DIR, "user_prompt.j2")
+
+
+class GenerationHandlerPromptsConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    system_prompt_template: Annotated[
+        str,
+        RJSFMetaTag.StringWidget.textarea(
+            rows=len(_SYSTEM_PROMPT_TEMPLATE.split("\n"))
+        ),
+    ] = Field(
+        default=_SYSTEM_PROMPT_TEMPLATE,
+        description="The system prompt for the source selection.",
+    )
+    user_prompt_template: Annotated[
+        str,
+        RJSFMetaTag.StringWidget.textarea(rows=len(_USER_PROMPT_TEMPLATE.split("\n"))),
+    ] = Field(
+        default=_USER_PROMPT_TEMPLATE,
+        description="The user prompt for the source selection.",
+    )

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/prompts/system_prompt.j2

@@ -0,0 +1,15 @@
+{# System prompt with common instructions #}
+You are an expert technical writer tasked with creating clear, concise, and well-organized content summaries.
+
+## Your Task
+Generate a comprehensive summary for a specific section based on the provided content. The section header is already defined - focus on creating the body content.
+
+## Guidelines
+- Synthesize the provided information into a coherent narrative
+- Organize content logically, using sub-sections if it improves clarity
+- Maintain a professional and informative tone
+- Be concise while preserving key information
+- If section-specific instructions are provided, follow them carefully
+
+{{ common_instruction }}
+

unique_toolkit/_common/experimental/write_up_agent/services/generation_handler/prompts/user_prompt.j2

@@ -0,0 +1,21 @@
+{# User prompt with section context and optional previous summary #}
+## Section Context
+You are writing content for the section titled: **"{{ section_name }}"**
+
+**IMPORTANT**: The section header is already provided in the final report. Do NOT include the section title (# {{ section_name }}) in your response. Start directly with the content.
+
+{% if group_instruction %}
+## Section-Specific Instructions
+{{ group_instruction }}
+
+{% endif %}
+{% if previous_summary %}
+## Context from Previous Batch
+The following is a summary from the previous batch of content for this section. Build upon this context when generating your summary:
+
+{{ previous_summary }}
+
+{% endif %}
+## Content to Summarize
+{{ content }}
+