Tools API

Contents

Tools API#

Auto-generated API documentation for tools.

ArXiv Tools#

ArXiv tools for searching and retrieving research papers.

Provides tools for: - Searching ArXiv papers by query - Fetching paper details by ID - Downloading and analyzing papers with MMORE RAG - Listing and querying analyzed papers

src.tools.arxiv_tools.create_download_and_analyze_tool(mmore_client)[source]#

Create a tool for downloading and analyzing ArXiv papers with MMORE.

Parameters:

mmore_client (MMOREClient) – MMORE client instance for uploading papers

Returns:

LangChain tool for downloading and analyzing papers

src.tools.arxiv_tools.create_ask_papers_tool(mmore_client)[source]#

Create a tool for asking questions about analyzed papers.

Parameters:

mmore_client (MMOREClient) – MMORE client instance for retrieving information

Returns:

LangChain tool for querying papers

src.tools.arxiv_tools.create_list_papers_tool(mmore_client)[source]#

Create a tool for listing analyzed papers.

Parameters:

mmore_client (MMOREClient) – MMORE client for accessing the knowledge base

Returns:

LangChain tool for listing papers

src.tools.arxiv_tools.create_arxiv_tools(mmore_client=None)[source]#

Create a list of ArXiv-related tools.

Parameters:

mmore_client (MMOREClient | None) – Optional MMORE client for analysis tools

Returns:

List of LangChain tools for ArXiv operations

Return type:

list

MMORE Client#

MMORE API Client for multimodal document retrieval.

Provides interface to mmore service for advanced RAG capabilities.

Targets the MMORE v2 API. v2 ingests asynchronously: POST /v1/files returns 202 with a job id and indexes on a background queue, rather than blocking until the document is searchable as v1 did. MMOREClient.upload_file() polls that job to completion so callers keep the synchronous semantics they rely on. See services/mmore/frozen-v1/ for the pre-v2 stack.

src.tools.mmore_client.empty_rag_enabled()[source]#

Whether the empty-RAG control condition is active.

In this mode the RAG tools are still created and offered to the agent, but retrieval and listing return nothing, so the benchmark can measure whether the mere availability of a retrieval tool changes agent behavior.

Read at call time, not cached: evaluation harnesses flip the variable after this module is imported.

Note this deliberately makes the client lie about the index. Tooling that needs the truth – anything that decides what to ingest, or asserts that documents are present – must refuse to run while it is set rather than trust the answer. See scripts/build_mmore_index.py.

src.tools.mmore_client.default_min_similarity()[source]#

Score floor for agent-facing document search.

Overridable via MMORE_MIN_SIMILARITY. The floor is worth tuning if an upstream bump shifts the similarity scale: too high and retrieval silently returns nothing while the service looks perfectly healthy.

src.tools.mmore_client.set_mmore_progress_callback(callback)[source]#

Set the progress callback for MMORE operations.

Parameters:

callback (Callable[[str, str], None] | None) – Callback function(step: str, message: str) or None to clear

src.tools.mmore_client.get_mmore_progress_callback()[source]#

Get the current MMORE progress callback.

exception src.tools.mmore_client.MMOREFileNotFoundError[source]#

Bases: FileNotFoundError

Exception raised when a file is not found for MMORE operations.

exception src.tools.mmore_client.MMOREDuplicateFileError(file_id, *args, **kwargs)[source]#

Bases: HTTPError

Raised when MMORE already holds a file with the requested fileId.

MMORE v2 answers a duplicate upload with 409. Subclasses requests.HTTPError so existing broad except requests.HTTPError handlers keep working.

__init__(file_id, *args, **kwargs)[source]#

Initialize RequestException with request and response objects.

exception src.tools.mmore_client.MMOREIndexingError[source]#

Bases: RuntimeError

Raised when a background indexing job reaches a terminal failure state.

exception src.tools.mmore_client.MMOREIndexingTimeoutError[source]#

Bases: TimeoutError

Raised when a background indexing job does not finish in time.

class src.tools.mmore_client.MMOREClient(base_url=None, *, index_timeout=None, poll_interval=None)[source]#

