curie-backend 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.
- curie_backend-0.1.0/PKG-INFO +116 -0
- curie_backend-0.1.0/README.md +98 -0
- curie_backend-0.1.0/curie_backend.egg-info/PKG-INFO +116 -0
- curie_backend-0.1.0/curie_backend.egg-info/SOURCES.txt +24 -0
- curie_backend-0.1.0/curie_backend.egg-info/dependency_links.txt +1 -0
- curie_backend-0.1.0/curie_backend.egg-info/entry_points.txt +2 -0
- curie_backend-0.1.0/curie_backend.egg-info/requires.txt +10 -0
- curie_backend-0.1.0/curie_backend.egg-info/top_level.txt +2 -0
- curie_backend-0.1.0/extension/backend/__init__.py +1 -0
- curie_backend-0.1.0/extension/backend/__main__.py +20 -0
- curie_backend-0.1.0/extension/backend/config.py +26 -0
- curie_backend-0.1.0/extension/backend/main.py +25 -0
- curie_backend-0.1.0/extension/backend/routes/__init__.py +1 -0
- curie_backend-0.1.0/extension/backend/routes/generation.py +99 -0
- curie_backend-0.1.0/extension/backend/schemas.py +83 -0
- curie_backend-0.1.0/extension/backend/services/__init__.py +1 -0
- curie_backend-0.1.0/extension/backend/services/database_service.py +184 -0
- curie_backend-0.1.0/extension/backend/services/model_service.py +133 -0
- curie_backend-0.1.0/ml/__init__.py +1 -0
- curie_backend-0.1.0/ml/finetune.py +854 -0
- curie_backend-0.1.0/ml/inference.py +154 -0
- curie_backend-0.1.0/ml/load_datasets.py +143 -0
- curie_backend-0.1.0/ml/model_loader.py +211 -0
- curie_backend-0.1.0/ml/prompting.py +94 -0
- curie_backend-0.1.0/pyproject.toml +31 -0
- curie_backend-0.1.0/setup.cfg +4 -0
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: curie-backend
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Text2SQL Assistant Backend
|
|
5
|
+
Classifier: Programming Language :: Python :: 3
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Requires-Dist: fastapi[standard]>=0.110.0
|
|
9
|
+
Requires-Dist: uvicorn>=0.28.0
|
|
10
|
+
Requires-Dist: torch>=2.1.0
|
|
11
|
+
Requires-Dist: transformers>=4.38.0
|
|
12
|
+
Requires-Dist: peft>=0.9.0
|
|
13
|
+
Requires-Dist: accelerate>=0.27.0
|
|
14
|
+
Requires-Dist: SQLAlchemy>=2.0.0
|
|
15
|
+
Requires-Dist: psycopg[binary]>=3.1.0
|
|
16
|
+
Requires-Dist: huggingface-hub>=0.21.0
|
|
17
|
+
Requires-Dist: python-dotenv>=1.0.0
|
|
18
|
+
|
|
19
|
+
# Project Curie - Text2SQL
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
## Project description
|
|
24
|
+
* **Idea:** Use natural language prompts to generate SQL queries
|
|
25
|
+
* **Front-end: VS Code Plugin**
|
|
26
|
+
* Integration of trained LLM into useful plugin with quick insights
|
|
27
|
+
* Establish connection with existing database
|
|
28
|
+
* Use natural language to describe queries to this database, which are generated by fine-tuned LLM
|
|
29
|
+
* Ability to run generated SQL code directly via plugin to see results immediately with no bothersome copying
|
|
30
|
+
|
|
31
|
+
* ...
|
|
32
|
+
|
|
33
|
+
### Simple example:
|
|
34
|
+

