docs/Agent/lab.py

import os
import yaml
import shutil
import hashlib
import subprocess
from datetime import datetime
from typing import Dict, Any, List, Optional
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("lab")

def create_experiment_with_plate(
    experiment_id: str,
    aim: str,
    project: str,
    researcher: str,
    plate: Dict[str, Any],
    sample_preparation: Optional[Dict[str, Any]] = None,
    downstream_application: Optional[Dict[str, Any]] = None,
    status: str = "in_progress"
) -> str:
    """
    Create a new experiment with plate layout information.
    
    Args:
        experiment_id: Unique identifier for the experiment (EXP-YYYYMMDD format)
        aim: Brief description of the experiment's goal
        project: Project this experiment belongs to
        researcher: GitHub username or name of the researcher
        plate: Dictionary with plate ID and layout information
        sample_preparation: Optional dictionary with sample preparation details
        downstream_application: Optional dictionary with downstream application details
        status: Current status of the experiment (default: in_progress)
        
    Returns:
        Path to the created experiment file
    """
    logger.info(f"Creating experiment {experiment_id} for {researcher}")
    
    # Create basic experiment structure
    experiment = {
        "experiment_id": experiment_id,
        "aim": aim,
        "project": project,
        "researcher": researcher,
        "status": status,
        "created": datetime.now().strftime("%Y-%m-%d"),
        "plate": plate,
        "tasks": []
    }
    
    # Add optional sections if provided
    if sample_preparation:
        experiment["sample_preparation"] = sample_preparation
    
    if downstream_application:
        experiment["downstream_application"] = downstream_application
        
    # Create data directory for this experiment if it doesn't exist
    data_dir = os.path.join("Data", experiment_id)
    os.makedirs(data_dir, exist_ok=True)
    
    # Save the experiment file
    experiment_file = f"Experiments/{experiment_id}.yaml"
    with open(experiment_file, "w") as f:
        yaml.dump(experiment, f, sort_keys=False)
    
    logger.info(f"Experiment file created: {experiment_file}")
    return experiment_file

def add_experiment_task(experiment_id: str, task_title: str) -> Dict[str, Any]:
    """
    Add a task to an experiment and create a corresponding GitHub Issue.
    
    Args:
        experiment_id: ID of the experiment to add task to
        task_title: Title of the task/Issue
        
    Returns:
        Dictionary with the created issue details
    """
    logger.info(f"Adding task '{task_title}' to experiment {experiment_id}")
    
    # Create GitHub Issue
    cmd = f'gh issue create --title "{task_title}" --body "Experiment: {experiment_id}\n\nTask for experiment {experiment_id}"'
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    
    if result.returncode != 0:
        logger.error(f"Failed to create GitHub Issue: {result.stderr}")
        raise Exception(f"Failed to create GitHub Issue: {result.stderr}")
    
    # Extract issue number from URL
    issue_url = result.stdout.strip()
    issue_number = issue_url.split("/")[-1]
    
    # Add issue to project board
    add_to_project_cmd = f'gh issue edit {issue_number} --add-project "Lab Tasks Board"'
    subprocess.run(add_to_project_cmd, shell=True)
    
    # Update experiment YAML
    experiment_file = f"Experiments/{experiment_id}.yaml"
    with open(experiment_file, "r") as f:
        experiment = yaml.safe_load(f)
    
    # Add task to tasks array
    task = {
        "id": issue_number,
        "title": task_title,
        "status": "open"
    }
    
    if "tasks" not in experiment:
        experiment["tasks"] = []
    
    experiment["tasks"].append(task)
    
    with open(experiment_file, "w") as f:
        yaml.dump(experiment, f, sort_keys=False)
    
    logger.info(f"Added task {issue_number} to experiment {experiment_id}")
    return task

