generate_function_descriptions.py

#!/usr/bin/env python3
"""
Function Description Generation Pipeline

This script generates AI-powered function descriptions using OpenAI API based on
extracted function information, following best practices for code summarization.
"""

import json
import sys
import time
import argparse
from typing import Dict, List, Any, Optional
from pathlib import Path
from datetime import datetime
from tqdm import tqdm
import multiprocessing as mp

try:
    from openai import OpenAI
except ImportError:
    print("❌ OpenAI library not found. Please install with: pip install openai")
    sys.exit(1)

# --- Multiprocessing helpers (module-level, picklable) ---
_mp_config: Optional[Dict[str, Any]] = None
_mp_client: Optional[OpenAI] = None  # type: ignore[name-defined]


def _init_worker(config: Dict[str, Any]) -> None:
    """Initializer for worker processes to set up a per-process OpenAI client."""
    global _mp_config, _mp_client
    _mp_config = config
    try:
        _mp_client = OpenAI(
            api_key=config['api_key'],
            base_url=config['base_url'],
            timeout=config.get('timeout', 30),
        )
    except Exception:
        # Defer error handling to the worker call where we can return a structured result
        _mp_client = None


def _build_prompt(function_data: Dict[str, Any]) -> str:
    """Build the same prompt as the instance method, but as a free function for workers."""
    function_name = function_data.get('name', 'unknown')
    function_signature = function_data.get('signature', '')
    function_body = function_data.get('body', '')
    parameters = function_data.get('parameters', [])
    complexity_score = function_data.get('complexity_score', 0)
    calls_made = function_data.get('calls_made', [])
    control_structures = function_data.get('control_structures', [])

    param_info: List[str] = []
    for param in parameters:
        param_desc = f"- {param.get('name', '')}"
        if param.get('type_annotation'):
            param_desc += f": {param['type_annotation']}"
        if param.get('default'):
            param_desc += f" = {param['default']}"
        param_desc += f" ({param.get('kind', 'param')})"
        param_info.append(param_desc)

    param_text = "\n".join(param_info) if param_info else "No parameters"

    prompt = f"""You are an expert code documentation assistant. Generate a clear, concise, and professional function description for the following function.

**FUNCTION TO ANALYZE:**
```python/typescript/javascript/...
{function_body}
```

**FUNCTION METADATA:**
- Name: {function_name}
- Signature: {function_signature}
- Complexity Score: {complexity_score}
- Control Structures: {', '.join(control_structures) if control_structures else 'None'}
- Function Calls: {', '.join(calls_made[:10]) if calls_made else 'None'}

**PARAMETERS:**
{param_text}

**INSTRUCTIONS:**
Generate a professional function description that includes:

1. **Purpose**: What does this function do? (1-2 sentences)
2. **Behavior**: How does it work? Explain each step of the function in detail
3. **Parameters**: Brief description of each parameter (if any)
4. **Returns**: What does it return? (if applicable)
5. **Game Context**: If this is game-related, explain the game development context

**FORMATTING REQUIREMENTS:**
- Use clear, professional language
- Keep the description concise but informative

**OUTPUT FORMAT:**
Provide ONLY the function description text, no additional formatting or prefixes."""
    return prompt


def _worker_generate_description(function_data: Dict[str, Any]) -> Dict[str, Any]:
    """Worker-safe description generation.

    Returns a dict with: { 'description': Optional[str], 'success': bool, 'tokens_used': int }
    """
    global _mp_client, _mp_config
    result: Dict[str, Any] = {'description': None, 'success': False, 'tokens_used': 0}

    try:
        if _mp_client is None or _mp_config is None:
            raise RuntimeError("Worker client not initialized")

        prompt = _build_prompt(function_data)

        max_retries = int(_mp_config.get('max_retries', 3))
        retry_delay = float(_mp_config.get('retry_delay', 1.0))

        for attempt in range(max_retries):
            try:
                response = _mp_client.chat.completions.create(
                    model=_mp_config['quick_model'],
                    messages=[
                        {"role": "system", "content": "You are an expert code documentation assistant specialized in GUI Application and game development. Generate clear, professional function descriptions."},
                        {"role": "user", "content": prompt},
                    ],
                    max_tokens=int(_mp_config.get('max_tokens', 500)),
                    temperature=float(_mp_config.get('temperature', 0.3)),
                    top_p=0.9,
                )

                description = response.choices[0].message.content.strip()
                if len(description) < 10:
                    raise ValueError("Generated description too short")

                tokens_used = 0
                if hasattr(response, 'usage') and getattr(response, 'usage') is not None:
                    # Some clients expose total_tokens; guard in case it's missing
                    tokens_used = getattr(response.usage, 'total_tokens', 0) or 0

                result.update({'description': description, 'success': True, 'tokens_used': tokens_used})
                return result
            except Exception as e:
                if attempt < max_retries - 1:
                    time.sleep(retry_delay * (attempt + 1))
                    continue
                else:
                    print(e)
                    raise
    except Exception:
        # On failure, keep defaults (success False, description None)
        return result


