Feature Request: Make PyTorch and sentence-transformers Optional Dependencies

# Feature Request: Make PyTorch and sentence-transformers Optional Dependencies

## Summary

Make PyTorch, sentence-transformers, and related ML dependencies optional to significantly reduce container image size and startup time when using cloud-based embedding providers like Amazon Bedrock, OpenAI, or Cohere via LiteLLM.

## Problem Statement

Currently, the registry container image includes heavy ML dependencies (PyTorch, sentence-transformers, scikit-learn) totaling approximately 1.5-2GB, even when users configure cloud-based embedding providers that don't require these libraries.

**Current Situation:**
- PyTorch: ~800MB-1GB
- sentence-transformers: ~100-200MB (plus model downloads)
- scikit-learn: ~50-100MB
- Total overhead: ~1.5-2GB+ for dependencies that may never be used

**Container Image Impact:**
```bash
# Current image size (with PyTorch)
REPOSITORY SIZE
mcp-gateway-registry 3.2GB

# Potential size after optimization (cloud embeddings only)
mcp-gateway-registry 1.5GB
```

**Startup Time Impact:**
- Loading PyTorch and sentence-transformers adds 5-15 seconds to container startup
- Model download (first run) can take 30-60 seconds for all-MiniLM-L6-v2
- Cloud-based embeddings (Bedrock/OpenAI) have <1 second initialization

**Use Cases Affected:**

1. **Cloud Embeddings Only** (most common in production):
- `EMBEDDINGS_PROVIDER=litellm`
- `EMBEDDINGS_MODEL_NAME=bedrock/amazon.titan-embed-text-v2`
- No need for PyTorch or sentence-transformers

2. **Mixed Deployments**:
- Development: Local sentence-transformers for offline testing
- Production: Amazon Bedrock for scalability and cost
- Currently forced to ship large image to both environments

3. **Cost Optimization**:
- Smaller images reduce ECR storage costs
- Faster pulls reduce ECS task startup time
- Lower bandwidth costs in CI/CD pipelines

## Proposed Solution

Use Python optional dependencies pattern to install PyTorch and sentence-transformers only when needed.

### Implementation Strategy

#### Option 1: Python Optional Dependencies (Recommended)

**Approach**: Define optional dependency groups in `pyproject.toml` and install based on `EMBEDDINGS_PROVIDER` environment variable.

**Benefits**:
- Standard Python packaging pattern
- Clear dependency management
- Works with `uv` package manager
- Easy to understand and maintain
- Supports multiple embedding provider "flavors"

**Implementation**:

**1. Update pyproject.toml**

```toml
[project]
name = "mcp-registry"
version = "0.1.0"
description = "A registry for MCP servers"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
"fastapi>=0.115.12",
"itsdangerous>=2.2.0",
"jinja2>=3.1.6",
"mcp>=1.9.3",
"pydantic>=2.11.3",
"pydantic-settings>=2.0.0",
"httpx>=0.27.0",
"python-dotenv>=1.1.0",
"python-multipart>=0.0.20",
"uvicorn[standard]>=0.34.2",
"faiss-cpu>=1.7.4", # Keep FAISS for vector search
"websockets>=15.0.1",
"bandit>=1.8.3",
"langchain-mcp-adapters>=0.0.11",
"langgraph>=0.4.3",
"langchain-aws>=0.2.23",
"pytz>=2025.2",
"strands-agents>=0.1.6",
"strands-agents-tools>=0.1.4",
"pyjwt>=2.10.1",
"typing-extensions>=4.8.0",
"httpcore[asyncio]>=1.0.9",
"pyyaml>=6.0.0",
"langchain-anthropic>=0.3.17",
"matplotlib>=3.10.5",
"psutil>=6.1.0",
"email-validator>=2.2.0",
"aiohttp>=3.8.0",
"rich>=13.0.0",
"requests>=2.31.0",
"cisco-ai-a2a-scanner @ git+https://github.com/cisco-ai-defense/a2a-scanner.git@main",
"awscli>=1.36.0",
"boto3>=1.35.0",
"opensearch-py>=2.4.0",
"litellm>=1.50.0",
# NOTE: PyTorch, sentence-transformers, and scikit-learn moved to optional dependencies
]

[project.optional-dependencies]
# Local embeddings support (sentence-transformers)
embeddings-local = [
"sentence-transformers>=3.0.0",
"torch>=1.6.0",
"scikit-learn>=1.3.0",
"huggingface-hub>=0.31.1",
]

# All embedding providers
embeddings-all = [
"sentence-transformers>=3.0.0",
"torch>=1.6.0",
"scikit-learn>=1.3.0",
"huggingface-hub>=0.31.1",
]

# Development dependencies
dev = [
"pytest>=8.0.0",
"pytest-asyncio>=0.23.0",
"pytest-cov>=4.1.0",
"pytest-mock>=3.12.0",
"pytest-xdist>=3.5.0",
"coverage[toml]>=7.4.0",
"httpx>=0.27.0",
"pytest-html>=4.1.1",
"pytest-json-report>=1.5.0",
"factory-boy>=3.3.0",
"faker>=24.0.0",
"freezegun>=1.4.0",
]

docs = [
"mkdocs>=1.5.0",
"mkdocs-material>=9.4.0",
"mkdocs-git-revision-date-localized-plugin>=1.2.0",
"mkdocs-minify-plugin>=0.7.0",
"pymdown-extensions>=10.0.0",
]
```

