Documentation
Chat
Implementation

Implementation

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

Refined Jupyter Agent Integration Plan

This plan outlines the steps to integrate the external AI agent (CodeAct/Calliope Data Agent) into JupyterLab, incorporating user feedback for enhanced configuration, security, and model selection flexibility.

I. Architecture Overview (v2)

The core components remain, but the configuration and session management are centralized around a PostgreSQL database accessible by the Agent.

  1. JupyterLab Frontend Extension (TypeScript/React): UI for chat, configuration viewing, RAG training, and notebook interaction. Fetches available models/datasources from the backend extension.
  2. JupyterLab Server Extension (Python/Tornado): Proxies requests to the Agent API. Handles Jupyter-specific aspects like user session identification (best effort) and retrieving the Agent API URL from Jupyter server config. It no longer manages agent-specific configurations like datasources or model defaults.
  3. External Agent (Dockerized Python/Flask): The core agent service.
    • Exposes REST API (/api/*) for chat, SQL, RAG, listing providers/models, listing datasources.
    • Exposes secured Admin REST API (/admin/*) for managing secrets and datasource configurations.
    • Reads secrets from its environment/secrets management.
    • Reads datasource configurations from the PostgreSQL DB.
    • Uses the PostgreSQL DB for persistent session management (PostgresSessionManager).
  4. PostgreSQL Database: Stores agent configurations (datasources) and user session history. Accessible only by the Agent container.
  5. Shared Volume: For file exchange between JupyterLab and the Agent.
  6. Vector Store (ChromaDB): Continues to store RAG context, likely persisted on the Shared Volume.

graph TD
A[JupyterLab UI (Frontend Ext)] -- REST API --> B(Jupyter Server Ext);
B -- REST API --> C{Agent Container (Flask API)};
C -- Reads/Writes --> D[(Shared Volume)];
A -- Reads/Writes --> D;
C -- Accesses --> E[PostgreSQL DB (Agent Config & Sessions)];
C -- Accesses --> F(Vector Store - ChromaDB on Shared Volume);
C -- Accesses --> G(External LLM/Tool APIs);
B -- Reads --> H(Jupyter Server Config - Agent URL);
I[Admin UI/CLI] -- Secured REST API --> C; subgraph Agent Internals
C -- Uses --> J(LLM Factory);
C -- Uses --> K(Session Manager - Postgres);
C -- Uses --> L(Config Loader - DB/Env);
C -- Uses --> M(RAG Store);
C -- Uses --> N(Secret Resolver);
end

II. Refined Implementation Details

  1. Configuration Management:
    • Agent Configuration (Datasources):
      • Storage: Stored in a dedicated table within the agent’s PostgreSQL database.
      • Management: CRUD operations via new, secured Agent Admin API endpoints (e.g., POST /admin/datasources, GET /admin/datasources, DELETE /admin/datasources/{id}).
      • Loading: The agent’s config loading mechanism (config.py needs refactoring) will query this DB table on startup or dynamically to get datasource definitions. It will need DB connection details from its environment variables.
      • Listing for UI: A public Agent API endpoint (GET /api/datasources) lists configured datasources (names, descriptions, dialects - NO connection details/secrets). The Jupyter Backend Extension proxies calls to this endpoint for the Frontend.
    • Secrets (API Keys / DB Passwords):
      • Storage: Managed via secured Agent Admin API endpoints (e.g., POST /admin/secrets, GET /admin/secrets (list names only), DELETE /admin/secrets/{name}). Secrets stored securely (e.g., encrypted in the DB or using a dedicated secrets manager accessible only by the agent).
      • Referencing: Connection strings/configs stored in the DB can reference secrets using a placeholder format (e.g., password={secret:POSTGRES_PASSWORD}).
      • Resolution: The agent resolves these placeholders internally just before use (e.g., when get_sql_database is called in tools.py), fetching the actual secret value from its secure storage.
    • Model Selection:
      • Agent Knowledge: The agent knows about all providers compiled into it (llm/providers/). It checks for necessary API keys (via its secret management) to determine which are active.
      • Listing for UI: New Agent API endpoint (GET /api/providers) lists available providers and the models the agent can potentially use for each (based on compiled code and active keys). Example response: {“providers”: [{“name”: “openai”, “active”: true, “models”: [“gpt-4o”, “gpt-4o-mini”]}, {“name”: “ollama”, “active”: true, “models”: [“gemma3:12b”, “qwen2.5-coder”]}, …]}. The Jupyter Backend Extension proxies calls to this endpoint.
      • Request Override: Agent API endpoints (/api/chat, /api/sql/ask) accept optional llm_provider (e.g., “openai”) and llm_model (e.g., “gpt-4o”) parameters in the request body.
      • Agent Logic: The agent’s LLMFactory (llm/factory.py) and potentially the agent invocation logic (server.py handlers) need modification to:
        1. Accept these optional provider/model parameters from the request.
        2. If provided, use them to instantiate the LLM for that specific request, overriding any agent-level defaults.
        3. If not provided, fall back to agent-level default models (which might still be loaded from its config DB or environment).
    • Jupyter Extension Config: Minimal - primarily stores the Agent API URL. Managed via standard Jupyter server configuration (jupyter_server_config.py or similar).
  2. Security: Addressed by the Admin API for secrets and the internal secret resolution mechanism within the agent. Admin endpoints require authentication (e.g., a bearer token configured via an agent environment variable AGENT_ADMIN_TOKEN).
  3. File Sharing (Shared Volume): No changes from the previous plan. Mount a shared volume to /data in the agent and an accessible path (e.g., /home/jovyan/agent_data) in JupyterLab.
  4. Feature Integration (UI & Notebook):
    • Chat Panel / Magic Command: Primary approach remains the same (dedicated panel, %%agent magic). UI fetches available models/datasources from the backend extension (which proxies to the agent) to populate dropdowns or provide context. When sending a request, the UI includes the selected model/provider override if the user chose one.
    • Jupyter AI Integration (/sql, @agent): This remains a potential future enhancement. It would require:
      • Creating a jupyter-ai “provider” or extension point.
      • Registering the slash commands/mention triggers.
      • Mapping jupyter-ai’s chat UI events and context to your agent’s API calls. This adds significant complexity and dependency on jupyter-ai’s architecture. Recommendation: Implement the dedicated panel/magic first.
  5. Multi-User Support & Session Management:
    • Persistence: Use PostgresSessionManager within the agent, pointing to the PostgreSQL DB.
    • Session ID:
      • The Jupyter Server Extension should try to obtain a stable identifier for the user/session. Investigate:
        • self.current_user in the Tornado handler (might give the Jupyter user).
        • Server session IDs (might be transient).
        • JupyterHub user information if running in Hub.
      • If a stable Jupyter-provided ID is found, use it.
      • Fallback: If no stable ID is reliably available, the Jupyter Server Extension manages a session cookie (e.g., jupyter-agent-session-id). It checks for this cookie on incoming requests. If present, it uses the cookie’s value. If not present, it generates a UUID, sets the cookie in the response, and uses the new UUID.
      • Pass the chosen session_id to the agent API.
    • Clear Session UI: Remains the same - button calls backend, which calls DELETE /api/sessions/{id} on the agent.
  6. RAG Interaction: No changes to the core flow. UI fetches available datasource_ids from the agent via the backend extension to populate the training target dropdown.
  7. Abstraction (SDK): Remains the same - recommended for frontend (handlers.ts), optional for backend (agent_client.py).

III. Project Structure (Conceptual)

jupyterlab-data-agent/
├── pyproject.toml        # Jupyter extension packaging
├── setup.py              # Jupyter extension packaging
├── src/                  # Frontend Extension (TypeScript/React)
│   ├── index.ts          # Entry point
│   ├── components/       # React components (ChatPanel, ConfigPanel, etc.)
│   ├── handlers.ts       # Frontend API client/SDK (talks to Backend Ext)
│   └── styles/           # CSS/Styling
├── data_agent_jupyter_extension/  # Backend Extension (Python)
│   ├── __init__.py
│   ├── handlers.py       # API Handlers (adapted from POC)
│   ├── config.py         # Extension configuration loading (adapted from POC)
│   └── agent_client.py   # (Optional) Client for Agent API
├── docker-compose.yml    # To run JupyterLab, Agent, optional DB
├── agent/                # Your agent code (can be separate repo or submodule)
│   ├── Dockerfile
│   ├── codeact_agent/    # Core agent library
│   ├── main.py           # Agent entry point
│   └── ... (agent files)
└── ... (other config files like .gitignore, README)

IV. Sequential Task List for Implementation

This list breaks down the implementation into granular tasks, ordered logically (Agent -> Backend -> Frontend).

Phase 1: Agent Enhancements

  1. Task A1: Database Setup
    • Goal: Add PostgreSQL support to the agent’s environment.
    • Implementation: Update the agent’s Dockerfile and docker-compose.yml to include a PostgreSQL service. Configure environment variables for the agent to connect to this DB (e.g., DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME).
    • Design: Use standard Postgres image. Network the agent and DB containers.
  2. Task A2: Implement PostgresSessionManager
    • Goal: Enable persistent session storage using the DB.
    • Implementation: Ensure psycopg2-binary is in agent requirements. Update session_manager.py:
      • Complete the PostgresSessionManager class (connection logic, table creation SQL, CRUD methods for sessions using Json for messages).
      • Modify SessionManagerFactory to read DB connection details from the config/environment and instantiate PostgresSessionManager when session_store is “postgres”.
      • Modify get_session_manager to load config and use the factory.
    • Design: Table schema: session_id (TEXT PK), messages (JSONB), updated_at (TIMESTAMP). Handle connection pooling/management if necessary for high load (though likely not needed initially).
  3. Task A3: Refactor Agent Configuration Loading
    • Goal: Load datasource configurations from the DB instead of only environment variables. Secrets remain loaded from env/Docker secrets.
    • Implementation: Modify config.py:
      • Define a DB schema for storing datasource configurations (e.g., datasources table: id (TEXT PK), name (TEXT), description (TEXT), dialect (TEXT), connection_details (JSONB) - store host, port, db, user without password here).
      • Update get_config():
        • Load DB connection details from environment variables.
        • Connect to the DB.
        • Query the datasources table.
        • Populate the config[‘sql_datasources’] dictionary from the DB results.
        • Continue loading secrets (API keys) and other settings (LLM provider defaults, ports) from environment variables as before. Merge DB config with env var config.
    • Design: Ensure DB query happens early in config loading. Handle DB connection errors gracefully (e.g., log error, potentially run with no datasources).
  4. Task A4: Implement Secret Management & Resolution
    • Goal: Securely store and use secrets (API keys, DB passwords) referenced in configurations.
    • Implementation:
      • Define DB schema for secrets (e.g., secrets table: name (TEXT PK), value (TEXT/BYTEA encrypted), updated_at (TIMESTAMP)). Choose an encryption method (e.g., cryptography library with a key derived from an agent env var SECRET_ENCRYPTION_KEY).
      • Create secured Agent Admin API endpoints (/admin/secrets) in server.py for CRUD operations on secrets. Use a decorator to check for an Authorization: Bearer <token> header matching AGENT_ADMIN_TOKEN from env vars.
      • Implement secret resolution logic: Create a function resolve_secrets(config_value: Any, secret_store: Dict) that recursively traverses config dictionaries/lists/strings, finds placeholders like {secret:NAME}, queries the (decrypted) secret value from the DB/cache, and replaces the placeholder.
      • Modify places where secrets are used (e.g., get_sql_database in tools.py, LLM provider instantiation in llm/factory.py) to call resolve_secrets on the relevant config part just before using it.
    • Design: Prioritize security. Avoid logging resolved secrets. Cache decrypted secrets in memory briefly if performance is critical, but be mindful of security implications.
  5. Task A5: Implement Agent Admin API for Datasources
    • Goal: Allow managing datasource configurations via a secured API.
    • Implementation: Create secured Agent Admin API endpoints (/admin/datasources) in server.py for CRUD operations on the datasources DB table. Secure using the same admin token mechanism as secrets.
    • Design: Ensure validation of incoming datasource configurations. Remember not to handle passwords here; use secret placeholders.
  6. Task A6: Implement Public API for Datasource/Provider Listing
    • Goal: Allow the UI to discover available datasources and models.
    • Implementation:
      • Create GET /api/datasources endpoint in server.py. Query the datasources table from the DB and return a list of basic info (id, name, description, dialect).
      • Create GET /api/providers endpoint in server.py.
        • Import the provider registry (llm/providers.get_provider_registry()).
        • Iterate through registered providers.
        • Check if required API keys/secrets are present (using the secret management system) to determine if a provider is active.
        • List default/known models for each provider (can be hardcoded initially based on llm/providers/*.py or made more dynamic).
        • Return the structured list as described in the plan.
    • Design: Cache provider/model info briefly if LLM checks are slow.
  7. Task A7: Implement Request-Level Model Override & Dynamic Graph Handling
    • Goal: Allow API callers to specify the LLM provider/model for a single request and ensure the agent graph uses the correct model for that session/request.
    • Implementation:
      • API Handlers: Modify Agent API handlers (/api/chat, /api/sql/ask in server.py) to accept optional llm_provider and llm_model in the JSON request body.
      • Agent Invocation: Pass these parameters down to the agent invocation logic (e.g., composable_agent.invoke, sql_agent.invoke). Add them to the initial state dictionary passed to the graph.
      • Agent Classes: Modify the agent classes (ComposableAgent, SQLAgent, CodeActAgent) invoke methods to accept these overrides in the initial state.
      • LLMFactory: Modify LLMFactory.get_llm to accept optional provider_override and model_override parameters. If provided, instantiate the specific provider/model, otherwise fall back to the standard config lookup. Ensure the factory still uses the centrally managed secrets for API keys.
      • Graph Instantiation/Model Binding:
        • Remove Global Graph Compilation: Instead of compiling the graph once on agent startup, the graph definition (create_graph method in agent classes) should define the structure, but the LLM instances used within the graph nodes should be determined at runtime based on the state passed into the node.
        • Node Modification: Modify graph nodes that use LLMs (e.g., call_model_node, planner_node, generate_and_call_llm, generate_summary_node) to:
          1. Retrieve llm_provider and llm_model overrides from the current state dictionary if present.
          2. Call LLMFactory.get_llm with these overrides (or None if not present) to get the correct LLM instance for that specific node execution.
          3. Use the obtained LLM instance for the invocation within the node.
        • State Passing: Ensure the initial llm_provider and llm_model overrides are correctly passed into the graph.invoke() call within the agent’s invoke method. LangGraph’s state mechanism will propagate this information to the nodes.
    • Design:
      • The graph structure remains static, but the components (specifically LLMs) used within the nodes become dynamic per invocation based on the state.
      • This avoids compiling a new graph per request, which would be inefficient.
      • Ensure overrides are validated against the list of available/active providers/models (potentially within the API handler or LLMFactory).
      • The checkpointer still manages the state persistence across steps for a given thread_id, regardless of the LLM used in each step.

Phase 2: Jupyter Backend Extension

  1. Task B1: Setup Project Structure
    • Goal: Create the basic file structure for a JupyterLab extension (frontend + backend).
    • Implementation: Use the JupyterLab extension cookiecutter or follow the extension tutorial. Set up pyproject.toml, setup.py, src/ for frontend, and a Python package (e.g., data_agent_jupyter_extension/) for the backend.
  2. Task B2: Configure Agent API URL
    • Goal: Allow Jupyter admins to configure where the backend extension finds the agent API.
    • Implementation: Use Jupyter’s server configuration system. Define a config traitlet (e.g., c.DataAgent.agent_api_url = “http://agent:5000”) in a config file (e.g., /etc/jupyter/jupyter_server_config.py). Adapt the POC config.py within the backend extension to load this setting.
  3. Task B3: Implement API Proxy Handlers
    • Goal: Create Tornado handlers in the backend extension to forward requests from the frontend to the agent API.
    • Implementation: Refactor the POC handlers.py. Create handlers for:
      • POST /data-agent/chat: Forwards request body (including optional llm_provider, llm_model) to Agent’s POST /api/chat. Includes session ID.
      • POST /data-agent/sql/ask: Forwards request body (including optional llm_provider, llm_model) to Agent’s POST /api/sql/ask. Includes session ID.
      • GET /data-agent/history/{thread_id}: Calls Agent’s GET /api/sessions/{thread_id}.
      • DELETE /data-agent/clear_session/{thread_id}: Calls Agent’s DELETE /api/sessions/{thread_id}.
      • GET /data-agent/providers: Calls Agent’s GET /api/providers.
      • GET /data-agent/datasources: Calls Agent’s GET /api/datasources.
      • (Optional) Handlers for proxying RAG training calls if UI controls this directly.
    • Design: Use tornado.httpclient.AsyncHTTPClient for making non-blocking requests to the agent API. Handle errors from the agent API gracefully. Read the agent_api_url from the loaded config. Implement session ID handling (get from request, pass to agent).
  4. Task B4: Implement Session ID Logic
    • Goal: Obtain the best possible session identifier to pass to the agent.
    • Implementation: In the base handler or relevant handlers:
      • Attempt to get user info (self.current_user).
      • Investigate other potential stable identifiers from self.request or Jupyter server context (this requires research into Jupyter Server internals/APIs).
      • If a stable ID is found, use it.
      • Fallback: If no stable ID is reliably available, the Jupyter Server Extension checks for a specific cookie (e.g., jupyter-agent-session-id) in the request. If the cookie exists, its value is used. If not, the handler generates a new UUID, uses it for the current request, and sets the cookie in the response (self.set_cookie(…)) for subsequent requests.
    • Design: Log clearly which type of session ID is being used (Jupyter user, cookie, etc.). Configure cookie properties (name, path, secure, httponly).

Phase 3: Jupyter Frontend Extension

  1. Task C1: Basic UI Setup
    • Goal: Set up the React/TypeScript environment and create a basic panel in JupyterLab.
    • Implementation: Follow JupyterLab extension tutorials. Use Lumino widgets and React. Create a main widget added to the JupyterLab shell (e.g., left or right sidebar).
  2. Task C2: Frontend API Client
    • Goal: Create a TypeScript module to interact with the backend extension’s API endpoints.
    • Implementation: Create handlers.ts (or apiClient.ts). Use JupyterLab’s requestAPI utility. Implement functions like sendMessage(text, provider?, model?), getHistory(), clearSession(), getProviders(), getDatasources(). Note: Session ID is now handled by the backend via cookies, so the frontend doesn’t need to manage/send it explicitly.
  3. Task C3: Build Chat Panel UI
    • Goal: Create the main chat interface component.
    • Implementation: Use React. Create components for:
      • Message list (displaying user and agent messages, rendering code blocks, potentially Plotly charts).
      • Input area (text input, send button, clear button).
      • Dropdowns (optional) for selecting provider/model (populated by getProviders()).
    • Design: Manage chat history state. Call API client functions on user actions (passing selected provider/model overrides). Handle loading states. Render Markdown in responses. Use a library like react-plotly.js to render Plotly JSON received from the agent (via backend).
  4. Task C4: Build Configuration/Status Panel UI
    • Goal: Display agent status and configuration info.
    • Implementation: Create a separate React component or tab.
      • Display the Agent API URL (from backend config, if exposed).
      • List available providers/models (fetched via getProviders()).
      • List configured datasources (fetched via getDatasources()).
      • (Optional) Add UI for RAG training (text areas, file upload, dropdown for datasource target, button to trigger backend training call).
      • (Optional) Add basic UI for interacting with Admin APIs if desired for non-production use (requires handling the admin token securely on the frontend, which is generally discouraged - better to use a separate admin tool/CLI).
    • Design: Fetch data on component mount. Display information clearly.
  5. Task C5: Implement Notebook Magic (%%agent)
    • Goal: Allow running agent queries directly from notebook cells.
    • Implementation:
      • Server-side approach (Simpler): Define the magic command in the backend extension. When executed, the kernel sends the cell content and potentially magic arguments (e.g., %%agent --provider openai --model gpt-4o) to the backend. The backend calls the agent API (similar to the chat handler), including any overrides. The backend streams/sends results back to the kernel, which displays them in the cell output.
    • Design: Start with the server-side approach. Parse magic arguments in the backend to get provider/model overrides. Ensure results (text, DataFrames as HTML, Plotly JSON) are rendered correctly using appropriate MIME types.
  6. Task C6: Styling and Finalization
    • Goal: Apply styling to make the UI consistent with JupyterLab.
    • Implementation: Use CSS or JupyterLab’s theme variables. Ensure responsiveness.
    • Design: Test in different themes (light/dark).

This detailed task list provides a clear path forward, starting with the necessary agent modifications and progressing through the backend and frontend layers. Remember to test each integration point thoroughly. Good luck!

Extending JupyterLab: Integrating External REST API AgentsLangGraph CodeAct: Building Agents with Executable Code