IMPLEMENTATION_PLAN.md

MediShield AI Document Classification — Implementation Plan

Workflow Diagram

                        ┌─────────────────────────────────┐
                        │        Uploaded Image            │
                        └──────────────┬──────────────────┘
                                       │
                        ┌──────────────▼──────────────────┐
                        │   Stage 1: Rules Engine          │
                        │   regex: ^bill_                  │
                        └──────────────┬──────────────────┘
                                       │
                     ┌─────────────────┴──────────────────┐
               bill_ match?                           no match
                     │                                     │
          ┌──────────▼──────────┐          ┌──────────────▼──────────────┐
          │  doc_type = "bill"   │          │   Stage 2: KYC OCR           │
          │  method   = "rules"  │          │   easyocr → keyword regex    │
          │  ✓ DONE              │          └──────────────┬──────────────┘
          └─────────────────────┘                         │
                                          ┌──────────────┴──────────────┐
                                    KYC match?                     no match
                                          │                              │
                             ┌────────────▼────────┐   ┌───────────────▼─────────────┐
                             │  doc_type = "kyc"   │   │  Stage 3: Gemini LLM         │
                             │  method   = "ocr"   │   │  gemma-4-31b-it              │
                             │  ✓ DONE             │   │  → Patient Bills             │
                             └─────────────────────┘   │  → Claim Forms               │
                                                        │  → Medical Reports           │
                                                        │  → Prescriptions             │
                                                        │  → Unknown                   │
                                                        └──────────────────────────────┘
                                                                     │
                                                        ┌────────────▼────────────────┐
                                                        │  doc_type = "image"          │
                                                        │  sub_type = <category>       │
                                                        │  method   = "llm"            │
                                                        │  ✓ DONE                      │
                                                        └─────────────────────────────┘

All stages emit @traceable spans → LangSmith (traces · tokens · latency)
All results served via FastAPI → Drag & Drop UI
Container deployed on Azure Container Apps via GitHub Actions CI/CD

Architecture Overview

┌─────────────────────────────────────────────────────┐
│                   Frontend UI                        │
│         Drag & Drop  ·  frontend/index.html          │
│  - Batch upload (all files in one POST)              │
│  - Concurrent server processing (asyncio.gather)     │
│  - Live progress bar + per-file status rows          │
│  - Color-coded badges: bill/kyc/image                │
└───────────────────┬─────────────────────────────────┘
                    │ POST /classify (multipart)
                    ▼
┌─────────────────────────────────────────────────────┐
│           FastAPI  ·  src/api.py  ·  Port 8000       │
│  POST /classify  ·  GET /health  ·  GET /metrics     │
│  asyncio.gather + run_in_executor (concurrent files) │
└────────────────────┬────────────────────────────────┘
                     │
          ┌──────────▼──────────┐
          │   src/classifier.py  │  (orchestrator)
          └──┬──────┬──────┬────┘
             │      │      │
    rules    │  ocr │  llm │
    engine   │      │      │
             ▼      ▼      ▼
┌────────────────────────────────────────────────────┐
│  src/monitoring.py  — LangSmith @traceable spans   │
│  trace_rules_engine · trace_kyc_ocr                │
│  trace_llm_classify · trace_classify               │
└────────────────────────────────────────────────────┘
                    │
         ┌──────────┴──────────┐
         ▼                     ▼
   LangSmith               Azure Monitor
   (traces/tokens)         (container logs)

Decision Rules

Condition	doc_type	method	Sent to LLM?
filename matches ^bill_ (regex)	bill	rules	No
OCR text contains KYC keywords	kyc	ocr	No
Everything else	image	llm	Yes

Final File Layout

