Enhance security and add vLLM model support with deployment improvements

.env.example

@@ -0,0 +1,60 @@
+# PremSQL Environment Configuration
+# Copy this file to .env and fill in your values
+
+# ============================================================
+# Security Configuration
+# ============================================================
+
+# MODE SELECTION:
+# - Development: PREMSQL_DJANGO_DEBUG=true (no token required)
+# - Production: PREMSQL_DJANGO_DEBUG=false (token REQUIRED)
+
+# For development (recommended for local testing):
+PREMSQL_DJANGO_DEBUG=true
+
+# For production, uncomment and set these:
+#PREMSQL_DJANGO_DEBUG=false
+#PREMSQL_API_TOKEN=<generate-with: python -c "import secrets; print(secrets.token_hex(32))">
+#PREMSQL_DJANGO_SECRET_KEY=<generate-with: python -c "import secrets; print(secrets.token_hex(32))">
+
+# Optional: Custom session name (auto-generated if not set)
+#PREMSQL_SESSION_NAME=my-session
+
+# ============================================================
+# Django Settings
+# ============================================================
+#PREMSQL_DJANGO_DEBUG=false
+#PREMSQL_ALLOWED_HOSTS=127.0.0.1,localhost
+#PREMSQL_ALLOWED_ORIGINS=http://127.0.0.1:8501,http://localhost:8501
+
+# ============================================================
+# LLM Provider Configuration
+# Choose ONE provider from the options below
+# ============================================================
+
+# Option 1: vLLM (self-hosted models with OpenAI-compatible API)
+# Start vLLM: vllm serve /path/to/model --port 8000
+#VLLM_BASE_URL=http://localhost:8000/v1
+#VLLM_MODEL_NAME=your-model-name
+
+# Option 2: Custom OpenAI-Compatible Service
+# For LM Studio, LocalAI, Text Generation WebUI, or any custom deployment
+#CUSTOM_BASE_URL=http://localhost:1234/v1
+#CUSTOM_MODEL_NAME=your-model-name
+#CUSTOM_API_KEY=your-key-if-required
+
+# Option 3: OpenAI (official API)
+# Get your API key from https://platform.openai.com/
+#OPENAI_API_KEY=sk-your-openai-api-key
+#OPENAI_MODEL_NAME=gpt-4o-mini
+
+# Option 4: PremAI
+# Get your API key from https://app.premai.io/
+#PREMAI_API_KEY=your-premai-api-key
+#PREMAI_PROJECT_ID=your-project-id
+
+# Option 5: Ollama (local models)
+# Install Ollama: https://ollama.com
+# Run: ollama pull llama3.2
+#OLLAMA_BASE_URL=http://127.0.0.1:11434
+#OLLAMA_MODEL_NAME=llama3.2

PR_DESCRIPTION.md

