AI Log Classification Engine

This project implements a pipeline that intercepts raw logs using Vector, enriches them with AI-predicted tags using XGBoost & TF-IDF, and forwards them to SigNoz for observability.

🏗 Architecture

The system follows a sidecar architecture pattern:

1. Log Source: Applications/VMs send logs via HTTP (JSON).
2. Vector (Collector): Intercepts the logs.
3. ML API (Sidecar): Vector pauses, sends the log text to the Python API.
4. AI Engine:
   • TF-IDF: Converts text to numbers.
   • XGBoost: Predicts source_type (e.g., Security) and content_type (e.g., Error).
5. SigNoz: Displays the logs with the new AI tags.

📂 Repository Structure

Directory/File	Description
ml_api/	The Flask API that runs inside Docker. Serves predictions.
training/	Source code to retrain the models from scratch.
vector/	Configuration for the Vector log collector.
models/	Binary model artifacts generated by training.
docker-compose.yaml	Orchestrates the entire pipeline. Starts the Vector and ML API containers.
app.py	(Optional) Script for local development/testing.

🧪 Methodology & Experimentation

The development process evolved through three distinct phases:

1. Phase 1 (Deep Learning): Sentence Transformer + Neural Network.
2. Phase 2 (The ML Hybrid): Sentence Transformer + XGBoost Classifier.
3. Phase 3 (Current - Distillation): Qwen Labeling + TF-IDF + XGBoost.

---

Phase 1: Pure Deep Learning

The initial attempt focused purely on semantic understanding using pre-trained transformer models.

• Architecture:
  • Embedding: all-MiniLM-L6-v2 to generate dense vectors.
  • Classifier: A fully connected Neural Network (PyTorch/TensorFlow).
• Outcome:
  • ✅ Pros: Accurate classification.
  • ❌ Cons: Extremely slow inference and difficult to tune. The Neural Network required massive data to converge.

---

Phase 2: The ML Hybrid

To improve training stability and interpretability, we replaced the Neural Network with XGBoost, but kept the heavy Sentence Transformer embeddings.

• Architecture:
  • Embedding: all-MiniLM-L6-v2
  • Classifier: XGBoost (Gradient Boosting).
• Why: XGBoost is generally faster and more robust than Neural Networks on tabular data (embeddings).
• The Bottleneck:
  • ❌ Latency: Slow embedding process.
  • ❌ Semantic Mismatch: Sentence Transformer is pre-trained on standard English (natural language). It focuses on the literal meaning of words, which causes poor performance on technical logs.

---

Phase 3: The "Student-Teacher" Distillation

To solve the latency bottleneck, we moved to a Knowledge Distillation strategy. We realized we didn't need heavy Sentence Transformer embeddings at runtime if we had good labels.

The "Teacher"

• Model: Qwen (LLM).
• Role: We use Qwen to "read" raw logs and generate high-quality ground-truth labels. It acts as the intelligent teacher.

The "Student" (Real-Time)

• Model: TF-IDF + XGBoost.
• Why: TF-IDF conversion takes lesser time.
• Result: The XGBoost model learns to mimic Qwen's high-quality labels using simple text features (TF-IDF).

---

Ablation Study & Performance Comparison

The table below summarizes the trade-offs across all three phases.

Phase	Embedding Strategy	Classifier	Verdict
Phase 1	Sentence Transformer	Neural Network	Too Slow.
Phase 2	Sentence Transformer	XGBoost	Bottlenecked. Embedding generation is too heavy.
Phase 3	TF-IDF (Sparse)	XGBoost	Current Methodology. Speedup with higher accuracy due to LLM labeling.

Setup & Prerequisites

Follow these steps to set up the pipeline on a new server or development machine.

1. Prerequisites

• Docker & Docker Compose (v2.0+)
• Python 3.9+ (For local training)
• Git

2. Installation

Clone the repository:

git clone http://<your-forgejo-url>/user/log-classification.git
cd log-classification

3. Environment Setup (Local Training)

cd training
python -m venv .venv
source .venv/bin/activate
pip install -r training/requirements.txt

Training the AI Models

1. The Training Workflow

The training logic is split into 5 sequential steps located in training/src/. This pipeline transforms raw unstructured logs into a trained XGBoost classifier.

Order	Script	Purpose
1	parse_logs.py	Ingestion: Reads raw data (JSON/CSV) and converts unstructured text into a standard format.
2	label.py	Applies heuristic labels (e.g., regex for "Security", "Network") to create initial training labels.
3	collapse.py	Deduplication: Groups similar labeled logs into cluster to prevent the model from overfitting on repetitive error loops.
4	augmentation.py	Balancing: Uses NLP techniques to generate synthetic data for rare log categories.
5	train.py	Training: Vectorizes the dataset (TF-IDF) and trains the final XGBoost classifiers.

2. Method A: Automated Training (Recommended)

We have provided a helper script to run the entire pipeline in the correct order. This is the preferred way to retrain the models.

Run the following command in your terminal:

./run_training.sh

Note: If you get a permission error, run chmod +x run_training.sh first.

What this script does:

• Sets set -e to ensure the process stops immediately if any step fails.
• Provides progress logs (e.g., [1/5] Parsing Logs...).
• Saves all artifacts to the models/ directory.

3. Method B: Manual Execution

If you need to debug a specific step, you can run the Python scripts individually from the root directory:

# 1. Parse raw input
python training/src/parse_logs.py

# 2. Label the data
python training/src/label.py

# 3. Collapse
python training/src/collapse.py

# 4. Augment (Generate synthetic samples)
python training/src/augmentation.py

# 5. Train Models
python training/src/train.py

4. Verification & Expected Output

Once the training script finishes, navigate to the models/ directory. You must see the following four files for the application to work:

File	Description
tfidf_vectorizer.joblib	The model that converts text to numbers.
xgb_content_model.joblib	Classifier for Content Type (Error, Info, Warning).
xgb_source_model.json	Classifier for Source Type (Security, Network, App).
xgb_label_maps.json	Mapping dictionary to convert numeric predictions back to text labels.

If these files exist, you are ready to proceed to Deployment.

# Docker Deployment

Use **Docker Compose** to run the ML API and Vector side-by-side.

### 1. Architecture

The `docker-compose.yaml` defines two services:

*   **`ml-api`:** Mounts the local `models/` folder so it can load the brains you just trained.
*   **`vector`:** Mounts `vector/vector.toml` to configure the log pipeline.

### 2. Start the Pipeline

Run the following command in the root directory:

```bash
# Build and start in background
docker-compose up -d --build

3. Verify Status

Check if both containers are running:

docker ps

You should see ml_api and vector-sidecar with status Up.

4. Debugging

If containers are crashing, check the logs:

# Check ML API logs
docker logs ml_api

# Check Vector logs
docker logs vector-sidecar

Testing the Pipeline

Test the system by sending logs manually via curl.

1. Direct ML API Test

Test if the Python logic works (bypassing Vector):

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"message": "[example-logs]"}'

Expected Response:

{
  "source_type": "[SOURCE]",
  "content_type": "[CONTENT]",
  "confidence": [CONFIDENCE]
}

2. Full Pipeline Test (Vector)

Test the full flow (Vector -> ML -> Vector -> Output):

NOW=$(date +%s%N)
curl -v -X POST http://localhost:5001/ -H "Content-Type: application/json" -d '{"resourceLogs":[{"resource":{},"scopeLogs":[{"scope":{},"logRecords":[{"body":{"stringValue":"Connection refused on port 80"},"attributes":[]}]}]}]}'