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

unique_toolkit/_common/experimental/write_up_agent/__init__.py

@@ -0,0 +1,22 @@
+"""
+Write-Up Agent: Template-driven DataFrame summarization and report generation.
+"""
+
+from unique_toolkit._common.experimental.write_up_agent.agent import WriteUpAgent
+from unique_toolkit._common.experimental.write_up_agent.config import (
+    WriteUpAgentConfig,
+)
+from unique_toolkit._common.experimental.write_up_agent.schemas import (
+    GroupData,
+    ProcessedGroup,
+)
+
+__all__ = [
+    # Main agent
+    "WriteUpAgent",
+    # Configuration
+    "WriteUpAgentConfig",
+    # Data schemas
+    "GroupData",
+    "ProcessedGroup",
+]

unique_toolkit/_common/experimental/write_up_agent/agent.py

@@ -0,0 +1,170 @@
+"""
+Write-Up Agent - Main pipeline orchestrator.
+"""
+
+import logging
+
+import pandas as pd
+
+from unique_toolkit._common.experimental.write_up_agent.config import (
+    WriteUpAgentConfig,
+)
+from unique_toolkit._common.experimental.write_up_agent.schemas import GroupData
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler import (
+    DataFrameHandler,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler.exceptions import (
+    DataFrameGroupingError,
+    DataFrameHandlerError,
+    DataFrameProcessingError,
+    DataFrameValidationError,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler import (
+    GenerationHandler,
+    GenerationHandlerError,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.template_handler import (
+    TemplateHandler,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.template_handler.exceptions import (
+    ColumnExtractionError,
+    TemplateHandlerError,
+    TemplateParsingError,
+    TemplateRenderingError,
+    TemplateStructureError,
+)
+from unique_toolkit.language_model.service import LanguageModelService
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class WriteUpAgent:
+    """
+    Main pipeline orchestrator for DataFrame summarization.
+
+    Orchestrates the complete pipeline:
+    1. Extract template info (grouping + columns)
+    2. Validate DataFrame
+    3. Create groups
+    4. Render each group
+    5. Process with LLM
+    6. Return results
+    """
+
+    def __init__(self, config: WriteUpAgentConfig):
+        """
+        Initialize WriteUpAgent.
+
+        Args:
+            config: Configuration with template and settings
+        """
+        self._config = config
+        self._template_handler = TemplateHandler(config.template)
+        self._dataframe_handler = DataFrameHandler()
+
+        # Create generation handler with injected renderer
+        def renderer(group_data: GroupData) -> str:
+            return self._template_handler.render_group(group_data)
+
+        # TODO [UN-16142]: Find a better way to inject the renderer
+        self._generation_handler = GenerationHandler(
+            self._config.generation_handler_config, renderer
+        )
+
+    def process(self, df: pd.DataFrame, llm_service: LanguageModelService) -> str:
+        """
+        Execute complete pipeline and generate final report.
+
+        Args:
+            df: pandas DataFrame to process
+            llm_service: LanguageModelService to use for generating summaries
+
+        Returns:
+            Final markdown report as a single string with all groups processed
+
+        Raises:
+            Various handler exceptions if processing fails
+
+        Example:
+            >>> config = WriteUpAgentConfig(template="...", max_rows_per_group=10)
+            >>> agent = WriteUpAgent(config)
+            >>> report = agent.process(df)
+            >>> print(report)
+        """
+        # TODO [UN-16142]: Add error handling for each step separately
+        try:
+            # Step 1: Extract template structure
+            _LOGGER.info("Extracting template structure...")
+            grouping_column = self._template_handler.get_grouping_column()
+            selected_columns = self._template_handler.get_selected_columns()
+            _LOGGER.info(f"Detected grouping column: {grouping_column}")
+            _LOGGER.info(f"Detected data columns: {selected_columns}")
+
+            # Step 2: Validate DataFrame
+            _LOGGER.info("Validating DataFrame columns...")
+            self._dataframe_handler.validate_columns(
+                df, grouping_column, selected_columns
+            )
+
+            # Step 3: Create groups
+            _LOGGER.info("Creating groups from DataFrame...")
+            groups = self._dataframe_handler.create_groups(
+                df, grouping_column, selected_columns
+            )
+            _LOGGER.info(f"Created {len(groups)} groups")
+
+            # Step 4: Process groups with GenerationHandler
+            _LOGGER.info("Processing groups with GenerationHandler...")
+            processed_groups = self._generation_handler.process_groups(
+                groups, grouping_column, llm_service
+            )
+            _LOGGER.info(f"Generation complete for {len(processed_groups)} groups")
+
+            # Step 5: Render final report with LLM responses
+            _LOGGER.info("Rendering final report...")
+
+            final_report = self._template_handler.render_all_groups(processed_groups)
+
+            _LOGGER.info(f"Report generated ({len(final_report)} characters)")
+
+            return final_report
+
+        except TemplateParsingError as e:
+            _LOGGER.error(f"Template parsing failed: {e}")
+            raise
+
+        except TemplateStructureError as e:
+            _LOGGER.error(f"Template structure invalid: {e}")
+            raise
+
+        except ColumnExtractionError as e:
+            _LOGGER.error(f"Column extraction failed: {e}")
+            raise
+
+        except DataFrameValidationError as e:
+            _LOGGER.error(f"DataFrame validation failed: {e}")
+            raise
+
+        except DataFrameGroupingError as e:
+            _LOGGER.error(f"DataFrame grouping failed: {e}")
+            raise
+
+        except DataFrameProcessingError as e:
+            _LOGGER.error(f"DataFrame processing failed: {e}")
+            raise
+
+        except GenerationHandlerError as e:
+            _LOGGER.error(f"Generation failed: {e}")
+            raise
+
+        except TemplateRenderingError as e:
+            _LOGGER.error(f"Final rendering failed: {e}")
+            raise
+
+        except (TemplateHandlerError, DataFrameHandlerError) as e:
+            _LOGGER.error(f"Handler error: {e}")
+            raise
+
+        except Exception as e:
+            _LOGGER.error(f"Unexpected error: {e}", exc_info=True)
+            raise

unique_toolkit/_common/experimental/write_up_agent/config.py

@@ -0,0 +1,42 @@
+from pydantic import BaseModel, Field, field_validator
+
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.config import (
+    GenerationHandlerConfig,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.template_handler import (
+    default_jinja_template_loader,
+)
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+
+
+class WriteUpAgentConfig(BaseModel):
+    """Configuration for the Write-Up Agent that generates summaries from DataFrame data.
+
+    The agent uses a Jinja template as the single source of truth for data structure.
+    The template is parsed to automatically detect grouping columns and data references.
+    """
+
+    model_config = get_configuration_dict()
+
+    # Template Configuration (single source of truth)
+    template: str = Field(
+        default_factory=default_jinja_template_loader,
+        description=(
+            "Jinja template string that defines the structure of the summary. "
+            "The template is parsed to automatically detect grouping columns and data references. "
+            "If not provided, loads the default Q&A template. "
+            "Example: '{% for g in groups %}## {{ g.section }}{% endfor %}'"
+        ),
+    )
+
+    generation_handler_config: GenerationHandlerConfig = Field(
+        default_factory=GenerationHandlerConfig,
+        description="Configuration for the generation handler.",
+    )
+
+    @field_validator("template")
+    @classmethod
+    def validate_template_not_empty(cls, v: str) -> str:
+        if not v.strip():
+            raise ValueError("Template must not be empty")
+        return v

unique_toolkit/_common/experimental/write_up_agent/examples/data.csv

@@ -0,0 +1,13 @@
+section,question,answer
+Introduction,What is the Write-Up Agent?,The Write-Up Agent is a tool that automatically generates summaries from structured DataFrame data using LLM technology.
+Introduction,Who should use this tool?,Data scientists and analysts who need to convert tabular data into readable reports.
+Introduction,What are the key benefits?,Automated report generation with customizable templates and intelligent summarization.
+Methods,How does the agent process data?,The agent groups data by sections and generates summaries for each group using LLM calls with adaptive batching.
+Methods,What is adaptive batching?,A technique that splits large groups into smaller batches to fit within token limits while maintaining context.
+Methods,Can I customize the output format?,Yes! You can provide custom Jinja templates to control the structure and style of the generated report.
+Results,What level of accuracy can I expect?,The agent leverages state-of-the-art LLMs to produce accurate and contextually relevant summaries.
+Results,How fast is the processing?,Processing speed depends on the LLM provider and the size of your dataset. Batching helps optimize performance.
+Results,Can it handle large datasets?,Yes! The agent automatically batches large groups to handle datasets of any size efficiently.
+Conclusion,Is this production-ready?,Yes! The agent includes robust error handling and type-safe operations using Pydantic schemas.
+Conclusion,Where can I find more examples?,Check the examples directory for additional use cases and custom template examples.
+

unique_toolkit/_common/experimental/write_up_agent/examples/example_usage.py

@@ -0,0 +1,78 @@
+"""
+Example: Using the Write-Up Agent to generate summaries from DataFrame data.
+"""
+
+# TODO [UN-16142]: Add example usage in tutorial instead of here
+
+import logging
+from pathlib import Path
+
+import pandas as pd
+
+from unique_toolkit._common.experimental.write_up_agent import (
+    WriteUpAgent,
+    WriteUpAgentConfig,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.generation_handler.config import (
+    GenerationHandlerConfig,
+)
+from unique_toolkit.app.unique_settings import UniqueSettings
+from unique_toolkit.language_model.service import LanguageModelService
+
+logging.basicConfig(level=logging.DEBUG)
+
+# Setup paths
+current_dir = Path(__file__).parent
+env_path = current_dir / "unique.env"
+data_path = current_dir / "data.csv"
+
+# Initialize SDK with your API keys
+_SETTINGS = UniqueSettings.from_env(env_file=env_path)
+_SETTINGS.init_sdk()
+
+# Configure the Write-Up Agent
+# Using default configuration which expects: section, question, answer columns
+write_up_agent_config = WriteUpAgentConfig(
+    generation_handler_config=GenerationHandlerConfig(
+        # Optional: Customize generation settings
+        # max_rows_per_batch=20,  # Max rows per batch (default: 20)
+        # max_tokens_per_batch=4000,  # Max tokens per batch (default: 4000)
+        # common_instruction="You are a technical writer...",  # Custom system prompt
+        # group_specific_instructions={
+        #     # IMPORTANT: Both column and value must be in snake_case
+        #     # DataFrame: Section="Introduction" → Key: "section:introduction"
+        #     "section:introduction": "Be welcoming and engaging",
+        #     "section:methods": "Be precise and technical"
+        # }
+    )
+)
+
+# Initialize the agent with LLM service
+write_up_agent = WriteUpAgent(
+    config=write_up_agent_config,
+)
+
+# Load your DataFrame
+# IMPORTANT: DataFrame must have columns: section, question, answer (otherwise adapt the template)
+df = pd.read_csv(data_path)
+
+print(f"Processing {len(df)} rows across {df['section'].nunique()} sections...")
+print(f"Columns in DataFrame: {list(df.columns)}")
+print()
+
+llm_service = LanguageModelService.from_settings(_SETTINGS)
+
+# Generate the report
+report = write_up_agent.process(df, llm_service=llm_service)
+
+# Display the result
+print("=" * 80)
+print("GENERATED REPORT")
+print("=" * 80)
+print(report)
+
+# Optional: Save to file
+output_path = current_dir / "report.md"
+output_path.write_text(report)
+print()
+print(f"Report saved to: {output_path}")

unique_toolkit/_common/experimental/write_up_agent/schemas.py

@@ -0,0 +1,36 @@
+"""Data schemas for the Write-Up Agent."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class GroupData(BaseModel):
+    """
+    Represents a group of rows from a DataFrame.
+
+    This is the core data structure passed between handlers in the pipeline.
+    """
+
+    group_key: str = Field(
+        ...,
+        description="The value of the grouping column for this group (e.g., 'Introduction', 'Methods')",
+    )
+
+    rows: list[dict[str, Any]] = Field(
+        ...,
+        description="List of row dictionaries containing the selected columns for this group",
+    )
+
+
+class ProcessedGroup(GroupData):
+    """
+    Represents a group after LLM processing.
+
+    Extends GroupData with the LLM-generated response.
+    """
+
+    llm_response: str = Field(
+        ...,
+        description="The LLM-generated summary/output for this group",
+    )

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

@@ -0,0 +1,13 @@
+"""Services for the write-up agent pipeline."""
+
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler import (
+    DataFrameHandler,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.template_handler import (
+    TemplateHandler,
+)
+
+__all__ = [
+    "DataFrameHandler",
+    "TemplateHandler",
+]

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

@@ -0,0 +1,19 @@
+"""DataFrame handler module."""
+
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler.exceptions import (
+    DataFrameGroupingError,
+    DataFrameHandlerError,
+    DataFrameProcessingError,
+    DataFrameValidationError,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler.service import (
+    DataFrameHandler,
+)
+
+__all__ = [
+    "DataFrameHandler",
+    "DataFrameHandlerError",
+    "DataFrameValidationError",
+    "DataFrameGroupingError",
+    "DataFrameProcessingError",
+]

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

@@ -0,0 +1,29 @@
+"""Exceptions for DataFrame handler operations."""
+
+
+class DataFrameHandlerError(Exception):
+    """Base exception for all DataFrame handler errors."""
+
+    pass
+
+
+class DataFrameValidationError(DataFrameHandlerError):
+    """Raised when DataFrame validation fails (e.g., missing columns)."""
+
+    def __init__(self, message: str, missing_columns: list[str] | None = None):
+        super().__init__(message)
+        self.missing_columns = missing_columns or []
+
+
+class DataFrameGroupingError(DataFrameHandlerError):
+    """Raised when DataFrame grouping operation fails."""
+
+    def __init__(self, message: str, grouping_column: str | None = None):
+        super().__init__(message)
+        self.grouping_column = grouping_column
+
+
+class DataFrameProcessingError(DataFrameHandlerError):
+    """Raised when general DataFrame processing fails."""
+
+    pass

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

@@ -0,0 +1,150 @@
+"""DataFrame handler service."""
+
+import pandas as pd
+
+from unique_toolkit._common.experimental.write_up_agent.schemas import GroupData
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler.exceptions import (
+    DataFrameGroupingError,
+    DataFrameProcessingError,
+    DataFrameValidationError,
+)
+from unique_toolkit._common.experimental.write_up_agent.services.dataframe_handler.utils import (
+    dataframe_to_dict_records,
+    normalize_column_names,
+    to_snake_case,
+)
+
+
+class DataFrameHandler:
+    """
+    Handles all DataFrame operations.
+
+    This handler automatically converts all column names to snake_case to ensure
+    compatibility with Jinja template syntax. For example:
+    - "My Column" becomes "my_column"
+    - "UserName" becomes "user_name"
+    - "column-name" becomes "column_name"
+
+    This normalization happens automatically during validation and grouping operations.
+
+    Responsibilities:
+    - Normalize column names to snake_case
+    - Validate DataFrame has required columns
+    - Create groups from DataFrame
+    """
+
+    def validate_columns(
+        self, df: pd.DataFrame, grouping_column: str, selected_columns: list[str]
+    ) -> None:
+        """
+        Validate DataFrame has required columns.
+
+        NOTE: Column names are automatically converted to snake_case before validation.
+        Ensure your template uses snake_case column references (e.g., {{ row.my_column }}).
+
+        Args:
+            df: pandas DataFrame to validate
+            grouping_column: Column to group by (should be in snake_case)
+            selected_columns: Columns that should exist (should be in snake_case)
+
+        Raises:
+            DataFrameValidationError: If columns are missing after normalization
+
+        Example:
+            >>> df = pd.DataFrame({"My Section": [1], "My Question": [2]})
+            >>> handler.validate_columns(df, "my_section", ["my_question"])
+            # Validation passes because "My Section" -> "my_section"
+        """
+        # Normalize DataFrame columns to snake_case
+        normalized_df = normalize_column_names(df)
+
+        required_columns = {grouping_column} | set(selected_columns)
+        missing_columns = required_columns - set(normalized_df.columns)
+
+        if missing_columns:
+            raise DataFrameValidationError(
+                f"DataFrame missing required columns after snake_case normalization: {sorted(missing_columns)}. "
+                f"Available columns: {sorted(normalized_df.columns)}",
+                missing_columns=sorted(missing_columns),
+            )
+
+    def create_groups(
+        self, df: pd.DataFrame, grouping_column: str, selected_columns: list[str]
+    ) -> list[GroupData]:
+        """
+        Create groups from DataFrame.
+
+        NOTE: Column names are automatically converted to snake_case.
+        Group values (group_key) are kept as-is from the DataFrame (NOT normalized).
+
+        The returned GroupData instances will have:
+        - snake_case column names in their rows
+        - Original group_key values from DataFrame
+
+        IMPORTANT: Groups are returned in the order of their first appearance in the DataFrame,
+        NOT sorted alphabetically. This preserves the logical flow of your data.
+
+        Args:
+            df: pandas DataFrame to group
+            grouping_column: Column to group by (should be in snake_case)
+            selected_columns: Columns to include in rows (should be in snake_case)
+
+        Returns:
+            List of GroupData instances in order of first appearance, each containing
+            group_key (in snake_case) and rows with snake_case columns
+
+        Raises:
+            DataFrameGroupingError: If grouping fails
+            DataFrameProcessingError: If data processing fails
+
+        Example:
+            >>> df = pd.DataFrame({
+            ...     "My Section": ["Intro", "Methods", "Results", "Intro"],
+            ...     "My Question": ["Q1", "Q2", "Q3", "Q4"],
+            ... })
+            >>> groups = handler.create_groups(df, "my_section", ["my_question"])
+            >>> [g.group_key for g in groups]
+            ['intro', 'methods', 'results']  # Values normalized to snake_case, order preserved
+        """
+        # Normalize column names to snake_case
+        normalized_df = normalize_column_names(df)
+
+        if grouping_column not in normalized_df.columns:
+            raise DataFrameGroupingError(
+                f"Grouping column '{grouping_column}' not found in normalized DataFrame. "
+                f"Available columns: {sorted(normalized_df.columns)}",
+                grouping_column=grouping_column,
+            )
+
+        try:
+            # Use sort=False to preserve the order of first appearance in the DataFrame
+            grouped = normalized_df.groupby(grouping_column, sort=False)
+        except Exception as e:
+            raise DataFrameGroupingError(
+                f"Failed to group DataFrame by '{grouping_column}': {e}",
+                grouping_column=grouping_column,
+            ) from e
+
+        results = []
+
+        try:
+            for group_key, group_df in grouped:
+                # Filter columns if specified
+                if selected_columns:
+                    cols_to_use = [c for c in selected_columns if c in group_df.columns]
+                    limited_df = group_df.loc[:, cols_to_use]
+                else:
+                    limited_df = group_df
+
+                # Convert to dict records
+                rows = dataframe_to_dict_records(limited_df)
+
+                # Normalize group_key value to snake_case for consistency with template syntax
+                normalized_group_key = to_snake_case(str(group_key))
+
+                # Create GroupData instance with proper typing
+                results.append(GroupData(group_key=normalized_group_key, rows=rows))
+        except Exception as e:
+            raise DataFrameProcessingError(f"Error processing grouped data: {e}") from e
+
+        return results