**2. Update Dockerfile.registry**

Create a multi-stage Dockerfile with build args to control which dependencies are installed:

```dockerfile
# Registry Dockerfile - optimized for cloud embeddings by default
FROM python:3.12-slim

ENV PYTHONUNBUFFERED=1 \
DEBIAN_FRONTEND=noninteractive

# Build argument to control embeddings dependencies
ARG EMBEDDINGS_PROVIDER="litellm"
ENV EMBEDDINGS_PROVIDER=$EMBEDDINGS_PROVIDER

# Install system dependencies including nginx with lua module and Node.js
RUN apt-get update && apt-get install -y --no-install-recommends \
nginx \
nginx-extras \
lua-cjson \
curl \
procps \
openssl \
git \
build-essential \
ca-certificates \
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
&& apt-get install -y nodejs \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Build argument for version (will be set at build time from git)
ARG BUILD_VERSION="1.0.0"
ENV BUILD_VERSION=$BUILD_VERSION

# Install uv and create virtual environment
RUN pip install uv && \
uv venv .venv --python 3.12

# Copy pyproject.toml and install base dependencies
COPY pyproject.toml /app/
COPY . /app/

# Install dependencies based on EMBEDDINGS_PROVIDER
RUN . .venv/bin/activate && \
if [ "$EMBEDDINGS_PROVIDER" = "sentence-transformers" ]; then \
echo "Installing with local embeddings support (PyTorch + sentence-transformers)..." && \
uv pip install -e ".[embeddings-local]"; \
else \
echo "Installing with cloud embeddings only (no PyTorch)..." && \
uv pip install -e .; \
fi

# Copy nginx configurations
COPY docker/nginx_rev_proxy_http_only.conf /app/docker/nginx_rev_proxy_http_only.conf
COPY docker/nginx_rev_proxy_http_and_https.conf /app/docker/nginx_rev_proxy_http_and_https.conf

# Build React frontend
WORKDIR /app/frontend
COPY frontend/package.json ./
RUN npm install --legacy-peer-deps
COPY frontend/ ./
RUN npm run build

# Return to app directory
WORKDIR /app

# Create logs directory
RUN mkdir -p /app/logs

# Expose ports for nginx (HTTP/HTTPS) and registry
EXPOSE 80 443 7860

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
CMD curl -f http://localhost:7860/health || exit 1

# Entrypoint script
COPY docker/registry-entrypoint.sh /app/registry-entrypoint.sh
RUN chmod +x /app/registry-entrypoint.sh

ENTRYPOINT ["/app/registry-entrypoint.sh"]
```

**3. Update registry/embeddings/client.py**

Add helpful error message when sentence-transformers is not installed:

```python
class SentenceTransformersClient(EmbeddingsClient):
"""Client for local sentence-transformers models."""

def _load_model(self) -> None:
"""Load the sentence-transformers model."""
if self._model is not None:
return

try:
from sentence_transformers import SentenceTransformer
except ImportError as e:
logger.error(
"sentence-transformers is not installed. "
"Install with: uv pip install -e '.[embeddings-local]' "
"or set EMBEDDINGS_PROVIDER=litellm to use cloud embeddings"
)
raise RuntimeError(
"sentence-transformers is not installed. "
"To use local embeddings, install with: uv pip install -e '.[embeddings-local]'. "
"Alternatively, configure cloud embeddings by setting EMBEDDINGS_PROVIDER=litellm "
"and EMBEDDINGS_MODEL_NAME to a supported model (e.g., 'bedrock/amazon.titan-embed-text-v2')"
) from e

# Rest of implementation...
```

**4. Update registry-entrypoint.sh**

Add runtime check and optional installation:

