data2/instruction_generation/extract_repo_functions.py

#!/usr/bin/env python3
"""
Multi-language Function Parsing Script
Scans code files in each repository, uses Qwen to parse dependencies and functions,
generates functions_with_context.csv
"""

import os
import sys
import json
import csv
import asyncio
import argparse
import hashlib
from pathlib import Path
from typing import List, Dict, Optional
from tqdm import tqdm
from dotenv import load_dotenv

# Load .env file (before importing other modules)
env_file = Path(__file__).parent / ".env"
if env_file.exists():
    load_dotenv(env_file)
elif (Path(__file__).parent.parent / ".env").exists():
    # If not in current directory, try loading from project root
    load_dotenv(Path(__file__).parent.parent / ".env")

# Add current directory to path (for importing schemas)
sys.path.insert(0, str(Path(__file__).parent))
# Add domain_code/src to path for reusing util functions
sys.path.insert(0, str(Path(__file__).parent.parent / "domain_code" / "src"))
from util import call_llm, init_logger, logger, CODE_EXTENSIONS
from schemas import FileParseResult

# Exclude markdown files (should not be parsed as code files)
PARSEABLE_CODE_EXTENSIONS = CODE_EXTENSIONS - {".md", ".markdown"}


# Default output filename (written back to repository directory)
CSV_FILENAME = "functions_with_context.csv"
SUMMARY_FILENAME = "README_SUMMARY.md"


def detect_language(file_path: Path) -> str:
    """
    Detect programming language based on file extension
    
    Args:
        file_path: File path
        
    Returns:
        Programming language name (e.g., python, cpp, java)
    """
    ext_map = {
        ".py": "python",
        ".ipynb": "python",
        ".java": "java",
        ".c": "c",
        ".cpp": "cpp",
        ".cc": "cpp",
        ".cxx": "cpp",
        ".h": "cpp",
        ".hpp": "cpp",
        ".hh": "cpp",
        ".F": "fortran",
        ".f90": "fortran",
        ".f": "fortran",
        ".f95": "fortran",
        ".r": "r",
        ".R": "r",
        ".m": "matlab",
        ".sh": "shell",
        ".bash": "shell",
        ".rs": "rust",
        ".go": "go",
        ".jl": "julia",
    }
    
    ext = file_path.suffix.lower()
    return ext_map.get(ext, ext.lstrip(".") if ext else "unknown")


def read_readme_summary(repo_dir: Path) -> Optional[str]:
    """
    Read README_SUMMARY.md content as project context
    
    Args:
        repo_dir: Repository root directory
        
    Returns:
        README summary text or None
    """
    summary_file = repo_dir / SUMMARY_FILENAME
    if not summary_file.exists():
        return None
    
    try:
        with open(summary_file, "r", encoding="utf-8", errors="ignore") as f:
            return f.read().strip()
    except Exception as e:
        logger.warning(f"Unable to read README summary file {summary_file}: {e}")
        return None


def find_code_files(repo_dir: Path, max_file_chars: int = 200000) -> List[Path]:
    """
    Find all code files in the repository (files covered by CODE_EXTENSIONS)
    
    Args:
        repo_dir: Repository root directory
        max_file_chars: Maximum file size (chars), files exceeding this are skipped
        
    Returns:
        List of code file paths
    """
    code_files = []
    
    for root, dirs, files in os.walk(repo_dir):
        # Skip hidden directories and common non-source directories
        dirs[:] = [d for d in dirs if not d.startswith(".") and d not in ["__pycache__", "node_modules", ".git"]]
        
        for file in files:
            file_path = Path(root) / file
            # Use PARSEABLE_CODE_EXTENSIONS to exclude markdown files
            if file_path.suffix.lower() in PARSEABLE_CODE_EXTENSIONS:
                # Check file size
                try:
                    size = file_path.stat().st_size
                    # Simple estimation: assume average 1 byte per char (UTF-8 encoding)
                    if size <= max_file_chars:
                        code_files.append(file_path)
                    else:
                        logger.debug(f"Skipping large file: {file_path} ({size} bytes)")
                except Exception as e:
                    logger.warning(f"Unable to get file size {file_path}: {e}")
    
    return sorted(code_files)