@@ -0,0 +1,201 @@
+## Summary
+
+This PR addresses security vulnerabilities (#40), fixes dependency conflicts (#37), and adds new features for better model support and deployment experience.
+
+---
+
+## Dependency Fixes (Issue #37)
+
+Fixed version conflicts reported during installation:
+
+| Conflict                                                | Cause                              | Resolution               |
+| ------------------------------------------------------- | ---------------------------------- | ------------------------ |
+| `httpx<0.29` vs `>=0.27` required by ollama/browser-use | fastapi 0.112 pinned old httpx     | Update fastapi >=0.115.0 |
+| `starlette>=0.41.3` vs 0.38.6                           | fastapi 0.112 pinned old starlette | Update fastapi >=0.115.0 |
+| `python==3.12 not compatible`                           | Version range `^3.10`              | Extend to `>=3.10,<3.13` |
+
+**Changes in pyproject.toml:**
+
+```toml
+# Before
+python = "^3.10"
+fastapi = "^0.112.0"
+
+# After
+python = ">=3.10,<3.13"  # Support 3.11, 3.12
+fastapi = ">=0.115.0"    # Brings httpx>=0.27, starlette>=0.41
+httpx = ">=0.27.0"       # Explicit for clarity
+starlette = ">=0.41.0"   # Explicit for clarity
+```
+
+---
+
+## Security Fixes (Issue #40)
+
+### Critical Severity
+- **RCE via eval()** → Replace with `ast.literal_eval()`
+- **SSRF** → `normalize_base_url()` restricts to loopback addresses
+- **Missing Authentication** → Token-based auth via `PREMSQL_API_TOKEN`
+
+### High Severity
+- **SQL Write Operations** → `enforce_read_only_sql()` blocks INSERT/UPDATE/DELETE/DROP
+- **Path Traversal** → Whitelist validation: `[A-Za-z0-9_-]{1,64}`
+- **Information Disclosure** → Removed sensitive fields from API responses
+- **Error Message Leakage** → `safe_error_message()` whitelist mechanism
+
+### Medium Severity
+- **Pickle RCE** → Add `weights_only=True` to `torch.load()`
+- **Swagger Exposure** → Restrict access based on DEBUG mode
+- **Process Management** → PID file-based precise process tracking
+- **SQL Injection** → Parameterized queries with `?` placeholders
+
+---
+
+## New Features
+
+### 1. LLM Provider Support
+
+Added new generators for self-hosted and custom LLM deployments:
+
+| Generator                           | Use Case                                                 |
+| ----------------------------------- | -------------------------------------------------------- |
+| `Text2SQLGeneratorVLLM`             | vLLM deployed models (auto Qwen3 thinking mode handling) |
+| `Text2SQLGeneratorOpenAICompatible` | Any OpenAI-compatible API (LM Studio, LocalAI, etc.)     |
+
+**Usage:**
+```python
+from premsql.generators import Text2SQLGeneratorVLLM
+
+generator = Text2SQLGeneratorVLLM(
+    model_name="/models/qwen",
+    base_url="http://localhost:8000/v1",
+    experiment_name="test", type="test"
+)
+```
+
+### 2. Easy Deployment Script
+
+Added `start_agent.py` for one-command AgentServer startup:
+
+```bash
+# Configure in .env, then:
+python start_agent.py
+```
+
+Auto-detects configured LLM provider from environment variables.
+
+### 3. Bug Fixes
+
+| Fix                     | Description                                                                    |
+| ----------------------- | ------------------------------------------------------------------------------ |
+| Plot image generation   | Fixed `plot_image=False` hardcoded in server_mode, now generates base64 images |
+| Matplotlib backend      | Added `Agg` backend for non-interactive server mode                            |
+| Session duplicate error | Auto-delete existing session when creating new one with same name              |
+| Session deletion memory | Recreate memory DB table after deletion to prevent OperationalError            |
+| API token propagation   | Fixed Django backend not passing API token to AgentServer                      |
+| Environment loading     | Added dotenv loading in Django manage.py and Streamlit main.py                 |
+
+### 4. UI Improvements
+
+- Session list now shows delete button for each session (no need to type session name)
+- Delete operation auto-refreshes the page
+- Simplified session creation form with placeholder hints
+
+---
+
+## Configuration
+
+### Environment Variables (.env)
+
+All configuration is optional for local development - tokens are auto-generated if not set.
+
+```bash
+# Security (auto-generated for local dev)
+#PREMSQL_API_TOKEN=your-token-here
+#PREMSQL_DJANGO_SECRET_KEY=your-secret-here
+
+# LLM Provider (choose one)
+VLLM_BASE_URL=http://localhost:8000/v1
+VLLM_MODEL_NAME=/models/your-model
+
+# Or custom OpenAI-compatible service
+CUSTOM_BASE_URL=http://localhost:1234/v1
+CUSTOM_MODEL_NAME=local-model
+
+# Or official OpenAI
+OPENAI_API_KEY=sk-your-key
+OPENAI_MODEL_NAME=gpt-4o-mini
+```
+
+### Quick Start
+
+```bash
+# 1. Copy and configure .env
+cp .env.example .env
+# Edit .env to set your LLM provider
+
+# 2. Start services
+python start_agent.py          # AgentServer on port 8100
+premsql launch all             # Django + Streamlit
+
+# 3. Open browser
+http://localhost:8501          # PremSQL Playground
+```
+
+---
+
+## Security Configuration
+
+### Development Mode (Recommended for Local Testing)
+
+Set `PREMSQL_DJANGO_DEBUG=true` to enable development mode:
+
+```bash
+# .env for development
+PREMSQL_DJANGO_DEBUG=true
+# No need to set PREMSQL_API_TOKEN - auto-generated
+```
+
+In development mode:
+- API token is auto-generated (printed in startup logs)
+- Authentication is skipped for all services
+- Suitable for local testing only
+
+### Production Mode (Required for Deployment)
+
+**IMPORTANT**: Production mode requires explicit token configuration:
+
+```bash
+# .env for production
+PREMSQL_DJANGO_DEBUG=false
+PREMSQL_API_TOKEN=<your-strong-random-token>  # REQUIRED!
+PREMSQL_DJANGO_SECRET_KEY=<your-strong-random-secret>
+```
+
+**How to generate secure tokens:**
+
+```bash
+# Using Python (recommended)
+python -c "import secrets; print(secrets.token_hex(32))"
+
+# Using OpenSSL
+openssl rand -hex 32
+
+# Using UUID
+python -c "import uuid; print(uuid.uuid4().hex)"
+```
+
+Example output: `a1b2c3d4e5f6...` (64 characters hex string)
+
+**Security behavior summary:**
+
+| Mode        | PREMSQL_API_TOKEN | Behavior                     |
+| ----------- | ----------------- | ---------------------------- |
+| Development | Not set           | Auto-generated, auth skipped |
+| Development | Set               | Use configured token         |
+| Production  | Not set           | **Service fails to start**   |
+| Production  | Set               | Enforce authentication       |
+
+---
+
+Resolves #37 #40 

README.md

@@ -131,7 +131,11 @@ plot = agent(
 You can launch the PremSQL Playground (as shown in the above video by adding these two additional lines after instantiating Agent)
 
 ```python
-agent_server = AgentServer(agent=agent, port={port})
+agent_server = AgentServer(
+    agent=agent,
+    port={port},
+    api_token=os.environ.get("PREMSQL_API_TOKEN")
+)
 agent_server.launch()
 ```
 
@@ -147,6 +151,24 @@ and on the second side of the terminal write:
 python start_agent.py
 ```
 
+### Security Defaults
+
+Recent hardening changes added safer defaults for Playground and AgentServer deployments:
+
+- Playground session registration only accepts loopback AgentServer URLs such as `http://127.0.0.1:8100`.
+- API authentication is supported through `PREMSQL_API_TOKEN`. When this variable is set, the Django backend, Playground clients, and FastAPI AgentServer all expect the same token.
+- Generated SQL is now restricted to read-only queries. `INSERT`, `UPDATE`, `DELETE`, `DROP`, and similar statements are rejected by default.
+- Direct raw SQL passthrough using backtick-wrapped prompts is disabled by default.
+- CSV imports now enforce upload limits to reduce accidental or malicious resource exhaustion.
+
+Recommended local environment variables:
+
+```bash
+export PREMSQL_API_TOKEN="change-this-token"
+export PREMSQL_DJANGO_SECRET_KEY="change-this-secret"
+export PREMSQL_ALLOWED_HOSTS="127.0.0.1,localhost"
+```
+
 ## 📦 Components Overview
 
 ### [Datasets](https://docs.premai.io/premsql/introduction)
@@ -447,13 +469,18 @@ In the above section you have see how we have defined our agent. You can deploy
 
 ```python
 # File name: start_agent_server.py
+import os
 from premsql.playground import AgentServer
 from premsql.agents import BaseLineAgent
 
 # Define your agent as shown above:
 agent = BaseLineAgent(...)
 
-agent_server = AgentServer(agent=agent, port={port})
+agent_server = AgentServer(
+    agent=agent,
+    port={port},
+    api_token=os.environ.get("PREMSQL_API_TOKEN")
+)
 agent_server.launch()
 ```
 
@@ -463,7 +490,7 @@ Now inside another terminal write:
 python start_agent_server.py
 ```
 
-This can be any python file name. This will run a fastapi server. You need to paste the deployed url and paste it inside `Register New Session` part of the UI. Below shows, how the basic backend architecture looks like on how Playground communicates with the server.
+This can be any python file name. This will run a fastapi server. You need to paste a loopback url such as `http://127.0.0.1:8100` inside `Register New Session` in the UI. Below shows, how the basic backend architecture looks like on how Playground communicates with the server.
 
 ![](/assets/agent_server.png)
 

premsql/agents/__init__.py

@@ -1,4 +1,11 @@
-from premsql.agents.baseline.main import BaseLineAgent
-from premsql.agents.memory import AgentInteractionMemory
+from importlib import import_module
 
-__all__ = ["BaseLineAgent", "AgentInteractionMemory"]
+__all__ = ["BaseLineAgent", "AgentInteractionMemory"]
+
+
+def __getattr__(name):
+    if name == "BaseLineAgent":
+        return import_module("premsql.agents.baseline.main").BaseLineAgent
+    if name == "AgentInteractionMemory":
+        return import_module("premsql.agents.memory").AgentInteractionMemory
+    raise AttributeError(f"module 'premsql.agents' has no attribute {name!r}")

premsql/agents/base.py

@@ -158,13 +158,14 @@ def __call__(
         question: str,
         input_dataframe: Optional[dict] = None,
         server_mode: Optional[bool] = False,
+        generate_plot_image: Optional[bool] = True,
     ) -> Union[ExitWorkerOutput, AgentOutput]:
         if server_mode:
             kwargs = self.route_worker_kwargs.get("plot", None)
             kwargs = (
-                {"plot_image": False}
+                {"plot_image": generate_plot_image}
                 if kwargs is None
-                else {**kwargs, "plot_image": False}
+                else {**kwargs, "plot_image": generate_plot_image}
             )
             self.route_worker_kwargs["plot"] = kwargs
 

premsql/agents/baseline/prompts.py

@@ -129,6 +129,7 @@
 a single line (string format).
 - Make sure the column names are correct and exists in the table
 - For column names which has a space with it, make sure you have put `` in that column name
+- Only generate read-only SQL queries. Never generate INSERT, UPDATE, DELETE, DROP, ALTER, or PRAGMA statements.
 
 # Database and Table Schema:
 {schemas}
@@ -145,6 +146,7 @@
 a single line (string format).
 - Make sure the column names are correct and exists in the table
 - For column names which has a space with it, make sure you have put `` in that column name
+- Only generate read-only SQL queries. Never generate INSERT, UPDATE, DELETE, DROP, ALTER, or PRAGMA statements.
 
 # Database and Table Schema:
 {schemas}