multimodal-ai/
├── src/
│   ├── rules_engine.py       # Step 1 ✅
│   ├── kyc_detector.py       # Step 2 ✅
│   ├── llm_classifier.py     # Step 3 ✅
│   ├── classifier.py         # Step 4 ✅
│   ├── api.py                # Step 5 ✅
│   └── monitoring.py         # Step 7 ✅
├── frontend/
│   └── index.html            # Step 6 ✅
├── tests/
│   ├── test_rules_engine.py  # 11 tests  ✅
│   ├── test_kyc_detector.py  # 21 tests  ✅
│   ├── test_llm_classifier.py # 32 tests ✅
│   ├── test_classifier.py    # 29 tests  ✅
│   ├── test_api.py           # 11 tests  ✅
│   └── test_monitoring.py    # 13 tests  ✅  (117 total)
├── infra/
│   ├── deploy.sh             # Step 9 ✅ Azure Container Apps
│   └── teardown.sh           # ✅
├── .github/workflows/
│   ├── ci.yml                # Step 10 ✅ test on every push
│   └── deploy.yml            # Step 10 ✅ deploy on merge to main
├── Dockerfile                # Step 8 ✅ two-stage build
├── .dockerignore             # Step 8 ✅
├── README.md                 # Step 10 ✅
├── pyproject.toml
└── .env.example

---

Steps

Phase 1 — Core Classification Engine

• Step 1 — Rules Engine (src/rules_engine.py)

  • Compiled regex re.compile(r"^bill_") — case-sensitive, anchored to start of filename
  • Strips directory prefix so full paths work (dataset/bill_x.png)
  • Returns RulesResult(filename, doc_type, send_to_llm)
  • Changed from plan: Used re.compile regex instead of str.startswith() as requested
  • ✅ 11/11 tests passing

• Step 2 — KYC Detector (src/kyc_detector.py)

  • 11 compiled regex patterns covering Aadhaar, PAN, Passport, Govt of India, DOB, 12-digit Aadhaar number, PAN card format
  • easyocr Reader is a lazy singleton — loaded once on first use, not at import time
  • reader is injectable (passed as parameter) so tests never load the real model
  • Returns KYCResult(filename, doc_type, send_to_llm, ocr_text)
  • ✅ 21/21 tests passing

• Step 3 — LLM Classifier (src/llm_classifier.py)

  • Sends image bytes + structured prompt to gemma-4-31b-it via google-genai
  • Prompt instructs model to return exactly one category name
  • _parse_category() does case-insensitive match + strips whitespace, falls back to "Unknown"
  • Captures input_tokens and output_tokens from response.usage_metadata
  • client is injectable for testing — zero live API calls in test suite
  • Returns LLMResult(filename, doc_type, sub_type, method, input_tokens, output_tokens, raw_response)
  • ✅ 32/32 tests passing

• Step 4 — Pipeline Orchestrator (src/classifier.py)

  • classify(filename, image_bytes, ocr_reader, llm_client) — single document
  • classify_dataset(dataset_dir, ...) — scans all PNGs in a directory
  • Returns ClassificationResult(filename, doc_type, sub_type, method, latency_ms, input_tokens, output_tokens)
  • Each stage emits a LangSmith trace span (added in Step 7)
  • ✅ 29/29 tests passing

---

Phase 2 — FastAPI Server

• Step 5 — API Server (src/api.py)
  • POST /classify — multipart file upload, returns JSON array
  • GET /health — liveness probe
  • GET /metrics — in-memory counters per method/doc_type/token usage
  • GET /docs — auto Swagger UI
  • Changed from plan: asyncio.gather + run_in_executor runs all uploaded files concurrently — bill_ files return in < 10 ms without waiting behind OCR/LLM calls
  • easyocr Reader and Gemini Client loaded once at startup via FastAPI lifespan
  • CORS middleware enabled for browser UI
  • ✅ 11/11 tests passing (patched at src.api.classify)

---

Phase 3 — Frontend UI

• Step 6 — Drag & Drop UI (frontend/index.html)
  • Self-contained single HTML file, no external dependencies
  • Drag & drop + click-to-browse, deduplicates files by name
  • Changed from plan (sequential → batch): Sends all files in ONE POST /classify — server processes concurrently so bill_ files don't wait behind slow OCR/LLM calls
  • Results table appears immediately with queued… rows; fills in as server responds
  • Live progress bar + Processing file N of M text
  • Color-coded badges: bill=blue, kyc=orange, image=green, rules=purple, ocr=red, llm=teal
  • All controls (classify, clear, remove buttons, drop zone) disabled during processing
  • Summary bar: counts per type + average latency
  • Error banner for API failures and unsupported file types

---

Phase 4 — Monitoring (LangSmith)

