sequel2sql/sequel2sql.py at main · SVijayB/sequel2sql · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
"""
Sequel2SQL - Main Application Entry Point

Runs a web-based chat interface for the SQL agent using Pydantic AI's built-in web UI.
The agent is connected to a PostgreSQL database running in Docker.

Usage:
    # Use default database (postgres)
    uv run python sequel2sql.py

    # Connect to a specific database
    DATABASE=california_schools_template uv run python sequel2sql.py

    # Or set in .env file:
    DATABASE=formula_1_template
"""

import os
import sys

from dotenv import load_dotenv

from src.agent.sqlagent import (
    SUPPORTED_MODELS,
    get_database_deps,
    webui_agent,
)
from src.webui import create_branded_web_app

# Load environment variables
load_dotenv()

# =============================================================================
# Configuration
# =============================================================================

DATABASE_NAME = os.getenv("DATABASE", "postgres")

print("=" * 70)
print("SEQUEL2SQL - Web Interface")
print("=" * 70)
print(f"\nDatabase: {DATABASE_NAME}")
print("To use a different database, set DATABASE environment variable")

# =============================================================================
# Connect to Database
# =============================================================================

try:
    print(f"Connecting to {DATABASE_NAME}...")
    deps = get_database_deps(DATABASE_NAME)
    print("Connected successfully!")
    print(f"   Tables available: {len(deps.database.table_names)}")
except Exception as e:
    print(f"\nFailed to connect to database '{DATABASE_NAME}': {e}")
    print("\nMake sure PostgreSQL Docker container is running:")
    print("   docker compose -f benchmark/docker-compose.yml up -d postgresql")
    sys.exit(1)

# =============================================================================
# Web Application
# =============================================================================

app = create_branded_web_app(
    agent=webui_agent,
    deps=deps,
    models={
        "Mistral Large Latest": SUPPORTED_MODELS["mistral"],
        "Google Gemini Flash": SUPPORTED_MODELS["google"],
    },
)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)