```bash
#!/bin/bash
set -e

echo "Registry entrypoint starting..."
echo "EMBEDDINGS_PROVIDER: ${EMBEDDINGS_PROVIDER:-sentence-transformers}"

# If using sentence-transformers but it's not installed, install it now
if [ "${EMBEDDINGS_PROVIDER}" = "sentence-transformers" ]; then
if ! python -c "import sentence_transformers" 2>/dev/null; then
echo "sentence-transformers not found, installing embeddings-local dependencies..."
. /app/.venv/bin/activate
uv pip install -e ".[embeddings-local]"
else
echo "sentence-transformers already installed"
fi
fi

# Continue with normal startup...
```

**5. Update build scripts**

Create separate build targets for different deployment scenarios:

```bash
#!/bin/bash
# scripts/build-docker-images.sh

# Build for cloud embeddings (default, smallest image)
docker build \
--build-arg EMBEDDINGS_PROVIDER=litellm \
-t mcp-gateway-registry:latest \
-t mcp-gateway-registry:cloud \
-f docker/Dockerfile.registry .

# Build for local embeddings (includes PyTorch)
docker build \
--build-arg EMBEDDINGS_PROVIDER=sentence-transformers \
-t mcp-gateway-registry:local \
-f docker/Dockerfile.registry .

# Build for both (largest image, all dependencies)
docker build \
--build-arg EMBEDDINGS_PROVIDER=all \
-t mcp-gateway-registry:full \
-f docker/Dockerfile.registry .
```

#### Option 2: Runtime Installation (Alternative)

**Approach**: Install PyTorch and sentence-transformers at container startup if `EMBEDDINGS_PROVIDER=sentence-transformers`.

**Benefits**:
- Single container image supports both modes
- Flexible runtime configuration
- No need to build multiple images

**Drawbacks**:
- Slower first startup when downloading packages
- Requires internet connectivity at runtime
- Package cache not part of image (repeated downloads)
- More complex entrypoint logic

**Implementation**: See entrypoint.sh changes in Option 1.

### Comparison: Option 1 vs Option 2

| Criteria | Option 1: Build-time | Option 2: Runtime Install |
|----------|---------------------|---------------------------|
| **Image Size** | Smallest (1.5GB cloud, 3.2GB local) | Medium (1.5GB base) |
| **Startup Time** | Fast (cloud), Medium (local) | Slow on first run (downloads) |
| **Flexibility** | Multiple images, clear intent | Single image, runtime choice |
| **Internet Required** | Build time only | Runtime (first use) |
| **Cache Efficiency** | Dependencies baked in | Must download each new container |
| **CI/CD Speed** | Fast (pre-built) | Slow (runtime install) |
| **Complexity** | Medium (multiple build targets) | Low (single Dockerfile) |
| **Recommended For** | Production deployments | Development/testing only |

**Recommendation**: Use **Option 1 (Build-time)** for production with separate image tags:
- `mcp-gateway-registry:latest` / `mcp-gateway-registry:cloud` - Cloud embeddings (Bedrock/OpenAI)
- `mcp-gateway-registry:local` - Local sentence-transformers
- `mcp-gateway-registry:full` - Both (for flexibility)

### Files Requiring Changes

#### Modified Files:

1. **pyproject.toml**
- Move `torch`, `sentence-transformers`, `scikit-learn`, `huggingface-hub` to `[project.optional-dependencies]`
- Create `embeddings-local` group for local embedding dependencies
- Create `embeddings-all` group for all embedding providers
- Estimated changes: 10 lines moved, 15 lines added

2. **docker/Dockerfile.registry**
- Add `ARG EMBEDDINGS_PROVIDER=litellm` build argument
- Conditional installation based on build arg
- Use `uv pip install -e ".[embeddings-local]"` when local embeddings needed
- Estimated changes: 20 lines modified

3. **docker/registry-entrypoint.sh**
- Add runtime check for sentence-transformers availability
- Optional runtime installation if provider is sentence-transformers
- Clear logging for which mode is active
- Estimated changes: 15 lines added

4. **registry/embeddings/client.py**
- Improve ImportError messages for sentence-transformers
- Suggest installation command or alternative configuration
- Estimated changes: 10 lines modified

#### New Files:

1. **scripts/build-docker-images.sh**
- Script to build multiple image variants
- Clear naming: `cloud`, `local`, `full`
- Documentation for each variant

2. **docs/embeddings-deployment.md**
- Guide for choosing embedding provider
- Docker build instructions for each scenario
- Environment variable configuration
- Cost and performance comparison