• Step 7 — LangSmith Integration (src/monitoring.py)
  • Four @traceable functions forming a parent/child span tree:
    • trace_classify — top-level chain span per document
    • trace_rules_engine — tool span for Stage 1
    • trace_kyc_ocr — tool span for Stage 2; records ocr_text_length not raw text (PII safety)
    • trace_llm_classify — llm span for Stage 3; records token breakdown
  • record_token_usage() extracts input/output/total_tokens from Gemini usage_metadata
  • Tracing is a no-op when LANGCHAIN_TRACING_V2 is not set — CI safe
  • Required env vars:

    LANGCHAIN_TRACING_V2=true
    LANGCHAIN_API_KEY=<key>
    LANGCHAIN_PROJECT=medishield-classification
    

  • ✅ 13/13 tests passing

---

Phase 5 — Docker

• Step 8 — Dockerfile
  • Two-stage build: uv builder → python:3.12-slim runtime
  • Installs OS libs for easyocr/opencv/weasyprint in runtime stage
  • Pre-downloads easyocr models at build time as appuser — container starts in ~10s not 60s
  • Runs as non-root appuser (with home dir so easyocr can write model cache)
  • HEALTHCHECK polls /health every 30s, 60s start period
  • 2 uvicorn workers for concurrency
  • Fix applied during build: Created home dir for appuser and set EASYOCR_MODULE_PATH to fix permission error on model cache write
  • ✅ Build verified, /classify tested inside container

---

Phase 6 — Azure Deployment

• Step 9 — Azure Container Apps (infra/deploy.sh)
  • Changed from plan: Azure instead of AWS (simpler setup, no separate load balancer, built-in HTTPS)
  • Provisions: Resource Group → ACR → Log Analytics → Container Apps Environment → Container App
  • Container App: 0.5 vCPU / 2 GB RAM, min 1 replica, max 3, public HTTPS ingress
  • Secrets (GOOGLE_API_KEY, LANGCHAIN_API_KEY) injected via Container Apps secret references
  • infra/teardown.sh for full cleanup
  • CI/CD via .github/workflows/deploy.yml:
    • Tests gate deploy (deploy only runs if tests pass)
    • az acr build builds in Azure cloud (no local Docker in CI)
    • az containerapp update rolling deploy
    • Smoke tests live /health endpoint post-deploy
    • OIDC login (no long-lived secrets in GitHub)

---

Phase 7 — Documentation

• Step 10 — README + CI (README.md, .github/workflows/ci.yml)
  • Professional README with ASCII architecture diagram, workflow diagram, full API reference, setup guide, deployment guide, test matrix, environment variable table
  • ci.yml runs all 117 tests on every push/PR — no real API keys needed

---

Build Order Summary

#	Deliverable	Test Gate	Status
1	Rules Engine	pytest tests/test_rules_engine.py — 11 passed	✅
2	KYC Detector	pytest tests/test_kyc_detector.py — 21 passed	✅
3	LLM Classifier	pytest tests/test_llm_classifier.py — 32 passed	✅
4	Orchestrator	pytest tests/test_classifier.py — 29 passed	✅
5	FastAPI Server	pytest tests/test_api.py — 11 passed + Swagger check	✅
6	Frontend UI	Batch POST, live progress, controls locked during processing	✅
7	LangSmith Monitoring	pytest tests/test_monitoring.py — 13 passed	✅
8	Docker	docker build + /classify tested inside container	✅
9	Azure Deploy	infra/deploy.sh + GitHub Actions CI/CD pipeline	✅
10	README + CI	ci.yml + deploy.yml + README.md	✅

Total: 117 tests · 10 steps · all complete ✅

Key Changes vs Original Plan

Area	Original Plan	What We Actually Built
Rules matching	str.startswith("bill_")	re.compile(r"^bill_") regex
API concurrency	Sequential file loop	asyncio.gather + run_in_executor
UI upload strategy	One request per file (sequential)	One batch request, server concurrent
Cloud provider	AWS ECS Fargate	Azure Container Apps
Metrics	OpenTelemetry + CloudWatch	LangSmith + Azure Monitor
Docker user	Root	Non-root appuser with home dir
easyocr models	Downloaded at runtime	Pre-baked into image at build time