Documentation
Chat
Extending JupyterLab: Integrating External REST API Agents

Extending JupyterLab: Integrating External REST API Agents

Calliope Integration: This component is integrated into the Calliope AI platform. Some features and configurations may differ from the upstream project.

This guide provides a comprehensive overview of how to extend JupyterLab by integrating an external agent running in a separate Docker container, exposing a REST API. The architecture and patterns described here are inspired by Vanna AI’s Retrieval-Augmented Generation (RAG) approach and the CodeAct paper’s executable Python code actions.

Table of Contents

Introduction

JupyterLab is a highly extensible environment for interactive computing. Integrating external agents, such as those inspired by Vanna AI and CodeAct, can significantly enhance its capabilities, especially in data analysis, natural language querying, and automated code execution.

Architecture Overview

The proposed architecture involves:

  • Frontend Extension: Built using TypeScript and React, integrated into JupyterLab’s UI.
  • Backend Server Extension: Python-based Jupyter server extension handling API requests from the frontend.
  • External Agent: A Dockerized REST API service inspired by Vanna AI and CodeAct, providing advanced data querying, code execution, and RAG-based SQL generation.

High-Level Diagram 

JupyterLab Frontend (React/TypeScript)
        |
        | HTTP Requests (REST API)
        v
Jupyter Server Extension (Python)
        |
        | HTTP Requests (REST API)
        v
Dockerized External Agent (REST API)
        |
        | Database Queries, Code Execution, RAG Operations
        v
Data Sources, Vector Stores, LLM Providers

JupyterLab Extension Points

JupyterLab provides several extension points:

  • Frontend Extensions: Modify or add UI components, menus, commands, and widgets.
  • Server Extensions: Extend backend functionality, handle API requests, and manage server-side logic.
  • Kernel Extensions: Extend kernel capabilities (less relevant for REST API integration).

Relevant Documentation

Integration Patterns

1. Frontend-Backend Communication

  • REST API Pattern: Frontend sends HTTP requests to the backend server extension, which proxies requests to the external agent.
  • WebSocket Pattern: For real-time interactions, WebSockets can be used (optional).

2. Backend-Agent Communication

  • REST API Calls: The backend server extension communicates with the external Dockerized agent via REST API calls.

3. State Management

  • Frontend State: Managed using React state or Redux.
  • Backend State: Managed using session management (e.g., in-memory, Redis, or database).

Step-by-Step Guide

Step 1: Setup Dockerized External Agent

  • Create a Docker container exposing REST API endpoints for:
    • Natural language to SQL generation (RAG-based).
    • Code execution (CodeAct-inspired).
    • Data querying and visualization.

Example Docker Compose:

version: '3.8'
services:
  agent:
    build: ./agent
    ports:
      - "5000:5000"
    environment:
      - OLLAMA_BASE_URL=http://ollama:11434
      - SQL_AGENT_ENABLED=true
      - SQL_DIALECT=postgres
      - SQL_CONNECTION_STRING=postgresql://user:pass@db:5432/mydb

Step 2: Create Jupyter Server Extension

  • Implement Python handlers extending APIHandler from jupyter_server.base.handlers.
  • Proxy requests to the external agent.

Example handler:

from jupyter_server.base.handlers import APIHandler
import tornado.web
import requests

class AgentProxyHandler(APIHandler):
    @tornado.web.authenticated
    async def post(self):
        body = self.get_json_body()
        response = requests.post("http://agent:5000/api/chat", json=body)
        self.finish(response.json())

Register handler:

def setup_handlers(web_app):
    host_pattern = ".*$"
    base_url = web_app.settings["base_url"]
    route = url_path_join(base_url, "agent", "chat")
    handlers = [(route, AgentProxyHandler)]
    web_app.add_handlers(host_pattern, handlers)

Step 3: Create JupyterLab Frontend Extension

  • Use TypeScript and React to build UI components.
  • Communicate with backend via REST API.

Example React component:

import { requestAPI } from './handler';

async function sendMessage(message: string) {
  const response = await requestAPI<any>('chat', {
    method: 'POST',
    body: JSON.stringify({ message })
  });
  console.log(response);
}

Step 4: Integrate UI Components

  • Add widgets, panels, or commands to JupyterLab UI.
  • Utilize JupyterLab’s UI components and extension points.

Example widget registration:

import { Widget } from '@lumino/widgets';
import { ICommandPalette } from '@jupyterlab/apputils';

export function activate(app, palette: ICommandPalette) {
  const widget = new Widget();
  widget.node.textContent = 'Agent Panel';
  widget.id = 'agent-panel';
  widget.title.label = 'Agent';
  widget.title.closable = true;

  app.shell.add(widget, 'left');

  palette.addItem({ command: 'agent:open', category: 'Agent' });
}

Best Practices

  • Security: Validate and sanitize all inputs. Securely handle API keys and sensitive data.
  • Error Handling: Implement robust error handling and logging.
  • Performance: Use asynchronous requests and caching where appropriate.
  • Testing: Write unit, integration, and end-to-end tests.

References

Docker Image Naming ConventionImplementation