feat: Add Amazon Q Agentic Chat integration support and build improvements (#959)

Author: TRNWWZ

README.md

@@ -7,6 +7,14 @@ These images come in two variants, CPU and GPU, and include deep learning framew
 Keras; popular Python packages like numpy, scikit-learn and pandas; and IDEs like Jupyter Lab. The distribution contains
 the _latest_ versions of all these packages _such that_ they are _mutually compatible_.
 
+Starting with v2.9.5+ and new v3.5+, the images include Amazon Q Agentic Chat integration for enhanced AI-powered development assistance in JupyterLab.
+
+### Amazon Q Agentic Chat Integration
+
+The images include pre-configured Amazon Q artifacts and shared web client libraries:
+- `/etc/web-client/libs/` - Shared JavaScript libraries (JSZip) for all web applications
+- `/etc/amazon-q-agentic-chat/artifacts/jupyterlab/` - Amazon Q server and client artifacts for JupyterLab
+
 This project follows semver (more on that below) and comes with a helper tool to automate new releases of the
 distribution.
 

assets/README.md

@@ -0,0 +1,46 @@
+# Assets
+
+This directory contains utility scripts and files used during the Docker image build process.
+
+## extract_amazon_q_agentic_chat_urls.py
+
+A Python script that extracts Amazon Q Agentic Chat artifact URLs from a manifest file for Linux x64 platform.
+
+### Usage
+```bash
+python extract_amazon_q_agentic_chat_urls.py <manifest_file> <version>
+```
+
+### Parameters
+- `manifest_file`: Path to the JSON manifest file containing artifact information
+- `version`: The server version to extract artifacts for
+
+### Output
+The script outputs environment variables for use in shell scripts:
+- `SERVERS_URL`: URL for the servers.zip artifact
+- `CLIENTS_URL`: URL for the clients.zip artifact
+
+## download_amazon_q_agentic_chat_artifacts.sh
+
+A modular shell script that downloads and extracts Amazon Q Agentic Chat artifacts for IDE integration.
+
+### Usage
+```bash
+bash download_amazon_q_agentic_chat_artifacts.sh <version> <target_dir> <ide_type>
+```
+
+### Parameters
+- `version`: Amazon Q server version (defaults to $FLARE_SERVER_VERSION_JL)
+- `target_dir`: Target directory for artifacts (defaults to /etc/amazon-q-agentic-chat/artifacts/jupyterlab)
+- `ide_type`: IDE type for logging (defaults to jupyterlab)
+
+### Features
+- Downloads JSZip library to shared web client location (/etc/web-client/libs/) for reuse across all web applications
+- Modular design supports future VSCode integration
+- Comprehensive error handling with retry logic
+- Automatic cleanup of temporary files
+
+### Directory Structure
+- `/etc/web-client/libs/` - Shared web client libraries (JSZip, etc.) for any web application
+- `/etc/amazon-q-agentic-chat/artifacts/jupyterlab/` - Amazon Q specific artifacts for JupyterLab
+- `/etc/amazon-q-agentic-chat/artifacts/vscode/` - Future Amazon Q artifacts for VSCode

assets/download_amazon_q_agentic_chat_artifacts.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+set -e
+
+# Download Amazon Q artifacts for IDE integration
+# Usage: download_amazon_q_artifacts.sh <version> <target_dir> <ide_type>
+# Example: download_amazon_q_artifacts.sh 1.25.0 /etc/amazon-q/artifacts/agentic-chat jupyterlab
+
+VERSION=${1:-$FLARE_SERVER_VERSION_JL}
+TARGET_DIR=${2:-"/etc/amazon-q-agentic-chat/artifacts/jupyterlab"}
+IDE_TYPE=${3:-"jupyterlab"}
+
+if [ -z "$VERSION" ]; then
+    echo "Error: Version not specified and FLARE_SERVER_VERSION_JL not set"
+    exit 1
+fi
+
+echo "Downloading Amazon Q artifacts for $IDE_TYPE (version: $VERSION)"
+
+# Create target directories
+sudo mkdir -p "$TARGET_DIR"
+
+# Download manifest and extract artifact URLs
+MANIFEST_URL="https://aws-toolkit-language-servers.amazonaws.com/qAgenticChatServer/0/manifest.json"
+curl -L --retry 3 --retry-delay 5 --fail "$MANIFEST_URL" -o "/tmp/manifest.json" || {
+    echo "Failed to download manifest"
+    exit 1
+}
+
+# Extract artifact URLs
+ARTIFACT_URLS=$(python3 /tmp/extract_amazon_q_agentic_chat_urls.py /tmp/manifest.json "$VERSION")
+if [ $? -ne 0 ] || [ -z "$ARTIFACT_URLS" ]; then
+    echo "Failed to extract Amazon Q artifact URLs"
+    exit 1
+fi
+
+eval "$ARTIFACT_URLS"
+
+# Download and extract servers.zip
+echo "Downloading servers.zip..."
+curl -L --retry 3 --retry-delay 5 --fail "$SERVERS_URL" -o "/tmp/servers.zip" || {
+    echo "Failed to download servers.zip"
+    exit 1
+}
+sudo unzip "/tmp/servers.zip" -d "$TARGET_DIR/servers" || {
+    echo "Failed to extract servers.zip"
+    exit 1
+}
+
+# Download and extract clients.zip
+echo "Downloading clients.zip..."
+curl -L --retry 3 --retry-delay 5 --fail "$CLIENTS_URL" -o "/tmp/clients.zip" || {
+    echo "Failed to download clients.zip"
+    exit 1
+}
+sudo unzip "/tmp/clients.zip" -d "$TARGET_DIR/clients" || {
+    echo "Failed to extract clients.zip"
+    exit 1
+}
+
+# Clean up temporary files
+rm -f /tmp/manifest.json /tmp/servers.zip /tmp/clients.zip
+
+echo "Amazon Q artifacts downloaded successfully to $TARGET_DIR"

assets/extract_amazon_q_agentic_chat_urls.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Extract Amazon Q artifact URLs from manifest for Linux x64 platform."""
+
+import json
+import sys
+
+
+def extract_urls(manifest_file, version, platform="linux", arch="x64"):
+    """Extract servers.zip and clients.zip URLs for specified platform/arch."""
+    try:
+        with open(manifest_file) as f:
+            manifest = json.load(f)
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Manifest file not found: {manifest_file}")
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid JSON in manifest file {manifest_file}: {str(e)}")
+
+    for ver in manifest["versions"]:
+        if ver["serverVersion"] == version:
+            for target in ver["targets"]:
+                if target["platform"] == platform and target.get("arch") == arch:
+                    servers_url = None
+                    clients_url = None
+
+                    for content in target["contents"]:
+                        if content["filename"] == "servers.zip":
+                            servers_url = content["url"]
+                        elif content["filename"] == "clients.zip":
+                            clients_url = content["url"]
+
+                    if servers_url is None or clients_url is None:
+                        raise ValueError(
+                            f"Required files (servers.zip/clients.zip) not found for version {version} {platform} {arch}"
+                        )
+
+                    return servers_url, clients_url
+
+    raise ValueError(f"Version {version} not found for {platform} {arch}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        print("Usage: extract_amazon_q_agentic_chat_urls.py <manifest_file> <version>")
+        sys.exit(1)
+
+    manifest_file, version = sys.argv[1], sys.argv[2]
+    servers_url, clients_url = extract_urls(manifest_file, version)
+
+    print(f"SERVERS_URL={servers_url}")
+    print(f"CLIENTS_URL={clients_url}")

src/main.py

@@ -131,6 +131,15 @@ def _copy_static_files(base_version_dir, new_version_dir, new_version_major, run
     if os.path.exists(aws_cli_key_path):
         shutil.copy2(aws_cli_key_path, new_version_dir)
 
+    # Copy Amazon Q agentic chat scripts from assets
+    q_extract_script_path = os.path.relpath(f"assets/extract_amazon_q_agentic_chat_urls.py")
+    if os.path.exists(q_extract_script_path):
+        shutil.copy2(q_extract_script_path, new_version_dir)
+
+    q_download_script_path = os.path.relpath(f"assets/download_amazon_q_agentic_chat_artifacts.sh")
+    if os.path.exists(q_download_script_path):
+        shutil.copy2(q_download_script_path, new_version_dir)
+
     if int(new_version_major) >= 1:
         # dirs directory doesn't exist for v0. It was introduced only for v1
         dirs_relative_path = os.path.relpath(f"{base_path}/dirs")
@@ -268,7 +277,7 @@ def _build_local_images(
             # Minimal patch build, use .patch Dockerfiles
             dockerfile = f"./Dockerfile-{image_type}.patch"
         else:
-            dockerfile="./Dockerfile"
+            dockerfile = "./Dockerfile"
         try:
             image, log_gen = _docker_client.images.build(
                 path=target_version_dir, dockerfile=dockerfile, rm=True, pull=True, buildargs=config["build_args"]

src/package_report.py

@@ -1,11 +1,11 @@
 import json
 import os
+import subprocess
 import warnings
 from datetime import datetime
 from itertools import islice
 
 import boto3
-import conda.cli.python_api
 from conda.models.match_spec import MatchSpec
 from condastats.cli import overall
 from dateutil.relativedelta import relativedelta
@@ -39,11 +39,18 @@ def _get_package_versions_in_upstream(target_packages_match_spec_out, target_ver
             continue
         channel = match_spec_out.get("channel").channel_name
         subdir_filter = "[subdir=" + match_spec_out.get("subdir") + "]"
-        search_result = conda.cli.python_api.run_command(
-            "search", channel + "::" + package + ">=" + str(package_version) + subdir_filter, "--json"
-        )
-        # Load the first result as json. The API sends a json string inside an array
-        package_metadata = json.loads(search_result[0])[package]
+        try:
+            search_result = subprocess.run(
+                ["conda", "search", channel + "::" + package + ">=" + str(package_version) + subdir_filter, "--json"],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            # Load the result as json
+            package_metadata = json.loads(search_result.stdout)[package]
+        except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e:
+            print(f"Error searching for package {package}: {str(e)}")
+            continue
         # Response is of the structure
         # { 'package_name': [{'url':<someurl>, 'dependencies': <List of dependencies>, 'version':
         # <version number>}, ..., {'url':<someurl>, 'dependencies': <List of dependencies>, 'version':
@@ -90,7 +97,7 @@ def _generate_staleness_report_per_image(
         package_string = (
             package
             if version_in_sagemaker_distribution == package_versions_in_upstream[package]
-            else "${\color{red}" + package + "}$"
+            else "${\\color{red}" + package + "}$"
         )
 
         if download_stats:
@@ -170,7 +177,7 @@ def _validate_new_package_size(new_package_total_size, target_total_size, image_
             + str(new_package_total_size_percent)
             + "%)"
         )
-        new_package_total_size_percent_string = "${\color{red}" + str(new_package_total_size_percent) + "}$"
+        new_package_total_size_percent_string = "${\\color{red}" + str(new_package_total_size_percent) + "}$"
 
     print(
         "The total size of newly introduced Python packages is "
@@ -276,10 +283,13 @@ def _generate_python_package_dependency_report(image_config, base_version_dir, t
     for package, version in new_packages.items():
         try:
             # Pull package metadata from conda-forge and dump into json file
-            search_result = conda.cli.python_api.run_command(
-                "search", "-c", "conda-forge", f"{package}=={version}", "--json"
+            search_result = subprocess.run(
+                ["conda", "search", "-c", "conda-forge", f"{package}=={version}", "--json"],
+                capture_output=True,
+                text=True,
+                check=True,
             )
-            package_metadata = json.loads(search_result[0])[package][0]
+            package_metadata = json.loads(search_result.stdout)[package][0]
             results[package] = {"version": package_metadata["version"], "depends": package_metadata["depends"]}
         except Exception as e:
             print(f"Error in report generation: {str(e)}")

src/utils.py

@@ -1,8 +1,8 @@
 import json
 import os
+import subprocess
 
-import conda.cli.python_api
-from conda.env.specs import RequirementsSpec
+from conda.env.specs.requirements import RequirementsSpec
 from conda.exceptions import PackagesNotFoundError
 from conda.models.match_spec import MatchSpec
 from semver import Version
@@ -48,7 +48,8 @@ def get_semver(version_str) -> Version:
     return version
 
 
-def read_env_file(file_path) -> RequirementsSpec:
+def read_env_file(file_path):
+    """Read environment file using conda's RequirementsSpec"""
     return RequirementsSpec(filename=file_path)
 
 
@@ -106,9 +107,15 @@ def pull_conda_package_metadata(image_config, image_artifact_dir):
         if str(match_spec_out).startswith("conda-forge"):
             # Pull package metadata from conda-forge and dump into json file
             try:
-                search_result = conda.cli.python_api.run_command("search", str(match_spec_out), "--json")
-                package_metadata = json.loads(search_result[0])[package][0]
+                search_result = subprocess.run(
+                    ["conda", "search", str(match_spec_out), "--json"], capture_output=True, text=True, check=True
+                )
+                package_metadata = json.loads(search_result.stdout)[package][0]
                 results[package] = {"version": package_metadata["version"], "size": package_metadata["size"]}
+            except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError, IndexError) as e:
+                print(
+                    f"Failed to pull package metadata for {package}, {match_spec_out} from conda-forge, ignore. Error: {str(e)}"
+                )
             except PackagesNotFoundError:
                 print(
                     f"Failed to pull package metadata for {package}, {match_spec_out} from conda-forge, ignore. Potentially this package is broken."

template/v2/Dockerfile

@@ -6,6 +6,12 @@ ARG ENV_IN_FILENAME
 ARG PINNED_ENV_IN_FILENAME
 ARG ARG_BASED_ENV_IN_FILENAME
 ARG IMAGE_VERSION
+
+# Amazon Q Agentic Chat version - update this default value when needed
+ARG FLARE_SERVER_VERSION_JL=1.25.0
+# IDE type for Amazon Q integration
+ARG AMAZON_Q_IDE_TYPE=jupyterlab
+
 LABEL "org.amazon.sagemaker-distribution.image.version"=$IMAGE_VERSION
 
 ARG AMZN_BASE="/opt/amazon/sagemaker"
@@ -49,6 +55,8 @@ ENV MAMBA_USER=$NB_USER
 ENV USER=$NB_USER
 
 COPY aws-cli-public-key.asc /tmp/
+COPY extract_amazon_q_agentic_chat_urls.py /tmp/
+COPY download_amazon_q_agentic_chat_artifacts.sh /tmp/
 
 RUN apt-get update && apt-get upgrade -y && \
     apt-get install -y --no-install-recommends sudo gettext-base wget curl unzip git rsync build-essential openssh-client nano cron less mandoc jq ca-certificates gnupg && \
@@ -73,7 +81,6 @@ RUN apt-get update && apt-get upgrade -y && \
     unzip q.zip && \
     Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \
     rm -rf q q.zip && \
-    : && \
     echo "source /usr/local/bin/_activate_current_env.sh" | tee --append /etc/profile && \
 # CodeEditor - create server, user data dirs
     mkdir -p /opt/amazon/sagemaker/sagemaker-code-editor-server-data /opt/amazon/sagemaker/sagemaker-code-editor-user-data \
@@ -121,6 +128,13 @@ RUN if [[ -z $ARG_BASED_ENV_IN_FILENAME ]] ; \
     find /opt/conda -name "yarn.lock" -type f -delete && \
     rm -rf /tmp/*.in && \
     sudo ln -s $(which python3) /usr/bin/python && \
+    # Download shared web client libraries
+    sudo mkdir -p /etc/web-client/libs && \
+    sudo curl -L --retry 3 --retry-delay 5 --fail "https://cdnjs.cloudflare.com/ajax/libs/jszip/3.10.1/jszip.min.js" -o "/etc/web-client/libs/jszip.min.js" || (echo "Failed to download JSZip library" && exit 1) && \
+    # Download Amazon Q Agentic Chat artifacts for JupyterLab integration
+    bash /tmp/download_amazon_q_agentic_chat_artifacts.sh $FLARE_SERVER_VERSION_JL /etc/amazon-q-agentic-chat/artifacts/$AMAZON_Q_IDE_TYPE $AMAZON_Q_IDE_TYPE && \
+    # Fix ownership for JupyterLab access
+    sudo chown -R $MAMBA_USER:$MAMBA_USER /etc/amazon-q-agentic-chat/ /etc/web-client/ && \
     # Update npm version
     npm i -g npm && \
     # Enforce to use `conda-forge` as only channel, by removing `defaults`