def read_code_file(file_path: Path) -> Optional[str]:
    """
    Read code file content
    
    Args:
        file_path: File path
        
    Returns:
        File content or None
    """
    try:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()
    except Exception as e:
        logger.warning(f"Unable to read file {file_path}: {e}")
        return None


def compute_file_hash(file_path: Path, content: str) -> str:
    """
    Compute SHA1 hash of file
    
    Args:
        file_path: File path
        content: File content
        
    Returns:
        SHA1 hash (hex string)
    """
    return hashlib.sha1(content.encode("utf-8")).hexdigest()


def compute_function_hash(repo_name: str, path: str, start_line: int, end_line: int, body: str) -> str:
    """
    Compute function hash (for deduplication)
    
    Args:
        repo_name: Repository name
        path: Relative file path
        start_line: Function start line number
        end_line: Function end line number
        body: Function body
        
    Returns:
        SHA1 hash (hex string)
    """
    key = f"{repo_name}:{path}:{start_line}:{end_line}:{body}"
    return hashlib.sha1(key.encode("utf-8")).hexdigest()


async def parse_code_file(
    file_path: Path,
    repo_dir: Path,
    project_context: str,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
) -> Optional[Dict]:
    """
    Use LLM to parse code file, extract dependencies and function information
    
    Args:
        file_path: Code file path
        repo_dir: Repository root directory
        project_context: Project context (README summary)
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        
    Returns:
        Parse result (dict) or None
    """
    # Read code content
    code_content = read_code_file(file_path)
    if not code_content:
        return None
    
    # Detect language
    language = detect_language(file_path)
    
    # Build relative path
    rel_path = str(file_path.relative_to(repo_dir))
    
    # Read prompt template
    prompt_template_path = Path(__file__).parent / "prompts" / "function_extract.txt"
    try:
        with open(prompt_template_path, "r", encoding="utf-8") as f:
            prompt_template = f.read()
    except Exception as e:
        logger.error(f"Unable to read prompt template: {e}")
        return None
    
    # Build prompt
    prompt = prompt_template.format(
        project_context=project_context or "(No project context)",
        file_path=rel_path,
        language=language,
        code_content=code_content,
    )
    
    # Call LLM
    messages = [{"role": "user", "content": prompt}]
    
    try:
        result = await call_llm(
            messages=messages,
            model=model,
            base_url=base_url,
            api_key=api_key,
            pydantic_object=FileParseResult,
            log_file=log_file,
        )
        
        if result is None:
            logger.warning(f"LLM call returned None, skipping file: {rel_path}")
            return None
        
        # If result is a string, try to parse JSON
        if isinstance(result, str):
            try:
                result = json.loads(result)
            except json.JSONDecodeError:
                logger.warning(f"Unable to parse JSON from LLM response: {result[:200]}")
                return None
        
        # Add file path (ensure consistency)
        if isinstance(result, dict):
            result["file_path"] = rel_path
            result["language"] = language
        
        return result
    except Exception as e:
        logger.error(f"LLM call failed (file: {rel_path}): {e}")
        return None


def extract_repo_name(repo_dir: Path) -> str:
    """
    Extract repository name from directory name (owner___repo -> owner/repo)
    
    Args:
        repo_dir: Repository root directory
        
    Returns:
        Repository name (owner/repo format)
    """
    dir_name = repo_dir.name
    return dir_name.replace("___", "/")