3. **.env.example updates**
- Add comments explaining EMBEDDINGS_PROVIDER options
- Document which dependencies are needed for each provider

### Docker Compose Configuration

Update `docker-compose.yml` to support both deployment modes:

```yaml
# Cloud embeddings (default - smallest image)
services:
registry-cloud:
build:
context: .
dockerfile: docker/Dockerfile.registry
args:
EMBEDDINGS_PROVIDER: litellm
environment:
- EMBEDDINGS_PROVIDER=litellm
- EMBEDDINGS_MODEL_NAME=bedrock/amazon.titan-embed-text-v2
- AWS_REGION=us-east-1
# ... rest of config

# Local embeddings (larger image with PyTorch)
registry-local:
build:
context: .
dockerfile: docker/Dockerfile.registry
args:
EMBEDDINGS_PROVIDER: sentence-transformers
environment:
- EMBEDDINGS_PROVIDER=sentence-transformers
- EMBEDDINGS_MODEL_NAME=all-MiniLM-L6-v2
# ... rest of config
```

### Terraform Updates

Update ECS task definitions to support different image variants:

```hcl
# terraform/aws-ecs/variables.tf

variable "embeddings_provider" {
description = "Embeddings provider: litellm (cloud) or sentence-transformers (local)"
type = string
default = "litellm"
validation {
condition = contains(["litellm", "sentence-transformers"], var.embeddings_provider)
error_message = "embeddings_provider must be 'litellm' or 'sentence-transformers'"
}
}

variable "registry_image_tag" {
description = "Registry container image tag (cloud, local, or full)"
type = string
default = "cloud"
}
```

### Testing Strategy

#### Unit Tests:

1. **Test Optional Import Handling**:
```python
def test_sentence_transformers_missing_import():
"""Test graceful handling when sentence-transformers not installed."""
with patch('builtins.__import__', side_effect=ImportError):
with pytest.raises(RuntimeError, match="sentence-transformers is not installed"):
client = SentenceTransformersClient(model_name="test")
client.encode(["test"])
```

2. **Test Provider Selection**:
```python
def test_create_embeddings_client_litellm():
"""Test creating LiteLLM client without sentence-transformers."""
client = create_embeddings_client(
provider="litellm",
model_name="bedrock/amazon.titan-embed-text-v2"
)
assert isinstance(client, LiteLLMClient)
```

#### Integration Tests:

1. **Build Cloud Image**:
```bash
docker build --build-arg EMBEDDINGS_PROVIDER=litellm -t test-cloud .
# Verify image size < 2GB
# Verify sentence-transformers NOT in pip list
```

2. **Build Local Image**:
```bash
docker build --build-arg EMBEDDINGS_PROVIDER=sentence-transformers -t test-local .
# Verify sentence-transformers in pip list
# Verify can load model
```

3. **Runtime Configuration Test**:
```bash
# Start with cloud image, try to use sentence-transformers (should fail gracefully)
docker run -e EMBEDDINGS_PROVIDER=sentence-transformers test-cloud
# Should show clear error message
```

#### Performance Tests:

1. **Image Size Comparison**:
```bash
docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}"
# Expected: cloud ~1.5GB, local ~3.2GB
```

2. **Startup Time Comparison**:
```bash
time docker run test-cloud /app/.venv/bin/python -c "import registry; print('ready')"
time docker run test-local /app/.venv/bin/python -c "import registry; print('ready')"
# Expected: cloud <2s, local 5-10s
```

### Migration Guide

For existing deployments:

#### Development/Testing Environments:
```bash
# Option 1: Use cloud embeddings (recommended)
export EMBEDDINGS_PROVIDER=litellm
export EMBEDDINGS_MODEL_NAME=bedrock/amazon.titan-embed-text-v2
docker-compose up -d

# Option 2: Continue using local embeddings (no change)
export EMBEDDINGS_PROVIDER=sentence-transformers
docker build --build-arg EMBEDDINGS_PROVIDER=sentence-transformers -t mcp-gateway-registry:local .
```

#### Production Deployments:
```bash
# 1. Update Terraform variables
embeddings_provider = "litellm"
registry_image_tag = "cloud"

# 2. Update ECS environment variables
EMBEDDINGS_PROVIDER=litellm
EMBEDDINGS_MODEL_NAME=bedrock/amazon.titan-embed-text-v2
EMBEDDINGS_MODEL_DIMENSIONS=1024

# 3. Deploy new task definition
terraform apply
```

### Documentation Updates

1. **README.md**:
- Add "Embeddings Configuration" section
- Explain cloud vs local embeddings
- Document build args and environment variables