|
|
35
|
+
|
|
36
|
+
## Project structure
|
|
37
|
+
|
|
38
|
+
```text
|
|
39
|
+
evaluation/ Benchmark setups for Text-to-SQL evaluation
|
|
40
|
+
data/ Local datasets and sample data
|
|
41
|
+
docs/ Project documentation, presentations, and research notes
|
|
42
|
+
experiments/ Ongoing fine-tuning experiments and run notes
|
|
43
|
+
extension/ VS Code extension and demo database
|
|
44
|
+
ml/ Training, fine-tuning, and inference code
|
|
45
|
+
ml/models/ Model metadata, configs, and model cards
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
## Project team and responsibilities
|
|
49
|
+
|
|
50
|
+
| Name | E-Mail | Identifier | Responsibilities / tasks |
|
|
51
|
+
|-------------------|--------------------------|------------|--------------------------|
|
|
52
|
+
| Tamer Lewis | tamer.lewis@tum.de | go68maj | TBD |
|
|
53
|
+
| Antonia Gilles | antonia.gilles@tum.de | go34cih | TBD |
|
|
54
|
+
| Leon Pfänder | leon.pfaender@tum.de | ge27fob | TBD |
|
|
55
|
+
| Christoph Klein | christoph.m.klein@tum.de | ge42low | TBD |
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
## Project setup
|
|
61
|
+
|
|
62
|
+
The project uses Python 3.12 and a pinned ML stack from `requirements.txt`.
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
python3.12 -m venv .venv
|
|
66
|
+
source .venv/bin/activate
|
|
67
|
+
python -m pip install --upgrade pip
|
|
68
|
+
python -m pip install -r requirements.txt
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
Verify that the installed dependencies are consistent:
|
|
72
|
+
|
|
73
|
+
```bash
|
|
74
|
+
python -m pip check
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### Git LFS for model weights
|
|
78
|
+
|
|
79
|
+
Git LFS stores large model binaries on the GitLab LFS server while Git tracks
|
|
80
|
+
small pointer files. This repository uses LFS for `*.safetensors` files.
|
|
81
|
+
|
|
82
|
+
Minimal setup:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
brew install git-lfs
|
|
86
|
+
git lfs install
|
|
87
|
+
git lfs pull
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
Add and push model weights through the normal Git workflow:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
git add .gitattributes ml/models/<model>/adapter_model.safetensors
|
|
94
|
+
git commit -m "Add model adapter"
|
|
95
|
+
git push
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
After cloning, Git LFS normally downloads tracked files automatically. If a model file is only an LFS pointer or is missing, run `git lfs pull`. Verify the available model files with `git lfs ls-files`.
|
|
99
|
+
|
|
100
|
+
LFS does not change how the benchmark is invoked:
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
python3 evaluation/spider/run.py \
|
|
104
|
+
--split dev \
|
|
105
|
+
--model-dir ml/models/<model>
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
### Spider
|
|
109
|
+
|
|
110
|
+
Download the Spider dataset and place it in the `data/` folder:
|
|
111
|
+
|
|
112
|
+
https://drive.google.com/file/d/1403EGqzIDoHMdQF4c9Bkyl7dZLZ5Wt6J/view
|
|
113
|
+
|
|
114
|
+
## Sources
|
|
115
|
+
|
|
116
|
+
TBD
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
# Project Curie - Text2SQL
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
## Project description
|
|
6
|
+
* **Idea:** Use natural language prompts to generate SQL queries
|
|
7
|
+
* **Front-end: VS Code Plugin**
|
|
8
|
+
* Integration of trained LLM into useful plugin with quick insights
|
|
9
|
+
* Establish connection with existing database
|
|
10
|
+
* Use natural language to describe queries to this database, which are generated by fine-tuned LLM
|
|
11
|
+
* Ability to run generated SQL code directly via plugin to see results immediately with no bothersome copying
|
|
12
|
+
|
|
13
|
+
* ...
|
|
14
|
+
|
|
15
|
+
### Simple example:
|
|
16
|
+