def record_data(exp_id: str, file_path: str, data_type: str) -> Dict[str, Any]:
    """
    Record data file for an experiment. This function:
    1. Computes checksum of the file
    2. Copies file to Data/{exp_id}/ if not already there
    3. Adds entry to the experiment's YAML data section
    
    Args:
        exp_id: Experiment ID 
        file_path: Path to the data file
        data_type: Type of data (e.g., qPCR, flow_cytometry, imaging)
        
    Returns:
        The data entry added to the experiment
    """
    logger.info(f"Recording {data_type} data for experiment {exp_id}")
    
    # Verify the experiment exists
    experiment_file = f"Experiments/{exp_id}.yaml"
    if not os.path.exists(experiment_file):
        raise FileNotFoundError(f"Experiment {exp_id} not found")
    
    # Make sure the data directory exists
    data_dir = os.path.join("Data", exp_id)
    os.makedirs(data_dir, exist_ok=True)
    
    # Compute file checksum
    with open(file_path, "rb") as f:
        file_data = f.read()
        sha256 = hashlib.sha256(file_data).hexdigest()
    
    # Get just the filename
    filename = os.path.basename(file_path)
    
    # Construct target path in the data directory
    target_path = os.path.join(data_dir, filename)
    
    # Copy file if it's not already in the data directory
    if os.path.abspath(file_path) != os.path.abspath(target_path):
        shutil.copy2(file_path, target_path)
        logger.info(f"Copied data file to {target_path}")
    
    # Update the experiment YAML
    with open(experiment_file, "r") as f:
        experiment = yaml.safe_load(f)
    
    # Create the data entry
    data_entry = {
        "path": f"Data/{exp_id}/{filename}",
        "type": data_type,
        "sha256": sha256,
        "added": datetime.now().strftime("%Y-%m-%d")
    }
    
    # Add to data array
    if "data" not in experiment:
        experiment["data"] = []
    
    experiment["data"].append(data_entry)
    
    # Save the updated experiment file
    with open(experiment_file, "w") as f:
        yaml.dump(experiment, f, sort_keys=False)
    
    # Commit the changes
    commit_cmd = f'git add "{target_path}" "{experiment_file}" && git commit -m "Data: add {data_type} result for {exp_id}"'
    subprocess.run(commit_cmd, shell=True)
    
    logger.info(f"Recorded data file {filename} for experiment {exp_id}")
    return data_entry

def close_task(issue_number: int) -> None:
    """
    Close a GitHub Issue and update the corresponding experiment task status.
    
    Args:
        issue_number: The GitHub Issue number to close
    """
    logger.info(f"Closing task (issue #{issue_number})")
    
    # Close the GitHub issue
    close_cmd = f'gh issue close {issue_number}'
    result = subprocess.run(close_cmd, shell=True)
    
    if result.returncode != 0:
        logger.error(f"Failed to close GitHub Issue #{issue_number}")
        return
    
    # Find which experiment this issue belongs to
    experiments_dir = "Experiments"
    for filename in os.listdir(experiments_dir):
        if not filename.endswith(".yaml"):
            continue
            
        file_path = os.path.join(experiments_dir, filename)
        with open(file_path, "r") as f:
            experiment = yaml.safe_load(f)
            
        if "tasks" not in experiment:
            continue
            
        # Check if this issue is in the tasks
        for i, task in enumerate(experiment["tasks"]):
            if str(task.get("id")) == str(issue_number):
                # Update task status
                experiment["tasks"][i]["status"] = "closed"
                
                # Save the updated experiment
                with open(file_path, "w") as f:
                    yaml.dump(experiment, f, sort_keys=False)
                    
                logger.info(f"Updated task status in experiment {experiment['experiment_id']}")
                
                # Commit the change
                commit_cmd = f'git add "{file_path}" && git commit -m "Task: close issue #{issue_number} for {experiment["experiment_id"]}"'
                subprocess.run(commit_cmd, shell=True)
                return

    logger.warning(f"Could not find task {issue_number} in any experiment")

