pgnode 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pgnode-0.1.0/PKG-INFO +223 -0
- pgnode-0.1.0/README.md +188 -0
- pgnode-0.1.0/app/__init__.py +1 -0
- pgnode-0.1.0/app/agent/__init__.py +1 -0
- pgnode-0.1.0/app/agent/agent.py +216 -0
- pgnode-0.1.0/app/agent/followup.py +127 -0
- pgnode-0.1.0/app/agent/intent_agent.py +21 -0
- pgnode-0.1.0/app/agent/memory.py +61 -0
- pgnode-0.1.0/app/agent/sql_metadata.py +101 -0
- pgnode-0.1.0/app/agent/sql_quality.py +117 -0
- pgnode-0.1.0/app/agent/sql_quoter.py +92 -0
- pgnode-0.1.0/app/agent/validator.py +72 -0
- pgnode-0.1.0/app/cli/__init__.py +1 -0
- pgnode-0.1.0/app/cli/diagnostics.py +314 -0
- pgnode-0.1.0/app/cli/formatting.py +319 -0
- pgnode-0.1.0/app/cli/history.py +76 -0
- pgnode-0.1.0/app/cli/main.py +801 -0
- pgnode-0.1.0/app/core/__init__.py +5 -0
- pgnode-0.1.0/app/core/config.py +157 -0
- pgnode-0.1.0/app/db/__init__.py +1 -0
- pgnode-0.1.0/app/db/connection.py +63 -0
- pgnode-0.1.0/app/db/query_executor.py +61 -0
- pgnode-0.1.0/app/db/schema_loader.py +53 -0
- pgnode-0.1.0/app/db/ssh_tunnel.py +82 -0
- pgnode-0.1.0/app/llm/__init__.py +1 -0
- pgnode-0.1.0/app/llm/intent_classifier.py +146 -0
- pgnode-0.1.0/app/llm/ollama_client.py +33 -0
- pgnode-0.1.0/app/llm/prompt_builder.py +250 -0
- pgnode-0.1.0/app/llm/sql_generator.py +63 -0
- pgnode-0.1.0/pgnode.egg-info/PKG-INFO +223 -0
- pgnode-0.1.0/pgnode.egg-info/SOURCES.txt +47 -0
- pgnode-0.1.0/pgnode.egg-info/dependency_links.txt +1 -0
- pgnode-0.1.0/pgnode.egg-info/entry_points.txt +2 -0
- pgnode-0.1.0/pgnode.egg-info/requires.txt +15 -0
- pgnode-0.1.0/pgnode.egg-info/top_level.txt +1 -0
- pgnode-0.1.0/pyproject.toml +77 -0
- pgnode-0.1.0/setup.cfg +4 -0
- pgnode-0.1.0/tests/test_cli_launch.py +50 -0
- pgnode-0.1.0/tests/test_config.py +91 -0
- pgnode-0.1.0/tests/test_diagnostics.py +49 -0
- pgnode-0.1.0/tests/test_followup.py +112 -0
- pgnode-0.1.0/tests/test_formatting.py +95 -0
- pgnode-0.1.0/tests/test_formatting_id.py +35 -0
- pgnode-0.1.0/tests/test_intent_classifier.py +110 -0
- pgnode-0.1.0/tests/test_prompt_builder.py +87 -0
- pgnode-0.1.0/tests/test_sql_quality.py +83 -0
- pgnode-0.1.0/tests/test_sql_quoter.py +65 -0
- pgnode-0.1.0/tests/test_ssh_tunnel.py +24 -0
- pgnode-0.1.0/tests/test_validator.py +48 -0
pgnode-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,223 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pgnode
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Terminal-native local AI agent for PostgreSQL databases.
|
|
5
|
+
Author: Chayan Mann
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/chayan-mann/pgnode
|
|
8
|
+
Project-URL: Repository, https://github.com/chayan-mann/pgnode
|
|
9
|
+
Project-URL: Changelog, https://github.com/chayan-mann/pgnode/blob/main/CHANGELOG.md
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Environment :: Console
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Topic :: Database
|
|
19
|
+
Requires-Python: >=3.11
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
Requires-Dist: ollama==0.3.3
|
|
22
|
+
Requires-Dist: psycopg[binary]>=3.2.0
|
|
23
|
+
Requires-Dist: python-dotenv==1.0.1
|
|
24
|
+
Requires-Dist: questionary>=2.0.1
|
|
25
|
+
Requires-Dist: rich>=13.7.1
|
|
26
|
+
Requires-Dist: SQLAlchemy>=2.0.44
|
|
27
|
+
Requires-Dist: sshtunnel>=0.4.0
|
|
28
|
+
Requires-Dist: typer>=0.16.0
|
|
29
|
+
Provides-Extra: dev
|
|
30
|
+
Requires-Dist: build>=1.2.0; extra == "dev"
|
|
31
|
+
Requires-Dist: mypy>=1.11.0; extra == "dev"
|
|
32
|
+
Requires-Dist: pytest>=8.3.0; extra == "dev"
|
|
33
|
+
Requires-Dist: ruff>=0.6.0; extra == "dev"
|
|
34
|
+
Requires-Dist: twine>=5.1.0; extra == "dev"
|
|
35
|
+
|
|
36
|
+
# **pgnode**
|
|
37
|
+
|
|
38
|
+
**Local LLM (Ollama) + PostgreSQL Agent**
|
|
39
|
+
|
|
40
|
+
An offline-first AI agent that converts natural language into **validated SQL queries** and executes them safely on your PostgreSQL database.
|
|
41
|
+
|
|
42
|
+
---
|
|
43
|
+
|
|
44
|
+
## Overview
|
|
45
|
+
|
|
46
|
+
**pgnode** is a local AI-powered database operator. It connects to your PostgreSQL instance and allows you to interact with your data using plain English while ensuring safety, control, and privacy.
|
|
47
|
+
|
|
48
|
+
* No external APIs
|
|
49
|
+
* No data leaves your system
|
|
50
|
+
* Fully local using Ollama
|
|
51
|
+
|
|
52
|
+
---
|
|
53
|
+
|
|
54
|
+
## Core Features
|
|
55
|
+
|
|
56
|
+
* Natural language → SQL conversion
|
|
57
|
+
* Safe query execution with validation layer
|
|
58
|
+
* Schema-aware query generation
|
|
59
|
+
* Works with existing PostgreSQL databases (pgAdmin compatible)
|
|
60
|
+
* CLI-first interface (fast and developer-friendly)
|
|
61
|
+
* Fully offline with local LLM
|
|
62
|
+
|
|
63
|
+
---
|
|
64
|
+
|
|
65
|
+
## Architecture
|
|
66
|
+
|
|
67
|
+
```
|
|
68
|
+
User Prompt
|
|
69
|
+
↓
|
|
70
|
+
Agent (planner)
|
|
71
|
+
↓
|
|
72
|
+
SQL Generator (LLM)
|
|
73
|
+
↓
|
|
74
|
+
Validator (safety layer)
|
|
75
|
+
↓
|
|
76
|
+
Query Executor (PostgreSQL)
|
|
77
|
+
↓
|
|
78
|
+
Response
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
---
|
|
82
|
+
|
|
83
|
+
## Tech Stack
|
|
84
|
+
|
|
85
|
+
### Core
|
|
86
|
+
|
|
87
|
+
* Python
|
|
88
|
+
* PostgreSQL
|
|
89
|
+
* Ollama (local LLM runtime)
|
|
90
|
+
|
|
91
|
+
### Libraries
|
|
92
|
+
|
|
93
|
+
* SQLAlchemy → DB interaction
|
|
94
|
+
* psycopg2 → PostgreSQL adapter
|
|
95
|
+
* Typer → CLI interface
|
|
96
|
+
* FastAPI (optional) → API layer
|
|
97
|
+
* LlamaIndex / FAISS (optional) → schema-aware retrieval
|
|
98
|
+
|
|
99
|
+
---
|
|
100
|
+
|
|
101
|
+
## Environment
|
|
102
|
+
|
|
103
|
+
Create `.env` in project root:
|
|
104
|
+
|
|
105
|
+
- `OLLAMA_HOST=http://127.0.0.1:11434`
|
|
106
|
+
- `DATABASE_URL=postgresql+psycopg2://user:pass@localhost:5432/dbname`
|
|
107
|
+
- `LLM_MODEL=deepseek-coder:6.7b` (optional)
|
|
108
|
+
|
|
109
|
+
## Run CLI
|
|
110
|
+
|
|
111
|
+
Install from source while developing:
|
|
112
|
+
|
|
113
|
+
```bash
|
|
114
|
+
pip install -e .
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
First-time setup:
|
|
118
|
+
|
|
119
|
+
```bash
|
|
120
|
+
pgnode connect
|
|
121
|
+
pgnode doctor
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
`connect` saves your database URL, Ollama host, and exact local model name to your user config. Environment variables still override saved config when present.
|
|
125
|
+
|
|
126
|
+
SSH tunnel databases are supported too. Choose `ssh` during `pgnode connect` or use flags:
|
|
127
|
+
|
|
128
|
+
```bash
|
|
129
|
+
pgnode connect \
|
|
130
|
+
--connection-type ssh \
|
|
131
|
+
--database-url "postgresql://user:pass@internal-db:5432/dbname" \
|
|
132
|
+
--ssh-host "bastion.example.com" \
|
|
133
|
+
--ssh-port 22 \
|
|
134
|
+
--ssh-user "ubuntu" \
|
|
135
|
+
--ssh-key-path "~/.ssh/id_rsa" \
|
|
136
|
+
--remote-host "127.0.0.1" \
|
|
137
|
+
--remote-port 5432 \
|
|
138
|
+
--local-port 0 \
|
|
139
|
+
--model "deepseek-coder-v2:16b"
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
Activate venv once:
|
|
143
|
+
|
|
144
|
+
```bash
|
|
145
|
+
source venv/bin/activate
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
Interactive conversation (context kept only in current session):
|
|
149
|
+
|
|
150
|
+
```bash
|
|
151
|
+
./pgnode run
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
or simply:
|
|
155
|
+
|
|
156
|
+
```bash
|
|
157
|
+
./pgnode
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
Useful chat commands:
|
|
161
|
+
|
|
162
|
+
- `/history` show recent turns
|
|
163
|
+
- `/clear` clear current session context
|
|
164
|
+
- `/exit` or `/quit` leave session
|
|
165
|
+
- Natural language meta-questions also work, e.g.:
|
|
166
|
+
- `what question did i ask you last`
|
|
167
|
+
- `which query did you execute last`
|
|
168
|
+
- `last result`
|
|
169
|
+
|
|
170
|
+
One-shot mode (no prior context):
|
|
171
|
+
|
|
172
|
+
```bash
|
|
173
|
+
./pgnode run "list all users with limit 5"
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
SQL-only generation (no execution):
|
|
177
|
+
|
|
178
|
+
```bash
|
|
179
|
+
./pgnode sql "top 5 customers by revenue last month"
|
|
180
|
+
```
|
|
181
|
+
|
|
182
|
+
Explain mode (SQL + short reasoning, no execution):
|
|
183
|
+
|
|
184
|
+
```bash
|
|
185
|
+
./pgnode explain "monthly revenue trend"
|
|
186
|
+
```
|
|
187
|
+
|
|
188
|
+
Schema helpers:
|
|
189
|
+
|
|
190
|
+
```bash
|
|
191
|
+
./pgnode tables
|
|
192
|
+
./pgnode describe Product
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
Environment and connectivity checks:
|
|
196
|
+
|
|
197
|
+
```bash
|
|
198
|
+
./pgnode config
|
|
199
|
+
./pgnode config-set --model "deepseek-coder-v2:16b"
|
|
200
|
+
./pgnode config-set --database-url "postgresql://user:pass@localhost:5432/dbname"
|
|
201
|
+
./pgnode config-set --connection-type ssh --ssh-host "bastion.example.com" --ssh-user "ubuntu" --ssh-key-path "~/.ssh/id_rsa" --remote-host "127.0.0.1" --remote-port 5432
|
|
202
|
+
./pgnode doctor
|
|
203
|
+
./pgnode models
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
Persistent local history:
|
|
207
|
+
|
|
208
|
+
```bash
|
|
209
|
+
./pgnode history
|
|
210
|
+
./pgnode rerun 12
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
Write behavior:
|
|
214
|
+
|
|
215
|
+
```bash
|
|
216
|
+
./pgnode run "update users set phone='999' where id=1"
|
|
217
|
+
```
|
|
218
|
+
|
|
219
|
+
`INSERT/UPDATE` now require confirmation by default. Use `--yes` to skip prompt.
|
|
220
|
+
|
|
221
|
+
```bash
|
|
222
|
+
./pgnode run --yes "update users set phone='999' where id=1"
|
|
223
|
+
```
|
pgnode-0.1.0/README.md
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
# **pgnode**
|
|
2
|
+
|
|
3
|
+
**Local LLM (Ollama) + PostgreSQL Agent**
|
|
4
|
+
|
|
5
|
+
An offline-first AI agent that converts natural language into **validated SQL queries** and executes them safely on your PostgreSQL database.
|
|
6
|
+
|
|
7
|
+
---
|
|
8
|
+
|
|
9
|
+
## Overview
|
|
10
|
+
|
|
11
|
+
**pgnode** is a local AI-powered database operator. It connects to your PostgreSQL instance and allows you to interact with your data using plain English while ensuring safety, control, and privacy.
|
|
12
|
+
|
|
13
|
+
* No external APIs
|
|
14
|
+
* No data leaves your system
|
|
15
|
+
* Fully local using Ollama
|
|
16
|
+
|
|
17
|
+
---
|
|
18
|
+
|
|
19
|
+
## Core Features
|
|
20
|
+
|
|
21
|
+
* Natural language → SQL conversion
|
|
22
|
+
* Safe query execution with validation layer
|
|
23
|
+
* Schema-aware query generation
|
|
24
|
+
* Works with existing PostgreSQL databases (pgAdmin compatible)
|
|
25
|
+
* CLI-first interface (fast and developer-friendly)
|
|
26
|
+
* Fully offline with local LLM
|
|
27
|
+
|
|
28
|
+
---
|
|
29
|
+
|
|
30
|
+
## Architecture
|
|
31
|
+
|
|
32
|
+
```
|
|
33
|
+
User Prompt
|
|
34
|
+
↓
|
|
35
|
+
Agent (planner)
|
|
36
|
+
↓
|
|
37
|
+
SQL Generator (LLM)
|
|
38
|
+
↓
|
|
39
|
+
Validator (safety layer)
|
|
40
|
+
↓
|
|
41
|
+
Query Executor (PostgreSQL)
|
|
42
|
+
↓
|
|
43
|
+
Response
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
---
|
|
47
|
+
|
|
48
|
+
## Tech Stack
|
|
49
|
+
|
|
50
|
+
### Core
|
|
51
|
+
|
|
52
|
+
* Python
|
|
53
|
+
* PostgreSQL
|
|
54
|
+
* Ollama (local LLM runtime)
|
|
55
|
+
|
|
56
|
+
### Libraries
|
|
57
|
+
|
|
58
|
+
* SQLAlchemy → DB interaction
|
|
59
|
+
* psycopg2 → PostgreSQL adapter
|
|
60
|
+
* Typer → CLI interface
|
|
61
|
+
* FastAPI (optional) → API layer
|
|
62
|
+
* LlamaIndex / FAISS (optional) → schema-aware retrieval
|
|
63
|
+
|
|
64
|
+
---
|
|
65
|
+
|
|
66
|
+
## Environment
|
|
67
|
+
|
|
68
|
+
Create `.env` in project root:
|
|
69
|
+
|
|
70
|
+
- `OLLAMA_HOST=http://127.0.0.1:11434`
|
|
71
|
+
- `DATABASE_URL=postgresql+psycopg2://user:pass@localhost:5432/dbname`
|
|
72
|
+
- `LLM_MODEL=deepseek-coder:6.7b` (optional)
|
|
73
|
+
|
|
74
|
+
## Run CLI
|
|
75
|
+
|
|
76
|
+
Install from source while developing:
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install -e .
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
First-time setup:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
pgnode connect
|
|
86
|
+
pgnode doctor
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
`connect` saves your database URL, Ollama host, and exact local model name to your user config. Environment variables still override saved config when present.
|
|
90
|
+
|
|
91
|
+
SSH tunnel databases are supported too. Choose `ssh` during `pgnode connect` or use flags:
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
pgnode connect \
|
|
95
|
+
--connection-type ssh \
|
|
96
|
+
--database-url "postgresql://user:pass@internal-db:5432/dbname" \
|
|
97
|
+
--ssh-host "bastion.example.com" \
|
|
98
|
+
--ssh-port 22 \
|
|
99
|
+
--ssh-user "ubuntu" \
|
|
100
|
+
--ssh-key-path "~/.ssh/id_rsa" \
|
|
101
|
+
--remote-host "127.0.0.1" \
|
|
102
|
+
--remote-port 5432 \
|
|
103
|
+
--local-port 0 \
|
|
104
|
+
--model "deepseek-coder-v2:16b"
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
Activate venv once:
|
|
108
|
+
|
|
109
|
+
```bash
|
|
110
|
+
source venv/bin/activate
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
Interactive conversation (context kept only in current session):
|
|
114
|
+
|
|
115
|
+
```bash
|
|
116
|
+
./pgnode run
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
or simply:
|
|
120
|
+
|
|
121
|
+
```bash
|
|
122
|
+
./pgnode
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
Useful chat commands:
|
|
126
|
+
|
|
127
|
+
- `/history` show recent turns
|
|
128
|
+
- `/clear` clear current session context
|
|
129
|
+
- `/exit` or `/quit` leave session
|
|
130
|
+
- Natural language meta-questions also work, e.g.:
|
|
131
|
+
- `what question did i ask you last`
|
|
132
|
+
- `which query did you execute last`
|
|
133
|
+
- `last result`
|
|
134
|
+
|
|
135
|
+
One-shot mode (no prior context):
|
|
136
|
+
|
|
137
|
+
```bash
|
|
138
|
+
./pgnode run "list all users with limit 5"
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
SQL-only generation (no execution):
|
|
142
|
+
|
|
143
|
+
```bash
|
|
144
|
+
./pgnode sql "top 5 customers by revenue last month"
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
Explain mode (SQL + short reasoning, no execution):
|
|
148
|
+
|
|
149
|
+
```bash
|
|
150
|
+
./pgnode explain "monthly revenue trend"
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
Schema helpers:
|
|
154
|
+
|
|
155
|
+
```bash
|
|
156
|
+
./pgnode tables
|
|
157
|
+
./pgnode describe Product
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
Environment and connectivity checks:
|
|
161
|
+
|
|
162
|
+
```bash
|
|
163
|
+
./pgnode config
|
|
164
|
+
./pgnode config-set --model "deepseek-coder-v2:16b"
|
|
165
|
+
./pgnode config-set --database-url "postgresql://user:pass@localhost:5432/dbname"
|
|
166
|
+
./pgnode config-set --connection-type ssh --ssh-host "bastion.example.com" --ssh-user "ubuntu" --ssh-key-path "~/.ssh/id_rsa" --remote-host "127.0.0.1" --remote-port 5432
|
|
167
|
+
./pgnode doctor
|
|
168
|
+
./pgnode models
|
|
169
|
+
```
|
|
170
|
+
|
|
171
|
+
Persistent local history:
|
|
172
|
+
|
|
173
|
+
```bash
|
|
174
|
+
./pgnode history
|
|
175
|
+
./pgnode rerun 12
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
Write behavior:
|
|
179
|
+
|
|
180
|
+
```bash
|
|
181
|
+
./pgnode run "update users set phone='999' where id=1"
|
|
182
|
+
```
|
|
183
|
+
|
|
184
|
+
`INSERT/UPDATE` now require confirmation by default. Use `--yes` to skip prompt.
|
|
185
|
+
|
|
186
|
+
```bash
|
|
187
|
+
./pgnode run --yes "update users set phone='999' where id=1"
|
|
188
|
+
```
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Application package."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Agent orchestration: NL to validated SQL execution."""
|
|
@@ -0,0 +1,216 @@
|
|
|
1
|
+
"""MVP agent: schema -> LLM SQL -> validate -> execute."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from dataclasses import dataclass
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
from app.agent.intent_agent import run_intent_agent
|
|
9
|
+
from app.agent.sql_metadata import extract_sql_metadata
|
|
10
|
+
from app.agent.sql_quality import detect_quality_issue
|
|
11
|
+
from app.agent.sql_quoter import quote_schema_identifiers
|
|
12
|
+
from app.agent.validator import validate_sql
|
|
13
|
+
from app.db.query_executor import run_query
|
|
14
|
+
from app.db.schema_loader import get_full_schema
|
|
15
|
+
from app.llm.prompt_builder import strip_sql_fences
|
|
16
|
+
from app.llm.sql_generator import generate_sql, repair_sql
|
|
17
|
+
|
|
18
|
+
_MAX_SQL_RETRIES = 1
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@dataclass
|
|
22
|
+
class AgentResult:
|
|
23
|
+
"""Outcome of a single agent turn."""
|
|
24
|
+
|
|
25
|
+
user_prompt: str
|
|
26
|
+
sql: str | None
|
|
27
|
+
validation_error: str | None
|
|
28
|
+
execution_result: list[dict[str, Any]] | str | None
|
|
29
|
+
llm_error: str | None = None
|
|
30
|
+
operation: str | None = None
|
|
31
|
+
table: str | None = None
|
|
32
|
+
changed_columns: list[str] | None = None
|
|
33
|
+
where_clause: str | None = None
|
|
34
|
+
out_of_scope: bool = False
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def run_agent(
|
|
38
|
+
user_prompt: str,
|
|
39
|
+
*,
|
|
40
|
+
dry_run: bool = False,
|
|
41
|
+
model: str | None = None,
|
|
42
|
+
conversation_context: str | None = None,
|
|
43
|
+
forced_sql: str | None = None,
|
|
44
|
+
) -> AgentResult:
|
|
45
|
+
"""
|
|
46
|
+
Load schema, generate SQL via local LLM, validate, optionally execute.
|
|
47
|
+
|
|
48
|
+
If dry_run is True, the query executor returns a dry-run message without executing.
|
|
49
|
+
"""
|
|
50
|
+
try:
|
|
51
|
+
schema = get_full_schema()
|
|
52
|
+
except RuntimeError as exc:
|
|
53
|
+
return AgentResult(
|
|
54
|
+
user_prompt=user_prompt,
|
|
55
|
+
sql=None,
|
|
56
|
+
validation_error=str(exc),
|
|
57
|
+
execution_result=None,
|
|
58
|
+
llm_error=None,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
try:
|
|
62
|
+
intent = run_intent_agent(
|
|
63
|
+
user_prompt,
|
|
64
|
+
schema,
|
|
65
|
+
model=model,
|
|
66
|
+
conversation_context=conversation_context,
|
|
67
|
+
)
|
|
68
|
+
except RuntimeError as exc:
|
|
69
|
+
return AgentResult(
|
|
70
|
+
user_prompt=user_prompt,
|
|
71
|
+
sql=None,
|
|
72
|
+
validation_error=None,
|
|
73
|
+
execution_result=None,
|
|
74
|
+
llm_error=str(exc),
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
if intent.intent == "out_of_scope":
|
|
78
|
+
return AgentResult(
|
|
79
|
+
user_prompt=user_prompt,
|
|
80
|
+
sql=None,
|
|
81
|
+
validation_error=None,
|
|
82
|
+
execution_result=intent.response,
|
|
83
|
+
out_of_scope=True,
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
if forced_sql is not None:
|
|
87
|
+
sql = quote_schema_identifiers(strip_sql_fences(forced_sql).strip(), schema)
|
|
88
|
+
else:
|
|
89
|
+
try:
|
|
90
|
+
raw_sql = generate_sql(
|
|
91
|
+
user_prompt,
|
|
92
|
+
schema,
|
|
93
|
+
model=model,
|
|
94
|
+
conversation_context=conversation_context,
|
|
95
|
+
)
|
|
96
|
+
except RuntimeError as exc:
|
|
97
|
+
return AgentResult(
|
|
98
|
+
user_prompt=user_prompt,
|
|
99
|
+
sql=None,
|
|
100
|
+
validation_error=None,
|
|
101
|
+
execution_result=None,
|
|
102
|
+
llm_error=str(exc),
|
|
103
|
+
)
|
|
104
|
+
|
|
105
|
+
sql = quote_schema_identifiers(strip_sql_fences(raw_sql).strip(), schema)
|
|
106
|
+
sql = _maybe_repair_sql_quality(
|
|
107
|
+
user_prompt=user_prompt,
|
|
108
|
+
sql=sql,
|
|
109
|
+
schema=schema,
|
|
110
|
+
model=model,
|
|
111
|
+
conversation_context=conversation_context,
|
|
112
|
+
)
|
|
113
|
+
metadata = extract_sql_metadata(sql)
|
|
114
|
+
|
|
115
|
+
validation_error = validate_sql(sql)
|
|
116
|
+
if validation_error is not None:
|
|
117
|
+
return AgentResult(
|
|
118
|
+
user_prompt=user_prompt,
|
|
119
|
+
sql=sql,
|
|
120
|
+
validation_error=validation_error,
|
|
121
|
+
execution_result=None,
|
|
122
|
+
operation=metadata.operation,
|
|
123
|
+
table=metadata.table,
|
|
124
|
+
changed_columns=metadata.changed_columns,
|
|
125
|
+
where_clause=metadata.where_clause,
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
execution_result = run_query(sql, dry_run=dry_run)
|
|
129
|
+
attempts = 0
|
|
130
|
+
while (
|
|
131
|
+
not dry_run
|
|
132
|
+
and isinstance(execution_result, str)
|
|
133
|
+
and execution_result.startswith("Query execution failed:")
|
|
134
|
+
and attempts < _MAX_SQL_RETRIES
|
|
135
|
+
):
|
|
136
|
+
attempts += 1
|
|
137
|
+
try:
|
|
138
|
+
repaired_sql = repair_sql(
|
|
139
|
+
user_request=user_prompt,
|
|
140
|
+
schema=schema,
|
|
141
|
+
previous_sql=sql,
|
|
142
|
+
db_error=execution_result,
|
|
143
|
+
model=model,
|
|
144
|
+
conversation_context=conversation_context,
|
|
145
|
+
).strip()
|
|
146
|
+
except RuntimeError as exc:
|
|
147
|
+
return AgentResult(
|
|
148
|
+
user_prompt=user_prompt,
|
|
149
|
+
sql=sql,
|
|
150
|
+
validation_error=None,
|
|
151
|
+
execution_result=execution_result,
|
|
152
|
+
llm_error=f"Repair attempt failed: {exc}",
|
|
153
|
+
operation=metadata.operation,
|
|
154
|
+
table=metadata.table,
|
|
155
|
+
changed_columns=metadata.changed_columns,
|
|
156
|
+
where_clause=metadata.where_clause,
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
repair_validation_error = validate_sql(repaired_sql)
|
|
160
|
+
if repair_validation_error is not None:
|
|
161
|
+
return AgentResult(
|
|
162
|
+
user_prompt=user_prompt,
|
|
163
|
+
sql=repaired_sql,
|
|
164
|
+
validation_error=repair_validation_error,
|
|
165
|
+
execution_result=execution_result,
|
|
166
|
+
operation=metadata.operation,
|
|
167
|
+
table=metadata.table,
|
|
168
|
+
changed_columns=metadata.changed_columns,
|
|
169
|
+
where_clause=metadata.where_clause,
|
|
170
|
+
)
|
|
171
|
+
|
|
172
|
+
sql = quote_schema_identifiers(repaired_sql, schema)
|
|
173
|
+
metadata = extract_sql_metadata(sql)
|
|
174
|
+
execution_result = run_query(sql, dry_run=False)
|
|
175
|
+
|
|
176
|
+
return AgentResult(
|
|
177
|
+
user_prompt=user_prompt,
|
|
178
|
+
sql=sql,
|
|
179
|
+
validation_error=None,
|
|
180
|
+
execution_result=execution_result,
|
|
181
|
+
operation=metadata.operation,
|
|
182
|
+
table=metadata.table,
|
|
183
|
+
changed_columns=metadata.changed_columns,
|
|
184
|
+
where_clause=metadata.where_clause,
|
|
185
|
+
)
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def _maybe_repair_sql_quality(
|
|
189
|
+
*,
|
|
190
|
+
user_prompt: str,
|
|
191
|
+
sql: str,
|
|
192
|
+
schema: dict[str, list[dict[str, str]]],
|
|
193
|
+
model: str | None,
|
|
194
|
+
conversation_context: str | None,
|
|
195
|
+
) -> str:
|
|
196
|
+
"""Run one quality-focused repair when generated SQL violates browse/lookup heuristics."""
|
|
197
|
+
quality_issue = detect_quality_issue(user_prompt, sql)
|
|
198
|
+
if quality_issue is None:
|
|
199
|
+
return sql
|
|
200
|
+
|
|
201
|
+
try:
|
|
202
|
+
repaired_sql = repair_sql(
|
|
203
|
+
user_request=user_prompt,
|
|
204
|
+
schema=schema,
|
|
205
|
+
previous_sql=sql,
|
|
206
|
+
db_error=f"Quality issue: {quality_issue}",
|
|
207
|
+
model=model,
|
|
208
|
+
conversation_context=conversation_context,
|
|
209
|
+
).strip()
|
|
210
|
+
except RuntimeError:
|
|
211
|
+
return sql
|
|
212
|
+
|
|
213
|
+
if validate_sql(repaired_sql) is not None:
|
|
214
|
+
return sql
|
|
215
|
+
|
|
216
|
+
return quote_schema_identifiers(repaired_sql, schema)
|