worklog_categorizer.py

import os
import logging
import google.generativeai as genai
from functools import lru_cache
from typing import List, Dict, Any, Optional, Tuple
import pandas as pd
from pathlib import Path
import time
from tqdm import tqdm
import re

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Initialize Gemini API
try:
    genai.configure(api_key=os.getenv("GEMINI_API_KEY","AIzaSyCunB1oTkxl7IINRMgQTVqIXKcFYw0Jqow"))
    model = genai.GenerativeModel("gemini-1.5-flash")
    logger.info("Gemini API initialized successfully")
except Exception as e:
    logger.error(f"Error initializing Gemini API: {e}")
    model = None

# Prompt for worklog categorization - modified for batch processing
BATCH_CATEGORIZATION_PROMPT = """
You are a technology skill categorizer. Analyze each worklog entry and assign a single technology category word that best represents the technical skill or technology involved.

Guidelines:
1. Respond with ONLY a single word (or hyphenated term if necessary) for each worklog
2. Focus on the core technology, framework, or skill
3. Be specific when the technology is clear (e.g., "React", "Python", "AWS")
4. Use broader categories when specific technology isn't clear (e.g., "Frontend", "Backend", "DevOps")
5. Prefer standard technology names over abbreviations
6. Don't include unnecessary adjectives or descriptions
7. Respond in a numbered list format matching the input worklogs

Examples:
Worklog 1: "fixing issue in next js application" → "NextJS"
Worklog 2: "Task issue fixing - next js application" → "NextJS"
Worklog 3: "Debugging Python script for data analysis" → "Python"
Worklog 4: "Creating responsive CSS layout" → "CSS"
Worklog 5: "Implementing REST API endpoints" → "Backend"

Here are the worklogs to categorize:
{worklogs}

For each worklog, respond with a numbered list containing only the category word for each entry:
1. [category for worklog 1]
2. [category for worklog 2]
...and so on
"""

def is_upskilling_issue(issue_text):
    """
    Check if an issue is related to upskilling using regex to match various formats.
    
    Args:
        issue_text: The issue text to check
        
    Returns:
        Boolean indicating if this is an upskilling issue
    """
    if not issue_text or not isinstance(issue_text, str):
        return False
        
    # Case insensitive search for "upskill" with potential variations
    # This will match: Upskilling, upskill, UPSKILLING, Up-skilling, Up skilling, etc.
    pattern = re.compile(r'up[-\s]?skill', re.IGNORECASE)
    return bool(pattern.search(issue_text))

def estimate_token_count(text: str) -> int:
    """
    Estimate token count for a given text string.
    
    This is an approximation based on GPT tokenization patterns:
    - Average of ~4 characters per token for English text
    - Spaces count as tokens
    - Special characters typically count as their own tokens
    
    Args:
        text: The text to estimate token count for
        
    Returns:
        Estimated token count
    """
    if not text:
        return 0
        
    # Count words (splitting by whitespace)
    words = len(text.split())
    
    # Count characters
    chars = len(text)
    
    # Count special tokens (punctuation, etc.)
    special_chars = len(re.findall(r'[^\w\s]', text))
    
    # Estimate based on a combination of factors
    # This formula is approximate and can be adjusted based on testing
    estimated_tokens = max(words, int(chars / 4) + special_chars)
    
    return estimated_tokens