def finish_experiment(exp_id: str) -> bool:
    """
    Finish an experiment by:
    1. Verifying all tasks are closed
    2. Checking data exists
    3. Setting status to completed
    4. Updating SESSION_LOG.md
    
    Args:
        exp_id: Experiment ID to finish
        
    Returns:
        True if experiment was successfully completed, False otherwise
    """
    logger.info(f"Attempting to finish experiment {exp_id}")
    
    # Verify the experiment exists
    experiment_file = f"Experiments/{exp_id}.yaml"
    if not os.path.exists(experiment_file):
        logger.error(f"Experiment {exp_id} not found")
        return False
    
    # Load the experiment
    with open(experiment_file, "r") as f:
        experiment = yaml.safe_load(f)
    
    # Check if all tasks are closed
    if "tasks" in experiment:
        open_tasks = [task for task in experiment["tasks"] if task.get("status") != "closed"]
        if open_tasks:
            task_ids = [task.get("id") for task in open_tasks]
            logger.warning(f"Cannot finish experiment: open tasks remain: {task_ids}")
            return False
    
    # Check if data exists
    if "data" not in experiment or not experiment["data"]:
        logger.warning(f"Cannot finish experiment: no data recorded")
        return False
    
    # Update status to completed
    experiment["status"] = "completed"
    experiment["completed"] = datetime.now().strftime("%Y-%m-%d")
    
    # Save the updated experiment
    with open(experiment_file, "w") as f:
        yaml.dump(experiment, f, sort_keys=False)
    
    # Update SESSION_LOG.md
    session_log_path = "SESSION_LOG.md"
    if not os.path.exists(session_log_path):
        with open(session_log_path, "w") as f:
            f.write("# Lab Session Log\n\n")
    
    with open(session_log_path, "a") as f:
        f.write(f"\n## {datetime.now().strftime('%Y-%m-%d')} - {experiment['researcher']}\n")
        f.write(f"- Completed experiment {exp_id}: {experiment['aim']}\n")
        f.write(f"- Data files: {len(experiment['data'])}\n")
    
    # Commit the changes
    commit_cmd = f'git add "{experiment_file}" "{session_log_path}" && git commit -m "Experiment: completed {exp_id}"'
    subprocess.run(commit_cmd, shell=True)
    
    logger.info(f"Successfully completed experiment {exp_id}")
    return True

def list_experiments(status_filter=None):
    """
    List all experiments, optionally filtered by status.
    
    Args:
        status_filter: Optional filter for experiment status ('in_progress', 'completed', etc.)
        
    Returns:
        List of experiment dictionaries
    """
    experiments = []
    experiments_dir = "Experiments"
    
    if not os.path.exists(experiments_dir):
        return []
        
    for filename in os.listdir(experiments_dir):
        if not filename.endswith(".yaml"):
            continue
            
        file_path = os.path.join(experiments_dir, filename)
        with open(file_path, "r") as f:
            experiment = yaml.safe_load(f)
            
        if status_filter is None or experiment.get("status") == status_filter:
            experiments.append(experiment)
            
    return experiments

def record_cli():
    """Entry point for lab-record CLI command"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Record data for an experiment")
    parser.add_argument("--exp", required=True, help="Experiment ID")
    parser.add_argument("--file", required=True, help="Path to data file")
    parser.add_argument("--type", required=True, help="Type of data")
    
    args = parser.parse_args()
    try:
        result = record_data(args.exp, args.file, args.type)
        print(f"Data recorded successfully: {result['path']}")
    except Exception as e:
        print(f"Error: {e}")
        exit(1)

def main():
    """Entry point for lab CLI command"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Lab Management CLI")
    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    
    # Record data command
    record_parser = subparsers.add_parser("record", help="Record data for an experiment")
    record_parser.add_argument("--exp", required=True, help="Experiment ID")
    record_parser.add_argument("--file", required=True, help="Path to data file")
    record_parser.add_argument("--type", required=True, help="Type of data")
    
    # Close task command
    close_parser = subparsers.add_parser("close-task", help="Close a task/issue")
    close_parser.add_argument("--issue", type=int, required=True, help="Issue number to close")
    
    # Finish experiment command
    finish_parser = subparsers.add_parser("finish", help="Mark an experiment as completed")
    finish_parser.add_argument("--exp", required=True, help="Experiment ID to finish")
    
    # List experiments command
    list_parser = subparsers.add_parser("list", help="List experiments")
    list_parser.add_argument("--status", help="Filter by status")
    
    # Initialize command (mostly for documentation, actual impl. is in init_extensions.py)
    init_parser = subparsers.add_parser("init-extensions", help="Initialize extensions and environment")
    
    # Parse arguments and execute
    args = parser.parse_args()
    
    if args.command == "record":
        record_data(args.exp, args.file, args.type)
    elif args.command == "close-task":
        close_task(args.issue)
    elif args.command == "finish":
        finish_experiment(args.exp)
    elif args.command == "list":
        experiments = list_experiments(args.status)
        for exp in experiments:
            print(f"{exp.get('experiment_id', 'No ID')} - {exp.get('aim', 'No description')} - {exp.get('status', 'unknown')}")
    elif args.command == "init-extensions":
        print("Please run 'lab-init-extensions' instead")
    else:
        parser.print_help()

# Command-line interface
if __name__ == "__main__":
    main()