2. **docs/embeddings.md**:
- Update with optional dependency information
- Add Docker build examples for each mode
- Cost and performance comparison table

3. **docs/deployment.md** (new):
- Production deployment recommendations
- Image size and performance metrics
- AWS Bedrock setup guide
- Local embeddings setup guide

4. **.env.example**:
```bash
# Embeddings Configuration
# Provider: 'litellm' (cloud, recommended for production) or 'sentence-transformers' (local)
EMBEDDINGS_PROVIDER=litellm

# For litellm provider (cloud embeddings):
EMBEDDINGS_MODEL_NAME=bedrock/amazon.titan-embed-text-v2
EMBEDDINGS_MODEL_DIMENSIONS=1024
EMBEDDINGS_AWS_REGION=us-east-1
# Note: PyTorch and sentence-transformers NOT required

# For sentence-transformers provider (local embeddings):
# EMBEDDINGS_PROVIDER=sentence-transformers
# EMBEDDINGS_MODEL_NAME=all-MiniLM-L6-v2
# EMBEDDINGS_MODEL_DIMENSIONS=384
# Note: Requires Docker build with --build-arg EMBEDDINGS_PROVIDER=sentence-transformers
```

### Benefits

**Image Size Reduction**:
- Cloud embeddings image: ~1.5GB (50% reduction from 3.2GB)
- Faster ECR push/pull operations
- Lower storage costs

**Startup Time Improvement**:
- Cloud embeddings startup: <2 seconds
- No model download on first run
- Faster ECS task replacement

**Cost Savings**:
- Reduced ECR storage costs (~$0.10/GB/month)
- Lower data transfer costs
- Faster CI/CD pipeline (smaller image builds)

**Flexibility**:
- Choose embedding provider per deployment
- Easy to test different configurations
- Support both cloud and on-premises deployments

### Success Criteria

- [ ] `pyproject.toml` updated with optional dependency groups
- [ ] Dockerfile supports build-time provider selection
- [ ] Cloud embeddings image size < 2GB
- [ ] Local embeddings image includes PyTorch and sentence-transformers
- [ ] Clear error messages when sentence-transformers not available
- [ ] Unit tests pass for both configurations
- [ ] Integration tests verify image builds and runtime behavior
- [ ] Documentation complete (README, embeddings.md, deployment guide)
- [ ] Terraform variables support embeddings_provider selection
- [ ] No regression in functionality for either mode

### Estimated Effort

**Development**: 1 week for experienced DevOps/Backend engineer
- Refactor pyproject.toml: 2 hours
- Update Dockerfile with build args: 4 hours
- Update entrypoint.sh: 2 hours
- Improve error messages: 2 hours
- Create build scripts: 4 hours
- Terraform updates: 4 hours

**Testing**: 3 days
- Unit tests: 1 day
- Integration tests: 1 day
- Performance benchmarking: 1 day

**Documentation**: 2 days
- Update existing docs: 1 day
- Create deployment guide: 1 day

**Total**: 2 weeks for complete implementation, testing, and documentation

### Future Enhancements

1. **CPU vs GPU PyTorch**: Separate optional dependency groups for `faiss-cpu` vs `faiss-gpu`
2. **Multi-arch Builds**: ARM64 support for AWS Graviton instances
3. **Layer Caching**: Optimize Dockerfile layers for faster rebuilds
4. **Slim Base Images**: Use `python:3.12-slim` variants for even smaller images
5. **Embedding Provider Plugins**: Dynamic loading of embedding provider modules

## Related Issues

- Embeddings abstraction implemented in `registry/embeddings/client.py`
- Configuration in `registry/core/config.py` (EMBEDDINGS_PROVIDER setting)
- Docker builds in `docker/Dockerfile.registry`

## Labels

`enhancement`, `docker`, `dependencies`, `embeddings`, `devops`, `cost-optimization`

## Acceptance Criteria

- [ ] PyTorch and sentence-transformers are optional dependencies
- [ ] Docker build with `--build-arg EMBEDDINGS_PROVIDER=litellm` produces image < 2GB
- [ ] Docker build with `--build-arg EMBEDDINGS_PROVIDER=sentence-transformers` includes PyTorch
- [ ] Clear error message when using sentence-transformers without dependencies
- [ ] Cloud embeddings (Bedrock/OpenAI) work without PyTorch installed
- [ ] Local embeddings work with PyTorch installed
- [ ] Unit and integration tests pass for both configurations
- [ ] Documentation updated with deployment options
- [ ] Terraform supports embeddings_provider variable
- [ ] No breaking changes for existing deployments