Bases: object

Client for interacting with MMORE RAG service API.

__init__(base_url=None, *, index_timeout=None, poll_interval=None)[source]#

Initialize MMORE client.

Parameters:
  • base_url (str | None) – Base URL of mmore service (defaults to MMORE_RAG_URL env var)

  • index_timeout (float | None) – Seconds to wait for a background indexing job before giving up (defaults to the MMORE_INDEX_TIMEOUT env var, then 900)

  • poll_interval (float | None) – Initial delay between job polls, in seconds. Grows geometrically up to 5s.

health_check()[source]#

Check if mmore service is healthy.

Probes GET /, which is the health route on MMORE’s live-retrieval app. That app does not expose /health – that belongs to the separate mmore rag generation server, which this project does not run.

Returns:

True if service is reachable, False otherwise

Return type:

bool

upload_file(file_path, file_id=None, original_name=None, *, wait=True)[source]#

Upload a file to mmore and, by default, block until it is indexed.

MMORE v2 accepts the upload with 202 and indexes it on a background queue, so the response arrives well before the document is searchable. This method polls GET /v1/jobs/{jobId} to preserve the synchronous “returns implies searchable” contract callers depend on.

Parameters:
  • file_path (str) – Path to the file to upload

  • file_id (str | None) – Optional unique identifier (defaults to filename)

  • original_name (str | None) – Optional original file name for display (uses file_path name if not provided)

  • wait (bool) – If False, return as soon as the upload is accepted, without waiting for indexing to finish.

Returns:

The terminal job snapshot, always carrying at least fileId and jobId. With wait=False, the immediate acceptance payload.

Raises:
Return type:

dict[str, Any]

get_job(job_id)[source]#

Fetch a snapshot of a background indexing job.

Parameters:

job_id (str) – The job identifier returned by an upload

Returns:

{jobId, fileId, filename, status, device, result, error}. On success result is {"chunks": <int>}.

Return type:

Job snapshot

Raises:

requests.HTTPError – If the job is unknown or expired (404). MMORE keeps job records in memory for two hours only.

replace_file(file_path, file_id, original_name=None, *, wait=True)[source]#

Replace an already-indexed file, reusing its fileId.

Uses PUT /v1/files/{fileId}, which avoids the delete-then-upload dance and the 409 that a plain re-upload would raise.

Parameters:
  • file_path (str) – Path to the replacement file

  • file_id (str) – ID of the existing document to replace

  • original_name (str | None) – Optional display name for progress messages

  • wait (bool) – If False, return without waiting for re-indexing to finish

Returns:

The terminal job snapshot (or the acceptance payload if wait=False)

Raises:
Return type:

dict[str, Any]

retrieve(query, file_ids=None, max_matches=5, min_similarity=0.0, collection_name=None)[source]#

Retrieve relevant documents for a query.

Parameters:
  • query (str) – Search query

  • file_ids (list[str] | None) – Optional list of file IDs to restrict search to

  • max_matches (int) – Maximum number of results to return

  • min_similarity (float) – Minimum similarity score threshold. Always sent explicitly, so MMORE’s own default never applies.

  • collection_name (str | None) – Optional Milvus collection override

Returns:

List of LangChain Document objects with retrieved content

Raises:

requests.HTTPError – If retrieval fails

Return type:

list[Document]

delete_file(file_id)[source]#

Delete a file from mmore index.

Parameters:

file_id (str) – ID of the file to delete

Returns:

Response from mmore API

Raises:

requests.HTTPError – If deletion fails

Return type:

dict[str, Any]

list_files(collection_name='my_docs', limit=None)[source]#

List all files stored in the mmore index.

Parameters:
  • collection_name (str) – Name of the Milvus collection (defaults to ‘my_docs’)

  • limit (int | None) – Optional cap on how many entries to return (MMORE defaults to 16000, and rejects anything above 100000)

Returns:

List of file IDs

Raises:

requests.HTTPError – If listing fails

Return type:

list[str]

HPC Tools#

HPC cluster job management tools for SLURM job submission and monitoring.

This module provides tools for interacting with HPC clusters via SSH, including job submission, status monitoring, cancellation, and output retrieval.

