src/utils/formatters.py

"""
Utility functions for formatting and extracting data.
"""


def gemini_text(res) -> str:
    """
    Extract text from Gemini API response.

    Args:
        res: Gemini API response object

    Returns:
        Concatenated text from all text parts
    """
    return "".join(part["text"] for part in res.content if part.get("type") == "text")


def extract_code(response: str) -> str:
    """
    Extract code from markdown code blocks.
    Handles ```python and ``` formats.

    Args:
        response: Text response potentially containing code blocks

    Returns:
        Extracted code string
    """
    # Try to find python code block first
    if "```python" in response:
        parts = response.split("```python", 1)
        if len(parts) > 1:
            code_part = parts[1].split("```", 1)
            if len(code_part) > 0:
                return code_part[0].strip()

    # Try generic code block
    elif "```" in response:
        parts = response.split("```", 1)
        if len(parts) > 1:
            code_part = parts[1].split("```", 1)
            if len(code_part) > 0:
                # Remove language identifier if present
                code = code_part[0].strip()
                # Remove first line if it's a language identifier
                lines = code.split("\n")
                if lines and lines[0].strip() in ["python", "py", "python3"]:
                    return "\n".join(lines[1:]).strip()
                return code

    # If no code blocks found, return as is
    return response.strip()


def format_data_descriptions(descriptions: dict) -> str:
    """
    Format data descriptions dictionary into readable string.

    Args:
        descriptions: Dict mapping filename to description

    Returns:
        Formatted string with file descriptions
    """
    if not descriptions:
        return "No data files analyzed yet."

    formatted_parts = []
    for filename, description in descriptions.items():
        formatted_parts.append(f"## File: {filename}\n{description}\n")

    return "\n".join(formatted_parts)


def format_plan(plan: list) -> str:
    """
    Format plan steps into readable string.

    Args:
        plan: List of PlanStep dictionaries

    Returns:
        Formatted plan string
    """
    if not plan:
        return "No plan steps yet."

    return "\n".join([f"{i + 1}. {step['description']}" for i, step in enumerate(plan)])