datagraph

<<<<<<< HEAD

datagraph

Data Model Dependency extraction using Databricks hosted foundation models

DataGraph AI

AI-Powered Data Model Dependency Extraction Hierarchical Multi-Tenant Platform for SQL, Power BI, Databricks, and Diagram Analysis

🎯 Overview

DataGraph AI automatically extracts dependencies from SQL databases (including Snowflake), Databricks Unity Catalog, BI tools, and diagram images using off-the-shelf foundation models (Meta Llama 3.1/3.3, Anthropic Claude Sonnet 4) and Vision AI (Claude 3.5 Sonnet Vision, GPT-4V) with sophisticated prompt engineering. It achieves 90-95% accuracy without the need for expensive fine-tuning.

Version 2 (Current) features hierarchical multi-tenant organization (Customer → Business Unit → Use Case), Databricks Unity Catalog live connectivity, ERD diagram visualization with Primary/Foreign Keys, and dual-server deployment architecture.

Key Insights

"Off-the-shelf foundation models with good engineering can match custom fine-tuned models—at a fraction of the cost and complexity."

"Vision AI makes legacy diagrams machine-readable—extract data models from screenshots, ERDs, and whiteboard photos with 90%+ accuracy."

✨ Key Features

🏢 Organization & Workflow

✅ Hierarchical Multi-Tenant: Customer → Business Unit → Use Case structure
✅ Lakebase Integration: PostgreSQL metadata storage for fast queries and history tracking
- Extractions are linked to Use Cases via foreign key (usecase_id)
- Each extraction record stores metadata (node/edge counts, confidence, model, timestamps)
- Actual extraction data files stored in Unity Catalog Volumes (hierarchical path)
- Database indexes enable fast queries by customer, BU, use case, source type, status
✅ Context-Aware Navigation: Sidebar tree navigation with breadcrumbs
✅ Extraction Management:
- Multiple extractions per use case (SQL, Power BI, Images, Databricks, Merged)
- History view filtered by customer/BU/use case
- Status tracking: uploading, processing, completed, failed

📊 Data Sources

✅ SQL Databases: Oracle, PostgreSQL, MySQL, SQL Server, Snowflake (.sql, .ddl files)
✅ Databricks Unity Catalog: Live connection to extract tables, views, functions, procedures, volumes, models, PK/FK constraints, and table lineage
✅ BI Tools: Power BI (.pbix, .json files) with automatic JSON extraction from .pbix archives
✅ Diagrams & Images: ERD diagrams, database schemas, Power BI screenshots (.png, .jpg, .jpeg, .pdf using Vision AI)
✅ SQL Dialect Patterns: CTEs, recursive CTEs, subqueries, triggers, foreign keys, materialized views, cross-schema dependencies, streams, tasks

AI & Models

✅ Foundation Model Integration: Llama 3.1 8B, Llama 3.3 70B, Claude Sonnet 4
✅ Vision AI: Claude 3.5 Sonnet Vision, GPT-4 Vision for diagram extraction
✅ Ensemble Strategy: Combine multiple models for improved accuracy
✅ Automatic Product Detection: Identifies source database/BI tool from files

🎨 Interface & Visualization

✅ Modern Web Interface: Next.js 16 with TypeScript, shadcn/ui, and ReactFlow
✅ ERD Diagram View: Entity-Relationship Diagrams with Primary Keys, Foreign Keys, and relationship arrows
✅ Interactive Graph View: Network visualization with multiple layouts (hierarchical, force-directed, radial)
✅ Export Capabilities: JSON, Cypher (Neo4j), CSV, PNG export
✅ Merge Extractions Within Use Case:
- Scope: Merge operates within a single use case only
- Constraints: Cannot merge two extractions of the same source type
- Supported combinations: SQL + Images, SQL + Power BI, Power BI + Images, Databricks + any other type
- Process: Select 2 completed extractions → Merge → Creates new extraction with source_type='merged'
- Storage: Merged result saved to use case's Volume path, metadata recorded in Lakebase
- Deduplication: Automatically handles duplicate nodes/edges across sources

🚀 Backend & Deployment

