Data Discovery AI

A consolidated framework for AI and ML components developed for the new portal.

🔧 Setup

Environment variables

To run the app locally:

1. Go to the root folder of this project.
2. Create a .env file. (or simply copy the .env.sample file)
3. Add the following (with your real keys):

API_KEY="your_actual_api_key_here"
OPENAI_API_KEY="your_actual_openai_api_key_here"
PROFILE="your_actual_environment_here"

Elasticsearch is required to process requests and store generated data. Also include:

ES_ENDPOINT="your_actual_elasticsearch_endpoint"
ES_API_KEY="your_actual_es_api_key"

🚀 Running the App

Option 1: Run with Docker (Recommended for Non-Developers)

1. Copy the content in .env.sampleto.env` and fill in your keys.
2. Run

   ./startServer.sh

   Then visit http://localhost:8000
3. Test App health at http://localhost:8000/api/v1/ml/health

Option 2: Run for Development (For Developers)

Conda (recommended for creating a virtual environment)

1. Install Conda (if not already installed):

   Follow the instructions at Conda Installation.

2. Create Conda virtual environment:

   conda env create -f environment.yml

Dependencies management

Poetry is used for dependency management, the pyproject.toml file is what is the most important, it will orchestrate the project and its dependencies.

You can update the file pyproject.toml for adding/removing dependencies by using

poetry add <pypi-dependency-name> # e.g poetry add numpy
poetry remove <pypi-dependency-name> # e.g. poetry remove numpy

You might want to update the poetry.lock file after manually modifying pyproject.toml with poetry lock command. To update all dependencies, use poetry update command.

Installation and Usage

1. Activate Conda virtual environment:

   conda activate data-discovery-ai

2. Install environment dependencies:

   # after cloning the repo with git clone command
   cd data-discovery-ai
   poetry install

Before Usage

FastAPI runs internal checks before making \process_record API calls. These checks include:

1. ✅ Required model resource files must be present in data_discovery_ai/resources/.
2. ✅ A valid OPENAI_API_KEY must be in .env unless you're in development environment.
3. ✅ If PROFILE=development, Ollama must be running locally at http://localhost:11434.

(Optional) Install Llama3 for Local Development

To use the Llama3 model locally without OpenAI:

1. Go to Ollama download page and download the version that matches your operating system (Windows, Linux, or macOS).

2. After installation, start Ollama either by launching the app or running the following command:

   ollama serve

3. Pull the "llama3" model used for local development：

   ollama pull llama3

4. (Optional) Consider installing Open WebUI to test Llama3 locally through a user-friendly interface:

   docker run -d --network=host -v open-webui:/app/backend/data -e PORT=8090 -e OLLAMA_BASE_URL=http://127.0.0.1:11434 --name open-webui ghcr.io/open-webui/open-webui:main

   Once the Open WebUI container is running, open your browser and go to: http://localhost:8090.

Run the FastAPI Server

Simply run:

python -m data_discovery_ai.server

Host will be run at http://localhost:8000.

Running Tests

Run all tests with:

poetry run python -m unittest discover -s tests

Code formatting

Run manual checks:

pre-commit run --all-files

Checks are also executed when you run git commit. The configurations for pre-commit hooks are defined in .pre-commit-config.yaml.

Git Commit Guide (Optional)

We are using gitmoji(OPTIONAL) with husky and commitlint. Here you have an example of the most used ones:

• 🎨 - Improving structure/format of the code.
• ⚡ - Improving performance.
• 🔥 - Removing code or files.
• 🐛 - Fixing a bug.
• 🚑 - Critical hotfix.
• ✨ - Introducing new features.
• 📝 - Adding or updating documentation.
• 🚀 - Deploying stuff.
• 💄 - Updating the UI and style files.
• 🎉 - Beginning a project.

Example of use: :wrench: add husky and commitlint config

Branching name

• hotfix/: for quickly fixing critical issues,
• usually/: with a temporary solution
• bugfix/: for fixing a bug
• feature/: for adding, removing or modifying a feature
• test/: for experimenting something which is not an issue
• wip/: for a work in progress

And add the issue id after an / followed with an explanation of the task.

Example of use: feature/5348-create-react-app

▶️ Using the API

Once the app is running, two routes are available:

Route	Description
GET /api/v1/ml/health	Health check
POST /api/v1/ml/process_record	One single point for calling AI models to process metadata record
DELETE /api/v1/ml/delete_doc	Deletes a document from the AI-related Elasticsearch index. Requires query parameter doc_id.

Health Check

GET /api/v1/ml/health always returns HTTP 200; readiness is in the body as status:

• STARTING: the server is up and Hugging Face models are downloading/loading in the background.
• UP: all components are ready.
• DOWN: a component failed, see components for details.

In Docker, Nginx serves this endpoint from /tmp/status/health.json, so the health check stays responsive while Python is busy (the same approach as data-access-service). Why Nginx is needed: startup work (Hugging Face model downloads, Elasticsearch setup) and CPU-bound inference can keep the event loop from answering in time, and the orchestrator would kill a container that fails its health check; a static file served by Nginx cannot be blocked by Python. Nginx also absorbs slow or large client connections in front of the single Uvicorn worker.

Uvicorn listens on 127.0.0.1:9000 (APP_HOST/APP_PORT, set in the Dockerfile) and Nginx is the only listener on port 8000, proxying everything except the health path. Both processes run as appuser: docker-entrypoint.sh starts Nginx in the background and then execs the container command, so CMD ["python", "-m", "data_discovery_ai.server"] is still PID 1 and still receives SIGTERM directly, as it did before Nginx was introduced. No supervisord or root process is involved. If Nginx dies, port 8000 stops answering and the orchestrator replaces the container, the same signal a dead Python process gives.

The health file only holds startup state (keyword_resources, models, elasticsearch): the application writes it at startup and when each background startup task finishes, and removes it during graceful shutdown, after which Nginx returns 404. The live llm check is not in the file, because the file is not rewritten when it changes; it is only reported by the FastAPI health route (when running Uvicorn directly) and enforced by process_record and delete_doc, which can therefore still return 503 while the file says UP.

If the application is killed without a graceful shutdown (e.g. OOM or SIGKILL), the last file stays in place until the container restarts and the application writes STARTING again. This is an accepted limitation, as in data-access-service.

process_record and delete_doc return HTTP 503 until the status is UP.

Elasticsearch setup runs in the background with a bounded retry policy. If those retries are exhausted, the elasticsearch component changes to DOWN and setup stops. The health endpoint remains available; correct the Elasticsearch configuration or service and restart the application to try again.

Example Request Body

{
    "selected_model":["description_formatting"],
    "uuid": "test",
    "title": "test title",
    "abstract": "test abstract"
}

Required Header

X-API-Key: your_api_key

(Must match the value of API_KEY specified in the environment variables).

AI Model Options

• selected_model: the AI models provided by data-discovery-ai. It should be a list of strings, which are the name of the AI task agents. Currently, four AI task agents available for distinctive tasks:
  • keyword_classification: predict keywords from AODN vocabularies based on metadata title and abstract with pretrained ML model.
  • delivery_classification: predict data delivery mode based on metadata title, abstract, and lineage with pretrained ML model.
  • description_formatting: reformatting long abstract into Markdown format based on metadata title and abstract with LLM model "gpt-4o-mini".
  • link_grouping: categorising links into four groups: ["Python Notebook", "Document", "Data Access", "Other"] based on metadata links.

🤖 Model Training (Optional)

Currently, two machine learning pipelines are available for training and evaluating models:

• keyword: keyword classification model, which is a Sequential model for multi-label classification task
• delivery: data delivery classification model, which is a self-learning model for binary classification task

How to Run

To run one of the pipelines (for example, the keyword one), you can use the following command in your terminal:

python -m data_discovery_ai.ml.pipeline --pipeline keyword --use_cached_raw False --start_from_preprocess True --model_name experimental

You can also use a shorter version:

 python -m data_discovery_ai.ml.pipeline -p keyword -r False -s True -n experimental

When Should I Re-Train?

If the raw data has changed (e.g., updated, cleaned, or expanded), you are recommended to re-train the model using the latest data. To do this, set:

--use_cached_raw False --start_from_preprocess True

As mentioned in Environment variables, ElasticSearch endpoint and API key are required to be set up in .env file.

What Happens When I Run a Pipeline?

Running a pipeline trains a Machine Learning model and saves several resource files for reuse — so you don’t have to reprocess data or retrain the model every time.

For cached purpose, raw data which contains all collections from OGCAPI is saved in data_discovery_ai/resources/raw_data.pkl

1. delivery pipeline

   Outputs are saved in: data_discovery_ai/resources/DataDeliveryModeFilter/

   File Name	Description
   filter_preprocessed.pkl	Preprocessed data used for training and testing
   development.pkl	Trained binary classification model

2. keyword pipeline

   Outputs are saved in: data_discovery_ai/resources/KeywordClassifier/

   File Name	Description
   keyword_sample.pkl	Preprocessed data used for training and testing
   keyword_label.pkl	Mapping between labels and internal IDs
   development.keras	Trained Keras model file (name set by --model_name)

These files are generated automatically during pipeline training and saved automatically after pipeline running. They are intended for reuse in subsequent runs to avoid retraining.

Accepted Model Names

The --model_name argument helps organise different versions of your model. Here's how they’re typically used:

Name	Purpose	When to Use
development	Active model development	For testing and iterating on new ideas
experimental	Try new techniques or tuning	For exploring new features or architectures
benchmark	Compare against the current baseline model	When validating improvements over a previous version
staging	Pre-production readiness	When testing full integration before final deployment
production	Final production model	Live version used in production APIs or systems

🛠 Tip: When working locally, use --model_name experimental to avoid overwriting files used in deployments.

Model Naming Guidelines

Each model name reflects a stage in the model lifecycle:

1. Development

   • Initial model design and prototyping
   • Reaches minimum performance targets with stable training

2. Experimental

   • Shows consistent performance improvements
   • Experiment logs and results are clearly documented

3. Benchmark

   • Outperforms the existing benchmark (usually a copy of the production model)
   • Validated using selected evaluation metrics

4. Staging

   • Successfully integrated with application components (e.g. APIs)
   • Ready for deployment, pending final checks

5. Production

   • Deployed in a live environment
   • Monitored continuously, supports user feedback and live data updates

Model Usage

In the configuration file data_discovery_ai/common/config-*.yaml, while * specify the environment for use. You can specify which model version each task should use. For example:

model:
  delivery_classification:
    pretrained_model: development

This means the agent handling the delivery_classification task will use the development version of the model.

Track Training Progress with MLflow

We use MLflow to track model training and performance over time (like hypermeters, accuracy, precision, etc.).

It is automatically started at the beginning of the pipeline. Once a pipeline starts, you can open the tracking dashboard in your browser: http://127.0.0.1:53000

You can change the model's training settings (like how long it trains or how fast it learns) by editing the trainer section in the file: data_discovery_ai/config/parameters.yaml

🛠️ Required Configuration Files

File	Description
data_discovery_ai/common/constants.py	Shared constants
data_discovery_ai/commom/parameters.yaml	Store parameter settings for ML models and AI agents.
data_discovery_ai/commom/config-*.yaml	Store parameter configuration for different environments.

📁 Project Structure

data_discovery_ai/
├── config/             # Common utilities and shared configurations/constants used across modules
├── enum/               # Enums to use
├── core/               # Core logic of the application such as API routes
├── agents/             # Task-specific agent modules using ML/AI/rule-based tools
├── ml/                 # Machine learning models: training, inference, evaluation logic
├── utils/              # Utility functions and helper scripts for various tasks
├── resources/          # Stored assets such as pretrained models, sample datasets, and other resources required for model inference
├── notebooks/          # Jupyter notebooks
├── tests/              # Unit tests for validating core components
│   ├── agents
│   ├── ml
│   └── utils
├── server.py             # FastAPI application entry point