How Athena built AI-enabled Jupyter Notebooks
2024-07-01
Felicia Chang
Athena is an enterprise-level AI data analyst. Used by Fortune 500 companies like AB InBev, Athena supports data teams across their various workflows.
The Challenge
For Athena to be efficient, the AI agent needed access to the data applications that analysts use in their day-to-day work.
The Athena team looked to Jamsocket for a solution to run Python code in a Jupyter notebook – one of the essential workflows they needed to support. The infrastructure needed to spin up Jupyter notebooks is non-trivial as each notebook requires an isolated environment to run user code.
The Solution
Using Jamsocket's customizable Jupyter notebook image, the Athena team was able to get the initial prototype working in their application in just a few hours.
We’re happy users. It works, nothing breaks, that’s what matters.
When a customer opens a Jupyter notebook, Athena’s application backend starts a Jupyter server inside a Jamsocket session backend.
The Jupyter server "serves" a Jupyter client that is embedded into Athena’s application interface, allowing users to write code and run cells.
When a user session ends, Jamsocket gracefully shuts down the Jupyter notebook and persists the data in Athena’s own S3 bucket.
Investing time into AI features instead of infrastructure
With this core integration handled, the Athena team was able to focus on building different ways users and their AI agent, Athena could collaborate on a Jupyter notebook.
For example, the AI agent can instantiate new cells or notebooks via custom frontend hooks to their Jupyter client.
To receive inline recommendations, the Athena team built a custom LLM that operates in the notebook environment using the Jupyter AI extension.
And since Jamsocket can spawn any Dockerized code, Athena built their Docker image to spawn the Jupyter server with their Python SDK, which manages access relevant customer data.
Smooth Sailing
The Athena team has been using Jamsocket for Jupyter notebooks for more than six months now. And they’ve found it to be “the most stable part of our stack.”
Since setting up Jamsocket, they were able to focus on what matters to their business – building the best AI data analyst.
“We haven’t had to think about Jamsocket, to be honest. That’s a good thing.”
If you’re interested in integrating Jupyter notebooks into your application, reach out to us or book a call to get setup with our open source notebook solution. To try spinning up a Jupyter notebook yourself, try it here.