AI as an API — Spam Classifier (FastAPI + TensorFlow + Astra DB)

Production-grade example of serving a text-classification model (spam vs ham) behind a high-performance REST API with caching and request logging backed by Apache Cassandra (DataStax Astra DB).

This repository contains:

• A FastAPI service exposing a Keras/TensorFlow model for text classification
• An optional mock model to run the API without training
• A complete data preparation and model training workflow
• A Cassandra/Astra DB-backed cache and request log, including a streaming endpoint
• Automated table sync and integration tests for API behavior

If you want a guided course with interactive labs, see the original workshop on DataStax Academy.

Contents

• Overview and Architecture
• Project Structure
• Requirements
• Setup
• Environment Configuration
• Train the Model
• Run the API
• API Endpoints
• Caching and Logging (Astra DB)
• Test Utilities
• Troubleshooting

Overview and Architecture

The service exposes a text-classifier to predict whether an input text is spam or ham. Requests are served through FastAPI. Predictions can be cached in Astra DB to reduce inference latency for repeated inputs. Each call is logged (per-caller, per-hour) and can be streamed back to clients.

Key components:

• Model layer (TensorFlow/Keras): loads model (.h5), tokenizer (JSON), and metadata (JSON)
• API layer (FastAPI): request/response validation (Pydantic), OpenAPI docs, endpoints
• Data layer (Astra DB/Cassandra via cqlengine): cache table, call-log table, schema sync
• Config layer (Pydantic BaseSettings): reads from .env; secrets are never returned by public endpoints

There are two API variants:

• Full API: api.main:app — includes caching and logging
• Minimal API: api.minimain:miniapp — simple prediction endpoints only, no DB integration

Project Structure

Key paths (non-exhaustive):

• api/
  • main.py — Full FastAPI app with caching, logging, streaming
  • minimain.py — Minimal FastAPI app (no DB)
  • AIModel.py — Keras/TensorFlow model wrapper
  • MockSpamAIModel.py — Mock classifier to run API without training artifacts
  • config.py — Settings management (Pydantic BaseSettings)
  • schema.py — Request/response models
  • dbio.py — DB I/O helpers (cache, recent call log, streaming JSON)
  • database/
    • db.py — Astra DB connection/session setup
    • models.py — cqlengine model definitions and keyspace wiring
  • tests/
    • aiModelTest.py — Sanity check for model loading/inference
    • apiRequestTest.py — Stress/consistency testing of API endpoints
    • paginationTest.py — Example of iterating DB results
• training/
  • dataset/spam-dataset.csv — Labeled dataset
  • prepared_dataset/ — Output location for preprocessed data
  • trained_model_v1/ — Output location for trained model/tokenizer/metadata
• prepareDataset.py — Step 1: dataset preprocessing
• trainModel.py — Step 2: model training and artifact export
• loadTestModel.py — Quick validation of trained artifacts
• requirements.txt — Python dependencies (pinned for TF 2.6 compatibility)

Requirements

• Python 3.8 or 3.9 recommended (TensorFlow 2.6 compatibility)
• pip (or pipx) and virtualenv tooling
• For full functionality: a DataStax Astra DB instance and a Secure Connect bundle
• For quick demos: you can run with a mock model and skip DB by setting MOCK_MODEL_CLASS=1

Note on versions: TensorFlow 2.6 is pinned together with protobuf==3.20.0 for compatibility (see requirements.txt).

Setup

1. Clone and create a virtual environment

• python -m venv .venv
• .venv\Scripts\activate (Windows) or source .venv/bin/activate (macOS/Linux)
• pip install -r requirements.txt

2. Environment file

• Copy .app-env.sample to .env and set values (see next section)
• If you plan to use Astra DB, download the Secure Connect bundle and place its path in .env

3. Optional: Train a model or enable the mock model

• To train, follow Train the Model
• To skip training, set MOCK_MODEL_CLASS="1" in .env and use the API right away (predictions will be fake but schema-compliant)

Environment Configuration

The API reads configuration from .env via Pydantic BaseSettings. Required keys:

Core

• API_NAME: name displayed at the root endpoint
• MODEL_VERSION: version label used in cache partitioning (e.g., v1)
• MODEL_DIRECTORY: path to model artifacts (e.g., training/trained_model_v1)
• CQLENG_ALLOW_SCHEMA_MANAGEMENT: set to "1" to allow cqlengine to create tables at startup
• MOCK_MODEL_CLASS: "0" to use the real model, "1" to use the mock model

Astra DB

• ASTRA_DB_KEYSPACE: keyspace for tables
• ASTRA_DB_SECURE_BUNDLE_PATH: path to the secure connect bundle zip
• ASTRA_DB_APPLICATION_TOKEN: application token with data access

Example (do not commit secrets):

API_NAME="Spam Classifier" MODEL_VERSION="v1" MODEL_DIRECTORY="training/trained_model_v1" CQLENG_ALLOW_SCHEMA_MANAGEMENT="1" MOCK_MODEL_CLASS="0"

