Skip to content

architecture.md

Latest commit

 

History

History
342 lines (266 loc) · 14.7 KB

architecture.md

File metadata and controls

342 lines (266 loc) · 14.7 KB
Raw

AgentML Architecture

System Overview

+------------------+       WebSocket        +------------------+
|                  | <--------------------> |                  |
|   React Frontend |       HTTP REST        |  FastAPI Backend  |
|   (Zustand Store)| <--------------------> |   (uvicorn)      |
|                  |                        |                  |
+------------------+                        +--------+---------+
                                                     |
                                  +------------------+------------------+
                                  |                  |                  |
                           +------+------+    +------+------+   +------+------+
                           |  LangGraph  |    |   Jupyter   |   |   SQLite    |
                           |  Agent      |    |   Kernel    |   |   Database  |
                           |  Pipeline   |    | (subprocess |   |             |
                           |             |    |  isolated)  |   |             |
                           +------+------+    +------+------+   +-------------+
                                  |                  |
                           +------+------+           |
                           |  LLM API   |    execute_code()
                           | (OpenAI /  |           |
                           |  Anthropic)|           v
                           +-------------+    Python execution
                                              (pandas, sklearn,
                                               matplotlib, etc.)

Agent Graph

The core pipeline is a LangGraph StateGraph compiled without a checkpointer. Session state is managed by our own session_manager (not LangGraph's MemorySaver) to avoid astream_events replaying all prior events on each call.

Graph Flow

START
  |
  v
entry_router (zero-cost, no LLM call)
  |
  +-- Chat mode --> general_agent --> tools --> general_agent --> END
  |                       |
  |                       +-- <stage> directive in autonomous mode
  |                               |
  |                               v
  |                         start_pipeline (extract problem_statement, init stage_history)
  |                               |
  +-- Pipeline mode -----> orchestrator
                               |
                               v
                    orchestrator_should_continue
                        |              |
                        v              v
                      route           END (complete or >10 calls)
                        |
                        v
                  [preprocessing | eda | feature_engineering | modeling | evaluation]
                        |
                        v
                  should_continue_agent
                     |           |
                     v           v
                   tools    hungry_review --> orchestrator (loop)
                     |
                     v
                  route_after_tools --> back to calling agent

Entry Router (_entry_router)

Zero-cost routing with no LLM call:

  • Chat mode: Always routes to the general agent. The user controls every action.
  • Autonomous mode: Routes to the orchestrator if mid-pipeline (current stage is a pipeline stage), otherwise to the general agent.

Orchestrator

Decides which pipeline stage to run next based on <stage> directives in its LLM response. Key behaviors:

  • Forward navigation: Follows the default order: preprocessing > eda > feature_engineering > modeling > evaluation > complete
  • Backward navigation: Can route to any earlier stage when evidence warrants it (poor metrics, data issues discovered later). Unmarking stages from the target onward in the ledger.
  • Iterative improvement loop: After evaluation, performance thresholds trigger automatic revisits:
    • Good (>85%): Route to complete
    • Moderate (65-85%): Try hyperparameter tuning, then different features
    • Poor (<65%): Route back to feature engineering, then preprocessing
    • Max 2 backward iterations per stage
  • "Now I Know Better" restart: After completing the full pipeline, the orchestrator can propose restarting with insights gained. This triggers a HITL gate for user approval.
  • Stage visit history: Tracks visit counts and sequence for backward navigation awareness
  • Completion guards: Multiple early-exit checks prevent the orchestrator from looping after <stage>complete</stage>

Stage Agents

Each pipeline stage has a dedicated agent with a specialized system prompt. All agents share:

  • BASE_INSTRUCTIONS: Output format rules (thinking/decision tags), error handling, quantitative output requirements
  • Dataset info and problem statement injection
  • Execution state from the ledger (variables, stage summaries, findings from previous stages)
  • Error loop detection with recovery guidance

Stage prompts (in agent/prompts/): preprocessing.py, eda.py, feature.py, modeling.py (A/B debate specialists), evaluation.py

General Agent

Handles all non-pipeline interaction:

  • In Chat mode: Runs code, answers questions, explores data. Never triggers the pipeline.
  • In Autonomous mode: Understands data thoroughly, asks clarifying questions about the target variable and problem type, then outputs <problem_statement> and <stage> to trigger the pipeline via the start_pipeline node.

Hungry Agent (Quality Review Gate)

A review-only agent (no tools) that fires between pipeline stages and the orchestrator. Key design:

  • Zero-cost pass-through for early stages (preprocessing, eda) — no LLM call
  • Active review for feature_engineering, modeling, and evaluation
  • Cross-references stage findings from previous stages to catch missed opportunities (e.g., EDA found recommended interactions but Feature Engineering didn't create them)
  • Adds tagged review feedback ([HUNGRY AGENT REVIEW]) to the message stream; the orchestrator can act on it
  • If the review LLM call fails, the pipeline continues uninterrupted

Tools Node

A LangGraph ToolNode that executes the execute_code tool. After execution, route_after_tools routes back to the calling agent (tracked via current_stage in state).

Human-in-the-Loop (HITL)

Architecture (agent/hitl.py, api/routes/ws.py)

HITL gates pause the agent pipeline and wait for user input via WebSocket:

  1. The orchestrator reaches a stage transition in autonomous mode
  2. check_hitl_gate() emits an hitl_request event with question and options
  3. wait_for_hitl() blocks on an asyncio.Event (up to 10 minutes)
  4. The frontend shows the HITL prompt; the user responds via hitl_response WebSocket message
  5. set_hitl_response() stores the response and triggers the event
  6. The orchestrator reads the response: approve (continue), modify (redo with feedback), or reject (redo stage)

Default Gates

Gate Triggered After Question
preprocessing_complete Preprocessing Ready to explore with EDA?
eda_complete EDA Proceed with feature engineering?
feature_engineering_complete Feature Engineering Ready to train models?
modeling_complete Modeling Proceed to evaluation?

Execution Engine

Kernel Process Isolation (executor/notebook.py)

FastAPI server  <-- mp.Queue (pipe) -->  _kernel_worker (child process)
                                              |  ZMQ
                                              v
                                         IPython kernel (grandchild)

The KernelManager and ZMQ client run in a child process so that libzmq C-level crashes only kill the child, not the FastAPI server. Communication uses multiprocessing.Queue (pipe), not ZMQ, between the server and the executor subprocess.

If the child crashes, the server detects it and spawns a new one. The ExecutionLedger replays variable creation code to recover state.

Error Loop Detection (common.py)

Counted on raw messages before windowing:

  • 2 consecutive errors: Recovery guidance injected into the system prompt ("try a fundamentally DIFFERENT approach")
  • 3 consecutive errors: Force return to orchestrator with failure summary
  • Kernel restarts: Tracked separately. 2 consecutive kernel restarts stop retrying.
  • Kernel fatal errors: Immediate stop, suggest new session

Message History Windowing (common.py)

Prevents token explosion on long conversations:

  • Triggers when message count exceeds 30
  • Keeps first 2 messages (original context) + last 20 messages (recent work)
  • Inserts a compact summary of dropped messages in between
  • Cut point adjusted to never split tool_call/ToolMessage pairs
  • After windowing, _sanitize_messages() fixes orphaned tool calls/results

Stage Findings and Ledger

ExecutionLedger (agent/ledger.py)

Tracks execution state outside the kernel:

  • ArtifactRecord: variable name, type, stage, creation code, shape
  • Stage summaries: Compact text summaries of completed stages
  • Completed stages: Set of stages that have been completed
  • Stage findings: Structured findings extracted from agent output (see below)

Stage Findings Extraction (agent/parser.py)

When agents output <stage_findings> tags, the content is extracted and stored in the ledger. These findings are:

  1. Injected into subsequent agent system prompts (so later stages know what earlier stages discovered)
  2. Cross-referenced by the Hungry Agent during quality reviews
  3. Persisted via ledger_json in SQLite

Ledger Context Injection

Every agent receives the current execution state in its system prompt:

  • Known variables with types and shapes
  • Stage summaries from completed stages
  • Findings from previous stages (for pipeline agents only)

Chat vs Agent Mode

Chat Mode (session.agent_mode == "chat")

  • Entry router always sends to the general agent
  • General agent's <stage> directives are ignored (no pipeline trigger)
  • User manually asks for EDA, model training, etc.
  • Mode is set from the frontend toggle and sent via the WS mode field

Agent Mode (session.agent_mode == "autonomous")

  • Entry router sends to orchestrator if mid-pipeline
  • General agent first understands data, asks clarifying questions, then triggers pipeline
  • HITL gates active at stage transitions
  • Full pipeline automation with human oversight

Session Management (utils/session_manager.py)

ChatSession

Holds per-session state:

  • executor: Jupyter kernel reference
  • graph: Compiled LangGraph instance
  • ledger: ExecutionLedger
  • message_history: Full LangChain message list
  • current_stage, stage_history, problem_statement
  • agent_mode: "chat" or "autonomous"
  • read_only: True when the session's CSV is no longer on disk
  • _emitter: EventEmitter reference for HITL access in graph nodes

Concurrency

  • Per-session asyncio.Lock prevents concurrent agent runs
  • If a user sends a message while the agent is processing, they get an error message instead of a race condition

Read-Only Sessions

When a session is restored from SQLite but its CSV file is not found on disk:

  • Session is marked read_only = True
  • A lightweight agent (run_readonly_agent) handles conversation — no tools, no code execution
  • The agent can discuss previous findings, explain decisions, and suggest next steps
  • If the user wants active analysis, they're told to re-upload their dataset

Persistence

  • Sessions persisted to SQLite after each agent run
  • Ledger serialized as JSON in the ledger_json column
  • On backend restart, sessions restored from SQLite with glob-based CSV file lookup

Database (db/database.py)

SQLite with WAL mode, context manager with retry:

Table Columns Purpose
sessions id, name, file_name, dataset_info, stage, ledger_json, timestamps Session metadata
messages session_id, role, parts_json, timestamp Conversation history
experiments session_id, model_name, metrics_json, stage Model training results

WebSocket (api/routes/ws.py)

Real-time streaming of agent events:

  • Heartbeat: 30-second ping to keep connections alive during long agent runs
  • Concurrency guard: Rejects new messages if the agent is already processing
  • Stage reset on reconnect: Prevents auto-resuming a pipeline when switching sessions
  • Agent task lifecycle: Background asyncio.Task for agent runs, cancelled on disconnect
  • HITL coordination: _hitl_events dict maps session IDs to asyncio.Event + response data

Frontend State Management

Zustand Store (store/chatStore.js)

State Purpose
settings LLM provider, API key, mode, persona (persisted to localStorage)
sessions Sidebar session list
sessionId, fileName, datasetInfo Current session info
currentStage, stageHistory Pipeline progress
messages Completed message list
streamingParts, isStreaming In-progress agent response
hitlPending Active HITL prompt
wsStatus WebSocket connection state

WebSocket Connection

  • Single connection per session with connection guard (skip if already CONNECTING or OPEN)
  • Automatic reconnection on unexpected disconnect (3-second delay)
  • Falls back to legacy HTTP chat if WebSocket connection fails

Data Flow

Upload Flow

User drops CSV
  --> POST /api/upload
  --> Validate (extension, size, parseable)
  --> Save to disk: {UPLOAD_DIR}/{session_id}_{filename}
  --> Create session (new kernel)
  --> Load dataset: df = pd.read_csv(...)
  --> Build dataset_info summary
  --> Persist to SQLite
  --> Return session_id to frontend

Chat Flow (No Upload)

User types message without session
  --> POST /api/sessions (create empty session)
  --> WebSocket connects to /ws/{session_id}
  --> General agent handles conversation (no kernel, no dataset)

Agent Run Flow

User sends message via WebSocket
  --> Save message to DB
  --> Build/reuse LangGraph graph
  --> Stream events via astream_events(version="v2")
      --> entry_router decides: general or orchestrator
      --> agent generates reasoning + code
      --> execute_code tool runs in subprocess-isolated kernel
      --> introspect variables, update ledger
      --> emit events (thinking, code, plots, tables)
      --> hungry_review checks quality (for late stages)
      --> orchestrator routes to next stage or END
  --> Save agent response to DB
  --> Persist session (including ledger)

Error Recovery Flow

Code execution fails (ZMQ error)
  --> Kernel restart via restart_kernel()
  --> Dataset (df) auto-reloaded
  --> KERNEL_RESTART signal returned to tool layer
  --> Tool replays artifacts from ledger
  --> Reports recovered/lost variables to agent
  --> Agent adapts with full awareness of state