✅ REST API Backend: FastAPI with Server-Sent Events (SSE) streaming and 4 workers
✅ Background Tasks: Asynchronous processing with polling for long extractions
✅ Metadata Storage: Databricks Lakebase (PostgreSQL) for structured metadata and project management
✅ File Storage: Unity Catalog Volumes for hierarchical file organization
✅ Production Ready: Databricks Apps deployment with dual-server architecture:
- Next.js server (port 8000, exposed)
- FastAPI server (port 8001, internal)
- Next.js proxies /api/* requests to FastAPI

Testing & Analytics

✅ Comprehensive Testing: 700+ synthetic schemas with ground truth
✅ Rich Analytics: Performance by database, domain, complexity, and object type
✅ Real-time Progress: Streaming test results with live metrics

📁 Project Structure

data-model-nel-archiecture/
│
├── 📖 README.md                       # This file
├── 📖 ARCHITECTURE_DIAGRAM.html      # Interactive system architecture diagram
├── 📖 DATABRICKS_DEPLOYMENT.md       # Databricks Apps deployment guide
├── app.yaml                          # Databricks Apps configuration (dual-server)
├── deploy_new.sh                     # Deployment script for Databricks Apps
│
├── api/                              # FastAPI Backend
│   ├── main.py                       # Thin ASGI entrypoint: `uvicorn api.main:app`
│   ├── app_factory.py                # `create_app()` + router registration
│   ├── core/                         # Cross-cutting concerns (config, wiring, identity, tracking)
│   ├── routers/                      # Modular FastAPI routers (HTTP layer)
│   ├── services/                     # Business logic helpers (non-HTTP)
│   ├── hierarchy_api.py              # Customer/BU/Usecase API endpoints (legacy router module)
│   ├── sql_extractor.py              # SQL dependency extraction
│   ├── sql_extractor_hierarchical.py # Hierarchical SQL extraction (>96K chars)
│   ├── json_extractor.py             # Power BI JSON extraction
│   ├── pbix_extractor.py             # Power BI .pbix extraction
│   ├── pbix_extractor_hierarchical.py # Hierarchical Power BI extraction
│   ├── vision_extractor.py           # Vision AI diagram extraction
│   ├── databricks_catalog_extractor.py # Unity Catalog live extraction
│   ├── model_adapter.py              # Unified AI model interface
│   ├── file_operations.py            # Volume/local file operations
│   ├── metadata_extractor.py         # Metadata detection
│   ├── jobs_api.py                   # Databricks Jobs API integration
│   ├── scripts/                      
│   │   ├── lakebase.py               # Lakebase (PostgreSQL) integration
│   │   ├── schema.sql                # PostgreSQL schema for metadata
│   │   └── [setup scripts]
│   ├── requirements.txt              # Python dependencies
│   └── 📖 README.md                  # API documentation
│
├── frontend/                         # Next.js Web Interface
│   ├── app/                          # App router pages
│   │   ├── page.tsx                  # Dashboard (homepage)
│   │   ├── layout.tsx                # Root layout with sidebar
│   │   ├── customers/                # Customer management
│   │   │   ├── new/                  # Create customer
│   │   │   └── detail/               # Customer details
│   │   ├── business-units/           # Business Unit management
│   │   │   ├── new/                  # Create BU
│   │   │   └── detail/               # BU details
│   │   ├── usecases/                 # Use Case management
│   │   │   ├── new/                  # Create use case
│   │   │   └── detail/               # Use case details
│   │   ├── upload/                   # File upload & Databricks connect
│   │   ├── extract/                  # Extraction progress & status
│   │   ├── visualize-erd/            # ERD diagram view (PK/FK)
│   │   ├── graph-view/               # Interactive graph visualization
│   │   ├── merge-extractions/        # Merge multiple extractions
│   │   ├── history/                  # Extraction history
│   │   ├── export/                   # Export (JSON, Cypher, CSV, PNG)
│   │   ├── run-tests/                # Batch testing interface
│   │   ├── test-results/             # Results visualization
│   │   ├── model-comparison/         # Model performance comparison
│   │   ├── admin/analytics/          # Admin analytics
│   │   ├── settings/                 # Application settings
│   │   ├── feedback/                 # User feedback
│   │   └── product-backlog/          # Feature requests
│   ├── components/                   # Reusable React components
│   │   ├── app-sidebar.tsx           # Hierarchical navigation sidebar
│   │   ├── erd-diagram.tsx           # ERD visualization components
│   │   ├── graph-visualizer.tsx      # Graph visualization
│   │   ├── backend-status-indicator.tsx # Health check indicator
│   │   └── ui/                       # shadcn/ui components
│   ├── contexts/                     # React Context providers
│   │   └── organization-context.tsx  # Customer/BU/Usecase state
│   ├── hooks/                        # Custom React hooks
│   ├── lib/                          # Utility functions
│   │   └── api-config.ts             # API fetch wrapper
│   ├── next.config.ts                # Next.js configuration (proxies)
│   ├── package.json                  # Node dependencies
│   └── 📖 README.md                  # Frontend documentation
│
├── data/                             # Local development data storage
│   ├── customers/                    # Hierarchical: {customer_id}/{bu_id}/{usecase_id}/
│   ├── sessions/                     # Extraction session files
│   └── training/                     # Training data
│       ├── sql_sources/              # SQL files by split (train/val/test)
│       ├── non_sql_sources/          # Power BI JSON files
│       └── image_sources/            # Diagram images
│
├── phase1_training_data/             # 🗄️ Synthetic Data Generation
│   ├── configs/                      # Configuration files
│   ├── data/                         # Generated schemas (700+)
│   ├── 00_load_training_data_BASIC_PHASE1_ONLY.py  # Data generator
│   ├── validate_training_data.py     # Validation script
│   └── 📖 README.md                  # Training data documentation
│
├── phase1_hierarchical_approach/     # 🗄️ Archived Fine-Tuning Experiments
│   ├── training/                     # 3-stage training scripts
│   ├── inference/                    # Hierarchical inference
│   └── 📖 README.md                  # Architecture documentation
│
└── docs/                             # 🗄️ Legacy Documentation
    ├── QUICKSTART.md
    ├── BLOG_POST_TWO_MONTH_JOURNEY.md
    └── [50+ detailed guides]

🚀 Quick Start (Local Development)

Prerequisites

Python 3.9+ with pip and venv
Node.js 18+ with npm
PostgreSQL 13+ (for Lakebase metadata storage)
Databricks Workspace with model serving endpoints (or API keys for Anthropic/OpenAI)

1. Clone Repository

git clone <repository-url>
cd data-model-nel-archiecture

2. Setup Lakebase (PostgreSQL)

Lakebase stores organizational metadata (customers, business units, use cases, extractions).

Option A: Local PostgreSQL

# Install PostgreSQL (if not already installed)
# macOS: brew install postgresql@14
# Ubuntu: sudo apt-get install postgresql-14
# Windows: Download from postgresql.org

# Start PostgreSQL service
# macOS: brew services start postgresql@14
# Ubuntu: sudo systemctl start postgresql
# Windows: Start via Services

# Create database and user
psql postgres
CREATE DATABASE data_lineage_metadata;
CREATE USER lakebase_user WITH PASSWORD 'lakebase_password';
GRANT ALL PRIVILEGES ON DATABASE data_lineage_metadata TO lakebase_user;
\q

# Run schema setup script
cd api/scripts
psql -U lakebase_user -d data_lineage_metadata -f schema.sql

Option B: Docker PostgreSQL (Easier)

# Start PostgreSQL in Docker
docker run --name lakebase-postgres \
  -e POSTGRES_DB=data_lineage_metadata \
  -e POSTGRES_USER=lakebase_user \
  -e POSTGRES_PASSWORD=lakebase_password \
  -p 5432:5432 \
  -d postgres:14

# Wait 10 seconds for PostgreSQL to start
sleep 10

# Run schema setup script
cd api/scripts
psql -h localhost -U lakebase_user -d data_lineage_metadata -f schema.sql
# Password: lakebase_password

3. Start Backend (API)

cd api

# Create virtual environment
python3 -m venv venv
source venv/bin/activate  # Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

# Configure environment (create .env file)
cat > .env << EOF
# Databricks Configuration
DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
DATABRICKS_TOKEN=your_token_here

# Model Endpoints
LLAMA_8B_ENDPOINT=https://.../serving-endpoints/llama-8b/invocations
LLAMA_70B_ENDPOINT=https://.../serving-endpoints/llama-70b/invocations
CLAUDE_SONNET_4_ENDPOINT=https://.../serving-endpoints/claude-sonnet-4/invocations

# Lakebase (PostgreSQL) Configuration
LAKEBASE_HOST=localhost
LAKEBASE_PORT=5432
LAKEBASE_DATABASE=data_lineage_metadata
LAKEBASE_USER=lakebase_user
LAKEBASE_PASSWORD=lakebase_password

# Data Paths (optional - defaults to <project_root>/data/)
LOCAL_TRAINING_DATA_PATH_SQL=../phase1_training_data/data/sql_sources
LOCAL_TRAINING_RESULTS_PATH_SQL=../phase1_training_data/data/test_results
EOF

# Start API server (entrypoint is stable)
cd ..
uvicorn api.main:app --reload --port 8000

API will be available at: http://localhost:8000

4. Start Frontend

cd frontend

# Install dependencies
npm install

# Configure environment
echo "NEXT_PUBLIC_API_URL=http://localhost:8000" > .env.local

# Start development server
npm run dev

Frontend will be available at: http://localhost:3000

5. Access the Application

Open your browser and navigate to:

Frontend: http://localhost:3000
API Docs: http://localhost:8000/docs

6. Quick Start Guide

Hierarchical Workflow (Customer → Business Unit → Use Case):

Dashboard → Click "Start Here" button
Create Customer → Enter customer name (e.g., "Acme Corp")
Add Business Unit → Enter BU name (e.g., "Finance Department")
Create Use Case → Enter use case name (e.g., "CRM Data Model") and select source type:
- sql_sources - SQL files
- power_bi - Power BI .pbix or JSON
- image_sources - Diagram images
- databricks_sources - Unity Catalog connection
Upload or Connect:
- File Upload: Drag & drop SQL, Power BI (.pbix), or diagram images
- Databricks Connect: Enter host, catalog, schema, SQL warehouse path
Extract Dependencies → Select AI model, monitor real-time progress
View Results:
- ERD Diagram → Entity-Relationship view with Primary/Foreign Keys
- Graph View → Interactive network visualization
- Export → JSON, Cypher, CSV, PNG

Alternative: One-Command Start (Both Backend & Frontend)

# From project root
./start.sh

This starts both services locally:

API: http://localhost:8000 (FastAPI, uvicorn api.main:app)
Frontend: http://localhost:3000 (Next.js dev server)

To stop everything cleanly:

./stop.sh

Local “Databricks Apps” simulation (optional)

If you want to run the same ports as Databricks Apps locally (Frontend :8000, API :8001):

./start_prod.sh
./stop.sh

start_prod.sh is intended for local validation; Databricks Apps deployment behavior is controlled by app.yaml.

🎨 User Workflows

New Hierarchical Workflow: Organize Your Data Models

Why Hierarchical? Organize extractions by customer, business unit, and use case for better project management, collaboration, and tracking across your organization.

Complete User Journey (4 Steps)

┌─────────────────────────────────────────────────────────────────────────┐
│ STEP 1: SETUP                                                            │
│ Create Customer → Business Unit → Use Case                              │
│                                                                           │
│ Customer: Your organization or client                                    │
│    └─ Business Unit: Department or team (e.g., Finance, Sales)         │
│         └─ Use Case: Specific project (e.g., "CRM Data Model")         │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│ STEP 2: UPLOAD                                                           │
│ Upload your source files (one type per use case)                        │
│                                                                           │
│ Source Types:                                                            │
│  • SQL Files (DDL, stored procedures, views)                            │
│  • Power BI Files (.pbix or .json files)                                │
│  • Diagrams & Images (ERD diagrams, screenshots)                        │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│ STEP 3: EXTRACT                                                          │
│ AI extracts dependencies automatically                                   │
│                                                                           │
│ • Select AI model (Llama 8B/70B, Claude Sonnet 4, Vision AI)           │
│ • Monitor real-time progress                                            │
│ • Background processing for large files                                 │
│ • Automatic save to Lakebase (metadata storage)                        │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│ STEP 4: VIEW GRAPH                                                       │
│ Explore your dependency graph                                            │
│                                                                           │
│ • Interactive visualization with multiple layouts                        │
│ • Table view with filtering                                             │
│ • Export to JSON, CSV, PNG, SVG, Neo4j Cypher                          │
│ • Merge multiple extractions within a use case                          │
└─────────────────────────────────────────────────────────────────────────┘

Detailed Workflow Steps

1. Setup Your Organization (One-Time)

a) Create Customer

Click "Start Here" or "New Customer" from Dashboard
Enter customer name (e.g., "Acme Corp")
Optional: Add industry, region, description

b) Create Business Unit

After creating customer, add business unit
Enter BU name (e.g., "Finance Analytics", "Sales Operations")
Optional: Add description

c) Create Use Case

After creating business unit, define your use case
Enter use case name (e.g., "CRM Data Model", "Sales Pipeline")
Select ONE source type: SQL, Power BI, or Diagrams
Optional: Add description

2. Upload Source Files

After creating a use case, you'll be automatically directed to upload:

SQL Files (.sql, .ddl):

Supported databases: Oracle, MySQL, PostgreSQL, SQL Server, Snowflake
Examples: DDL scripts, stored procedures, views, functions, triggers
Hierarchical extraction for large SQL files (>96K characters)

Power BI Files (.pbix, .json):

Upload .pbix files directly OR pre-extracted ReportLayout.json files
System automatically extracts JSON from .pbix files (using zipfile)
Detects tables, measures, relationships, DAX formulas, calculated columns
Hierarchical extraction for complex Power BI models

Diagrams & Images (.png, .jpg, .jpeg, .pdf):

Examples: ERD diagrams, database schemas, Power BI screenshots, whiteboard photos
Uses Vision AI (Claude 3.5 Sonnet Vision or GPT-4 Vision)
Confidence scoring from model (no hardcoded values)

3. Extract Dependencies

Model Selection:

For SQL/Power BI: Llama 8B, Llama 70B, or Claude Sonnet 4
For Diagrams: Claude 3.5 Sonnet Vision or GPT-4 Vision

Extraction Process:

Real-time progress updates (file processing, AI extraction, saving)
Background processing for large files (>96K characters)
Hierarchical extraction for very large files (chunks + merge)
Automatic save to Lakebase metadata storage
Results stored in Unity Catalog Volumes

4. View & Analyze Results

Dependency Graph:

Table View: Browse all nodes and edges with filtering
Graph View: Interactive visualization
- Force-directed layout (shows relationships)
- Hierarchical layout (top-down structure)
- Circular layout (radial design)
Zoom, pan, and click nodes for details

Export Options:

JSON (complete graph data)
CSV (tabular format)
PNG/SVG (visual export)
Neo4j Cypher (import to graph database)

Merge Extractions (Optional):

Navigate: /merge-extractions?usecase={usecase_id}
Select: Choose 2 completed extractions from the same use case
Constraint: Must be different source types (e.g., SQL + Images, not SQL + SQL)
Supported combinations: SQL + Power BI, SQL + Images, Power BI + Images, Databricks + any
Result: Creates new extraction record with source_type='merged'
Storage: Saved to use case's hierarchical Volume path: /Volumes/.../customers/{c}/{b}/{u}/results/merged_*.json
Metadata: Lakebase stores references to source extractions in metadata field
Deduplication: Automatically merges nodes (by name) and edges (by source→target→type)

5. Browse & Manage Projects

Sidebar Navigation:

Browse all customers → business units → use cases
Quick access to any extraction
View extraction history per use case
Track progress across your organization

Extraction History:

View all extractions across all use cases
Filter by customer, business unit, source type
Compare extractions over time
Access previous results anytime

Testing Workflow: Evaluate Models

Run Batch Tests
- Select models to test (1-4 models)
- Choose data split: train, val, or test
- Set sample size (default: 20 schemas)
- Monitor streaming results
Analyze Results
- Overview: Aggregate F1 scores, test counts, timing
- By Database: Performance per SQL dialect
- By Domain: Accuracy across business domains
- By Complexity: Simple vs. complex schema handling
- Object Types: Detection rates for tables, views, functions, etc.
- Dependencies: Accuracy per relationship type
- Errors: Debug failed tests
Compare Models
- Side-by-side performance comparison
- Cost analysis
- Response time benchmarks

📊 Performance Metrics

Current Results (Foundation Models)

Text Extraction (SQL/JSON)

Model	Overall F1	Node F1	Edge F1	Avg Time	Cost (1K schemas)
Meta Llama 3.1 8B	90%	93%	87%	5s	$26
Meta Llama 3.3 70B	93%	95%	91%	8s	$260
Claude Sonnet 4	95%	97%	93%	10s	$780

Vision Extraction (Diagrams/Images)

Model	Node Detection	Edge Detection	Overall Accuracy	Avg Time	Cost (1K diagrams)
Claude 3.5 Sonnet Vision	95%	85-90%	90-92%	8-10s	$10-20
GPT-4 Vision (Turbo)	95%	85-90%	90-92%	5-8s	$15-30

Production Readiness

Metric	Target	Current (Llama 70B)	Status
Overall F1	≥ 93%	93%	✅ Ready
Node F1	≥ 95%	95%	✅ Ready
Edge F1	≥ 90%	91%	✅ Ready
Response Time	< 15s	8s	✅ Ready

🎯 Supported Products & Features

SQL Databases

Database	Supported Features
Oracle	DDL, PL/SQL, procedures, functions, triggers, packages
PostgreSQL	DDL, PL/pgSQL, functions, triggers, materialized views
MySQL	DDL, procedures, views, functions, triggers
SQL Server	T-SQL, procedures, views, functions, triggers
Snowflake	DDL, SQL, procedures, functions, streams, tasks, transient tables, external tables, clustering

SQL Patterns Detected

CTEs (Common Table Expressions): WITH clause queries
Recursive CTEs: Self-referencing CTEs with UNION ALL
Subqueries: Nested SELECT statements (correlated and non-correlated)
Triggers: BEFORE/AFTER/INSTEAD OF triggers with actions
Foreign Keys: Table relationships and constraints
Materialized Views: Cached views with refresh logic
Function/Procedure Calls: Cross-object invocations
Cross-schema Dependencies: References across database schemas
Index Definitions: Performance indexes on tables
Sequences: Auto-increment generators
Streams (Snowflake): Change data capture on tables
Tasks (Snowflake): Scheduled SQL execution with dependencies
Transient/External Tables (Snowflake): Temporary and external data sources
Clustering Keys (Snowflake): Table clustering for performance

BI Tools

Tool	Supported Features
Power BI	Datasets, tables, columns, measures (DAX), relationships, visuals, reports, pages

Diagrams & Images (Vision AI)

Source Type	Supported Features
ERD Diagrams	Tables, columns, primary keys, foreign keys, relationships, cardinality (1-to-many, etc.)
Database Schemas	Entity boxes, attribute lists, relationship lines, key indicators (PK/FK)
Power BI Screenshots	Visual layouts, dataset connections, measure definitions
Whiteboard Photos	Hand-drawn diagrams, annotations, relationship arrows

Supported Formats: .png, .jpg, .jpeg, .pdf

Expected Accuracy:

Table/Entity Detection: 95%
Column/Attribute Extraction: 90%
Relationship Detection: 85-90%
Cardinality Recognition: 85%

💰 Cost Comparison

Fine-Tuning vs Foundation Models

Fine-Tuning Approach (Archived):

Training: $2,000-3,000 initial
Training time: 15-18 hours
Ongoing costs: $2,000-3,000/month for iterations
Inference: ~$0.01 per schema
Total (1K schemas): ~$2,010-3,010

Foundation Model Approach (Current):

Training: $0
Setup time: <1 hour
Ongoing costs: $0 (no training)
Inference: ~$0.26-0.78 per schema
Total (1K schemas): ~$260-780

Savings: 60-90% cost reduction + 100x faster iteration

🔧 Configuration

API Configuration

Create api/.env (optional - see api/env.template for all options):

# Databricks Configuration
DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
DATABRICKS_TOKEN=your_token_here

# Model Endpoints
LLAMA_8B_ENDPOINT=https://.../serving-endpoints/llama-8b/invocations
LLAMA_70B_ENDPOINT=https://.../serving-endpoints/llama-70b/invocations
CLAUDE_SONNET_4_ENDPOINT=https://.../serving-endpoints/claude-sonnet-4/invocations

# Vision AI (Optional - for diagram extraction)
# Use Databricks endpoint or Anthropic API key
# ANTHROPIC_API_KEY=your_anthropic_key_here
# OPENAI_API_KEY=your_openai_key_here  # For GPT-4V

# Metadata Storage (Optional - for Lakebase integration)
# LAKEBASE_INSTANCE_NAME=data-lineage-metadata  # Auto-configured in Databricks Apps

# Data Paths (optional - defaults to <project_root>/data/...)
# TRAINING_DATA_PATH_SQL=/custom/path/sql_sources
# TRAINING_DATA_PATH_IMAGE=/custom/path/image_sources
# TRAINING_RESULTS_PATH_SQL=/custom/path/results
# EXTRACTION_PATH_SQL=/custom/path/extracts

📝 Note: Data paths now default to <project_root>/data/ for local development. See docs/PATH_CONFIGURATION.md for complete details.

Frontend Configuration

Create frontend/.env.local:

NEXT_PUBLIC_API_URL=http://localhost:8000

🚢 Databricks Deployment

Quick Deployment (3 Steps)

This repo deploys as a single Databricks App (dual server):

Frontend: Next.js on port 8000 (exposed)
Backend: FastAPI on port 8001 (internal)
Frontend proxies /api/* to the backend based on app.yaml

Deploy (recommended)

./deploy_new.sh

What deploy_new.sh does:

(Optional) initializes/validates Lakebase schema if configured in app.yaml
bundles the code (includes api/core/, api/routers/, api/services/)
uploads to your workspace under /Workspace/Users/<you>/<APP_NAME>
runs databricks apps deploy <APP_NAME> --source-code-path <workspace-path>

Prerequisites

Databricks CLI authenticated (if auth expires: databricks auth login --host <workspace-url>)
npm available (for frontend deps during deployment)
python3 available (for local prep steps)

For environment variables (UC/Lakebase/etc.), edit app.yaml.

Authentication

Frontend: Databricks Apps OAuth (automatic user login)
Backend: OAuth token validation + Service Principal for models

See: AUTHENTICATION_SETUP.md for complete guide

Deployment Checklist

Set Up Unity Catalog - Create catalog, schema, and volumes
Configure Model Endpoints - Deploy foundation model serving endpoints
Upload Training Data - Upload 700 schemas to Unity Catalog volume
Deploy Application - Execute ./deploy_new.sh
Verify Deployment - Test API and frontend
Configure CORS - Update allowed origins if needed
Configure Monitoring - Set up alerts and tracking

See: DATABRICKS_DEPLOYMENT.md for detailed instructions

🧪 Synthetic Test Data

Generation

The system includes 700+ synthetic SQL schemas:

5 Databases: Oracle, PostgreSQL, MySQL, SQL Server, Snowflake
10 Domains: Ecommerce, Healthcare, Finance, Logistics, Manufacturing, Retail, Insurance, Education, Telecommunications, Real Estate
3 Splits: Train (70%), Val (15%), Test (15%)

Schema Characteristics

Each schema includes:

Tables: 5-8 base tables with foreign keys
Views: 2-3 views with complex queries
Functions: 1-2 reusable functions
Procedures: 1-2 stored procedures
Triggers: 1 audit/logging trigger

Complexity Distribution:

Simple: ≤12 objects
Medium: 13-16 objects
Complex: 17-20 objects
Very Complex: >20 objects

🛠️ Development

Tech Stack

Backend:

Python 3.9+
FastAPI (REST API with SSE streaming)
OpenAI SDK (Databricks vision endpoint compatibility)
Anthropic SDK (Claude Vision API)
asyncpg (PostgreSQL async driver for Lakebase)
Requests (HTTP client)
Pydantic (data validation)
Pillow (image processing)

Frontend:

TypeScript 5
Next.js 16 (App Router)
React 19
Tailwind CSS
shadcn/ui (Radix UI)
Recharts (visualizations)
ReactFlow (graph visualization)

Infrastructure:

Databricks (model serving)
Unity Catalog Volumes (persistent storage)
Databricks Lakebase (PostgreSQL for metadata)
Databricks Apps (hosting with OAuth)

Adding Features

Backend (API):

Add endpoint in the appropriate router under api/routers/ (or api/hierarchy_api.py for hierarchy endpoints)
If you add a new router module, include it in api/app_factory.py
Keep api/main.py minimal (it should only expose app = create_app())
Test with FastAPI docs: http://localhost:8000/docs
Update API documentation in api/README.md

Frontend:

Create page in frontend/app/your-feature/page.tsx
Add to sidebar in frontend/components/app-sidebar.tsx
Use existing UI components from frontend/components/ui/
Test in browser at http://localhost:3000

🐛 Troubleshooting

Common Issues

API won't start:

Check Python version: python --version (requires 3.9+)
Activate virtual environment: source venv/bin/activate
Install dependencies: pip install -r requirements.txt
Verify .env file exists with valid tokens

Frontend won't start:

Check Node.js version: node --version (requires 18+)
Clear cache: rm -rf node_modules .next && npm install
Verify .env.local has correct API URL

CORS errors:

Ensure API CORS is configured for http://localhost:3000
Check API is running on port 8000
Verify NEXT_PUBLIC_API_URL in .env.local

Model endpoint errors:

Verify Databricks token is valid
Check endpoint URLs are correct in .env
Test endpoint directly: curl -X POST <endpoint-url> -H "Authorization: Bearer <token>"

File size errors:

Llama models: Max 96K characters
Claude Sonnet 4: Max 150K characters
Split large files into smaller chunks
Remove unnecessary comments/whitespace

Vision AI errors:

Ensure ANTHROPIC_API_KEY or OPENAI_API_KEY is set in .env for direct API access
For Databricks endpoints, verify vision model is deployed
Check image format: .png, .jpg, .jpeg, .pdf only
Image size limit: 5MB per image (API limitation)
For poor results, try higher quality images or clearer diagrams

For detailed troubleshooting, see:

📚 Documentation

Architecture & Design

ARCHITECTURE_DIAGRAM.md - 🆕 Complete system architecture, workflows, and diagrams
METADATA_EXTRACTION_ARCHITECTURE.md - AI-first metadata extraction approach

Getting Started

QUICKSTART.md - 5-minute quick start guide
SETUP_GUIDE.md - Complete local setup instructions
RUN_LOCALLY.md - Local development setup

Deployment

DATABRICKS_DEPLOYMENT.md - Deploy to Databricks Apps
LAKEBASE_DEPLOYMENT_GUIDE.md - Lakebase (PostgreSQL) metadata storage setup
AUTHENTICATION_SETUP.md - Configure authentication
SECURITY_SETUP.md - Security best practices

Component Documentation

api/README.md - FastAPI backend documentation
frontend/README.md - Next.js frontend documentation
phase1_training_data/README.md - Data generation documentation
PROJECTS_AND_CONTEXT_MANAGEMENT.md - Project organization and metadata management

Vision AI Documentation

VISION_EXTRACTION_QUICKSTART.md - Quick start for diagram extraction
VISION_MODEL_COMPARISON.md - Model selection guide
MERGE_EXTRACTIONS_GUIDE.md - Combining SQL + Image + JSON results

Project Journey

BLOG_POST_TWO_MONTH_JOURNEY.md - Complete project journey, learnings, and insights

📈 Roadmap & Future Enhancements

Recently Completed ✅

✅ Hierarchical Multi-Tenant Workflow - 🆕 Customer → Business Unit → Use Case organization
✅ Lakebase Integration - PostgreSQL-based metadata storage for hierarchical project management
✅ Streamlined Navigation - 🆕 Context-aware UI with sidebar tree, workflow-aligned menus
✅ Power BI Support - Direct .pbix and .json file upload with hierarchical extraction
✅ Vision AI Extraction - Extract from ERD diagrams and screenshots (Claude 3.5 Sonnet Vision, GPT-4V)
✅ Merge Extractions - Combine SQL + Power BI + Images into unified graphs
✅ Background Tasks - Asynchronous processing with polling for long extractions
✅ Snowflake Support - Full support for Snowflake SQL including streams, tasks, transient tables, clustering
✅ Unity Catalog Volumes - Production-ready persistent storage with hierarchical paths

Suggested Future Directions

Chain-of-Thought (COT) Prompting
- Guide models through step-by-step extraction
- Potential 2-5% accuracy improvement
- Trade-off: Higher latency and token costs
Hybrid Symbolic + LLM Approach
- Use deterministic parsers for simple patterns
- Reserve LLMs for complex ambiguous cases
- Potential 50-70% cost reduction
Advanced Caching
- Cache results for common schemas
- Reduce redundant API calls
- Improve response times
Additional Data Sources
- Tableau metadata
- Looker models
- dbt models
- Teradata
- SAP HANA
Enhanced Visualizations
- Animated dependency flows
- Impact analysis (what-if scenarios)
- Export to diagram formats (Mermaid, PlantUML)
- Visual differentiation by source (SQL vs Image vs JSON)
Vision AI Enhancements
- OCR preprocessing for low-quality images
- Multi-page diagram support
- Automatic diagram type detection
- Confidence-based diagram quality warnings

See BLOG_POST_TWO_MONTH_JOURNEY.md for detailed future directions.

🤝 Contributing

Fork the repository
Create a feature branch: git checkout -b feature/your-feature
Make changes and test thoroughly
Submit a pull request with detailed description

📄 License

📞 Support

For questions or issues:

Check documentation files (listed above)
Review API documentation at http://localhost:8000/docs
Check logs in console output
Contact project maintainers

Version: 3.0.0
Last Updated: November 30, 2024
Status: Production Ready ✅

Architecture: Hierarchical Multi-Tenant + Foundation Models + Vision AI
New in 3.0: Customer/BU/Use Case hierarchy, Lakebase metadata storage, Streamlined workflow, Power BI support (.pbix and .json)
Supported Databases: Oracle, PostgreSQL, MySQL, SQL Server, Snowflake
Supported BI Tools: Power BI (.pbix, .json)
Supported Diagram Formats: PNG, JPG, PDF (via Vision AI)
Previous Approaches Explored: Hierarchical Fine-Tuning (archived) <<<<<<< HEAD

407f04c (Datagraph release 1) =======

datagraph

5d8df17 (first commit)

Name		Name	Last commit message	Last commit date
Latest commit History 1 Commit
api		api
data/training		data/training
docs		docs
frontend		frontend
scripts		scripts
test_security_samples		test_security_samples
.DS_Store		.DS_Store
ARCHITECTURE_DIAGRAM.html		ARCHITECTURE_DIAGRAM.html
LICENSE		LICENSE
README.md		README.md
USER_GUIDE_ARCHITECTURE.html		USER_GUIDE_ARCHITECTURE.html
app.yaml		app.yaml
deploy_new.sh		deploy_new.sh
set-token.sh		set-token.sh
start.sh		start.sh
stop.sh		stop.sh

Folders and files

Latest commit

History

Repository files navigation

datagraph

Data Model Dependency extraction using Databricks hosted foundation models

DataGraph AI

🎯 Overview

Key Insights

✨ Key Features

🏢 Organization & Workflow

📊 Data Sources

AI & Models

🎨 Interface & Visualization

🚀 Backend & Deployment

Testing & Analytics

📁 Project Structure

🚀 Quick Start (Local Development)

Prerequisites

1. Clone Repository

2. Setup Lakebase (PostgreSQL)

3. Start Backend (API)

4. Start Frontend

5. Access the Application

6. Quick Start Guide

Alternative: One-Command Start (Both Backend & Frontend)

Local “Databricks Apps” simulation (optional)

🎨 User Workflows

New Hierarchical Workflow: Organize Your Data Models

Complete User Journey (4 Steps)

Detailed Workflow Steps

1. Setup Your Organization (One-Time)

2. Upload Source Files

3. Extract Dependencies

4. View & Analyze Results

5. Browse & Manage Projects

Testing Workflow: Evaluate Models

📊 Performance Metrics

Current Results (Foundation Models)

Text Extraction (SQL/JSON)

Vision Extraction (Diagrams/Images)

Production Readiness

🎯 Supported Products & Features

SQL Databases

SQL Patterns Detected

BI Tools

Diagrams & Images (Vision AI)

💰 Cost Comparison

Fine-Tuning vs Foundation Models

🔧 Configuration

API Configuration

Frontend Configuration

🚢 Databricks Deployment

Quick Deployment (3 Steps)

Deploy (recommended)

Prerequisites

Authentication

Deployment Checklist

🧪 Synthetic Test Data

Generation

Schema Characteristics

🛠️ Development

Tech Stack

Adding Features

🐛 Troubleshooting

Common Issues

📚 Documentation

Architecture & Design

Getting Started

Deployment

Component Documentation

Vision AI Documentation

Project Journey

📈 Roadmap & Future Enhancements

Recently Completed ✅

Suggested Future Directions

🤝 Contributing

📄 License

📞 Support

datagraph