feat: add router endpoints with tests

Author: hspedro

.github/workflows/test.yml

@@ -36,12 +36,11 @@ jobs:
       run: poetry install
 
     - name: Run tests
-      run: poetry run pytest
+      run: make test
 
     - name: Check code quality
       run: |
-        poetry run black --check babeltron/
-        poetry run isort --check-only babeltron/
+        make lint
 
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v3

.gitignore

@@ -169,3 +169,6 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# Models
+models/*

Dockerfile

@@ -0,0 +1,37 @@
+# Base image with Python
+FROM python:3.10-slim AS base
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Poetry
+RUN curl -sSL https://install.python-poetry.org | python3 - && \
+    ln -s /root/.local/bin/poetry /usr/local/bin/
+
+# Copy Poetry configuration files
+COPY README.md pyproject.toml poetry.lock* ./
+
+# Copy application code
+COPY babeltron/ ./babeltron/
+
+# Configure Poetry to not use virtualenvs in Docker
+RUN poetry config virtualenvs.create false
+
+# Install dependencies
+RUN poetry install --without dev --no-interaction --no-ansi
+
+# Set environment variables
+ENV MODEL_PATH=/models
+ENV PYTHONPATH=/app
+
+# Expose the port the app runs on
+EXPOSE 8000
+
+# Command to run the application
+CMD ["uvicorn", "babeltron.app.main:app", "--host", "0.0.0.0", "--port", "8000"]

Makefile

@@ -1,4 +1,7 @@
-.PHONY: check-poetry install test lint format help system-deps coverage coverage-html
+.PHONY: check-poetry install test lint format help system-deps coverage coverage-html download-model download-model-small download-model-medium download-model-large serve serve-prod docker-build docker-run docker-compose-up docker-compose-down
+
+# Define model path variable with default value, can be overridden by environment
+MODEL_PATH ?= ./models
 
 # Extract target descriptions from comments
 help: ## Show this help message
@@ -40,11 +43,10 @@ install: check-poetry system-deps ## Install project dependencies
 
 test: check-poetry ## Run tests
 	@echo "Running tests..."
-	@poetry run pytest
+	@poetry run pytest tests/unit
 
 lint: check-poetry ## Run linters
 	@echo "Running linters..."
-	@poetry run flake8 babeltron
 	@poetry run isort --check babeltron
 	@poetry run black --check babeltron
 
@@ -55,8 +57,77 @@ format: check-poetry ## Format code
 
 coverage: check-poetry ## Run tests with coverage report
 	@echo "Running tests with coverage..."
-	@poetry run pytest --cov=babeltron --cov-report=term-missing
+	@poetry run pytest tests/unit --cov=babeltron --cov-report=term
 
 coverage-html: check-poetry ## Generate HTML coverage report
 	@echo "Generating HTML coverage report..."
-	@poetry run pytest --cov=babeltron --cov-report=html
+	@poetry run pytest tests/unit --cov=babeltron --cov-report=html
+	@echo "HTML coverage report generated in htmlcov/ directory"
+	@echo "Open htmlcov/index.html in your browser to view the report"
+
+# Model download commands
+download-model: download-model-small ## Download the default (small) translation model
+
+download-model-small: check-poetry ## Download small translation model (418M parameters, ~1GB)
+	@echo "Downloading small translation model (418M parameters)..."
+	@poetry run python -m babeltron.scripts.download_models --size 418M --output-dir $(MODEL_PATH)
+
+download-model-medium: check-poetry ## Download medium translation model (1.2B parameters, ~2.5GB)
+	@echo "Downloading medium translation model (1.2B parameters)..."
+	@poetry run python -m babeltron.scripts.download_models --size 1.2B --output-dir $(MODEL_PATH)
+
+download-model-large: check-poetry ## Download large translation model (12B parameters, ~24GB)
+	@echo "Downloading large translation model (12B parameters)..."
+	@poetry run python -m babeltron.scripts.download_models --size 12B --output-dir $(MODEL_PATH)
+
+# Add these commands to your Makefile
+serve: check-poetry ## Run the API server locally
+	@echo "Starting API server on http://localhost:8000..."
+	@poetry run uvicorn babeltron.app.main:app --reload --host 0.0.0.0 --port 8000
+
+serve-prod: check-poetry ## Run the API server in production mode (no reload)
+	@echo "Starting API server in production mode on http://localhost:8000..."
+	@poetry run uvicorn babeltron.app.main:app --host 0.0.0.0 --port 8000
+
+# Docker commands
+docker-build: ## Build Docker image
+	@echo "Building Docker image..."
+	@docker build -t babeltron:latest .
+
+docker-run: ## Run Docker container with model volume mount
+	@echo "Checking for model files..."
+	@if [ ! -d "$(MODEL_PATH)" ] || [ -z "$(shell ls -A $(MODEL_PATH) 2>/dev/null)" ]; then \
+		echo "No model files found in $(MODEL_PATH) directory."; \
+		read -p "Do you want to download the small model now? (y/n) " answer; \
+		if [ "$$answer" = "y" ]; then \
+			mkdir -p $(MODEL_PATH); \
+			echo "Downloading small model..."; \
+			poetry run python -m babeltron.scripts.download_models --size 418M --output-dir $(MODEL_PATH); \
+		else \
+			echo "Model download skipped. Container may not work properly."; \
+		fi; \
+	fi
+	@echo "Running Docker container..."
+	@docker run -p 8000:8000 -v $(shell pwd)/$(MODEL_PATH):/models babeltron:latest
+
+docker-up: ## Build and start services with docker-compose
+	@echo "Checking for model files..."
+	@if [ ! -d "$(MODEL_PATH)" ] || [ -z "$(shell ls -A $(MODEL_PATH) 2>/dev/null)" ]; then \
+		echo "No model files found in $(MODEL_PATH) directory."; \
+		read -p "Do you want to download the small model now? (y/n) " answer; \
+		if [ "$$answer" = "y" ]; then \
+			mkdir -p $(MODEL_PATH); \
+			echo "Downloading small model..."; \
+			poetry run python -m babeltron.scripts.download_models --size 418M --output-dir $(MODEL_PATH); \
+		else \
+			echo "Model download skipped. Container may not work properly."; \
+		fi; \
+	fi
+	@echo "Building and starting services with docker-compose..."
+	@docker-compose up -d --build
+	@echo "Services started successfully. API available at http://localhost:8000"
+	@echo "API documentation available at http://localhost:8000/docs"
+
+docker-down:
+	@echo "Stopping docker-compose services..."
+	@docker-compose down

README.md

@@ -3,6 +3,7 @@
 [![Tests](https://github.com/hspedro/babeltron/actions/workflows/test.yml/badge.svg)](https://github.com/hspedro/babeltron/actions/workflows/test.yml)
 [![PyPI version](https://badge.fury.io/py/babeltron.svg)](https://badge.fury.io/py/babeltron)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Coverage](https://img.shields.io/badge/coverage-90%25-brightgreen.svg)](https://github.com/hspedro/babeltron/actions/workflows/test.yml)
 
 A Python-based REST API that leverages single multilingual models like mBERT to
 provide efficient text translation services. Babeltron exposes a simple interface
@@ -68,6 +69,106 @@ The HTML report will be generated in the `htmlcov` directory.
 
 The project uses a `.coveragerc` file to configure coverage settings. This ensures consistent coverage reporting across different environments.
 
+## Downloading Translation Models
+
+Babeltron requires a translation model to function. You can download models of different sizes depending on your needs and hardware constraints:
+
+```bash
+# Download the small model (418M parameters, ~1GB disk space)
+make download-model
+
+# Or download medium model (1.2B parameters, ~2.5GB disk space)
+make download-model-medium
+
+# Or download large model (12B parameters, ~24GB disk space)
+make download-model-large
+```
+
+### Model Size Considerations
+
+- **Small (418M)**: ~1GB disk space, less memory required, faster but less accurate
+- **Medium (1.2B)**: ~2.5GB disk space, moderate memory requirements
+- **Large (12B)**: ~24GB disk space, requires significant RAM/GPU memory
+
+Choose based on your hardware constraints and translation quality requirements.
+
+## Running the API Server
+
+After installing dependencies and downloading a model, you can run the API server locally:
+
+```bash
+# Run the server in development mode with auto-reload
+make serve
+
+# Or run in production mode (no auto-reload)
+make serve-prod
+```
+
+The API will be available at http://localhost:8000.
+
+### API Usage Examples
+
+Once the server is running, you can use the translation API:
+
+```bash
+# Translate text from English to Spanish
+curl -X POST "http://localhost:8000/translate" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "text": "Hello, how are you?",
+    "src_lang": "en",
+    "tgt_lang": "es"
+  }'
+
+# Response:
+# {"translation":"Hola, ¿cómo estás?"}
+```
+
+You can also access the interactive API documentation at http://localhost:8000/docs.
+
+## API Documentation
+
+Babeltron provides interactive API documentation:
+
+- **Swagger UI**: Available at http://localhost:8000/docs when the server is running
+- **ReDoc**: Available at http://localhost:8000/redoc for an alternative documentation view
+
+These interactive documentation pages allow you to:
+- Explore all available endpoints
+- See request and response schemas
+- Test the API directly from your browser
+- View detailed descriptions of each endpoint and parameter
+
+## Running with Docker
+
+Babeltron can be run as a Docker container, which simplifies deployment and isolates dependencies.
+
+### Building and Running with Docker
+
+```bash
+# Start services with Docker Compose
+make docker-up
+```
+
+The API will be available at http://localhost:8000.
+
+### Stopping Docker Services
+
+```bash
+# Stop services
+make docker-down
+```
+
+### Docker Volume Mounts
+
+The Docker setup mounts the local `./models` directory to `/models` inside the container. This allows you to:
+
+1. Reuse downloaded models between container restarts
+2. Use different model sizes without rebuilding the image
+3. Persist models even if the container is removed
+
+If no models are found when starting the container, you'll be prompted to download the small model automatically.
+
 ## License
 
 MIT License

test/__init__.py babeltron/app/__init__.pytest/__init__.py renamed to babeltron/app/__init__.py

@@ -0,0 +1,47 @@
+from importlib.metadata import version
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from babeltron.app.utils import include_routers
+
+try:
+    __version__ = version("babeltron")
+except ImportError:
+    __version__ = "0.1.0-dev"
+
+app = FastAPI(
+    title="Babeltron Translation API",
+    description="API for machine translation using NLLB models",
+    version="0.1.0",
+    contact={
+        "name": "Your Name",
+        "url": "https://your-website.com",
+        "email": "your-email@example.com",
+    },
+    license_info={
+        "name": "MIT",
+        "url": "https://opensource.org/licenses/MIT",
+    },
+    docs_url="/docs",
+    redoc_url="/redoc",
+    openapi_url="/openapi.json",
+)
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+# Include all routers
+include_routers(app)
+
+# This allows running the app directly with uvicorn when this file is executed
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

babeltron/app/main.py

@@ -0,0 +1,67 @@
+from typing import Optional
+
+from fastapi import APIRouter, status
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+from babeltron.app.main import __version__
+from babeltron.app.routers.translate import model, tokenizer
+
+router = APIRouter()
+
+
+class HealthResponse(BaseModel):
+    status: str
+    model_loaded: bool
+    version: Optional[str] = None
+
+
+@router.get("/healthcheck", summary="Healthcheck")
+@router.get(
+    "/healthz",
+    summary="Check API health",
+    description="Returns the health status of the API and whether the translation model is loaded",
+    response_model=HealthResponse,
+    tags=["Control"],
+)
+async def healthcheck():
+    return {"status": "ok", "model_loaded": model is not None, "version": __version__}
+
+
+class ReadinessResponse(BaseModel):
+    status: str
+    version: Optional[str] = None
+    error: Optional[str] = None
+
+
+@router.get("/readiness", summary="Readiness Probe")
+@router.get(
+    "/readyz",
+    summary="Check API Readiness",
+    description="Returns the readiness status of the API. Able to process requests.",
+    response_model=ReadinessResponse,
+    tags=["Control"],
+)
+async def readiness():
+    try:
+        if model is None or tokenizer is None:
+            return JSONResponse(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                content={
+                    "status": "not ready",
+                    "error": "Model or tokenizer not loaded",
+                    "version": __version__,
+                },
+            )
+
+        test_sentence = "hello"
+        tokenizer.src_lang = "en"
+        encoded_text = tokenizer(test_sentence, return_tensors="pt")
+        _ = model.generate(**encoded_text)
+
+        return {"status": "ready", "version": __version__}
+    except Exception as e:
+        return JSONResponse(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            content={"status": "not ready", "error": str(e), "version": __version__},
+        )