def categorize_worklog_batch(worklogs: List[str]) -> List[str]:
    """
    Categorize multiple worklog entries with a single API call.
    
    Args:
        worklogs: List of worklog texts to categorize
        
    Returns:
        List of categories corresponding to each worklog
    """
    if not worklogs or model is None:
        return ["Unknown"] * len(worklogs)
    
    # Format worklogs as a numbered list for the prompt
    formatted_worklogs = "\n".join([f"{i+1}. {worklog}" for i, worklog in enumerate(worklogs)])
    prompt = BATCH_CATEGORIZATION_PROMPT.format(worklogs=formatted_worklogs)
    
    # Estimate token usage
    worklogs_token_count = sum(estimate_token_count(w) for w in worklogs)
    prompt_token_count = estimate_token_count(prompt)
    total_tokens = prompt_token_count
    
    logger.info(f"Sending batch with {len(worklogs)} worklogs (~{worklogs_token_count} worklog tokens, ~{total_tokens} total tokens)")
    
    try:
        response = model.generate_content(prompt)
        response_text = response.text.strip()
        
        logger.info(f"Response received: {response_text}")
        
        # Parse numbered response - looking for patterns like "1. Python", "2. JavaScript", etc.
        categories = []
        
        # First, try to match numbered lines (1. Category)
        number_pattern = re.compile(r'^\s*(\d+)\.\s*(.+?)$', re.MULTILINE)
        matches = number_pattern.findall(response_text)
        
        if matches:
            # Sort by the number to maintain order
            sorted_matches = sorted(matches, key=lambda x: int(x[0]))
            categories = [match[1].strip() for match in sorted_matches]
        else:
            # Fallback: try to split by lines
            lines = [line.strip() for line in response_text.split('\n') if line.strip()]
            categories = [line.split('.')[-1].strip() if '.' in line else line for line in lines]
        
        # Ensure we have the right number of categories
        if len(categories) != len(worklogs):
            logger.warning(f"Mismatch between number of worklogs ({len(worklogs)}) and categories ({len(categories)})")
            
            # Pad with "Unknown" if we have too few categories
            if len(categories) < len(worklogs):
                categories.extend(["Unknown"] * (len(worklogs) - len(categories)))
            # Truncate if we have too many categories
            else:
                categories = categories[:len(worklogs)]
        
        # Ensure each category is a single word
        for i, category in enumerate(categories):
            if len(category.split()) > 1 and "-" not in category:
                logger.warning(f"Response '{category}' contains multiple words, taking first word")
                categories[i] = category.split()[0]
                
        # Log the results for verification
        for i, (worklog, category) in enumerate(zip(worklogs, categories)):
            logger.info(f"Worklog {i+1}: '{worklog[:50]}{'...' if len(worklog) > 50 else ''}' → '{category}'")
            
        return categories
    except Exception as e:
        logger.error(f"Error categorizing worklog batch: {e}")
        return ["Unknown"] * len(worklogs)

def batch_process_worklogs(worklogs: List[str], batch_size: int = 10, 
                           pause_seconds: int = 5, show_progress: bool = True) -> List[str]:
    """
    Process multiple worklog entries in batches with pauses to avoid rate limits.
    Using 10 queries at a time with 5 seconds rest between batches.
    
    Args:
        worklogs: List of worklog texts to categorize
        batch_size: Number of worklogs to process in each batch (default: 10)
        pause_seconds: Seconds to pause between batches (default: 5)
        show_progress: Whether to show a progress bar
        
    Returns:
        List of categories corresponding to each worklog
    """
    results = []
    total_worklogs = len(worklogs)
    
    # Create batches
    batches = [worklogs[i:i + batch_size] for i in range(0, total_worklogs, batch_size)]
    
    # Process each batch with progress indication
    progress_bar = tqdm(total=total_worklogs, desc="Categorizing worklogs") if show_progress else None
    
    for i, batch in enumerate(batches):
        # Process current batch
        logger.info(f"Processing batch {i+1}/{len(batches)} with {len(batch)} worklogs")
        batch_results = categorize_worklog_batch(batch)
        results.extend(batch_results)
        
        # Update progress
        if progress_bar:
            progress_bar.update(len(batch))
        
        # Pause between batches (except after the last batch)
        if i < len(batches) - 1 and pause_seconds > 0:
            logger.info(f"Pausing for {pause_seconds}s before next batch. Processed {len(results)}/{total_worklogs} worklogs")
            if show_progress:
                for s in range(pause_seconds):
                    progress_bar.set_description(f"Waiting {pause_seconds-s}s before next batch")
                    time.sleep(1)
                progress_bar.set_description("Categorizing worklogs")
            else:
                time.sleep(pause_seconds)
    
    if progress_bar:
        progress_bar.close()
    
    logger.info(f"Completed processing {total_worklogs} worklogs")
    return results

