datascience-agent 0.8.0

AI Agent with dynamic planning and persistent Jupyter kernel execution for data analysis

DSAgent

An AI-powered autonomous agent for data science with persistent Jupyter kernel execution, session management, and conversational interface.

    ____  _____  ___                    __
   / __ \/ ___/ /   | ____ ____  ____  / /_
  / / / /\__ \ / /| |/ __ `/ _ \/ __ \/ __/
 / /_/ /___/ // ___ / /_/ /  __/ / / / /_
/_____//____//_/  |_\__, /\___/_/ /_/\__/
                   /____/

Features

• Conversational Interface: Interactive chat with persistent context and sessions
• Dynamic Planning: Agent creates and follows plans with step tracking
• Persistent Execution: Code runs in a Jupyter kernel with variable persistence across messages
• Session Management: Save and resume conversations with full kernel state
• Multi-Provider LLM: Supports OpenAI, Anthropic, Google, Ollama via LiteLLM
• MCP Tools: Connect to external tools (web search, databases, etc.) via Model Context Protocol
• Human-in-the-Loop: Configurable checkpoints for plan and code approval
• Notebook Generation: Automatically generates clean, runnable Jupyter notebooks

Installation

pip install datascience-agent

With optional features:

pip install "datascience-agent[api]"   # FastAPI server support
pip install "datascience-agent[mcp]"   # MCP tools support

For development:

git clone https://github.com/nmlemus/dsagent
cd dsagent
uv sync --all-extras

Docker

# Run API server
docker run -d -p 8000:8000 \
  -e OPENAI_API_KEY=sk-your-key \
  nmlemus/dsagent:latest

# Run interactive CLI
docker run -it \
  -e OPENAI_API_KEY=sk-your-key \
  nmlemus/dsagent:latest \
  dsagent chat

For Docker deployment details, see docs/DOCKER.md.

Quick Start

1. Setup (First Time)

Run the setup wizard to configure your LLM provider:

dsagent init

This will:

• Ask for your LLM provider (OpenAI, Anthropic, Google, local, etc.)
• Store your API key securely in ~/.dsagent/.env
• Automatically select a default model based on provider:
  • OpenAI → gpt-4o
  • Anthropic → claude-sonnet-4-5
  • Google → gemini/gemini-2.5-flash
  • Local → ollama/llama3
• Optionally configure MCP tools (web search, etc.)

To use a different model, edit ~/.dsagent/.env or use the --model flag:

dsagent --model gpt-4o-mini

2. Start Chatting

dsagent

This starts an interactive session where you can:

• Chat naturally with the agent
• Execute Python code with persistent variables
• Analyze data files
• Generate visualizations
• Resume previous sessions

3. One-Shot Tasks

For batch processing or scripts:

dsagent run "Analyze sales trends" --data ./sales.csv

CLI Commands

Command	Description
dsagent	Start interactive chat (default)
dsagent chat	Same as above, with explicit options
dsagent run "task"	Execute a one-shot task
dsagent init	Setup wizard for configuration
dsagent mcp list	List configured MCP servers
dsagent mcp add <template>	Add an MCP server

Examples

# Interactive chat with specific model
dsagent --model claude-sonnet-4-5

# One-shot analysis
dsagent run "Find patterns in this data" --data ./dataset.csv

# Resume a previous session
dsagent --session abc123

# With MCP tools (web search)
dsagent --mcp-config ~/.dsagent/mcp.yaml

# Human-in-the-loop mode
dsagent --hitl plan

For complete CLI documentation, see docs/CLI.md.

Python API

DSAgent provides two agents for different use cases:

ConversationalAgent (Interactive)

For building chat interfaces and interactive applications:

from dsagent import ConversationalAgent, ConversationalAgentConfig

config = ConversationalAgentConfig(model="gpt-4o")
agent = ConversationalAgent(config)
agent.start()

# Chat with persistent context
response = agent.chat("Load the iris dataset")
print(response.content)

response = agent.chat("Train a classifier on it")
print(response.content)  # Has access to previous variables

agent.shutdown()

PlannerAgent (Batch)

For one-shot tasks and automated pipelines:

from dsagent import PlannerAgent

with PlannerAgent(model="gpt-4o", data="./data.csv") as agent:
    result = agent.run("Analyze this dataset and create visualizations")
    print(result.answer)
    print(f"Notebook: {result.notebook_path}")

For complete API documentation, see docs/PYTHON_API.md.

Supported Models

DSAgent uses LiteLLM to support 100+ LLM providers:

Provider	Models	API Key
OpenAI	gpt-4o, o1, o3-mini	OPENAI_API_KEY
Anthropic	claude-sonnet-4-5, claude-opus-4	ANTHROPIC_API_KEY
Google	gemini-2.5-pro, gemini-2.5-flash	GOOGLE_API_KEY
DeepSeek	deepseek/deepseek-r1	DEEPSEEK_API_KEY
Ollama	ollama/llama3.2	None (local)

For detailed model setup, see docs/MODELS.md.

MCP Tools

Connect to external tools via the Model Context Protocol:

# Add web search capability
dsagent mcp add brave-search

# Use it in chat
dsagent --mcp-config ~/.dsagent/mcp.yaml

Available templates: brave-search, filesystem, github, memory, fetch, bigquery

For MCP configuration details, see docs/MCP.md.

Session Management

Sessions persist your conversation history and kernel state:

# List sessions
dsagent chat
> /sessions

# Resume a session
dsagent --session <session-id>

# Export session to notebook
> /export myanalysis.ipynb

Output Structure

Each run creates organized output:

workspace/
└── runs/{run_id}/
    ├── data/           # Input data (copied)
    ├── notebooks/      # Generated Jupyter notebooks
    ├── artifacts/      # Charts, models, exports
    └── logs/           # Execution logs

Included Libraries

DSAgent comes with essential data science libraries pre-installed:

Category	Libraries
Core	numpy, pandas, scipy
DataFrames	polars, pyarrow
Visualization	matplotlib, seaborn, plotly
Machine Learning	scikit-learn, xgboost, lightgbm, pycaret
Feature Selection	boruta
Statistics	statsmodels

Documentation

• CLI Reference - Complete command-line options
• Python API - Detailed API documentation
• Model Configuration - LLM provider setup
• MCP Tools - External tools integration
• Docker Guide - Container deployment

License

MIT