class FunctionDescriptionGenerator:
    """AI-powered function description generator using OpenAI"""
    
    def __init__(self, config_file: str = "openai_config.json"):
        """Initialize the description generator"""
        self.config_file = config_file
        self.client = None
        self.config = {}
        self.load_openai_config()
        self.setup_openai_client()
        
        # Statistics tracking
        self.stats = {
            'total_functions': 0,
            'successful_descriptions': 0,
            'failed_descriptions': 0,
            'api_calls': 0,
            'total_tokens_used': 0,
            'processing_time': 0
        }
    
    def load_openai_config(self):
        """Load OpenAI configuration from file"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                self.config = json.load(f)
            
            required_keys = ['api_key', 'base_url', 'model']
            for key in required_keys:
                if key not in self.config:
                    raise ValueError(f"Missing required configuration key: {key}")
            
            print(f"✅ Loaded OpenAI configuration from {self.config_file}")
            print(f"   Model: {self.config['quick_model']}")
            print(f"   Base URL: {self.config['base_url']}")
            
        except FileNotFoundError:
            print(f"❌ Configuration file {self.config_file} not found!")
            print("Creating template configuration file...")
            self.create_config_template()
            sys.exit(1)
        except Exception as e:
            print(f"❌ Error loading configuration: {e}")
            sys.exit(1)
    
    def create_config_template(self):
        """Create a template configuration file"""
        template_config = {
            "api_key": "your-openai-api-key-here",
            "base_url": "https://api.openai.com/v1",
            "model": "gpt-4",
            "max_tokens": 500,
            "temperature": 0.3,
            "timeout": 30,
            "max_retries": 3,
            "retry_delay": 1.0
        }
        
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(template_config, f, indent=2)
            print(f"📄 Created template configuration file: {self.config_file}")
            print("Please edit the file with your OpenAI API credentials.")
        except Exception as e:
            print(f"❌ Error creating template: {e}")
    
    def setup_openai_client(self):
        """Setup OpenAI client with custom configuration"""
        try:
            self.client = OpenAI(
                api_key=self.config['api_key'],
                base_url=self.config['base_url'],
                timeout=self.config.get('timeout', 30)
            )
            
            # Test the connection with a simple call
            print("🔍 Testing OpenAI API connection...")
            test_response = self.client.chat.completions.create(
                model=self.config['quick_model'],
                messages=[{"role": "user", "content": "Test connection. Reply with 'OK'."}],
                max_tokens=5,
                temperature=0
            )
            
            if test_response.choices[0].message.content.strip().upper() == "OK":
                print("✅ OpenAI API connection successful!")
            else:
                print("⚠️  OpenAI API connection test completed but unexpected response.")
                
        except Exception as e:
            print(f"❌ Error setting up OpenAI client: {e}")
            print("Please check your API key, base URL, and network connection.")
            sys.exit(1)
    
    def create_function_description_prompt(self, function_data: Dict[str, Any]) -> str:
        """Create an optimized prompt for function description generation"""
        
        # Extract key information
        function_name = function_data.get('name', 'unknown')
        function_signature = function_data.get('signature', '')
        function_body = function_data.get('body', '')
        parameters = function_data.get('parameters', [])
        complexity_score = function_data.get('complexity_score', 0)
        calls_made = function_data.get('calls_made', [])
        control_structures = function_data.get('control_structures', [])
        
        # Build parameter description
        param_info = []
        for param in parameters:
            param_desc = f"- {param['name']}"
            if param.get('type_annotation'):
                param_desc += f": {param['type_annotation']}"
            if param.get('default'):
                param_desc += f" = {param['default']}"
            param_desc += f" ({param['kind']})"
            param_info.append(param_desc)
        
        param_text = "\n".join(param_info) if param_info else "No parameters"
        
        # Create comprehensive prompt following best practices
        prompt = f"""You are an expert code documentation assistant. Generate a clear, concise, and professional function description for the following function.

