# tommyma3/ai-data-scientist-agent architecture

> Architecture diagram of the GitHub repository tommyma3/ai-data-scientist-agent, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-30). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/tommyma3/ai-data-scientist-agent
- Repository: https://github.com/tommyma3/ai-data-scientist-agent

## Overview

The product is a local-LLM agent that iteratively generates tabular ML solutions, executes them under a timeout, and ranks candidates by validation MSE. The main path runs from CLI configuration through orchestration, then draft/debug/improve generation, sandbox execution and journal updates, with best and good code saved between steps. Dataset previews, model inference, and execution results are distinct stages; the journal and text utilities support the search. Major documented features are represented, though generated submissions are produced by candidate code rather than a sampled repository module. Edges reflect sampled calls or README-documented behavior; unsampled internal details are not inferred.

Main entry points: [`src/ai_data_scientist_agent/runner.py`](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/runner.py) (Run orchestrator), [`src/ai_data_scientist_agent/agent.py`](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/agent.py) (Stage selector), [`src/ai_data_scientist_agent/interpreter.py`](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/interpreter.py) (Sandbox interpreter), [`src/ai_data_scientist_agent/cli.py`](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/cli.py) (CLI configuration)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entry and orchestration"]
  node_cli["CLI configuration<br/>[cli.py]"]
  node_runner["Run orchestrator<br/>[runner.py]"]
  node_config["Run settings<br/>[config.py]"]
end

subgraph group_search["Solution search"]
  node_preview["CSV preview<br/>[data_preview.py]"]
  node_agent["Stage selector<br/>[agent.py]"]
  node_llm["Local LLM client<br/>[llm.py]"]
  node_text["Code extraction<br/>[text_utils.py]"]
end

subgraph group_execution["Execution and evaluation"]
  node_interpreter["Sandbox interpreter<br/>[interpreter.py]"]
end

subgraph group_state["Search state and artifacts"]
  node_journal[("Candidate journal<br/>[models.py]")]
  node_persistence["Solution saver<br/>[persistence.py]"]
end

node_user(("CLI user"))
node_candidate["Generated candidate"]
node_dataset["CSV dataset"]
node_submission["Submission CSV"]
node_artifacts["Saved solutions"]

node_user -->|"sets options"| node_cli
node_cli -->|"builds config"| node_config
node_cli -->|"starts run"| node_runner
node_runner -->|"reads settings"| node_config
node_runner -->|"steps search"| node_agent
node_runner -->|"creates client"| node_llm
node_runner -->|"creates sandbox"| node_interpreter
node_runner -->|"creates journal"| node_journal
node_runner -->|"saves run"| node_persistence
node_agent -->|"builds context"| node_preview
node_preview -->|"reads CSVs"| node_dataset
node_agent -->|"requests generation"| node_llm
node_agent -->|"extracts code"| node_text
node_agent -->|"selects and appends"| node_journal
node_agent -->|"executes code"| node_interpreter
node_interpreter -->|"runs script"| node_candidate
node_candidate -->|"writes predictions"| node_submission
node_interpreter -->|"returns output"| node_agent
node_persistence -->|"selects solutions"| node_journal
node_persistence -->|"writes code"| node_artifacts

click node_cli "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/cli.py"
click node_runner "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/runner.py"
click node_config "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/config.py"
click node_preview "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/data_preview.py"
click node_agent "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/agent.py"
click node_llm "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/llm.py"
click node_text "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/text_utils.py"
click node_interpreter "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/interpreter.py"
click node_journal "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/models.py"
click node_persistence "https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/persistence.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_cli,node_runner,node_config,node_user toneBlue
class node_preview,node_agent,node_llm,node_text toneAmber
class node_interpreter toneMint
class node_journal,node_persistence toneRose
class node_candidate,node_dataset,node_submission,node_artifacts toneIndigo
```

## Components

### Entry and orchestration

- [CLI configuration](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/cli.py): `src/ai_data_scientist_agent/cli.py`
- [Run orchestrator](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/runner.py): `src/ai_data_scientist_agent/runner.py`
- [Run settings](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/config.py): `src/ai_data_scientist_agent/config.py`

### Solution search

- [CSV preview](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/data_preview.py): `src/ai_data_scientist_agent/data_preview.py`
- [Stage selector](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/agent.py): `src/ai_data_scientist_agent/agent.py`
- [Local LLM client](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/llm.py): `src/ai_data_scientist_agent/llm.py`
- [Code extraction](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/text_utils.py): `src/ai_data_scientist_agent/text_utils.py`

### Execution and evaluation

- [Sandbox interpreter](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/interpreter.py): `src/ai_data_scientist_agent/interpreter.py`

### Search state and artifacts

- [Candidate journal](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/models.py): `src/ai_data_scientist_agent/models.py`
- [Solution saver](https://github.com/tommyma3/ai-data-scientist-agent/blob/main/src/ai_data_scientist_agent/persistence.py): `src/ai_data_scientist_agent/persistence.py`

### Other components

- **CLI user**
- **Generated candidate**
- **CSV dataset**
- **Submission CSV**
- **Saved solutions**

## Connections

- CLI user → CLI configuration: sets options (evidence: `README.md`)
- CLI configuration → Run settings: builds config (evidence: `src/ai_data_scientist_agent/cli.py`)
- CLI configuration → Run orchestrator: starts run (evidence: `src/ai_data_scientist_agent/cli.py`)
- Run orchestrator → Run settings: reads settings (evidence: `src/ai_data_scientist_agent/runner.py`)
- Run orchestrator → Stage selector: steps search (evidence: `src/ai_data_scientist_agent/runner.py`)
- Run orchestrator → Local LLM client: creates client (evidence: `src/ai_data_scientist_agent/runner.py`)
- Run orchestrator → Sandbox interpreter: creates sandbox (evidence: `src/ai_data_scientist_agent/runner.py`)
- Run orchestrator → Candidate journal: creates journal (evidence: `src/ai_data_scientist_agent/runner.py`)
- Run orchestrator → Solution saver: saves run (evidence: `src/ai_data_scientist_agent/runner.py`)
- Stage selector → CSV preview: builds context (evidence: `src/ai_data_scientist_agent/agent.py`)
- CSV preview → CSV dataset: reads CSVs (evidence: `src/ai_data_scientist_agent/data_preview.py`)
- Stage selector → Local LLM client: requests generation (evidence: `src/ai_data_scientist_agent/agent.py`)
- Stage selector → Code extraction: extracts code (evidence: `src/ai_data_scientist_agent/agent.py`)
- Stage selector → Candidate journal: selects and appends (evidence: `src/ai_data_scientist_agent/agent.py`)
- Stage selector → Sandbox interpreter: executes code (evidence: `src/ai_data_scientist_agent/agent.py`)
- Sandbox interpreter → Generated candidate: runs script (evidence: `README.md`)
- Generated candidate → Submission CSV: writes predictions (evidence: `README.md`)
- Sandbox interpreter → Stage selector: returns output (evidence: `src/ai_data_scientist_agent/agent.py`)
- Solution saver → Candidate journal: selects solutions (evidence: `src/ai_data_scientist_agent/persistence.py`)
- Solution saver → Saved solutions: writes code (evidence: `src/ai_data_scientist_agent/persistence.py`)