src.tools.hpc.set_current_session_id(session_id)[source]#

Set the current session ID for credential isolation.

This should be called by Streamlit before invoking any HPC tools.

Parameters:

session_id (str | None) – Unique session identifier

src.tools.hpc.set_submit_budget(k)[source]#

Set the per-trial cap on full training-job submissions (None disables).

Called once by the evaluation harness for training_code_* runs; the cap is a cost guardrail enforced inside submit_slurm_job. Dryrun submissions (the cluster dryrun_* job) are never counted.

src.tools.hpc.reset_submit_count(session_id)[source]#

Reset (and free) the submitted-jobs counter and job-id list for a session.

The harness calls this at trial start and teardown so each trial gets a fresh budget and _submit_counts / _submit_job_ids do not grow across a long grid run. Call cancel_session_jobs before this if a dead trial’s jobs need cancelling – this clears the recorded ids.

src.tools.hpc.cancel_session_jobs(session_id, host_alias=None)[source]#

Best-effort scancel of every full-run job submitted in session_id.

Called by the eval harness at trial teardown: a trial that died (e.g. hit the agent recursion limit) must not leave PENDING/RUNNING SLURM jobs running unwatched on the cluster, wasting GPU and clobbering the moving seed_N wandb alias. scancel on an already-finished job is a harmless no-op, so this is safe to call after every trial. Never raises – a teardown must not be derailed by an SSH hiccup.

Returns the list of job ids it attempted to cancel (for logging / tests).

src.tools.hpc.set_hpc_progress_callback(callback)[source]#

Set the progress callback for HPC operations.

This should be called by Streamlit before invoking HPC tools.

Parameters:

callback (Any | None) – Callback function(step: str, message: str) or None to clear

src.tools.hpc.get_hpc_progress_callback()[source]#

Get the current HPC progress callback.

src.tools.hpc.validate_hostname(hostname)[source]#

Validate hostname format to prevent injection attacks.

Parameters:

hostname (str) – The hostname to validate

Returns:

True if valid, False otherwise

Return type:

bool

src.tools.hpc.validate_username(username)[source]#

Validate username format.

Parameters:

username (str) – The username to validate

Returns:

True if valid, False otherwise

Return type:

bool

src.tools.hpc.set_ssh_credentials(host=None, user=None, password=None, port=22, session_id=None)[source]#

Set SSH credentials for a specific session.

Parameters:
  • host (str | None) – Hostname (e.g., “hpc.example.com”)

  • user (str | None) – Username

  • password (str | None) – Password

  • port (int) – SSH port (default: 22)

  • session_id (str | None) – Session ID for isolation (uses current if not provided)

Returns:

Tuple of (success, message)

Return type:

tuple[bool, str]

src.tools.hpc.clear_ssh_credentials(session_id=None)[source]#

Clear stored SSH credentials for a session.

Parameters:

session_id (str | None) – Session ID (uses current if not provided)

src.tools.hpc.get_ssh_credentials(session_id=None)[source]#

Get SSH credentials for a session.

Parameters:

session_id (str | None) – Session ID (uses current if not provided)

Returns:

Credentials dict or None if not set/expired

Return type:

dict[str, Any] | None

src.tools.hpc.get_ssh_credentials_secure(session_id=None)[source]#

Securely get SSH credentials and clear from memory after use.

This context manager ensures the decrypted password is cleared from the returned dictionary after the block exits.

Parameters:

session_id (str | None) – Session ID (uses current if not provided)

Yields:

dict or None – Credentials dict or None if not set/expired

Example:

with get_ssh_credentials_secure() as creds:
    if creds:
        connection = HPCConnection(..., password=creds["password"])
# Password is cleared from creds after this point
src.tools.hpc.cleanup_expired_credentials()[source]#

Remove all expired credentials from storage.

Returns:

Number of expired sessions cleaned up

Return type:

int

src.tools.hpc.create_hpc_connection(host_alias)[source]#

Create an HPC connection using session credentials or SSH config.

Parameters:

host_alias (str) – Host alias for SSH config fallback

Returns:

Direct SSH or managed-broker connection instance

Return type:

HPCConnectionLike

CLI Tools#

CLI command execution tools for running local command-line applications.

Connection Tools#

SSH connection helper for transferring SLURM scripts to HPC and submitting jobs.

Uses Fabric (https://www.fabfile.org/) for robust SSH connection management. This module reads your SSH config file and provides utilities to: 1. Transfer generated SLURM scripts to a remote HPC cluster 2. Submit jobs via sbatch 3. Monitor job status 4. Download results

Usage:

from src.tools.connection import HPCConnection

# Initialize connection (reads ~/.ssh/config) hpc = HPCConnection(host_alias=”mycluster”)

# Copy SLURM script and submit job job_id = hpc.submit_job(slurm_file=”outputs/train_cgan_cnn_2d_beams2d_seed1.slurm”) print(f”Job submitted with ID: {job_id}”)

# Check job status status = hpc.get_job_status(job_id) print(status)

class src.tools.connection.HPCConnectionLike(*args, **kwargs)[source]#

Bases: Protocol

Operations used by EngiAI regardless of the SSH transport.

host_alias: str#
auth_mode: str#
get_remote_directory()[source]#
ensure_remote_dir(remote_dir)[source]#
put_file(local_file, remote_path)[source]#
get_file(remote_file, local_path)[source]#
test_slurm_script(remote_path)[source]#
remove_remote_file(remote_path)[source]#
submit_job(slurm_file, remote_dir='~/slurm_jobs')[source]#
get_job_status(job_id)[source]#
get_job_final_state(job_id)[source]#
cancel_job(job_id)[source]#
cancel_jobs(job_ids)[source]#
get_job_output(job_id, remote_dir='~/slurm_jobs', local_dir='outputs')[source]#
get_job_diagnostics(job_id, remote_dir='~/slurm_jobs')[source]#
list_active_jobs()[source]#
close()[source]#
__init__(*args, **kwargs)#
class src.tools.connection.HPCBrokerConnection(broker_url, token, host_alias='euler', progress_callback=None, timeout=120)[source]#

Bases: object

HPC connection backed by the native, key-holding host broker.

__init__(broker_url, token, host_alias='euler', progress_callback=None, timeout=120)[source]#
get_remote_directory()[source]#
ensure_remote_dir(remote_dir)[source]#
put_file(local_file, remote_path)[source]#
get_file(remote_file, local_path)[source]#
test_slurm_script(remote_path)[source]#
remove_remote_file(remote_path)[source]#
submit_job(slurm_file, remote_dir='~/slurm_jobs')[source]#
get_job_status(job_id)[source]#
get_job_final_state(job_id)[source]#
cancel_job(job_id)[source]#
cancel_jobs(job_ids)[source]#
get_job_output(job_id, remote_dir='~/slurm_jobs', local_dir='outputs')[source]#
get_job_diagnostics(job_id, remote_dir='~/slurm_jobs')[source]#
list_active_jobs()[source]#
close()[source]#

The broker owns SSH process lifetimes; there is no client socket to close.

class src.tools.connection.HPCConnection(host_alias='', host=None, user=None, password=None, port=22, progress_callback=None)[source]#

Bases: object

SSH connection to HPC cluster for job submission using Fabric.

__init__(host_alias='', host=None, user=None, password=None, port=22, progress_callback=None)[source]#

Initialize HPC connection.

Supports two authentication modes: 1. SSH config mode (default): Uses ~/.ssh/config with SSH keys 2. Password mode: Uses explicit username/password credentials

Parameters:
  • host_alias (str) – Host alias from ~/.ssh/config (e.g., “mycluster”)

  • host (str | None) – Explicit hostname (e.g., “hpc.example.com”). If provided, uses password auth.

  • user (str | None) – Username for password authentication

  • password (str | None) – Password for authentication

  • port (int) – SSH port (default: 22)

  • progress_callback (Callable[[str, str], None] | None) – Optional callback for progress updates (step, message)

run_command(command)[source]#

Run a command over SSH.

Parameters:

command (str) – Command to execute on remote

Returns:

Command output

Return type:

str

get_remote_directory()[source]#

Return the Euler working directory for a connection smoke test.

ensure_remote_dir(remote_dir)[source]#

Create an allowed working directory on the remote host.

test_slurm_script(remote_path)[source]#

Ask SLURM to validate a staged script without submitting it.

remove_remote_file(remote_path)[source]#

Remove a staged remote file.

put_file(local_file, remote_path)[source]#

Copy a file to remote using SFTP.

Parameters:
  • local_file (str) – Path to local file

  • remote_path (str) – Destination path on remote (can be relative to home)

get_file(remote_file, local_path)[source]#

Download a file from remote using SFTP.

Parameters:
  • remote_file (str) – Path to remote file

  • local_path (str) – Destination path on local

submit_job(slurm_file, remote_dir='~/slurm_jobs')[source]#

Transfer SLURM script and submit job to HPC.

Parameters:
  • slurm_file (str) – Path to local SLURM script

  • remote_dir (str) – Directory on remote to store scripts (default: ~/slurm_jobs)

Returns:

Job ID returned by sbatch

Return type:

str

close()[source]#

Close the underlying SSH connection, releasing the socket.

get_job_status(job_id)[source]#

Get status of a submitted job.

Parameters:

job_id (str) – SLURM job ID

Returns:

Job status output

Return type:

str

get_job_final_state(job_id)[source]#

Return a finished job’s terminal (state, exit_code) via sacct.

squeue only lists pending/running jobs, so a finished job – whether it COMPLETED or FAILED – simply vanishes from it; callers that only poll get_job_status cannot tell success from failure. sacct records the real final State and ExitCode (e.g. a crashed training run that exits non-zero shows FAILED), which is what failure detection needs.

state is "" when sacct has no record yet – accounting can lag a few seconds behind squeue – so callers should retry briefly on empty.

cancel_job(job_id)[source]#

Cancel a running job.

Parameters:

job_id (str) – SLURM job ID

cancel_jobs(job_ids)[source]#

Cancel one or more jobs in one SSH round-trip.

list_active_jobs()[source]#

List active jobs owned by the SSH account.

get_job_diagnostics(job_id, remote_dir='~/slurm_jobs')[source]#

Return the accounting line and recent stdout/stderr for a job.

get_job_output(job_id, remote_dir='~/slurm_jobs', local_dir='outputs')[source]#

Download job output files.

Parameters:
  • job_id (str) – SLURM job ID

  • remote_dir (str) – Remote directory where SLURM outputs are stored

  • local_dir (str) – Local directory to download to

Search Tools#

Search tools for web search capabilities.

class src.tools.search.WebSearchInput(*, query)[source]#

Bases: BaseModel

Input accepted by the provider-neutral web-search tool.

query: str#
model_config = {}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

src.tools.search.create_search_tool(max_results=5)[source]#

Create a web-search tool using the configured provider.

Parameters:

max_results (int) – Maximum number of search results to return

Returns:

Provider-neutral LangChain search tool

Return type:

BaseTool

Web Crawler#

Web crawler for extracting multiple pages from documentation sites.

This module crawls websites starting from a root URL and extracts all pages within the same domain. Each page is then uploaded to MMORE for HTML processing.

class src.tools.web_crawler.WebCrawler(max_pages=100, max_depth=3, same_domain_only=True, timeout=30)[source]#

Bases: object

Crawls websites and extracts HTML pages for indexing.

__init__(max_pages=100, max_depth=3, same_domain_only=True, timeout=30)[source]#

Initialize web crawler.

Parameters:
  • max_pages (int) – Maximum number of pages to crawl

  • max_depth (int) – Maximum depth to crawl from root URL

  • same_domain_only (bool) – Only crawl pages on the same domain

  • timeout (int) – HTTP request timeout in seconds

normalize_url(url)[source]#

Normalize URL by removing fragments and trailing slashes.

Parameters:

url (str) – URL to normalize

Returns:

Normalized URL

Return type:

str

is_same_domain(url1, url2)[source]#

Check if two URLs are on the same domain.

Parameters:
  • url1 (str) – First URL

  • url2 (str) – Second URL

Returns:

True if same domain, False otherwise

Return type:

bool

is_valid_url(url, base_domain)[source]#

Check if URL should be crawled.

Parameters:
  • url (str) – URL to check

  • base_domain (str) – Base domain for crawling

Returns:

True if URL should be crawled, False otherwise

Return type:

bool

Extract all valid links from HTML content.

Parameters:
  • html_content (str) – HTML content as string

  • base_url (str) – Base URL for resolving relative links

Returns:

List of absolute URLs

Return type:

list[str]

crawl(start_url)[source]#

Crawl website starting from a URL.

Yields pages as they are crawled for immediate processing.

Parameters:

start_url (str) – URL to start crawling from

Yields:

Dictionary with page information – - url: Page URL - content: HTML content - title: Page title (if available) - depth: Crawl depth

get_stats()[source]#

Get crawling statistics.

Returns:

Dictionary with crawling stats

Return type:

dict[str, Any]

Algorithms#

Algorithm type definitions for EngiOpt integration.

This is the SINGLE FILE to update when adding a new algorithm type: 1. Add to AlgorithmId Literal type 2. Add to SUPPORTED_ALGORITHMS mapping with properties

Everything else (algorithm lists, registries) is derived automatically.

src.tools.algorithms.canonical_algorithm(script_stem)[source]#

Map a candidate script’s filename stem to its canonical engiopt algorithm.

write_training_script versions candidate filenames (<algo>_candidate_<timestamp>_…), so a stem is not itself a valid engiopt algorithm id. The prompt instructs the agent to name candidates <algo>_candidate.py, so the canonical id is the SUPPORTED_ALGORITHMS entry the stem starts with. Both validate_training_script (to pick the contract/family) and generate_training_slurm (to pass --algo so the trained wandb artifact is the canonical name evaluate_model downloads) derive the algorithm this way, so the filename stem is the single source of truth across the local + cluster halves of the workflow.

Returns None when no supported algorithm matches; callers decide whether to fall back to a default or surface an actionable error.

EngiOpt Tools#

WandB model download tool for engineering design models.

This tool allows downloading pre-trained generative models from WandB (Weights & Biases) that can be used for inverse design tasks. The models are trained using the engiopt library and are available on the engibench WandB project.

class src.tools.engiopt.HPCInputs(algorithm: str, problem_id: str, epochs: int, slurm_gpus: str, slurm_time: str, slurm_cpus_per_task: str, slurm_mem_per_cpu: str)[source]#

Bases: object

algorithm: str#
problem_id: str#
epochs: int#
slurm_gpus: str#
slurm_time: str#
slurm_cpus_per_task: str#
slurm_mem_per_cpu: str#
__init__(algorithm, problem_id, epochs, slurm_gpus, slurm_time, slurm_cpus_per_task, slurm_mem_per_cpu)#
class src.tools.engiopt.HPCContext(algorithm: str, problem_id: str, epochs: int, gpu_count: int, total_hours: float, cpu_count: int, mem_value: float)[source]#

Bases: object

algorithm: str#
problem_id: str#
epochs: int#
gpu_count: int#
total_hours: float#
cpu_count: int#
mem_value: float#
__init__(algorithm, problem_id, epochs, gpu_count, total_hours, cpu_count, mem_value)#
class src.tools.engiopt.TrainingConfig(algorithm: str = 'cgan_cnn_2d', epochs: int = 200, seed: int = 1, wandb_entity: str | None = None, problem_id: Literal['beams2d', 'photonics2d'] = 'beams2d', gpus: int | None = None, time_hours: float | None = None)[source]#

Bases: object

algorithm: str = 'cgan_cnn_2d'#
epochs: int = 200#
seed: int = 1#
wandb_entity: str | None = None#
problem_id: Literal['beams2d', 'photonics2d'] = 'beams2d'#
gpus: int | None = None#
time_hours: float | None = None#
__init__(algorithm='cgan_cnn_2d', epochs=200, seed=1, wandb_entity=None, problem_id='beams2d', gpus=None, time_hours=None)#

EngiBench Tools#

EngiBench tools for engineering design optimization.

EngiBench (https://engibench.ethz.ch) is a library for benchmarking engineering design problems. It provides access to various optimization problems like beam design, airfoils, truss structures, etc.

These tools allow LLM agents to interact with EngiBench simulators.

src.tools.engibench.get_problem_class(problem_type)[source]#

Get the problem class for a given problem type.

src.tools.engibench.set_session_id(session_id)[source]#

Set the session ID for state isolation in batch evaluations.

Parameters:

session_id (str) – Unique identifier for this evaluation session/example

Example

>>> set_session_id("example_42")  # In evaluation framework
>>> optimize_design(...)  # State is isolated to this session
src.tools.engibench.get_current_session_id()[source]#

Get the current session ID, or ‘default’ if not set.

Returns:

Session ID string (defaults to ‘default’ for interactive use)

Return type:

str

src.tools.engibench.get_problem_state(problem_type)[source]#

Get the state dictionary for a given problem type in the current session.

This function is session-aware and prevents state leakage between evaluation examples. Each session gets its own isolated state dictionary for storing designs only.

Parameters:

problem_type (str) – Type of problem (e.g., ‘beams2d’)

Returns:

last_design, initial_design

Return type:

State dictionary with keys

src.tools.engibench.get_unified_last_design(problem_type)[source]#

Get the last design for a given problem type.

src.tools.engibench.set_unified_last_design(problem_type, design)[source]#

Store the last design for a given problem type.

src.tools.engibench.get_initial_design(problem_type)[source]#

Get the initial design (before optimization) for a given problem type.

src.tools.engibench.set_initial_design(problem_type, design)[source]#

Store the initial design (before optimization) for a given problem type.

src.tools.engibench.clear_session_state(session_id=None)[source]#

Clear state for a specific session to prevent memory leaks.

Parameters:

session_id (str | None) – Session ID to clear, or None to clear current session

Example

>>> clear_session_state("example_42")  # Clear specific session
>>> clear_session_state()  # Clear current session

STL Export#

STL export utilities for beam design arrays.

Converts 2D topology designs (numpy arrays) into watertight 3D STL meshes using xeltofab’s marching-cubes extruder. Suitable for 3D printing or CAD.

Job Monitor#

Job monitoring and notification system for HPC SLURM jobs.

This module provides tools for monitoring job status and getting notifications when jobs complete. It supports: 1. Email notifications via SLURM (configured in job script) 2. Polling-based monitoring with status change detection 3. Background monitoring for long-running jobs

Problem Registry#

Problem type definitions and registry for EngiBench integration.

This is the SINGLE FILE to update when adding a new problem type: 1. Add to ProblemId Literal type 2. Add the engibench import path to _PROBLEM_IMPORTS

Everything else (SUPPORTED_PROBLEMS, PROBLEM_CLASSES, the Beams2D / Photonics2D re-exports) is derived automatically.

The engibench problem classes are imported lazily, on first lookup. Importing engibench.problems.beams2d.v1 eagerly costs ~1.7 s (it drags in seaborn -> scipy.stats -> IPython/ipywidgets, and datasets -> pandas/pyarrow), and this module sits on the import path of every agent via src.utils.prompts -> src.tools – including agents that never touch engibench (arxiv, rag, search, cli). Deferring the class import keeps SUPPORTED_PROBLEMS (derived from the Literal, not from the classes) free.

src.tools.problems.harness_problem_id(problem_id)[source]#

Return the harness-injected problem_id if set, else problem_id.

The evaluation harness sets TRAINING_CODE_PROBLEM_ID per training-code run (from the registered benchmark, e.g. training_code_photonics2d -> photonics2d) so the whole pipeline – SLURM training, the local dry-run, contract validation, and model evaluation – targets that problem without the prompt having to name it. Unset -> the tool/agent-supplied value (default beams2d), so behaviour is unchanged when the harness does not set it.

The override is stripped and validated against SUPPORTED_PROBLEMS; an empty/whitespace or unsupported value (e.g. a typo in a manual export) is ignored with a warning and the supplied problem_id is kept, rather than silently binding the pipeline to a bogus problem.