|
|
17
|
+
|
|
18
|
+
## Project structure
|
|
19
|
+
|
|
20
|
+
```text
|
|
21
|
+
evaluation/ Benchmark setups for Text-to-SQL evaluation
|
|
22
|
+
data/ Local datasets and sample data
|
|
23
|
+
docs/ Project documentation, presentations, and research notes
|
|
24
|
+
experiments/ Ongoing fine-tuning experiments and run notes
|
|
25
|
+
extension/ VS Code extension and demo database
|
|
26
|
+
ml/ Training, fine-tuning, and inference code
|
|
27
|
+
ml/models/ Model metadata, configs, and model cards
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Project team and responsibilities
|
|
31
|
+
|
|
32
|
+
| Name | E-Mail | Identifier | Responsibilities / tasks |
|
|
33
|
+
|-------------------|--------------------------|------------|--------------------------|
|
|
34
|
+
| Tamer Lewis | tamer.lewis@tum.de | go68maj | TBD |
|
|
35
|
+
| Antonia Gilles | antonia.gilles@tum.de | go34cih | TBD |
|
|
36
|
+
| Leon Pfänder | leon.pfaender@tum.de | ge27fob | TBD |
|
|
37
|
+
| Christoph Klein | christoph.m.klein@tum.de | ge42low | TBD |
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
## Project setup
|
|
43
|
+
|
|
44
|
+
The project uses Python 3.12 and a pinned ML stack from `requirements.txt`.
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
python3.12 -m venv .venv
|
|
48
|
+
source .venv/bin/activate
|
|
49
|
+
python -m pip install --upgrade pip
|
|
50
|
+
python -m pip install -r requirements.txt
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
Verify that the installed dependencies are consistent:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
python -m pip check
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
### Git LFS for model weights
|
|
60
|
+
|
|
61
|
+
Git LFS stores large model binaries on the GitLab LFS server while Git tracks
|
|
62
|
+
small pointer files. This repository uses LFS for `*.safetensors` files.
|
|
63
|
+
|
|
64
|
+
Minimal setup:
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
brew install git-lfs
|
|
68
|
+
git lfs install
|
|
69
|
+
git lfs pull
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Add and push model weights through the normal Git workflow:
|
|
73
|
+
|
|
74
|
+
```bash
|
|
75
|
+
git add .gitattributes ml/models/<model>/adapter_model.safetensors
|
|
76
|
+
git commit -m "Add model adapter"
|
|
77
|
+
git push
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
After cloning, Git LFS normally downloads tracked files automatically. If a model file is only an LFS pointer or is missing, run `git lfs pull`. Verify the available model files with `git lfs ls-files`.
|
|
81
|
+
|
|
82
|
+
LFS does not change how the benchmark is invoked:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
python3 evaluation/spider/run.py \
|
|
86
|
+
--split dev \
|
|
87
|
+
--model-dir ml/models/<model>
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
### Spider
|
|
91
|
+
|
|
92
|
+
Download the Spider dataset and place it in the `data/` folder:
|
|
93
|
+
|
|
94
|
+
https://drive.google.com/file/d/1403EGqzIDoHMdQF4c9Bkyl7dZLZ5Wt6J/view
|
|
95
|
+
|
|
96
|
+
## Sources
|
|
97
|
+
|
|
98
|
+
TBD
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: curie-backend
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Text2SQL Assistant Backend
|
|
5
|
+
Classifier: Programming Language :: Python :: 3
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Requires-Dist: fastapi[standard]>=0.110.0
|
|
9
|
+
Requires-Dist: uvicorn>=0.28.0
|
|
10
|
+
Requires-Dist: torch>=2.1.0
|
|
11
|
+
Requires-Dist: transformers>=4.38.0
|
|
12
|
+
Requires-Dist: peft>=0.9.0
|
|
13
|
+
Requires-Dist: accelerate>=0.27.0
|
|
14
|
+
Requires-Dist: SQLAlchemy>=2.0.0
|
|
15
|
+
Requires-Dist: psycopg[binary]>=3.1.0
|
|
16
|
+
Requires-Dist: huggingface-hub>=0.21.0
|
|
17
|
+
Requires-Dist: python-dotenv>=1.0.0
|
|
18
|
+
|
|
19
|
+
# Project Curie - Text2SQL
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
## Project description
|
|
24
|
+
* **Idea:** Use natural language prompts to generate SQL queries
|
|
25
|
+
* **Front-end: VS Code Plugin**
|
|
26
|
+
* Integration of trained LLM into useful plugin with quick insights
|
|
27
|
+
* Establish connection with existing database
|
|
28
|
+
* Use natural language to describe queries to this database, which are generated by fine-tuned LLM
|
|
29
|
+
* Ability to run generated SQL code directly via plugin to see results immediately with no bothersome copying
|
|
30
|
+
|
|
31
|
+
* ...
|
|
32
|
+
|
|
33
|
+
### Simple example:
|
|
34
|
+

