Marimo Flow 🌊

Like marimo algae drifting in crystal waters, your data flows and evolves – each cell a living sphere of computation, gently touching others, creating ripples of reactive change. In this digital ocean, data streams like currents, models grow like organic formations, and insights emerge naturally from the depths. Let your ML experiments flow freely, tracked and nurtured, as nature intended.

What is marimo-flow? 🚀

An agentic scientific-computing platform for physics-informed ML.
A reactive multi-agent team orchestrates PINA / PINN workflows
end-to-end over a pydantic-graph state machine — backed by MLflow for
tracing and persistence, and exposed through marimo's chat UI plus
optional A2A and AG-UI ASGI servers.

Describe a PDE in natural language; the team composes the Problem,
designs the network, wires the solver, trains it, and grades the run —
every handoff typed, logged to MLflow, and indexed in a DuckDB
provenance store. A classic marimo_flow.core API is still there for
when you'd rather drive PINA by hand. Composition-first throughout: no
hardcoded PDE factories, plus a Docker deployment story for CPU /
NVIDIA / Intel GPUs.

Architecture 🏗️

Natural language goes in; a pydantic-ai multi-agent team turns it into
a composed, trained, tracked PINN. PINA, MLflow, and marimo are the three
backends the team drives:

   Natural language  ("Solve a 1D Poisson equation on [0,1] with u(0)=u(1)=0")
          │
          ▼
 ┌─────────────────────────────────────────────────────────┐
 │  pydantic-ai  ·  multi-agent team                        │
 │  Triage → Route → Problem · Model · Solver · Training     │
 │         → Validation → MLflow · Notebook                 │
 │  (pydantic-graph state machine · typed spec handoffs)    │
 └─────────────────────────────────────────────────────────┘
          │                    │                    │
          ▼                    ▼                    ▼
       ┌──────┐            ┌────────┐           ┌────────┐
       │ PINA │            │ MLflow │           │ marimo │
       └──────┘            └────────┘           └────────┘
    compose + train      runs · metrics       reactive .py
    PINNs (PyTorch)      registry · tracing    notebooks · chat
    Plotly 2D/3D viz     DuckDB provenance     MCP-wired UI

A classic marimo_flow.core API drives PINA by hand when you'd rather
skip the agents — both paths write to the same MLflow backend.

Stack 🧱

🧠 PINA — composition-first (no hardcoded PDE factories)

The core idea: agents emit typed EquationSpec + SubdomainSpec +
ConditionSpec and the composer compiles a pina.Problem subclass at
runtime via sympy.lambdify + pina.operator.grad/laplacian. There
is no ProblemKind enum — any PDE sympy can express is reachable.

Code lives in src/marimo_flow/core/ (PINA
solvers + training + viz3d), src/marimo_flow/control/
(MPC), and src/marimo_flow/agents/services/
(composer, mesh-domain, design aggregator).

Features ✨

• 🧑‍🚀 Multi-agent PINA team — free-form intent → typed TaskSpec
  → composed ProblemSpec / ModelSpec / SolverPlan → trained
  surrogate → validation verdict, over a 9-node pydantic-graph state
  machine (see docs/agents.md).
• 🔌 Three transports — in-notebook mo.ui.chat, an A2A ASGI
  server, and an AG-UI ASGI server, all wrapping the same lead agent.
• 🧩 Provider-agnostic LLMs — per-role model specs resolved through
  pydantic-ai's infer_model; defaults target Ollama Cloud, any
  provider in the catalogue works via config.yaml / env vars.
• 🔬 MLflow tracing + DuckDB provenance — every run, model, and
  metric tracked under data/mlflow/{db,artifacts}/; typed specs,
  decisions, handoffs, and lineage mirrored into a 16-table DuckDB
  store. Lightning checkpoints land inside the active run.
• 🧠 Composition-first PINA — agents compose PDEs from primitives;
  no hardcoded ProblemKind (see below).
• 📓 Reactive notebooks — Git-friendly .py notebooks with
  automatic dependency tracking.
• 🤖 MCP-powered AI — marimo, mlflow, and context7 MCP
  servers wired up; live library docs without leaving your notebook.
• 🐳 Multi-platform Docker — CPU, CUDA, Intel XPU images on GHCR.

Quick Start 🏃‍♂️

Talk to the agent team

The fastest path is the in-notebook chat. The lead agent dispatches to
the specialist team and streams its answer back:

from marimo_flow.agents import lead_chat, FlowDeps
import marimo as mo

deps = FlowDeps()  # resolves per-role models + MLflow URI from config / .env / env vars
chat = mo.ui.chat(
    lead_chat(deps=deps),
    prompts=["Solve a 1D Poisson equation on [0,1] with u(0)=u(1)=0 using a PINN."],
)
chat