**FUNCTION TO ANALYZE:**
```python/typescript/javascript/...
{function_body}
```

**FUNCTION METADATA:**
- Name: {function_name}
- Signature: {function_signature}
- Complexity Score: {complexity_score}
- Control Structures: {', '.join(control_structures) if control_structures else 'None'}
- Function Calls: {', '.join(calls_made[:10]) if calls_made else 'None'}

**PARAMETERS:**
{param_text}

**INSTRUCTIONS:**
Generate a professional function description that includes:

1. **Purpose**: What does this function do? (1-2 sentences)
2. **Behavior**: How does it work? Explain each step of the function in detail
3. **Parameters**: Brief description of each parameter (if any)
4. **Returns**: What does it return? (if applicable)
5. **Game Context**: If this is game-related, explain the game development context

**FORMATTING REQUIREMENTS:**
- Use clear, professional language
- Keep the description concise but informative

**OUTPUT FORMAT:**
Provide ONLY the function description text, no additional formatting or prefixes."""

        return prompt
    
    def generate_description(self, function_data: Dict[str, Any]) -> Optional[str]:
        """Generate description for a single function"""
        try:
            # Create the prompt
            prompt = self.create_function_description_prompt(function_data)
            
            # Make API call with retry logic
            max_retries = self.config.get('max_retries', 3)
            retry_delay = self.config.get('retry_delay', 1.0)
            
            for attempt in range(max_retries):
                try:
                    response = self.client.chat.completions.create(
                        model=self.config['quick_model'],
                        messages=[
                            {
                                "role": "system", 
                                "content": "You are an expert code documentation assistant specialized in GUI Application and game development. Generate clear, professional function descriptions."
                            },
                            {"role": "user", "content": prompt}
                        ],
                        max_tokens=self.config.get('max_tokens', 500),
                        temperature=self.config.get('temperature', 0.3),
                        top_p=0.9,
                        # frequency_penalty=0.1,
                        # presence_penalty=0.1
                    )
                    
                    # Update statistics
                    self.stats['api_calls'] += 1
                    if hasattr(response, 'usage') and response.usage:
                        self.stats['total_tokens_used'] += response.usage.total_tokens if response.usage.total_tokens else 0
                    
                    # Extract and clean the description
                    description = response.choices[0].message.content.strip()
                    
                    # Basic validation
                    if len(description) < 10:
                        raise ValueError("Generated description too short")
                    
                    self.stats['successful_descriptions'] += 1
                    return description
                    
                except Exception as e:
                    if attempt < max_retries - 1:
                        print(f"   Retry {attempt + 1}/{max_retries} for function {function_data.get('name', 'unknown')}: {e}")
                        time.sleep(retry_delay * (attempt + 1))  # Exponential backoff
                        continue
                    else:
                        print(e)
                        raise e
            
        except Exception as e:
            print(f"❌ Error generating description for {function_data.get('name', 'unknown')}: {e}")
            self.stats['failed_descriptions'] += 1
            return None
    
    def process_functions_file(self, 
                             input_file: str = "extracted_functions.json",
                             output_file: str = "extracted_functions_with_comments.json",
                             max_functions: Optional[int] = None,
                             batch_size: int = 10) -> Dict[str, Any]:
        """Process all functions in the extracted functions file"""
        
        print(f"🚀 Starting function description generation...")
        print(f"📂 Input file: {input_file}")
        print(f"📄 Output file: {output_file}")
        
        start_time = time.time()
        
        # Load extracted functions data
        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                functions_data = json.load(f)
            print(f"✅ Loaded function data from {input_file}")
        except Exception as e:
            print(f"❌ Error loading functions file: {e}")
            return {}
        
        # Extract all functions from all files
        all_functions = []
        for file_info in functions_data.get('files', []):
            for func in file_info.get('functions', []):
                # Add file context
                func_with_context = func.copy()
                func_with_context['file_path'] = file_info.get('file_path', '')
                func_with_context['repository'] = file_info.get('repository', '')
                func_with_context['relative_path'] = file_info.get('relative_path', '')
                all_functions.append(func_with_context)
        
        if max_functions:
            all_functions = all_functions[:max_functions]
        
        self.stats['total_functions'] = len(all_functions)
        print(f"📊 Processing {len(all_functions)} functions...")
        
        # Process functions in batches
        processed_functions = []

        if batch_size <= 1:
            # Serial path for batch_size of 1 to reduce overhead
            for i in tqdm(range(0, len(all_functions), 1), desc="Generating descriptions"):
                func = all_functions[i]
                description = self.generate_description(func)

                func_with_desc = func.copy()
                func_with_desc['ai_generated_description'] = description
                func_with_desc['description_generated_at'] = datetime.now().isoformat()
                func_with_desc['description_model'] = self.config['quick_model']

                if description:
                    func_with_desc['description_length'] = len(description)
                    func_with_desc['description_word_count'] = len(description.split())
                else:
                    func_with_desc['description_length'] = 0
                    func_with_desc['description_word_count'] = 0

                processed_functions.append(func_with_desc)
        else:
            # Parallel path: one process per function within each batch
            ctx = mp.get_context("spawn")
            with ctx.Pool(processes=batch_size, initializer=_init_worker, initargs=(self.config,)) as pool:
                for i in tqdm(range(0, len(all_functions), batch_size), desc="Generating descriptions"):
                    batch = all_functions[i:i + batch_size]

                    # Run the batch concurrently
                    results: List[Dict[str, Any]] = pool.map(_worker_generate_description, batch)

                    # Aggregate results and stats
                    for func, out in zip(batch, results):
                        description = out.get('description') if out.get('success') else None

                        # Update stats similar to generate_description
                        if out.get('success'):
                            self.stats['successful_descriptions'] += 1
                            self.stats['api_calls'] += 1
                            self.stats['total_tokens_used'] += int(out.get('tokens_used', 0))
                        else:
                            self.stats['failed_descriptions'] += 1

                        # Attach to function record
                        func_with_desc = func.copy()
                        func_with_desc['ai_generated_description'] = description
                        func_with_desc['description_generated_at'] = datetime.now().isoformat()
                        func_with_desc['description_model'] = self.config['quick_model']

                        if description:
                            func_with_desc['description_length'] = len(description)
                            func_with_desc['description_word_count'] = len(description.split())
                        else:
                            func_with_desc['description_length'] = 0
                            func_with_desc['description_word_count'] = 0

                        processed_functions.append(func_with_desc)

                    # Delay between batches to reduce rate limiting
                    if i + batch_size < len(all_functions):
                        time.sleep(0.5)
        
        # Compile enhanced results
        enhanced_data = {
            'generation_metadata': {
                'timestamp': datetime.now().isoformat(),
                'input_file': input_file,
                'model_used': self.config['quick_model'],
                'base_url': self.config['base_url'],
                'total_functions_processed': len(processed_functions),
                'successful_descriptions': self.stats['successful_descriptions'],
                'failed_descriptions': self.stats['failed_descriptions'],
                'api_calls_made': self.stats['api_calls'],
                'total_tokens_used': self.stats['total_tokens_used'],
                'processing_time_seconds': time.time() - start_time,
                'success_rate': (self.stats['successful_descriptions'] / len(processed_functions) * 100) if processed_functions else 0
            },
            'original_metadata': functions_data.get('extraction_metadata', {}),
            'enhanced_functions': processed_functions
        }
        
        # Save enhanced data
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(enhanced_data, f, indent=2, ensure_ascii=False)
            print(f"✅ Enhanced function data saved to {output_file}")
        except Exception as e:
            print(f"❌ Error saving enhanced data: {e}")
        
        # Update processing time
        self.stats['processing_time'] = time.time() - start_time
        
        # Print summary
        self.print_generation_summary()
        
        return enhanced_data
    
    def print_generation_summary(self):
        """Print comprehensive generation summary"""
        print("\n" + "="*80)
        print("📊 FUNCTION DESCRIPTION GENERATION SUMMARY")
        print("="*80)
        
        print(f"\n📈 PROCESSING STATISTICS:")
        print(f"Total Functions: {self.stats['total_functions']}")
        print(f"Successful Descriptions: {self.stats['successful_descriptions']}")
        print(f"Failed Descriptions: {self.stats['failed_descriptions']}")
        print(f"Success Rate: {(self.stats['successful_descriptions'] / self.stats['total_functions'] * 100) if self.stats['total_functions'] > 0 else 0:.1f}%")
        
        print(f"\n🔧 API USAGE:")
        print(f"API Calls Made: {self.stats['api_calls']}")
        print(f"Total Tokens Used: {self.stats['total_tokens_used']:,}")
        print(f"Processing Time: {self.stats['processing_time']:.1f} seconds")
        
        if self.stats['successful_descriptions'] > 0:
            avg_tokens_per_function = self.stats['total_tokens_used'] / self.stats['successful_descriptions']
            print(f"Average Tokens per Function: {avg_tokens_per_function:.1f}")
        
        print(f"\n⚙️  CONFIGURATION:")
        print(f"Model: {self.config['quick_model']}")
        print(f"Base URL: {self.config['base_url']}")
        print(f"Max Tokens: {self.config.get('max_tokens', 500)}")
        print(f"Temperature: {self.config.get('temperature', 0.3)}")
        
        print("="*80)
    
    def analyze_generated_descriptions(self, output_file: str = "extracted_functions_with_comments.json"):
        """Analyze the quality and characteristics of generated descriptions"""
        try:
            with open(output_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            functions = data.get('enhanced_functions', [])
            
            if not functions:
                print("No functions found for analysis.")
                return
            
            # Analyze description characteristics
            descriptions = [f.get('ai_generated_description', '') for f in functions if f.get('ai_generated_description')]
            
            if not descriptions:
                print("No descriptions found for analysis.")
                return
            
            desc_lengths = [len(desc) for desc in descriptions]
            word_counts = [len(desc.split()) for desc in descriptions]
            
            print(f"\n📊 DESCRIPTION QUALITY ANALYSIS:")
            print(f"Functions with Descriptions: {len(descriptions)}")
            print(f"Average Description Length: {sum(desc_lengths) / len(desc_lengths):.1f} characters")
            print(f"Average Word Count: {sum(word_counts) / len(word_counts):.1f} words")
            print(f"Shortest Description: {min(desc_lengths)} characters")
            print(f"Longest Description: {max(desc_lengths)} characters")
            
            # Sample descriptions
            print(f"\n📝 SAMPLE DESCRIPTIONS:")
            for i, func in enumerate(functions[:3]):
                if func.get('ai_generated_description'):
                    print(f"\n{i+1}. Function: {func.get('name', 'unknown')}")
                    print(f"   Description: {func['ai_generated_description'][:200]}...")
            
        except Exception as e:
            print(f"❌ Error analyzing descriptions: {e}")


def main():
    """Main function"""
    parser = argparse.ArgumentParser(description="Generate AI-powered function descriptions")
    parser.add_argument("--input-file", default="Jsons/extracted_functions.json",
                       help="Input JSON file with extracted functions")
    parser.add_argument("--output-file", default="Jsons/extracted_functions_with_comments.json",
                       help="Output JSON file for enhanced functions")
    parser.add_argument("--config-file", default="openai_config.json",
                       help="OpenAI configuration file")
    parser.add_argument("--max-functions", type=int,
                       help="Maximum number of functions to process")
    parser.add_argument("--batch-size", type=int, default=16,
                       help="Number of functions to process in each batch")
    parser.add_argument("--analyze", action="store_true",
                       help="Analyze existing generated descriptions")
    parser.add_argument("--test-config", action="store_true",
                       help="Test OpenAI configuration and exit")
    
    args = parser.parse_args()
    
    # Create generator
    generator = FunctionDescriptionGenerator(config_file=args.config_file)
    
    if args.test_config:
        print("✅ OpenAI configuration test completed successfully!")
        return
    
    if args.analyze:
        generator.analyze_generated_descriptions(args.output_file)
        return
    
    # Generate descriptions
    print(f"🚀 Starting function description generation...")
    
    result = generator.process_functions_file(
        input_file=args.input_file,
        output_file=args.output_file,
        max_functions=args.max_functions,
        batch_size=args.batch_size
    )
    
    print(f"\n✅ Function description generation complete!")
    print(f"📄 Enhanced data saved to: {args.output_file}")

if __name__ == "__main__":
    main()