DEV Community

novatechnolab
novatechnolab

Posted on

Bringing AI to Jupyter Notebooks: Meet jupyter-ext-ai

Bringing AI to Jupyter Notebooks: Meet jupyter-ext-ai

Are you tired of switching tabs to ask ChatGPT or Claude for code help while working in Jupyter Notebooks?

Meet jupyter-ext-ai, a new open-source IPython magical extension that brings OpenAI, Google Gemini, and Anthropic Claude right into your cells!

πŸš€ Why jupyter-ext-ai?

While there are many AI coding tools, I wanted a lightweight, seamless way to chat with different LLM providers directly from a Jupyter Notebook cell. Whether you want to quickly ask a syntax question or have an LLM write a complex data processing function for you directly in the notebookβ€”this extension has you covered.

✨ Features

  • Multiple Providers: Switch seamlessly between OpenAI (gpt-4o-mini, gpt-4o), Google Gemini, and Anthropic Claude.
  • One-line prompts: Use %ai <prompt> for quick questions.
  • Cell magic: Use %%ai [flags] where the cell body acts as the prompt for longer instructions.
  • Context memory: Maintains conversation history with %ai_history and %ai_reset to clear it.
  • Easy Config: View and update your provider, model, and temperature settings on the fly using %ai_config.

πŸ“¦ Getting Started

Installation is as simple as:

pip install jupyter-ext-ai
Enter fullscreen mode Exit fullscreen mode

Set up your API keys as environment variables in your terminal before launching Jupyter:

export OPENAI_API_KEY="sk-..."
export JUPYTER_AI_PROVIDER="openai" # default
export JUPYTER_AI_MODEL="gpt-4o-mini" # default
Enter fullscreen mode Exit fullscreen mode

Or configure Gemini or Anthropic Claude similarly!

πŸͺ„ Quick Demo

Once you're in a Jupyter notebook, load the extension:

%load_ext jupyter_ext_ai
Enter fullscreen mode Exit fullscreen mode

Ask a quick question

%ai Write a regex in Python to validate an email address
Enter fullscreen mode Exit fullscreen mode

Or use the cell magic for more complex tasks!

%%ai --provider gemini --model gemini-2.0-flash
Write a Python function to fetch historical stock data using yfinance, 
calculate a 50-day moving average, and plot it using matplotlib.
Enter fullscreen mode Exit fullscreen mode

If you ever want to check what models are supported for a provider, just run:

%ai_models claude
Enter fullscreen mode Exit fullscreen mode

🀝 Contributing & Source Code

The extension is open-source and available on GitHub! We'd love your feedback, feature requests, or pull requests.

Check it out on GitHub: novatechnolab/jupyter-ext-ai and PyPI.

Let me know what you think in the comments! πŸ‘‡

Top comments (0)