Defaults target Ollama Cloud (:cloud tags via a local Ollama
endpoint). Copy config.yaml.example to
config.yaml to point any role at OpenAI, Anthropic, Groq, … — auth
comes from each provider's standard env var.

Prefer a terminal? The CLI runs the same graph:

uv run marimo-flow solve "Solve the Burgers equation with a small PINN"
uv run marimo-flow solve -m lead=anthropic:claude-sonnet-4-6 "..."   # override a role
uv run marimo-flow config-show        # print the resolved models + URIs
uv run marimo-flow lab                # open examples/lab.py in marimo

Or expose the team as a standalone ASGI server:

uv run python -m marimo_flow.agents.server.a2a     # A2A protocol   → :8000
uv run python -m marimo_flow.agents.server.ag_ui   # AG-UI protocol → :8001

With Docker (Recommended)

git clone https://github.com/synapticore-io/marimo-flow.git
cd marimo-flow
docker compose -f docker/docker-compose.yaml up --build -d

# Marimo:  http://localhost:2718
# MLflow:  http://localhost:5000

Image Variants

GPU compose files: docker-compose.cuda.yaml (requires nvidia-docker)
and docker-compose.xpu.yaml (requires Intel GPU drivers).

Local Development

→ See SETUP.md — bare-metal uv path, MLflow + marimo
processes, MCP setup table, troubleshooting. Five minutes from clone
to running notebook.

Example Notebooks 📚

All notebooks live in examples/ and open with
uv run marimo edit examples/<file>.py.

Project Structure 📁

marimo-flow/
├── examples/                     # Demo notebooks
├── src/marimo_flow/
│   ├── core/                     # PINA solvers, training, plotly viz3d
│   ├── control/                  # Rolling-horizon MPC (scipy SLSQP)
│   └── agents/                   # Multi-agent team (pydantic-graph + MLflow)
│       ├── nodes/                # 9 graph nodes
│       ├── schemas/              # Typed Pydantic specs (ProblemSpec, …)
│       ├── toolsets/             # FunctionToolset per role
│       └── services/             # composer, mesh_domain, design,
│                                 #   provenance (DuckDB, 16 tables)
├── tests/                        # 232 passing, 0 xfailed
├── docker/                       # Dockerfiles + compose (CPU/CUDA/XPU)
├── docs/                         # Project documentation (see docs/INDEX.md)
└── data/mlflow/                  # MLflow storage (db + artifacts)

Two Workflows

Both write to the same MLflow backend (data/mlflow/). The two
packages do not depend on each other — pick whichever matches the task.

PINA Multi-Agent Team 🧑‍🚀

The team drives PINA workflows end-to-end — it doesn't just answer
questions. Free-form intent → typed TaskSpec → composed ProblemSpec
→ trained surrogate → validation verdict, all logged to MLflow + a
DuckDB provenance store.

TriageNode → RouteNode ─┬─ ProblemNode
                        ├─ ModelNode
                        ├─ SolverNode
                        ├─ TrainingNode
                        ├─ ValidationNode → (accept/retry/escalate/reject)
                        ├─ MLflowNode
                        └─ NotebookNode

Nine graph nodes (pydantic-graph), one toolset + skill per role,
provider-agnostic LLM config, and three transport options (in-notebook
chat, A2A server, AG-UI server). See the Quick Start
for the chat / CLI / server entry points.

→ Full architecture, role list, provenance schema, and config docs in
docs/agents.md.

MCP Integration 🔌

marimo and AI-assisted IDEs share MCP servers for live documentation
and notebook operations. Pre-configured in .marimo.toml (in-notebook)
and .vscode/mcp.json (VS Code / Claude Code).

→ Full configuration reference in docs/mcp-setup.md.

Claude Code Integration 🤖

marimo-flow is built to be driven from Claude Code.
Two integration points work together:

• marimo MCP server — start marimo with --mcp and Claude Code can
  introspect the running notebook: list cells, read variables/DataFrames,
  surface cell errors, and inspect DuckDB/SQL connections.

  marimo edit examples/ --mcp --no-token --headless --port 2718
  

• marimo-pair — the marimo-claude plugin's pair-programming skill.
  It attaches to that same live kernel so Claude Code can execute code,
  create cells, and build the notebook as an artifact interactively —
  not just read it. Ask Claude to "start a marimo notebook" or "work in
  the active marimo session" and it drives the kernel directly.

Together they turn a marimo notebook into a shared workspace: Claude
reads runtime state over MCP and writes cells back through marimo-pair,
with MLflow + the DuckDB provenance store recording everything the agent
team produces alongside.

Documentation 📚

Contributing 🤝

We welcome contributions! See CONTRIBUTING.md for
development setup, code standards, and the PR process.

# Quick path
git checkout -b my-feature
uv run pytest                         # 232 passing
uv run ruff format . && uv run ruff check --fix .
# open a PR

License 📄

MIT — see LICENSE.

Built with ❤️ using marimo, MLflow, and PINA.