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:
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.
- src.tools.mmore_client.get_mmore_progress_callback()[source]#
Get the current MMORE progress callback.
- exception src.tools.mmore_client.MMOREFileNotFoundError[source]#
Bases:
FileNotFoundErrorException raised when a file is not found for MMORE operations.
- exception src.tools.mmore_client.MMOREDuplicateFileError(file_id, *args, **kwargs)[source]#
Bases:
HTTPErrorRaised when MMORE already holds a file with the requested
fileId.MMORE v2 answers a duplicate upload with
409. Subclassesrequests.HTTPErrorso existing broadexcept requests.HTTPErrorhandlers keep working.
- exception src.tools.mmore_client.MMOREIndexingError[source]#
Bases:
RuntimeErrorRaised when a background indexing job reaches a terminal failure state.
- exception src.tools.mmore_client.MMOREIndexingTimeoutError[source]#
Bases:
TimeoutErrorRaised 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:
objectClient 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 separatemmore raggeneration server, which this project does not run.- Returns:
True if service is reachable, False otherwise
- Return type:
- 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
202and indexes it on a background queue, so the response arrives well before the document is searchable. This method pollsGET /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
fileIdandjobId. Withwait=False, the immediate acceptance payload.- Raises:
MMOREFileNotFoundError – The local file does not exist
MMOREDuplicateFileError – MMORE already holds this fileId (409)
MMOREIndexingError – The indexing job failed
MMOREIndexingTimeoutError – The job did not finish within index_timeout
requests.HTTPError – Any other non-2xx response
- Return type:
- 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 successresultis{"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:
- Returns:
The terminal job snapshot (or the acceptance payload if wait=False)
- Raises:
MMOREFileNotFoundError – The local file does not exist
requests.HTTPError – If MMORE holds no file with this ID (404)
MMOREIndexingError – The indexing job failed
MMOREIndexingTimeoutError – The job did not finish within index_timeout
- Return type:
- 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]
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 (
Nonedisables).Called once by the evaluation harness for
training_code_*runs; the cap is a cost guardrail enforced insidesubmit_slurm_job. Dryrun submissions (the clusterdryrun_*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_idsdo not grow across a long grid run. Callcancel_session_jobsbefore 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
scancelof every full-run job submitted insession_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_Nwandb alias.scancelon 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.validate_hostname(hostname)[source]#
Validate hostname format to prevent injection attacks.
- 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:
- Returns:
Tuple of (success, message)
- Return type:
- 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_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:
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:
ProtocolOperations used by EngiAI regardless of the SSH transport.
- __init__(*args, **kwargs)#
- class src.tools.connection.HPCBrokerConnection(broker_url, token, host_alias='euler', progress_callback=None, timeout=120)[source]#
Bases:
objectHPC connection backed by the native, key-holding host broker.
- class src.tools.connection.HPCConnection(host_alias='', host=None, user=None, password=None, port=22, progress_callback=None)[source]#
Bases:
objectSSH 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)
- test_slurm_script(remote_path)[source]#
Ask SLURM to validate a staged script without submitting it.
- submit_job(slurm_file, remote_dir='~/slurm_jobs')[source]#
Transfer SLURM script and submit job to HPC.
- get_job_final_state(job_id)[source]#
Return a finished job’s terminal
(state, exit_code)viasacct.squeueonly lists pending/running jobs, so a finished job – whether it COMPLETED or FAILED – simply vanishes from it; callers that only pollget_job_statuscannot tell success from failure.sacctrecords the real finalStateandExitCode(e.g. a crashed training run that exits non-zero showsFAILED), which is what failure detection needs.stateis""when sacct has no record yet – accounting can lag a few seconds behindsqueue– so callers should retry briefly on empty.
Search Tools#
Search tools for web search capabilities.
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:
objectCrawls websites and extracts HTML pages for indexing.
- __init__(max_pages=100, max_depth=3, same_domain_only=True, timeout=30)[source]#
Initialize web crawler.
- 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
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_scriptversions 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 theSUPPORTED_ALGORITHMSentry the stem starts with. Bothvalidate_training_script(to pick the contract/family) andgenerate_training_slurm(to pass--algoso the trained wandb artifact is the canonical nameevaluate_modeldownloads) derive the algorithm this way, so the filename stem is the single source of truth across the local + cluster halves of the workflow.Returns
Nonewhen 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- __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- __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- __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:
- 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_idif set, elseproblem_id.The evaluation harness sets
TRAINING_CODE_PROBLEM_IDper 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 (defaultbeams2d), 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 suppliedproblem_idis kept, rather than silently binding the pipeline to a bogus problem.