core/lab_notebook.py

import json
import time
import uuid
from typing import Dict, List, Any, Optional
from datetime import datetime
import os

class ExperimentEntry:
    def __init__(self, 
                 lab_name: str, 
                 experiment_name: str, 
                 parameters: Dict[str, Any], 
                 results: Dict[str, Any], 
                 user: str = "User",
                 tags: List[str] = None):
        self.entry_id = str(uuid.uuid4())
        self.timestamp = datetime.now().isoformat()
        self.lab_name = lab_name
        self.experiment_name = experiment_name
        self.parameters = parameters
        self.results = results
        self.user = user
        self.tags = tags or []

    def to_dict(self) -> Dict[str, Any]:
        return {
            "entry_id": self.entry_id,
            "timestamp": self.timestamp,
            "lab_name": self.lab_name,
            "experiment_name": self.experiment_name,
            "parameters": self.parameters,
            "results": self.results,
            "user": self.user,
            "tags": self.tags
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        entry = cls(
            lab_name=data["lab_name"],
            experiment_name=data["experiment_name"],
            parameters=data["parameters"],
            results=data["results"],
            user=data.get("user", "User"),
            tags=data.get("tags", [])
        )
        entry.entry_id = data.get("entry_id", str(uuid.uuid4()))
        entry.timestamp = data.get("timestamp", datetime.now().isoformat())
        return entry

class LabNotebook:
    def __init__(self, storage_path: str = "lab_notebook.json"):
        self.storage_path = storage_path
        self.entries: List[ExperimentEntry] = []
        self._load_from_disk()

    def add_entry(self, 
                  lab_name: str, 
                  experiment_name: str, 
                  parameters: Dict[str, Any], 
                  results: Dict[str, Any], 
                  tags: List[str] = None) -> str:
        """
        Record a new experiment entry.
        Returns the entry_id.
        """
        entry = ExperimentEntry(lab_name, experiment_name, parameters, results, tags=tags)
        self.entries.append(entry)
        self._save_to_disk()
        return entry.entry_id

    def search_entries(self, query: str = None, tag: str = None, lab: str = None) -> List[ExperimentEntry]:
        """
        Search for entries matching criteria.
        """
        results = self.entries
        
        if lab:
            results = [e for e in results if e.lab_name == lab]
            
        if tag:
            results = [e for e in results if tag in e.tags]
            
        if query:
            q = query.lower()
            results = [e for e in results if (q in e.experiment_name.lower() or 
                                            q in str(e.parameters).lower() or 
                                            q in str(e.results).lower())]
            
        return results

    def get_entry(self, entry_id: str) -> Optional[ExperimentEntry]:
        for entry in self.entries:
            if entry.entry_id == entry_id:
                return entry
        return None

    def export_report(self, entry_ids: List[str], output_file: str = "report.md"):
        """
        Export selected entries to a markdown report.
        """
        selected_entries = [self.get_entry(eid) for eid in entry_ids if self.get_entry(eid)]
        
        with open(output_file, 'w') as f:
            f.write(f"# Lab Notebook Report\n")
            f.write(f"Generated on: {datetime.now().isoformat()}\n\n")
            
            for entry in selected_entries:
                f.write(f"## {entry.experiment_name}\n")
                f.write(f"**ID**: `{entry.entry_id}`\n")
                f.write(f"**Lab**: {entry.lab_name} | **Time**: {entry.timestamp}\n")
                f.write(f"**User**: {entry.user}\n")
                f.write(f"**Tags**: {', '.join(entry.tags)}\n\n")
                
                f.write("### Parameters\n")
                f.write("```json\n")
                f.write(json.dumps(entry.parameters, indent=2))
                f.write("\n```\n\n")
                
                f.write("### Results\n")
                f.write("```json\n")
                f.write(json.dumps(entry.results, indent=2))
                f.write("\n```\n")
                f.write("---\n\n")
                
        return output_file

    def _save_to_disk(self):
        data = [e.to_dict() for e in self.entries]
        with open(self.storage_path, 'w') as f:
            json.dump(data, f, indent=2)

    def _load_from_disk(self):
        if not os.path.exists(self.storage_path):
            return
            
        try:
            with open(self.storage_path, 'r') as f:
                data = json.load(f)
                self.entries = [ExperimentEntry.from_dict(d) for d in data]
        except (json.JSONDecodeError, IOError):
            # Backup corrupt file if needed, detailed logging, etc.
            # For now, just start fresh or warn
            print(f"Warning: Could not load notebook from {self.storage_path}")
            self.entries = []

# Global default instance
default_notebook = LabNotebook()