ASTRA_DB_APPLICATION_TOKEN="AstraCS:..." ASTRA_DB_SECURE_BUNDLE_PATH="./secure-connect-classifier-db.zip" ASTRA_DB_KEYSPACE="classifier_space"

Train the Model

If you want real predictions, you must train and export artifacts into training/trained_model_v1/:

1. Prepare dataset

• python prepareDataset.py -v # -v for verbose (optional) This generates training/prepared_dataset/spam_training_data.pickle.

2. Train model

• python trainModel.py This writes three files into training/trained_model_v1/:
• spam_model.h5 — Keras model
• spam_tokenizer.json — tokenizer
• spam_metadata.json — metadata (max_seq_length, label legends, etc.)

3. Validate artifacts (optional)

• python loadTestModel.py "Your test sentence here"

If you prefer to skip training, set MOCK_MODEL_CLASS="1" in .env to use the mock classifier.

Run the API

From the repository root, with your virtual environment active:

Full API (with DB) — includes caching, call logging, and the streaming endpoint:

• uvicorn api.main:app --reload

Minimal API (no DB) — simple predictions only:

• uvicorn api.minimain:miniapp --reload

Docs and testing UI:

• Swagger UI: http://localhost:8000/docs
• ReDoc: http://localhost:8000/redoc

API Endpoints

All responses are validated with Pydantic models.

• GET / — Basic API info

  • Returns non-secret settings and runtime information (caller_id, started_at)

• POST /prediction — Classify a single text

  • Body: { text: string, echo_input?: bool=false, skip_cache?: bool=false }
  • Response: { input?, prediction: {label: score, ...}, top: {label, value}, from_cache: bool }
  • If skip_cache=false, the API reads from cache first; else it computes and stores the result

• GET /prediction — Same as above through query params

  • Example: /prediction?text=Hello%20world&echo_input=true&skip_cache=true

• POST /predictions — Classify multiple texts

  • Body: { texts: string[], echo_input?: bool=true, skip_cache?: bool=false }
  • Efficiently mixes cached and non-cached inputs, merges results, and returns them in the original order

• GET /recent_log — Stream recent call log (current hour) for the caller

  • Streams a JSON array of entries: [{ index, input, called_at }, ...]
  • Implemented with StreamingResponse to avoid holding large responses in memory

Example requests

• Single (POST):

  • curl -X POST "http://localhost:8000/prediction" -H "Content-Type: application/json" -d "{"text":"URGENT! You have WON a prize", "echo_input": true}"

• Single (GET):

  • curl "http://localhost:8000/prediction?text=Hello&echo_input=true&skip_cache=true"

• Multiple (POST):

  • curl -X POST "http://localhost:8000/predictions" -H "Content-Type: application/json" -d "{"texts":["Hello","WIN a FREE IPHONE"]}"

• Stream recent log:

  • curl "http://localhost:8000/recent_log"

Caching and Logging (Astra DB)

Two cqlengine models are used:

• SpamCacheItem (table spam_cache_items): caches top label/value and full prediction map by (model_version, input)
• SpamCallItem (table spam_calls_per_caller): logs calls by (caller_id, called_hour, called_at)

On startup, the API:

• Connects to Astra DB using the Secure Connect bundle and application token
• Syncs tables via cassandra.cqlengine.management.sync_table

Notes:

• Ensure CQLENG_ALLOW_SCHEMA_MANAGEMENT="1" in .env so cqlengine can create tables
• Keyspace must already exist (ASTRA_DB_KEYSPACE)
• readCachedPrediction requires MODEL_VERSION to match your stored items

Test Utilities

• api/tests/apiRequestTest.py

  • Starts by fetching reference scores, then repeatedly tests /prediction and /predictions
  • Intended to run with the API active at http://localhost:8000
  • Usage: python api/tests/apiRequestTest.py 50 # optional iterations

• api/tests/aiModelTest.py

  • Loads the model and runs a small prediction example, printing structured outputs

• api/tests/paginationTest.py

  • Demonstrates iterating results via cqlengine (assumes data exists for given keys)

Troubleshooting

• TensorFlow or protobuf errors

  • This repo pins tensorflow==2.6.0 and protobuf==3.20.0 to avoid incompatibilities
  • Use Python 3.8/3.9 for best compatibility with TF 2.6

• Model files missing

  • Train the model (prepareDataset.py, trainModel.py) or set MOCK_MODEL_CLASS="1"

• Astra DB connection issues

  • Verify ASTRA_DB_SECURE_BUNDLE_PATH points to the correct zip
  • Ensure ASTRA_DB_APPLICATION_TOKEN is valid and has permissions
  • Confirm ASTRA_DB_KEYSPACE exists

• Tables not found / schema errors

  • Ensure CQLENG_ALLOW_SCHEMA_MANAGEMENT="1"
  • Startup sync_table runs automatically; check logs on startup

• Import errors when running uvicorn

  • Run uvicorn from the repository root so that the api package is importable

• Windows path quoting

  • If your secure bundle path contains spaces, wrap it in quotes in .env

License

This project is provided for educational and demonstration purposes as part of the “AI as an API” workshop. Review third-party licenses for dependencies listed in requirements.txt.