def process_dataframe(df: pd.DataFrame, worklog_column: str = "Worklog", 
                      issue_column: str = "Issue", default_category: str = "N/A",
                      batch_size: int = 10, pause_seconds: int = 5,
                      show_progress: bool = True) -> pd.DataFrame:
    """
    Add a new column with technology categories to a dataframe.
    Only categorizes worklogs associated with upskilling issues.
    Processes 10 worklogs at a time with 5-second pauses between batches.
    
    Args:
        df: Pandas DataFrame containing worklog data
        worklog_column: Name of the column containing worklog text
        issue_column: Name of the column containing issue text
        default_category: Default value for non-upskilling worklogs
        batch_size: Number of worklogs to process in each batch (default: 10) 
        pause_seconds: Seconds to pause between batches (default: 5)
        show_progress: Whether to show a progress bar
        
    Returns:
        DataFrame with an additional 'TechCategory' column
    """
    # Initialize TechCategory column with default value
    df["TechCategory"] = default_category
    
    # Check if required columns exist
    if worklog_column not in df.columns:
        logger.error(f"Column '{worklog_column}' not found in DataFrame")
        return df
        
    if issue_column not in df.columns:
        logger.error(f"Column '{issue_column}' not found in DataFrame")
        return df
    
    # Filter for upskilling issues
    upskilling_mask = df[issue_column].apply(is_upskilling_issue)
    upskilling_rows = df[upskilling_mask].copy()
    
    logger.info(f"Found {len(upskilling_rows)} rows with upskilling issues out of {len(df)} total rows")
    
    if upskilling_rows.empty:
        logger.info("No upskilling issues found, returning dataframe with default category values")
        return df
    
    # Extract unique non-null worklog entries from upskilling issues
    unique_worklogs = upskilling_rows[worklog_column].dropna().unique().tolist()
    
    # Calculate total estimated tokens
    total_estimated_tokens = sum(estimate_token_count(worklog) for worklog in unique_worklogs)
    
    logger.info(f"Processing {len(unique_worklogs)} unique upskilling worklog entries with approximately {total_estimated_tokens} tokens")
    
    # Create a mapping of worklog text to category
    if unique_worklogs:
        categories = batch_process_worklogs(
            unique_worklogs, 
            batch_size=batch_size, 
            pause_seconds=pause_seconds,
            show_progress=show_progress
        )
        worklog_to_category = dict(zip(unique_worklogs, categories))
    else:
        worklog_to_category = {}
    
    # Apply categorization only to upskilling worklog entries
    df.loc[upskilling_mask, "TechCategory"] = df.loc[upskilling_mask, worklog_column].apply(
        lambda x: worklog_to_category.get(x, default_category) if pd.notna(x) else default_category
    )
    
    # Count the number of actually categorized entries
    categorized_count = len(df[df["TechCategory"] != default_category])
    logger.info(f"Successfully categorized {categorized_count} worklog entries")
    
    return df

def process_csv_file(
    csv_path: str, 
    worklog_column: str = "Worklog", 
    issue_column: str = "Issue",
    default_category: str = "N/A",
    output_path: Optional[str] = None,
    overwrite: bool = False,
    batch_size: int = 10,
    pause_seconds: int = 5
) -> str:
    """
    Process a CSV file to add technology categories based on worklog entries.
    Only categorizes worklogs associated with upskilling issues.
    Processes 10 worklogs at a time with 5-second pauses between batches.
    
    Args:
        csv_path: Path to the CSV file to process
        worklog_column: Name of the column containing worklog text
        issue_column: Name of the column containing issue text
        default_category: Default value for non-upskilling worklogs
        output_path: Path to save the processed file (if None, creates a new file with '_categorized' suffix)
        overwrite: If True, overwrite the original file
        batch_size: Number of worklogs to process in each batch (default: 10)
        pause_seconds: Seconds to pause between batches (default: 5)
        
    Returns:
        Path to the saved CSV file
    """
    try:
        # Check if file exists
        if not Path(csv_path).exists():
            logger.error(f"CSV file not found: {csv_path}")
            return ""
            
        # Read CSV
        logger.info(f"Reading CSV file: {csv_path}")
        df = pd.read_csv(csv_path)
        
        # Process dataframe
        processed_df = process_dataframe(
            df, 
            worklog_column=worklog_column,
            issue_column=issue_column,
            default_category=default_category,
            batch_size=batch_size,
            pause_seconds=pause_seconds
        )
        
        # Determine output path
        if overwrite:
            save_path = csv_path
        elif output_path:
            save_path = output_path
        else:
            # Create new filename with _categorized suffix
            path_obj = Path(csv_path)
            save_path = str(path_obj.with_stem(f"{path_obj.stem}_categorized"))
        
        # Save processed dataframe
        processed_df.to_csv(save_path, index=False)
        logger.info(f"Saved categorized CSV to: {save_path}")
        
        return save_path
        
    except Exception as e:
        logger.error(f"Error processing CSV file: {e}")
        return ""