IPychat — An AI extension for IPython

02 January 2025 · python · llm Tweet

Recently, I've been exploring geospatial data about Paris. Although I've seen others work with this type of data and know basic terms like shapefiles and polygons, this is my first time working with it myself. I've been looking up terms and going through docs to learn how various geospatial data libraries work.

IPython and Jupyter were my tools of choice here—they're pleasant to work with, and built exactly for this type of data exploration. However, after using Cursor at work to ask LLMs questions about libraries, generate boilerplate functions, or brainstorm ideas, I really wanted to have the same experience within IPython to avoid context switching between my terminal and browser.

So, I spent last weekend building an extension that lets me use IPython like Cursor! Here's a quick demo where I load the Titanic dataset into a dataframe and ask a LLM, "what insights can I draw from this df".

Referencing variables by name works here because the extension finds the relevant context—in this case, the column names, shape, first five rows of this dataframe—and sends it to the selected LLM along with the IPython history. This ensures the LLM's response is relevant to the question, exactly like Cursor.

Since all major LLMs can return markdown, I used rich to render the response nicely in the terminal. ✨

The extension lets you prompt GPT-4o, Claude 3.5 Sonnet, and Gemini. I'm also planning to add support for local models to compare their responses with Claude 3.5 Sonnet, which, in my experience, generates the best responses for coding questions.

You can find the code on GitHub.

Installation

You can install ipychat using pip:

  $ pip install ipychat

Configuration

Based on the model you want to use, either set OPENAI_API_KEY, or ANTHROPIC_API_KEY, or both environment variables. You can also run the ipychat config command to configure ipychat interactively.

Usage

You can start IPython with the extension already loaded by simply running the ipychat command:

  $ ipychat
  Welcome to ipychat! Use %ask to chat with gpt-4o.
  You can change models using %models.

  In [1]:

Alternatively, you can load the extension in a IPython session:

  In [1]: %load_ext ipychat

You can then start asking the model any question using the %ask magic.

  In [1]: %ask how can i download the titanic dataset using seaborn

If you want, you can change the current model using the %models magic.

Note: Placing question marks directly after a word doesn't work right now because it's a default IPython feature used to display an object's docs, and I didn't want to override it.

---

That's it! Let me know if you find ipychat useful or have any suggestions.