Source code for langgraph_compare.experiment

import os
import sqlite3
from dataclasses import dataclass
from typing import Optional
from langgraph.checkpoint.sqlite import SqliteSaver

[docs] @dataclass class ExperimentPaths: """ Class that holds and manages all paths for an experiment. :param name: Name of the experiment that will be used for folder structure. :type name: str :param base_dir: Base directory where all experiments are stored, defaults to "experiments". :type base_dir: str """ name: str base_dir: str = "experiments" _connection: Optional[sqlite3.Connection] = None _memory: Optional[SqliteSaver] = None @property def database(self) -> str: """ Returns path to the main SQLite database. :return: Full path to the SQLite database file. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.database 'experiments/test/db/test.sqlite' """ return os.path.join(self.base_dir, self.name, "db", f"{self.name}.sqlite") @property def connection(self) -> sqlite3.Connection: """ Returns or creates SQLite connection with appropriate settings. :return: SQLite connection object :rtype: sqlite3.Connection **Example:** >>> paths = ExperimentPaths("test") >>> conn = paths.connection # Creates new connection if none exists >>> conn2 = paths.connection # Returns existing connection """ if self._connection is None: self._connection = sqlite3.connect(self.database, check_same_thread=False) return self._connection @property def memory(self) -> SqliteSaver: """ Returns or creates SqliteSaver instance for the database. :return: SqliteSaver instance :rtype: SqliteSaver **Example:** >>> paths = ExperimentPaths("test") >>> saver = paths.memory # Creates new SqliteSaver if none exists >>> saver2 = paths.memory # Returns existing SqliteSaver """ if self._memory is None: self._memory = SqliteSaver(self.connection) return self._memory @property def json_dir(self) -> str: """ Returns path to the json directory. :return: Full path to the JSON directory. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.json_dir 'experiments/test/json' """ return os.path.join(self.base_dir, self.name, "json") @property def csv_dir(self) -> str: """ Returns path to the csv directory. :return: Full path to the CSV directory. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.csv_dir 'experiments/test/csv' """ return os.path.join(self.base_dir, self.name, "csv") @property def img_dir(self) -> str: """ Returns path to the img directory. :return: Full path to the images directory. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.img_dir 'experiments/test/img' """ return os.path.join(self.base_dir, self.name, "img") @property def reports_dir(self) -> str: """ Returns path to the reports directory. :return: Full path to the reports directory. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.reports_dir 'experiments/test/reports' """ return os.path.join(self.base_dir, self.name, "reports")
[docs] def get_csv_path(self, filename: str = "csv_output.csv") -> str: """ Returns full path for a CSV file. :param filename: Name of the CSV file. :type filename: str :return: Full path to the CSV file. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.get_csv_path("data.csv") 'experiments/test/csv/data.csv' """ return os.path.join(self.csv_dir, filename)
[docs] def get_img_path(self, filename: Optional[str] = None) -> str: """ Returns full path for an image file. :param filename: Name of the image file, if None uses the default from the calling function :type filename: Optional[str] :return: Full path to the image file. :rtype: str **Example:** >>> paths = ExperimentPaths("test") >>> paths.get_img_path("plot.png") 'experiments/test/img/plot.png' >>> # When used in functions, will use that function's default >>> generate_prefix_tree(event_log, paths.get_img_path()) # Will use 'tree.png' >>> generate_dfg(event_log, paths.get_img_path()) # Will use 'dfg.png' """ if filename is None: return self.img_dir return os.path.join(self.img_dir, filename)
def _create_folder_structure(folder_name: str, base_dir: str = "experiments") -> None: """ Create a folder structure with db, img, json, csv, and reports subfolders. :param folder_name: Name of the main folder to be created inside base directory. :type folder_name: str :param base_dir: Base directory where the experiment folder will be created, defaults to "experiments". :type base_dir: str **Example:** >>> _create_folder_structure("test_directory", base_dir="custom_experiments") Successfully created 'custom_experiments/test_directory' with subfolders: db, img, json, csv, reports """ # Create base directory if it doesn't exist if not os.path.exists(base_dir): os.makedirs(base_dir) # Full path for the main folder full_path = os.path.join(base_dir, folder_name) # Check if folder exists if os.path.exists(full_path): raise FileExistsError(f"Error: Folder '{full_path}' already exists!") try: # Create main folder inside base directory os.makedirs(full_path) # Create all subfolders basic_subfolders = ['db', 'img', 'json', 'csv', 'reports'] for subfolder in basic_subfolders: subfolder_path = os.path.join(full_path, subfolder) os.makedirs(subfolder_path) print(f"Successfully created '{full_path}' with subfolders: {', '.join(basic_subfolders)}") except OSError as error: print(f"Error creating folder structure: {error}")
[docs] def create_experiment(name: str, base_dir: str = "experiments") -> ExperimentPaths: """ Main function to set up experiment and return paths. This is the main entry point for creating a new experiment structure. :param name: Name of the experiment to be created. :type name: str :param base_dir: Base directory where the experiment folder will be created, defaults to "experiments". :type base_dir: str :return: ExperimentPaths object containing all relevant paths. :rtype: ExperimentPaths **Example:** >>> paths = create_experiment("my_experiment", base_dir="custom_experiments") Creating new experiment... Successfully created 'custom_experiments/my_experiment' with subfolders: db, img, json, csv, reports Experiment 'my_experiment' created successfully! Database path: custom_experiments/my_experiment/db/my_experiment.sqlite JSON directory: custom_experiments/my_experiment/json CSV directory: custom_experiments/my_experiment/csv Image directory: custom_experiments/my_experiment/img Reports directory: custom_experiments/my_experiment/reports """ print("\nCreating new experiment...") _create_folder_structure(name, base_dir) paths = ExperimentPaths(name, base_dir) print(f"\nExperiment '{paths.name}' created successfully!") print(f"Database path: {paths.database}") print(f"JSON directory: {paths.json_dir}") print(f"CSV directory: {paths.csv_dir}") print(f"Image directory: {paths.img_dir}") print(f"Reports directory: {paths.reports_dir}") return paths