diff --git a/cli/olaf/src/olaf/cli/main.py b/cli/olaf/src/olaf/cli/main.py index e14aa8f..0e29e90 100644 --- a/cli/olaf/src/olaf/cli/main.py +++ b/cli/olaf/src/olaf/cli/main.py @@ -9,6 +9,7 @@ from .datasets_cli import datasets_app from .run_cli import run_app from .config_cli import config_app +from .utils_cli import utils_app # <-- Import the utils app # Main OLAF application app = typer.Typer( name="olaf", @@ -21,6 +22,7 @@ app.add_typer(datasets_app, name="datasets") app.add_typer(run_app, name="run") app.add_typer(config_app, name="config") # <-- Register the new config app +app.add_typer(utils_app, name="utils") # <-- Register the utils app def main(): diff --git a/cli/olaf/src/olaf/cli/run_cli.py b/cli/olaf/src/olaf/cli/run_cli.py index 34966ea..5be2be3 100644 --- a/cli/olaf/src/olaf/cli/run_cli.py +++ b/cli/olaf/src/olaf/cli/run_cli.py @@ -5,28 +5,29 @@ from pathlib import Path from typing import List, Tuple, cast, Optional import subprocess +import json +from datetime import datetime import typer from rich.console import Console from rich.prompt import Prompt, IntPrompt from dotenv import load_dotenv +from olaf.config import DEFAULT_AGENT_DIR, ENV_FILE, OLAF_HOME + PACKAGE_ROOT = Path(__file__).resolve().parent.parent PACKAGE_AGENTS_DIR = PACKAGE_ROOT / "agents" PACKAGE_DATASETS_DIR = PACKAGE_ROOT / "datasets" PACKAGE_AUTO_METRICS_DIR = PACKAGE_ROOT / "auto_metrics" -# Define static in-container paths for primary and reference datasets SANDBOX_DATA_PATH = "/workspace/dataset.h5ad" SANDBOX_REF_DATA_PATH = "/workspace/reference.h5ad" - def _prompt_for_file( console: Console, user_dir: Path, package_dir: Path, extension: str, prompt_title: str ) -> Path: """ Generic helper to find files in both user and package directories and prompt for a selection. - User files take priority over package files with the same name. """ console.print(f"[bold]Select {prompt_title}:[/bold]") found_files = [] @@ -116,7 +117,6 @@ def main_run_callback( force_refresh: bool = typer.Option(False, "--force-refresh", help="Force refresh/rebuild of the sandbox environment."), ): # --- Heavy imports are deferred to here --- - from olaf.config import DEFAULT_AGENT_DIR, ENV_FILE from olaf.agents.AgentSystem import AgentSystem from olaf.core.io_helpers import collect_resources from olaf.core.sandbox_management import init_docker, init_singularity_exec @@ -182,7 +182,6 @@ def main_run_callback( if app_context.reference_dataset_path: app_context.resources.append((app_context.reference_dataset_path, SANDBOX_REF_DATA_PATH)) - # Build the analysis context string, including the reference dataset if it exists analysis_context_str = f"Primary dataset path: **{SANDBOX_DATA_PATH}**\n" if app_context.reference_dataset_path: analysis_context_str += f"Reference dataset path: **{SANDBOX_REF_DATA_PATH}**\n" @@ -194,9 +193,9 @@ def main_run_callback( def _setup_and_run_session(context: AppContext, history: list, is_auto: bool, max_turns: int, benchmark_modules: Optional[List[Path]] = None): """Helper to start, run, and stop the sandbox session.""" - # --- Heavy imports needed for the session are deferred to here --- from olaf.execution.runner import run_agent_session, SandboxManager from olaf.agents.AgentSystem import AgentSystem + from olaf.core.io_helpers import save_chat_history_as_json, save_chat_history_as_notebook sandbox_manager = cast(SandboxManager, context.sandbox_manager) console = context.console @@ -206,7 +205,6 @@ def _setup_and_run_session(context: AppContext, history: list, is_auto: bool, ma details = context.sandbox_details dataset_path = cast(Path, context.dataset_path) if details["is_exec_mode"] and hasattr(sandbox_manager, "set_data"): - # Pass all resources, including the reference dataset, for bind mounting all_resources = [(dataset_path, SANDBOX_DATA_PATH)] + context.resources sandbox_manager.set_data(all_resources) if not sandbox_manager.start_container(): @@ -215,9 +213,7 @@ def _setup_and_run_session(context: AppContext, history: list, is_auto: bool, ma try: if not details["is_exec_mode"]: - # Copy primary dataset details["copy_cmd"](str(dataset_path), f"{details['handle']}:{SANDBOX_DATA_PATH}") - # Copy all other resources, including the reference dataset for hp, cp in context.resources: details["copy_cmd"](str(hp), f"{details['handle']}:{cp}") @@ -237,6 +233,26 @@ def _setup_and_run_session(context: AppContext, history: list, is_auto: bool, ma finally: console.print("[cyan]Stopping sandbox...[/cyan]") sandbox_manager.stop_container() + if not is_auto: + if Prompt.ask("\n[bold]Do you want to save the chat history?[/bold]", choices=["y", "n"], default="y").lower() == 'y': + log_dir = OLAF_HOME / "runs" / "chat_logs" + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + + # --- NEW: Prompt for save format --- + save_format = Prompt.ask("Save format", choices=["json", "notebook"], default="notebook") + file_extension = ".ipynb" if save_format == "notebook" else ".json" + + default_path = log_dir / f"interactive_chat_{timestamp}{file_extension}" + save_path_str = Prompt.ask( + "Enter the save path for the log", + default=str(default_path) + ) + save_path = Path(save_path_str).expanduser() + + if save_format == "notebook": + save_chat_history_as_notebook(console, history, save_path) + else: + save_chat_history_as_json(console, history, save_path) @run_app.command("interactive") def run_interactive(ctx: typer.Context): diff --git a/cli/olaf/src/olaf/cli/utils_cli.py b/cli/olaf/src/olaf/cli/utils_cli.py new file mode 100644 index 0000000..b83ef0b --- /dev/null +++ b/cli/olaf/src/olaf/cli/utils_cli.py @@ -0,0 +1,134 @@ +# olaf/cli/utils_cli.py +import json +import re +from pathlib import Path +from typing import Optional + +import typer +from rich.console import Console + +# Import from the central config to know where chat logs are stored by default +from olaf.config import OLAF_HOME + +utils_app = typer.Typer( + name="utils", + help="Utility commands for managing OLAF artifacts like chat logs.", + no_args_is_help=True +) + +console = Console() +LOG_DIR = OLAF_HOME / "runs" / "chat_logs" + +def _convert_history_to_notebook(history_path: Path, output_path: Path): + """ + Parses an OLAF chat log and converts it into a Jupyter Notebook (.ipynb). + """ + try: + with open(history_path, 'r', encoding='utf-8') as f: + history = json.load(f) + except (FileNotFoundError, json.JSONDecodeError) as e: + console.print(f"[bold red]Error: Could not read or parse the history file at {history_path}.[/bold red]\n{e}") + raise typer.Exit(1) + + notebook = { + "cells": [], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" # This can be made more dynamic if needed + } + }, + "nbformat": 4, + "nbformat_minor": 5 + } + + # This regex specifically finds Python code blocks + code_block_re = re.compile(r"```python\n(.*?)\n```", re.DOTALL) + + for message in history: + role = message.get("role") + content = message.get("content", "") + + # We are primarily interested in the agent's responses + if role and "assistant" in role: + # Split the content by code blocks to interleave markdown and code + parts = code_block_re.split(content) + + for i, part in enumerate(parts): + part = part.strip() + if not part: + continue + + # Odd-indexed parts are the code blocks captured by the regex + if i % 2 == 1: + cell = { + "cell_type": "code", + "execution_count": None, + "metadata": {}, + "outputs": [], + "source": part + } + notebook["cells"].append(cell) + # Even-indexed parts are the explanatory text outside the code blocks + else: + cell = { + "cell_type": "markdown", + "metadata": {}, + "source": part + } + notebook["cells"].append(cell) + + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, 'w', encoding='utf-8') as f: + json.dump(notebook, f, indent=2) + console.print(f"[bold green]✓ Successfully converted chat log to notebook:[/bold green] {output_path}") + except Exception as e: + console.print(f"[bold red]Error writing notebook file: {e}[/bold red]") + raise typer.Exit(1) + +@utils_app.command("convert-to-notebook") +def convert_to_notebook( + chat_log: Path = typer.Argument( + ..., + help="Path to the interactive chat log JSON file to convert.", + exists=True, + readable=True, + resolve_path=True, + ), + output: Optional[Path] = typer.Option( + None, + "--output", + "-o", + help="Path to save the output .ipynb file. Defaults to the same name as the input file.", + writable=True, + resolve_path=True, + ), +): + """ + Converts an OLAF interactive chat log into an executable Jupyter Notebook. + + This command parses the JSON log file, extracts all Python code blocks generated + by the assistant, and arranges them into code cells. The explanatory text + between code blocks is converted into markdown cells, creating a clean, + reproducible protocol of the agent session. + """ + if not chat_log.name.startswith("interactive_chat_"): + console.print(f"[yellow]Warning: The input file '{chat_log.name}' does not look like a standard OLAF chat log.[/yellow]") + + output_path = output + if output_path is None: + # Default to the same name as the input file, but with an .ipynb extension + output_path = chat_log.with_suffix(".ipynb") + + # Ensure the output path has the correct extension + if output_path.suffix != ".ipynb": + output_path = output_path.with_suffix(".ipynb") + + console.print(f"Converting [cyan]{chat_log.name}[/cyan] to Jupyter Notebook...") + _convert_history_to_notebook(chat_log, output_path) diff --git a/cli/olaf/src/olaf/core/io_helpers.py b/cli/olaf/src/olaf/core/io_helpers.py index 2b9508f..91154dd 100644 --- a/cli/olaf/src/olaf/core/io_helpers.py +++ b/cli/olaf/src/olaf/core/io_helpers.py @@ -180,4 +180,54 @@ def format_execute_response(resp: dict, output_dir) -> str: img_paths.append(str(fname)) if img_paths: lines.append("Saved images: " + ", ".join(img_paths)) - return "\n".join(lines) \ No newline at end of file + return "\n".join(lines) + +def save_chat_history_as_json(console: Console, history: list, file_path: Path): + """Saves the interactive chat history to a user-specified JSON file.""" + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + with open(file_path, 'w', encoding='utf-8') as f: + json.dump(history, f, indent=2) + console.print(f"\n[bold green]✓ Chat history saved to:[/bold green] {file_path}") + except Exception as e: + console.print(f"\n[bold red]Error saving chat history: {e}[/bold red]") + +def save_chat_history_as_notebook(console: Console, history: list, file_path: Path): + """Parses an OLAF chat log and converts it into a Jupyter Notebook (.ipynb).""" + notebook = { + "cells": [], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { "name": "python", "version": "3.11" } + }, + "nbformat": 4, + "nbformat_minor": 5 + } + code_block_re = re.compile(r"```python\n(.*?)\n```", re.DOTALL) + + for message in history: + if message.get("role") and "assistant" in message.get("role", ""): + content = message.get("content", "") + parts = code_block_re.split(content) + for i, part in enumerate(parts): + part = part.strip() + if not part: continue + + if i % 2 == 1: # Code block + cell = {"cell_type": "code", "execution_count": None, "metadata": {}, "outputs": [], "source": part} + notebook["cells"].append(cell) + else: # Markdown + cell = {"cell_type": "markdown", "metadata": {}, "source": part} + notebook["cells"].append(cell) + + try: + file_path.parent.mkdir(parents=True, exist_ok=True) + with open(file_path, 'w', encoding='utf-8') as f: + json.dump(notebook, f, indent=2) + console.print(f"\n[bold green]✓ Notebook saved to:[/bold green] {file_path}") + except Exception as e: + console.print(f"\n[bold red]Error saving notebook: {e}[/bold red]") \ No newline at end of file