async def process_single_repo(
    repo_dir: Path,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
    max_file_chars: int = 200000,
    max_concurrency: int = 8,
    overwrite: bool = False,
) -> Dict[str, any]:
    """
    Process function parsing for a single repository
    
    Args:
        repo_dir: Repository root directory
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        max_file_chars: Maximum file size (chars)
        max_concurrency: Maximum concurrency
        overwrite: Whether to overwrite existing CSV file
        
    Returns:
        Processing result dictionary
    """
    repo_name = repo_dir.name
    csv_file = repo_dir / CSV_FILENAME
    
    # Check if CSV file already exists
    if csv_file.exists() and not overwrite:
        return {
            "repo": repo_name,
            "status": "skipped",
            "reason": "CSV file already exists",
        }
    
    # Read README summary as project context
    project_context = read_readme_summary(repo_dir)
    if not project_context:
        logger.warning(f"Repository {repo_name} has no README_SUMMARY.md, skipping")
        return {
            "repo": repo_name,
            "status": "no_summary",
            "reason": "README_SUMMARY.md not found",
        }
    
    # Find code files
    code_files = find_code_files(repo_dir, max_file_chars=max_file_chars)
    if not code_files:
        return {
            "repo": repo_name,
            "status": "no_code",
            "reason": "No code files found",
        }
    
    logger.info(f"Repository {repo_name}: found {len(code_files)} code files")
    
    # Parse all code files
    semaphore = asyncio.Semaphore(max_concurrency)
    
    async def parse_with_semaphore(file_path: Path):
        async with semaphore:
            return await parse_code_file(
                file_path=file_path,
                repo_dir=repo_dir,
                project_context=project_context,
                base_url=base_url,
                model=model,
                api_key=api_key,
                log_file=log_file,
            )
    
    # Parse all files concurrently
    tasks = [parse_with_semaphore(file_path) for file_path in code_files]
    parse_results = []
    
    for task in tqdm(asyncio.as_completed(tasks), total=len(tasks), desc=f"Parsing {repo_name}", leave=False):
        result = await task
        if result:
            parse_results.append(result)
    
    if not parse_results:
        return {
            "repo": repo_name,
            "status": "parse_failed",
            "reason": "All files failed to parse",
        }
    
    # Generate CSV file
    repo_name_normalized = extract_repo_name(repo_dir)
    
    # CSV fields
    fieldnames = [
        "repo_name",
        "readme_summary_path",
        "readme_summary_text",
        "path",
        "language",
        "dependencies",
        "function_name",
        "function_start_line",
        "function_end_line",
        "function_body",
        "doc_start_line",
        "doc_end_line",
        "file_size_bytes",
        "file_sha1",
        "function_hash",
        "ds_source",
    ]
    
    # Write CSV
    try:
        with open(csv_file, "w", encoding="utf-8", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            
            function_count = 0
            for parse_result in parse_results:
                file_path = parse_result["file_path"]
                language = parse_result["language"]
                dependencies = parse_result.get("dependencies", [])
                functions = parse_result.get("functions", [])
                
                # Read file content (for hash and size calculation)
                full_file_path = repo_dir / file_path
                file_content = read_code_file(full_file_path)
                file_size = len(file_content.encode("utf-8")) if file_content else 0
                file_sha1 = compute_file_hash(full_file_path, file_content) if file_content else ""
                
                # Write a row for each function
                for func in functions:
                    function_name = func.get("function_name", "")
                    function_start_line = func.get("function_start_line", 0)
                    function_end_line = func.get("function_end_line", 0)
                    function_body = func.get("function_body", "")
                    doc_start_line = func.get("doc_start_line")
                    doc_end_line = func.get("doc_end_line")
                    
                    function_hash = compute_function_hash(
                        repo_name_normalized,
                        file_path,
                        function_start_line,
                        function_end_line,
                        function_body,
                    )
                    
                    # Truncate project_context (if too long)
                    context_text = project_context[:5000] if len(project_context) > 5000 else project_context
                    
                    row = {
                        "repo_name": repo_name_normalized,
                        "readme_summary_path": SUMMARY_FILENAME,
                        "readme_summary_text": context_text,
                        "path": file_path,
                        "language": language,
                        "dependencies": json.dumps(dependencies, ensure_ascii=False),
                        "function_name": function_name,
                        "function_start_line": function_start_line,
                        "function_end_line": function_end_line,
                        "function_body": function_body,
                        "doc_start_line": doc_start_line if doc_start_line else "",
                        "doc_end_line": doc_end_line if doc_end_line else "",
                        "file_size_bytes": file_size,
                        "file_sha1": file_sha1,
                        "function_hash": function_hash,
                        "ds_source": "repos_filtered",
                    }
                    
                    writer.writerow(row)
                    function_count += 1
        
        logger.info(f"Repository {repo_name}: wrote {function_count} functions to {csv_file}")
        
        return {
            "repo": repo_name,
            "status": "success",
            "csv_file": str(csv_file),
            "file_count": len(code_files),
            "function_count": function_count,
        }
    except Exception as e:
        logger.error(f"Unable to write CSV file {csv_file}: {e}")
        return {
            "repo": repo_name,
            "status": "write_failed",
            "reason": str(e),
        }


async def process_all_repos(
    repos_dir: Path,
    base_url: str,
    model: str,
    api_key: str,
    log_file: str,
    max_file_chars: int = 200000,
    max_concurrency: int = 8,
    overwrite: bool = False,
) -> List[Dict]:
    """
    Process function parsing for all repositories
    
    Args:
        repos_dir: Repository root directory
        base_url: LLM API base URL
        model: Model name
        api_key: API key
        log_file: Log file path
        max_file_chars: Maximum file size (chars)
        max_concurrency: Maximum concurrency
        overwrite: Whether to overwrite existing CSV files
        
    Returns:
        List of processing results for all repositories
    """
    # Get all repository directories
    repo_dirs = [
        d for d in repos_dir.iterdir() 
        if d.is_dir() and not d.name.startswith(".")
    ]
    repo_dirs.sort()
    
    logger.info(f"Found {len(repo_dirs)} repositories, starting processing...")
    
    # Process each repository sequentially (concurrency is controlled at file level)
    results = []
    
    for repo_dir in tqdm(repo_dirs, desc="Processing repos"):
        result = await process_single_repo(
            repo_dir=repo_dir,
            base_url=base_url,
            model=model,
            api_key=api_key,
            log_file=log_file,
            max_file_chars=max_file_chars,
            max_concurrency=max_concurrency,
            overwrite=overwrite,
        )
        results.append(result)
    
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Multi-language Function Parsing Tool")
    parser.add_argument(
        "--repos_dir",
        type=str,
        default="/home/weifengsun/tangou1/domain_code/src/workdir/repos_filtered",
        help="Repository root directory path",
    )
    parser.add_argument(
        "--base_url",
        type=str,
        default=os.getenv("OPENAI_BASE_URL", "http://localhost:8000/v1"),
        help="LLM API base URL (default: http://localhost:8000/v1)",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="Qwen3",
        help="Model name (default: Qwen3)",
    )
    parser.add_argument(
        "--api_key_env",
        type=str,
        default="OPENAI_API_KEY",
        help="API key environment variable name (default: OPENAI_API_KEY)",
    )
    parser.add_argument(
        "--max_concurrency",
        type=int,
        default=8,
        help="Maximum concurrency (default: 8)",
    )
    parser.add_argument(
        "--max_file_chars",
        type=int,
        default=200000,
        help="Maximum file size in chars (default: 200000)",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing CSV files",
    )
    parser.add_argument(
        "--log_file",
        type=str,
        default="instruction_generation/workdir/logs/extract.log",
        help="Log file path",
    )
    
    args = parser.parse_args()
    
    # Initialize logger
    init_logger(args.log_file, level="INFO")
    
    # Get API key
    api_key = os.getenv(args.api_key_env, "none")
    
    # Process all repositories
    repos_dir = Path(args.repos_dir)
    if not repos_dir.exists():
        logger.error(f"Repository directory does not exist: {repos_dir}")
        sys.exit(1)
    
    # Create log directory
    log_file_path = Path(args.log_file)
    log_file_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Run main logic
    results = asyncio.run(
        process_all_repos(
            repos_dir=repos_dir,
            base_url=args.base_url,
            model=args.model,
            api_key=api_key,
            log_file=str(log_file_path),
            max_file_chars=args.max_file_chars,
            max_concurrency=args.max_concurrency,
            overwrite=args.overwrite,
        )
    )
    
    # Statistics
    status_counts = {}
    total_functions = 0
    for result in results:
        status = result["status"]
        status_counts[status] = status_counts.get(status, 0) + 1
        if "function_count" in result:
            total_functions += result["function_count"]
    
    logger.info("\n" + "=" * 80)
    logger.info("Processing complete!")
    logger.info("=" * 80)
    logger.info(f"Total: {len(results)} repositories")
    logger.info(f"Total: {total_functions} functions")
    for status, count in status_counts.items():
        logger.info(f"  {status}: {count}")
    logger.info("=" * 80)