|
|
35
|
+
|
|
36
|
+
## Project structure
|
|
37
|
+
|
|
38
|
+
```text
|
|
39
|
+
evaluation/ Benchmark setups for Text-to-SQL evaluation
|
|
40
|
+
data/ Local datasets and sample data
|
|
41
|
+
docs/ Project documentation, presentations, and research notes
|
|
42
|
+
experiments/ Ongoing fine-tuning experiments and run notes
|
|
43
|
+
extension/ VS Code extension and demo database
|
|
44
|
+
ml/ Training, fine-tuning, and inference code
|
|
45
|
+
ml/models/ Model metadata, configs, and model cards
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
## Project team and responsibilities
|
|
49
|
+
|
|
50
|
+
| Name | E-Mail | Identifier | Responsibilities / tasks |
|
|
51
|
+
|-------------------|--------------------------|------------|--------------------------|
|
|
52
|
+
| Tamer Lewis | tamer.lewis@tum.de | go68maj | TBD |
|
|
53
|
+
| Antonia Gilles | antonia.gilles@tum.de | go34cih | TBD |
|
|
54
|
+
| Leon Pfänder | leon.pfaender@tum.de | ge27fob | TBD |
|
|
55
|
+
| Christoph Klein | christoph.m.klein@tum.de | ge42low | TBD |
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
## Project setup
|
|
61
|
+
|
|
62
|
+
The project uses Python 3.12 and a pinned ML stack from `requirements.txt`.
|
|
63
|
+
|
|
64
|
+
```bash
|
|
65
|
+
python3.12 -m venv .venv
|
|
66
|
+
source .venv/bin/activate
|
|
67
|
+
python -m pip install --upgrade pip
|
|
68
|
+
python -m pip install -r requirements.txt
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
Verify that the installed dependencies are consistent:
|
|
72
|
+
|
|
73
|
+
```bash
|
|
74
|
+
python -m pip check
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
### Git LFS for model weights
|
|
78
|
+
|
|
79
|
+
Git LFS stores large model binaries on the GitLab LFS server while Git tracks
|
|
80
|
+
small pointer files. This repository uses LFS for `*.safetensors` files.
|
|
81
|
+
|
|
82
|
+
Minimal setup:
|
|
83
|
+
|
|
84
|
+
```bash
|
|
85
|
+
brew install git-lfs
|
|
86
|
+
git lfs install
|
|
87
|
+
git lfs pull
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
Add and push model weights through the normal Git workflow:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
git add .gitattributes ml/models/<model>/adapter_model.safetensors
|
|
94
|
+
git commit -m "Add model adapter"
|
|
95
|
+
git push
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
After cloning, Git LFS normally downloads tracked files automatically. If a model file is only an LFS pointer or is missing, run `git lfs pull`. Verify the available model files with `git lfs ls-files`.
|
|
99
|
+
|
|
100
|
+
LFS does not change how the benchmark is invoked:
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
python3 evaluation/spider/run.py \
|
|
104
|
+
--split dev \
|
|
105
|
+
--model-dir ml/models/<model>
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
### Spider
|
|
109
|
+
|
|
110
|
+
Download the Spider dataset and place it in the `data/` folder:
|
|
111
|
+
|
|
112
|
+
https://drive.google.com/file/d/1403EGqzIDoHMdQF4c9Bkyl7dZLZ5Wt6J/view
|
|
113
|
+
|
|
114
|
+
## Sources
|
|
115
|
+
|
|
116
|
+
TBD
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
curie_backend.egg-info/PKG-INFO
|
|
4
|
+
curie_backend.egg-info/SOURCES.txt
|
|
5
|
+
curie_backend.egg-info/dependency_links.txt
|
|
6
|
+
curie_backend.egg-info/entry_points.txt
|
|
7
|
+
curie_backend.egg-info/requires.txt
|
|
8
|
+
curie_backend.egg-info/top_level.txt
|
|
9
|
+
extension/backend/__init__.py
|
|
10
|
+
extension/backend/__main__.py
|
|
11
|
+
extension/backend/config.py
|
|
12
|
+
extension/backend/main.py
|
|
13
|
+
extension/backend/schemas.py
|
|
14
|
+
extension/backend/routes/__init__.py
|
|
15
|
+
extension/backend/routes/generation.py
|
|
16
|
+
extension/backend/services/__init__.py
|
|
17
|
+
extension/backend/services/database_service.py
|
|
18
|
+
extension/backend/services/model_service.py
|
|
19
|
+
ml/__init__.py
|
|
20
|
+
ml/finetune.py
|
|
21
|
+
ml/inference.py
|
|
22
|
+
ml/load_datasets.py
|
|
23
|
+
ml/model_loader.py
|
|
24
|
+
ml/prompting.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Backend package."""
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import uvicorn
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def main() -> None:
|
|
6
|
+
parser = argparse.ArgumentParser(description="Curie Backend Server")
|
|
7
|
+
parser.add_argument("--host", default="127.0.0.1", help="Host address to bind to")
|
|
8
|
+
parser.add_argument("--port", type=int, default=8000, help="Port to listen on")
|
|
9
|
+
args = parser.parse_args()
|
|
10
|
+
|
|
11
|
+
uvicorn.run(
|
|
12
|
+
"extension.backend.main:app",
|
|
13
|
+
host=args.host,
|
|
14
|
+
port=args.port,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
if __name__ == "__main__":
|
|
19
|
+
main()
|
|
20
|
+
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
|
|
4
|
+
from ml.inference import GenerationConfig
|
|
5
|
+
from ml.model_loader import Device
|
|
6
|
+
|
|
7
|
+
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
8
|
+
CACHE_DIR = Path.home() / ".cache" / "curie"
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass(frozen=True)
|
|
12
|
+
class Settings:
|
|
13
|
+
project_root: Path = PROJECT_ROOT
|
|
14
|
+
models_dir: Path = CACHE_DIR / "models"
|
|
15
|
+
model_dir: Path = CACHE_DIR / "models" / "llama-3.2-1b-spider-lora-3000ex"
|
|
16
|
+
prompt_path: Path = PROJECT_ROOT / "ml" / "prompts" / "test_prompt.txt"
|
|
17
|
+
device: Device = "auto"
|
|
18
|
+
generation: GenerationConfig = GenerationConfig(
|
|
19
|
+
max_new_tokens=128,
|
|
20
|
+
temperature=0.0,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
settings = Settings()
|
|
26
|
+
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
from collections.abc import AsyncIterator
|
|
2
|
+
from contextlib import asynccontextmanager
|
|
3
|
+
|
|
4
|
+
from fastapi import FastAPI
|
|
5
|
+
|
|
6
|
+
from extension.backend.routes.generation import router
|
|
7
|
+
from extension.backend.services.model_service import ModelService
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@asynccontextmanager
|
|
11
|
+
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
12
|
+
service = ModelService()
|
|
13
|
+
service.load()
|
|
14
|
+
|
|
15
|
+
app.state.model_service = service
|
|
16
|
+
yield
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
app = FastAPI(title="Curie Backend", lifespan=lifespan)
|
|
20
|
+
app.include_router(router, prefix="/api")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@app.get("/health")
|
|
24
|
+
def health() -> dict[str, str]:
|
|
25
|
+
return {"status": "ok"}
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Backend route package."""
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
from fastapi import APIRouter, HTTPException, Request
|
|
2
|
+
from extension.backend.schemas import (
|
|
3
|
+
ExecuteRequest,
|
|
4
|
+
ExecuteResponse,
|
|
5
|
+
GenerationRequest,
|
|
6
|
+
GenerationResponse,
|
|
7
|
+
SchemaRequest,
|
|
8
|
+
SchemaResponse,
|
|
9
|
+
AvailableModelResponse,
|
|
10
|
+
ModelsResponse,
|
|
11
|
+
SelectModelRequest,
|
|
12
|
+
)
|
|
13
|
+
from extension.backend.services.database_service import (
|
|
14
|
+
DatabaseConnectionError,
|
|
15
|
+
DatabaseService,
|
|
16
|
+
InvalidSqlError,
|
|
17
|
+
QueryExecutionError,
|
|
18
|
+
)
|
|
19
|
+
from extension.backend.services.model_service import ModelService
|
|
20
|
+
|
|
21
|
+
router = APIRouter()
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@router.get("/models", response_model=ModelsResponse)
|
|
25
|
+
def list_models(request: Request) -> ModelsResponse:
|
|
26
|
+
model_service: ModelService = request.app.state.model_service
|
|
27
|
+
|
|
28
|
+
return ModelsResponse(
|
|
29
|
+
models=[
|
|
30
|
+
AvailableModelResponse(
|
|
31
|
+
id=model.id,
|
|
32
|
+
name=model.name,
|
|
33
|
+
is_active=model.is_active,
|
|
34
|
+
)
|
|
35
|
+
for model in model_service.list_models()
|
|
36
|
+
]
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
@router.post("/models/select", response_model=AvailableModelResponse)
|
|
41
|
+
def select_model(body: SelectModelRequest, request: Request) -> AvailableModelResponse:
|
|
42
|
+
model_service: ModelService = request.app.state.model_service
|
|
43
|
+
|
|
44
|
+
try:
|
|
45
|
+
model = model_service.select_model(body.model_id)
|
|
46
|
+
except ValueError as error:
|
|
47
|
+
raise HTTPException(status_code=404, detail=str(error)) from error
|
|
48
|
+
|
|
49
|
+
return AvailableModelResponse(
|
|
50
|
+
id=model.id,
|
|
51
|
+
name=model.name,
|
|
52
|
+
is_active=model.is_active,
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@router.post("/generate", response_model=GenerationResponse)
|
|
57
|
+
def generate_sql(body: GenerationRequest, request: Request) -> GenerationResponse:
|
|
58
|
+
model_service: ModelService = request.app.state.model_service
|
|
59
|
+
|
|
60
|
+
if body.database is not None:
|
|
61
|
+
database_service = DatabaseService()
|
|
62
|
+
try:
|
|
63
|
+
schema = database_service.build_schema_text(body.database)
|
|
64
|
+
except DatabaseConnectionError as error:
|
|
65
|
+
raise HTTPException(status_code=400, detail=str(error)) from error
|
|
66
|
+
|
|
67
|
+
sql = model_service.generate_sql(body.question, schema)
|
|
68
|
+
return GenerationResponse(sql=sql)
|
|
69
|
+
|
|
70
|
+
if body.database_schema is None:
|
|
71
|
+
raise HTTPException(status_code=422, detail="Either schema or database must be provided")
|
|
72
|
+
|
|
73
|
+
sql = model_service.generate_sql(body.question, body.database_schema)
|
|
74
|
+
|
|
75
|
+
return GenerationResponse(sql=sql)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
@router.post("/execute", response_model=ExecuteResponse)
|
|
79
|
+
def execute_sql(body: ExecuteRequest) -> ExecuteResponse:
|
|
80
|
+
database_service = DatabaseService()
|
|
81
|
+
try:
|
|
82
|
+
result = database_service.execute(body.sql, body.database, body.read_only)
|
|
83
|
+
except InvalidSqlError as error:
|
|
84
|
+
raise HTTPException(status_code=400, detail=str(error)) from error
|
|
85
|
+
except DatabaseConnectionError as error:
|
|
86
|
+
raise HTTPException(status_code=400, detail=str(error)) from error
|
|
87
|
+
except QueryExecutionError as error:
|
|
88
|
+
raise HTTPException(status_code=400, detail=str(error)) from error
|
|
89
|
+
|
|
90
|
+
return ExecuteResponse(**result)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
@router.post("/schema", response_model=SchemaResponse)
|
|
94
|
+
def inspect_schema(body: SchemaRequest) -> SchemaResponse:
|
|
95
|
+
database_service = DatabaseService()
|
|
96
|
+
try:
|
|
97
|
+
return database_service.inspect_schema(body.database)
|
|
98
|
+
except DatabaseConnectionError as error:
|
|
99
|
+
raise HTTPException(status_code=400, detail=str(error)) from error
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
from typing import Any, Literal
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class PostgreSQLDatabaseConfig(BaseModel):
|
|
7
|
+
type: Literal["postgresql"]
|
|
8
|
+
host: str = Field(min_length=1)
|
|
9
|
+
port: int = Field(default=5432, ge=1, le=65535)
|
|
10
|
+
database: str = Field(min_length=1)
|
|
11
|
+
username: str = Field(min_length=1)
|
|
12
|
+
password: str
|
|
13
|
+
sslmode: Literal["disable", "prefer", "require"] = "prefer"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class GenerationRequest(BaseModel):
|
|
17
|
+
model_config = ConfigDict(populate_by_name=True)
|
|
18
|
+
|
|
19
|
+
question: str = Field(min_length=1)
|
|
20
|
+
database_schema: str | None = Field(default=None, alias="schema", min_length=1)
|
|
21
|
+
database: PostgreSQLDatabaseConfig | None = None
|
|
22
|
+
|
|
23
|
+
@model_validator(mode="after")
|
|
24
|
+
def require_schema_or_database(self) -> "GenerationRequest":
|
|
25
|
+
if self.database_schema is None and self.database is None:
|
|
26
|
+
raise ValueError("Either schema or database must be provided")
|
|
27
|
+
return self
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class GenerationResponse(BaseModel):
|
|
31
|
+
sql: str
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class DatabaseColumn(BaseModel):
|
|
35
|
+
name: str
|
|
36
|
+
type: str
|
|
37
|
+
primary_key: bool
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class DatabaseForeignKey(BaseModel):
|
|
41
|
+
columns: list[str]
|
|
42
|
+
referred_table: str
|
|
43
|
+
referred_columns: list[str]
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class DatabaseTable(BaseModel):
|
|
47
|
+
name: str
|
|
48
|
+
columns: list[DatabaseColumn]
|
|
49
|
+
foreign_keys: list[DatabaseForeignKey]
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
class SchemaRequest(BaseModel):
|
|
53
|
+
database: PostgreSQLDatabaseConfig
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class SchemaResponse(BaseModel):
|
|
57
|
+
tables: list[DatabaseTable]
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class ExecuteRequest(BaseModel):
|
|
61
|
+
sql: str = Field(min_length=1)
|
|
62
|
+
database: PostgreSQLDatabaseConfig
|
|
63
|
+
read_only: bool = True
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class ExecuteResponse(BaseModel):
|
|
67
|
+
columns: list[str]
|
|
68
|
+
rows: list[dict[str, Any]]
|
|
69
|
+
row_count: int
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class AvailableModelResponse(BaseModel):
|
|
73
|
+
id: str
|
|
74
|
+
name: str
|
|
75
|
+
is_active: bool
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class ModelsResponse(BaseModel):
|
|
79
|
+
models: list[AvailableModelResponse]
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class SelectModelRequest(BaseModel):
|
|
83
|
+
model_id: str = Field(min_length=1)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Backend service layer."""
|