src/utilities.py

from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

def save_notebook(notebook):
    from pathlib import Path
    import os,shutil

    folder_path = Path("/tmp/Notebook")

    if folder_path.exists() and folder_path.is_dir():
        shutil.rmtree(folder_path)

    folder_path.mkdir(parents=True, exist_ok=True)

    if notebook:
        file_path = notebook.name
        with open(os.path.join(folder_path,file_path), "wb") as file:
                file.write(notebook.getbuffer())

def safe_concurrent_batch(chain, inputs, max_workers=2):
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_input = {executor.submit(chain.invoke, inp): inp for inp in inputs}
        
        for future in tqdm(as_completed(future_to_input), total=len(inputs), desc="Processing"):
            inp = future_to_input[future]
            try:
                output = future.result()
                results.append({
                    "input": inp,
                    "output": output,
                    "status": "success"
                })
            except Exception as e:
                results.append({
                    "input": inp,
                    "output": None,
                    "status": f"failed: {str(e)}"
                })
    return results