fyrefly 0.3.0__tar.gz → 0.3.2__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.
- fyrefly-0.3.2/.github/workflows/publish.yml +29 -0
- fyrefly-0.3.2/.gitignore +7 -0
- fyrefly-0.3.2/PKG-INFO +155 -0
- fyrefly-0.3.2/README.md +126 -0
- {fyrefly-0.3.0 → fyrefly-0.3.2}/pyproject.toml +11 -4
- fyrefly-0.3.2/src/fyrefly/__init__.py +14 -0
- fyrefly-0.3.2/src/fyrefly/ai.py +327 -0
- {fyrefly-0.3.0 → fyrefly-0.3.2}/src/fyrefly/core.py +17 -0
- {fyrefly-0.3.0 → fyrefly-0.3.2}/src/fyrefly/dataops.py +51 -1
- fyrefly-0.3.2/src/fyrefly/viz.py +132 -0
- fyrefly-0.3.2/tests/test_ai.py +78 -0
- {fyrefly-0.3.0 → fyrefly-0.3.2}/tests/test_core.py +8 -1
- {fyrefly-0.3.0 → fyrefly-0.3.2}/tests/test_dataops.py +15 -2
- fyrefly-0.3.2/tests/test_viz.py +66 -0
- fyrefly-0.3.0/PKG-INFO +0 -96
- fyrefly-0.3.0/README.md +0 -69
- fyrefly-0.3.0/setup.cfg +0 -4
- fyrefly-0.3.0/src/fyrefly/__init__.py +0 -7
- fyrefly-0.3.0/src/fyrefly.egg-info/PKG-INFO +0 -96
- fyrefly-0.3.0/src/fyrefly.egg-info/SOURCES.txt +0 -13
- fyrefly-0.3.0/src/fyrefly.egg-info/dependency_links.txt +0 -1
- fyrefly-0.3.0/src/fyrefly.egg-info/requires.txt +0 -6
- fyrefly-0.3.0/src/fyrefly.egg-info/top_level.txt +0 -1
- {fyrefly-0.3.0 → fyrefly-0.3.2}/LICENSE +0 -0
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
build-and-publish:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
permissions:
|
|
11
|
+
id-token: write # required for PyPI Trusted Publishing (no API token needed)
|
|
12
|
+
|
|
13
|
+
steps:
|
|
14
|
+
- name: Check out the repository
|
|
15
|
+
uses: actions/checkout@v4
|
|
16
|
+
|
|
17
|
+
- name: Set up Python
|
|
18
|
+
uses: actions/setup-python@v5
|
|
19
|
+
with:
|
|
20
|
+
python-version: "3.x"
|
|
21
|
+
|
|
22
|
+
- name: Install build tool
|
|
23
|
+
run: pip install build
|
|
24
|
+
|
|
25
|
+
- name: Build sdist and wheel
|
|
26
|
+
run: python -m build
|
|
27
|
+
|
|
28
|
+
- name: Publish to PyPI
|
|
29
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
fyrefly-0.3.2/.gitignore
ADDED
fyrefly-0.3.2/PKG-INFO
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: fyrefly
|
|
3
|
+
Version: 0.3.2
|
|
4
|
+
Summary: A quant-based library to compute mathematical operations, run SQL queries, get AI-powered insights, and visualize your data.
|
|
5
|
+
Project-URL: Homepage, https://github.com/Ninad077/fyrefly
|
|
6
|
+
Project-URL: Repository, https://github.com/Ninad077/fyrefly
|
|
7
|
+
Author-email: Ninad Mandavkar <ninadmandavkar28@gmail.com>
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: ai,arithmetic,math,pandas,quant,sql,utility,visualization
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Requires-Python: >=3.9
|
|
20
|
+
Requires-Dist: duckdb
|
|
21
|
+
Requires-Dist: matplotlib
|
|
22
|
+
Requires-Dist: openpyxl
|
|
23
|
+
Requires-Dist: pandas
|
|
24
|
+
Requires-Dist: requests
|
|
25
|
+
Requires-Dist: seaborn
|
|
26
|
+
Provides-Extra: dev
|
|
27
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
|
|
30
|
+
# Fyrefly
|
|
31
|
+
|
|
32
|
+
A quant-based library to compute mathematical operations, run SQL queries, get AI-powered insights, and visualize your data.
|
|
33
|
+
|
|
34
|
+
## Installation
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install fyrefly
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
## Usage — Math functions
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
from fyrefly import add, mul, div, mod, diff
|
|
44
|
+
|
|
45
|
+
add(5, 4, 6, 10) # 25
|
|
46
|
+
mul(5, 4, 6, 10) # 1200
|
|
47
|
+
div(100, 5, 2) # 10.0
|
|
48
|
+
mod(17, 5) # 2
|
|
49
|
+
diff(10, 3, 2) # 5
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Each function accepts any number of arguments:
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
add(1) # 1
|
|
56
|
+
add(1, 2, 3, 4, 5) # 15
|
|
57
|
+
mul() # 1 (identity)
|
|
58
|
+
diff(10, 3) # 7
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Usage — Data functions
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
from fyrefly import load, loadh, loadc, loads, sql, xtract, clean
|
|
65
|
+
|
|
66
|
+
c = load("data.csv") # loads a CSV/Excel/Google Sheet, previews it
|
|
67
|
+
loadh(c) # preview column headers
|
|
68
|
+
loadc(c) # preview record count
|
|
69
|
+
loads(c) # preview schema (columns + datatypes)
|
|
70
|
+
|
|
71
|
+
c = clean(c) # dedupe, strip whitespace, standardize column names
|
|
72
|
+
|
|
73
|
+
d = sql("select * from c where age > 30") # run ANY SQL query on loaded data
|
|
74
|
+
xtract(d) # export the result to CSV
|
|
75
|
+
xtract(sql("select * from c"), filename="all.csv") # chained form also works
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
**Note:** for Google Sheets, the sheet must be shared as "Anyone with the link – Viewer" for `load()` to access it.
|
|
79
|
+
|
|
80
|
+
## Usage — AI functions (require an API key)
|
|
81
|
+
|
|
82
|
+
```python
|
|
83
|
+
from fyrefly import ask, insights
|
|
84
|
+
|
|
85
|
+
c = load("sales.csv")
|
|
86
|
+
|
|
87
|
+
# Ask questions in plain English — converted to SQL and run automatically
|
|
88
|
+
ask(c, "what were total sales by region last quarter?", api_key="AIza...")
|
|
89
|
+
|
|
90
|
+
# Query across multiple datasets (e.g. a join) — scales to any count
|
|
91
|
+
ask((c, d), "who are our top customers by total sales?", api_key="AIza...")
|
|
92
|
+
ask((c, d, e), "...", api_key="AIza...")
|
|
93
|
+
|
|
94
|
+
# Get a plain-English summary of a dataset
|
|
95
|
+
insights(c, api_key="AIza...")
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
The provider (Anthropic, OpenAI, Gemini, or any OpenAI-compatible API like Groq/Mistral) is **auto-detected from your API key's format** — you usually don't need to specify it. API keys can also be set via environment variables instead of passing them directly: `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `GEMINI_API_KEY`.
|
|
99
|
+
|
|
100
|
+
⚠️ **Never commit an API key into a script you push to GitHub.** Use an environment variable instead:
|
|
101
|
+
```bash
|
|
102
|
+
export GEMINI_API_KEY="AIza..."
|
|
103
|
+
```
|
|
104
|
+
```python
|
|
105
|
+
ask(c, "what were total sales by region?") # picked up automatically from the env var
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
## Usage — Visualization (no AI, no API key needed)
|
|
109
|
+
|
|
110
|
+
```python
|
|
111
|
+
from fyrefly import viz
|
|
112
|
+
|
|
113
|
+
viz(c) # auto -> correlation heatmap
|
|
114
|
+
viz(c, x="department") # auto -> bar chart of counts
|
|
115
|
+
viz(c, x="age") # auto -> histogram
|
|
116
|
+
viz(c, kind="scatter", x="age", y="salary")
|
|
117
|
+
viz(c, kind="box", x="department", y="salary")
|
|
118
|
+
viz(c, kind="pie", x="department")
|
|
119
|
+
viz(c, kind="line", x="date", y="revenue")
|
|
120
|
+
viz(c, kind="pairplot", hue="department")
|
|
121
|
+
viz(c, kind="bar", x="department", y="salary", save="chart.png") # also saves a file
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
Supported chart types: `line`, `bar`, `barh`, `scatter`, `hist`, `box`, `violin`, `kde`, `pie`, `heatmap`, `pairplot`, `area`, plus `"auto"` which picks a sensible chart based on your data.
|
|
125
|
+
|
|
126
|
+
## Functions
|
|
127
|
+
|
|
128
|
+
| Function | Description | Needs API key? |
|
|
129
|
+
|---|---|---|
|
|
130
|
+
| `add(*nos)` | Adds all given numbers | No |
|
|
131
|
+
| `mul(*nos)` | Multiplies all given numbers | No |
|
|
132
|
+
| `div(*nos)` | Divides left to right: first ÷ second ÷ third ... | No |
|
|
133
|
+
| `mod(*nos)` | Finds remainder left to right: first % second % third ... | No |
|
|
134
|
+
| `diff(*nos)` | Subtracts left to right: first − second − third ... | No |
|
|
135
|
+
| `load(path)` | Loads a CSV/Excel/Google Sheet into a DataFrame and previews it | No |
|
|
136
|
+
| `loadh(df)` | Previews just the column headers | No |
|
|
137
|
+
| `loadc(df)` | Previews the record count | No |
|
|
138
|
+
| `loads(df)` | Previews the schema (columns + datatypes) | No |
|
|
139
|
+
| `sql(query)` | Runs any SQL query against a loaded DataFrame, by variable name | No |
|
|
140
|
+
| `xtract(df, filename=None)` | Exports a DataFrame to CSV in the current directory | No |
|
|
141
|
+
| `clean(df)` | Dedupes, strips whitespace, standardizes column names | No |
|
|
142
|
+
| `ask(data, question, api_key=...)` | Converts a plain-English question into SQL and runs it | **Yes** |
|
|
143
|
+
| `insights(df, api_key=...)` | Generates a plain-English summary of a dataset | **Yes** |
|
|
144
|
+
| `viz(df, kind="auto", ...)` | Visualizes a DataFrame — 12+ chart types | No |
|
|
145
|
+
|
|
146
|
+
## Development
|
|
147
|
+
|
|
148
|
+
```bash
|
|
149
|
+
pip install -e ".[dev]"
|
|
150
|
+
pytest
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
## License
|
|
154
|
+
|
|
155
|
+
MIT
|
fyrefly-0.3.2/README.md
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
# Fyrefly
|
|
2
|
+
|
|
3
|
+
A quant-based library to compute mathematical operations, run SQL queries, get AI-powered insights, and visualize your data.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install fyrefly
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Usage — Math functions
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from fyrefly import add, mul, div, mod, diff
|
|
15
|
+
|
|
16
|
+
add(5, 4, 6, 10) # 25
|
|
17
|
+
mul(5, 4, 6, 10) # 1200
|
|
18
|
+
div(100, 5, 2) # 10.0
|
|
19
|
+
mod(17, 5) # 2
|
|
20
|
+
diff(10, 3, 2) # 5
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
Each function accepts any number of arguments:
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
add(1) # 1
|
|
27
|
+
add(1, 2, 3, 4, 5) # 15
|
|
28
|
+
mul() # 1 (identity)
|
|
29
|
+
diff(10, 3) # 7
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
## Usage — Data functions
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from fyrefly import load, loadh, loadc, loads, sql, xtract, clean
|
|
36
|
+
|
|
37
|
+
c = load("data.csv") # loads a CSV/Excel/Google Sheet, previews it
|
|
38
|
+
loadh(c) # preview column headers
|
|
39
|
+
loadc(c) # preview record count
|
|
40
|
+
loads(c) # preview schema (columns + datatypes)
|
|
41
|
+
|
|
42
|
+
c = clean(c) # dedupe, strip whitespace, standardize column names
|
|
43
|
+
|
|
44
|
+
d = sql("select * from c where age > 30") # run ANY SQL query on loaded data
|
|
45
|
+
xtract(d) # export the result to CSV
|
|
46
|
+
xtract(sql("select * from c"), filename="all.csv") # chained form also works
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
**Note:** for Google Sheets, the sheet must be shared as "Anyone with the link – Viewer" for `load()` to access it.
|
|
50
|
+
|
|
51
|
+
## Usage — AI functions (require an API key)
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from fyrefly import ask, insights
|
|
55
|
+
|
|
56
|
+
c = load("sales.csv")
|
|
57
|
+
|
|
58
|
+
# Ask questions in plain English — converted to SQL and run automatically
|
|
59
|
+
ask(c, "what were total sales by region last quarter?", api_key="AIza...")
|
|
60
|
+
|
|
61
|
+
# Query across multiple datasets (e.g. a join) — scales to any count
|
|
62
|
+
ask((c, d), "who are our top customers by total sales?", api_key="AIza...")
|
|
63
|
+
ask((c, d, e), "...", api_key="AIza...")
|
|
64
|
+
|
|
65
|
+
# Get a plain-English summary of a dataset
|
|
66
|
+
insights(c, api_key="AIza...")
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
The provider (Anthropic, OpenAI, Gemini, or any OpenAI-compatible API like Groq/Mistral) is **auto-detected from your API key's format** — you usually don't need to specify it. API keys can also be set via environment variables instead of passing them directly: `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `GEMINI_API_KEY`.
|
|
70
|
+
|
|
71
|
+
⚠️ **Never commit an API key into a script you push to GitHub.** Use an environment variable instead:
|
|
72
|
+
```bash
|
|
73
|
+
export GEMINI_API_KEY="AIza..."
|
|
74
|
+
```
|
|
75
|
+
```python
|
|
76
|
+
ask(c, "what were total sales by region?") # picked up automatically from the env var
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
## Usage — Visualization (no AI, no API key needed)
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
from fyrefly import viz
|
|
83
|
+
|
|
84
|
+
viz(c) # auto -> correlation heatmap
|
|
85
|
+
viz(c, x="department") # auto -> bar chart of counts
|
|
86
|
+
viz(c, x="age") # auto -> histogram
|
|
87
|
+
viz(c, kind="scatter", x="age", y="salary")
|
|
88
|
+
viz(c, kind="box", x="department", y="salary")
|
|
89
|
+
viz(c, kind="pie", x="department")
|
|
90
|
+
viz(c, kind="line", x="date", y="revenue")
|
|
91
|
+
viz(c, kind="pairplot", hue="department")
|
|
92
|
+
viz(c, kind="bar", x="department", y="salary", save="chart.png") # also saves a file
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
Supported chart types: `line`, `bar`, `barh`, `scatter`, `hist`, `box`, `violin`, `kde`, `pie`, `heatmap`, `pairplot`, `area`, plus `"auto"` which picks a sensible chart based on your data.
|
|
96
|
+
|
|
97
|
+
## Functions
|
|
98
|
+
|
|
99
|
+
| Function | Description | Needs API key? |
|
|
100
|
+
|---|---|---|
|
|
101
|
+
| `add(*nos)` | Adds all given numbers | No |
|
|
102
|
+
| `mul(*nos)` | Multiplies all given numbers | No |
|
|
103
|
+
| `div(*nos)` | Divides left to right: first ÷ second ÷ third ... | No |
|
|
104
|
+
| `mod(*nos)` | Finds remainder left to right: first % second % third ... | No |
|
|
105
|
+
| `diff(*nos)` | Subtracts left to right: first − second − third ... | No |
|
|
106
|
+
| `load(path)` | Loads a CSV/Excel/Google Sheet into a DataFrame and previews it | No |
|
|
107
|
+
| `loadh(df)` | Previews just the column headers | No |
|
|
108
|
+
| `loadc(df)` | Previews the record count | No |
|
|
109
|
+
| `loads(df)` | Previews the schema (columns + datatypes) | No |
|
|
110
|
+
| `sql(query)` | Runs any SQL query against a loaded DataFrame, by variable name | No |
|
|
111
|
+
| `xtract(df, filename=None)` | Exports a DataFrame to CSV in the current directory | No |
|
|
112
|
+
| `clean(df)` | Dedupes, strips whitespace, standardizes column names | No |
|
|
113
|
+
| `ask(data, question, api_key=...)` | Converts a plain-English question into SQL and runs it | **Yes** |
|
|
114
|
+
| `insights(df, api_key=...)` | Generates a plain-English summary of a dataset | **Yes** |
|
|
115
|
+
| `viz(df, kind="auto", ...)` | Visualizes a DataFrame — 12+ chart types | No |
|
|
116
|
+
|
|
117
|
+
## Development
|
|
118
|
+
|
|
119
|
+
```bash
|
|
120
|
+
pip install -e ".[dev]"
|
|
121
|
+
pytest
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
## License
|
|
125
|
+
|
|
126
|
+
MIT
|
|
@@ -1,18 +1,25 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
1
5
|
[project]
|
|
2
6
|
name = "fyrefly"
|
|
3
|
-
version = "0.3.
|
|
4
|
-
description = "A quant-based library to compute mathematical operations
|
|
7
|
+
version = "0.3.2"
|
|
8
|
+
description = "A quant-based library to compute mathematical operations, run SQL queries, get AI-powered insights, and visualize your data."
|
|
5
9
|
readme = "README.md"
|
|
6
10
|
requires-python = ">=3.9"
|
|
7
11
|
license = { text = "MIT" }
|
|
8
12
|
authors = [
|
|
9
|
-
{ name = "
|
|
13
|
+
{ name = "Ninad Mandavkar", email = "ninadmandavkar28@gmail.com" }
|
|
10
14
|
]
|
|
11
|
-
keywords = ["math", "arithmetic", "quant", "utility", "sql", "pandas"]
|
|
15
|
+
keywords = ["math", "arithmetic", "quant", "utility", "sql", "pandas", "ai", "visualization"]
|
|
12
16
|
dependencies = [
|
|
13
17
|
"pandas",
|
|
14
18
|
"duckdb",
|
|
15
19
|
"openpyxl",
|
|
20
|
+
"requests",
|
|
21
|
+
"matplotlib",
|
|
22
|
+
"seaborn",
|
|
16
23
|
]
|
|
17
24
|
classifiers = [
|
|
18
25
|
"Development Status :: 3 - Alpha",
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
"""A quant-based library to compute mathematical operations, run SQL queries, get AI-powered insights, and visualize your data."""
|
|
2
|
+
|
|
3
|
+
from .core import add, mul, div, mod, diff
|
|
4
|
+
from .dataops import load, loadh, loadc, loads, sql, xtract, clean
|
|
5
|
+
from .ai import ask, insights
|
|
6
|
+
from .viz import viz
|
|
7
|
+
|
|
8
|
+
__all__ = [
|
|
9
|
+
"add", "mul", "div", "mod", "diff",
|
|
10
|
+
"load", "loadh", "loadc", "loads", "sql", "xtract", "clean",
|
|
11
|
+
"ask", "insights",
|
|
12
|
+
"viz",
|
|
13
|
+
]
|
|
14
|
+
__version__ = "0.3.2"
|
|
@@ -0,0 +1,327 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import inspect
|
|
3
|
+
import pandas as pd
|
|
4
|
+
import duckdb
|
|
5
|
+
import requests
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def _describe_schema(dataframes):
|
|
9
|
+
"""Builds a plain-text schema description of all loaded DataFrames for the LLM prompt."""
|
|
10
|
+
lines = []
|
|
11
|
+
for name, df in dataframes.items():
|
|
12
|
+
cols = ", ".join(f"{col} ({dtype})" for col, dtype in df.dtypes.items())
|
|
13
|
+
lines.append(f"Table '{name}': {cols}")
|
|
14
|
+
return "\n".join(lines)
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _clean_sql(text):
|
|
18
|
+
"""Strips markdown code fences and extra whitespace/labels from a model's raw response."""
|
|
19
|
+
text = text.strip()
|
|
20
|
+
if text.startswith("```"):
|
|
21
|
+
text = text.strip("`").strip()
|
|
22
|
+
lines = text.split("\n")
|
|
23
|
+
if lines and lines[0].strip().lower() in ("sql", "sqlite", "duckdb"):
|
|
24
|
+
lines = lines[1:]
|
|
25
|
+
text = "\n".join(lines).strip()
|
|
26
|
+
# Some models prefix with "SQL:" or similar
|
|
27
|
+
for prefix in ("SQL:", "sql:", "Query:", "query:"):
|
|
28
|
+
if text.startswith(prefix):
|
|
29
|
+
text = text[len(prefix):].strip()
|
|
30
|
+
return text.rstrip(";").strip()
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _build_prompt(question, schema_description):
|
|
34
|
+
return (
|
|
35
|
+
"You are a SQL generator. Given the table schema(s) below, write a single "
|
|
36
|
+
"valid SQL query (DuckDB dialect) that answers the user's question. "
|
|
37
|
+
"Return ONLY the raw SQL query — no explanation, no markdown formatting, no backticks.\n\n"
|
|
38
|
+
f"Schema:\n{schema_description}\n\n"
|
|
39
|
+
f"Question: {question}\n\n"
|
|
40
|
+
"SQL:"
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _detect_provider(api_key):
|
|
45
|
+
"""
|
|
46
|
+
Guesses the provider from the API key's format. Order matters —
|
|
47
|
+
check the most specific prefixes before more generic ones.
|
|
48
|
+
"""
|
|
49
|
+
if api_key.startswith("sk-ant-"):
|
|
50
|
+
return "anthropic", None
|
|
51
|
+
if api_key.startswith("AIza"):
|
|
52
|
+
return "gemini", None
|
|
53
|
+
if api_key.startswith("gsk_"):
|
|
54
|
+
return "openai_compatible", "https://api.groq.com/openai/v1"
|
|
55
|
+
if api_key.startswith("sk-"):
|
|
56
|
+
return "openai", None
|
|
57
|
+
raise ValueError(
|
|
58
|
+
"Couldn't auto-detect the provider from this API key's format. "
|
|
59
|
+
"Please specify it explicitly, e.g. ask(question, provider=\"openai\", api_key=\"...\")"
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def _call_anthropic(prompt, api_key, model):
|
|
64
|
+
model = model or "claude-haiku-4-5-20251001"
|
|
65
|
+
response = requests.post(
|
|
66
|
+
"https://api.anthropic.com/v1/messages",
|
|
67
|
+
headers={
|
|
68
|
+
"x-api-key": api_key,
|
|
69
|
+
"anthropic-version": "2023-06-01",
|
|
70
|
+
"content-type": "application/json",
|
|
71
|
+
},
|
|
72
|
+
json={
|
|
73
|
+
"model": model,
|
|
74
|
+
"max_tokens": 300,
|
|
75
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
76
|
+
},
|
|
77
|
+
timeout=30,
|
|
78
|
+
)
|
|
79
|
+
response.raise_for_status()
|
|
80
|
+
data = response.json()
|
|
81
|
+
return data["content"][0]["text"]
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _call_openai(prompt, api_key, model, base_url=None):
|
|
85
|
+
model = model or "gpt-4o-mini"
|
|
86
|
+
url = f"{(base_url or 'https://api.openai.com/v1').rstrip('/')}/chat/completions"
|
|
87
|
+
response = requests.post(
|
|
88
|
+
url,
|
|
89
|
+
headers={
|
|
90
|
+
"Authorization": f"Bearer {api_key}",
|
|
91
|
+
"Content-Type": "application/json",
|
|
92
|
+
},
|
|
93
|
+
json={
|
|
94
|
+
"model": model,
|
|
95
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
96
|
+
"max_tokens": 300,
|
|
97
|
+
},
|
|
98
|
+
timeout=30,
|
|
99
|
+
)
|
|
100
|
+
response.raise_for_status()
|
|
101
|
+
data = response.json()
|
|
102
|
+
return data["choices"][0]["message"]["content"]
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def _call_gemini(prompt, api_key, model):
|
|
106
|
+
model = model or "gemini-1.5-flash"
|
|
107
|
+
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
|
|
108
|
+
response = requests.post(
|
|
109
|
+
url,
|
|
110
|
+
json={"contents": [{"parts": [{"text": prompt}]}]},
|
|
111
|
+
timeout=30,
|
|
112
|
+
)
|
|
113
|
+
response.raise_for_status()
|
|
114
|
+
data = response.json()
|
|
115
|
+
return data["candidates"][0]["content"]["parts"][0]["text"]
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def _call_llm(prompt, provider, api_key, model, base_url):
|
|
119
|
+
"""Dispatches a prompt to the given provider and returns the raw text response."""
|
|
120
|
+
provider = provider.lower()
|
|
121
|
+
|
|
122
|
+
if provider == "anthropic":
|
|
123
|
+
return _call_anthropic(prompt, api_key, model)
|
|
124
|
+
elif provider == "openai":
|
|
125
|
+
return _call_openai(prompt, api_key, model)
|
|
126
|
+
elif provider in ("gemini", "google"):
|
|
127
|
+
return _call_gemini(prompt, api_key, model)
|
|
128
|
+
elif provider == "openai_compatible":
|
|
129
|
+
if not base_url:
|
|
130
|
+
raise ValueError("provider='openai_compatible' requires base_url=\"...\"")
|
|
131
|
+
if not model:
|
|
132
|
+
raise ValueError("provider='openai_compatible' requires model=\"...\"")
|
|
133
|
+
return _call_openai(prompt, api_key, model, base_url=base_url)
|
|
134
|
+
else:
|
|
135
|
+
raise ValueError(
|
|
136
|
+
f"Unsupported provider '{provider}'. "
|
|
137
|
+
"Supported: anthropic, openai, gemini, openai_compatible"
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def _resolve_provider_and_key(api_key=None, provider=None, base_url=None):
|
|
142
|
+
"""
|
|
143
|
+
Resolves the API key (from argument or environment variable) and,
|
|
144
|
+
if provider wasn't specified, auto-detects it from the key's format.
|
|
145
|
+
Shared by ask() and insights().
|
|
146
|
+
"""
|
|
147
|
+
if not api_key:
|
|
148
|
+
for env_var in ("ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY"):
|
|
149
|
+
if os.environ.get(env_var):
|
|
150
|
+
api_key = os.environ[env_var]
|
|
151
|
+
break
|
|
152
|
+
if not api_key:
|
|
153
|
+
raise ValueError(
|
|
154
|
+
"No API key found. Pass api_key=\"...\" or set one of: "
|
|
155
|
+
"ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY"
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
if not provider:
|
|
159
|
+
detected_provider, detected_base_url = _detect_provider(api_key)
|
|
160
|
+
provider = detected_provider
|
|
161
|
+
base_url = base_url or detected_base_url
|
|
162
|
+
print(f"Detected provider: {provider}")
|
|
163
|
+
|
|
164
|
+
return provider, api_key, base_url
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def _generate_sql(question, schema_description, provider, api_key, model, base_url):
|
|
168
|
+
prompt = _build_prompt(question, schema_description)
|
|
169
|
+
raw = _call_llm(prompt, provider, api_key, model, base_url)
|
|
170
|
+
return _clean_sql(raw)
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
def _describe_dataframe_stats(df):
|
|
174
|
+
"""Builds a plain-text statistical summary of a DataFrame for the LLM prompt."""
|
|
175
|
+
lines = [f"Shape: {df.shape[0]} rows, {df.shape[1]} columns"]
|
|
176
|
+
lines.append("Columns and dtypes: " + ", ".join(f"{c} ({t})" for c, t in df.dtypes.items()))
|
|
177
|
+
|
|
178
|
+
missing = df.isnull().sum()
|
|
179
|
+
missing = missing[missing > 0]
|
|
180
|
+
if len(missing):
|
|
181
|
+
lines.append("Missing values: " + ", ".join(f"{c}: {v}" for c, v in missing.items()))
|
|
182
|
+
else:
|
|
183
|
+
lines.append("Missing values: none")
|
|
184
|
+
|
|
185
|
+
numeric_df = df.select_dtypes(include="number")
|
|
186
|
+
if not numeric_df.empty:
|
|
187
|
+
lines.append("Numeric summary:\n" + numeric_df.describe().to_string())
|
|
188
|
+
|
|
189
|
+
return "\n".join(lines)
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def _resolve_dataframes(data, frame):
|
|
193
|
+
"""
|
|
194
|
+
Resolves one or more DataFrames into a {name: DataFrame} mapping for
|
|
195
|
+
registering with DuckDB, recovering each DataFrame's original variable
|
|
196
|
+
name from the caller's scope where possible.
|
|
197
|
+
|
|
198
|
+
Accepts:
|
|
199
|
+
- a single DataFrame: ask(c, "...")
|
|
200
|
+
- a tuple/list of DataFrames, any length "up to n": ask((c, d, e), "...")
|
|
201
|
+
- a dict of {name: DataFrame} for fully explicit naming
|
|
202
|
+
"""
|
|
203
|
+
if isinstance(data, dict):
|
|
204
|
+
for name, val in data.items():
|
|
205
|
+
if not isinstance(val, pd.DataFrame):
|
|
206
|
+
raise ValueError(f"data['{name}'] is not a pandas DataFrame.")
|
|
207
|
+
return data
|
|
208
|
+
|
|
209
|
+
if isinstance(data, pd.DataFrame):
|
|
210
|
+
items = [data]
|
|
211
|
+
elif isinstance(data, (tuple, list)):
|
|
212
|
+
if not data:
|
|
213
|
+
raise ValueError("No DataFrames provided.")
|
|
214
|
+
items = list(data)
|
|
215
|
+
for i, val in enumerate(items):
|
|
216
|
+
if not isinstance(val, pd.DataFrame):
|
|
217
|
+
raise ValueError(f"Item {i} in the tuple/list is not a pandas DataFrame.")
|
|
218
|
+
else:
|
|
219
|
+
raise ValueError(
|
|
220
|
+
"First argument must be a pandas DataFrame, a tuple/list of "
|
|
221
|
+
"DataFrames (e.g. (c, d)), or a dict of {name: DataFrame}."
|
|
222
|
+
)
|
|
223
|
+
|
|
224
|
+
# Recover each DataFrame's original variable name from the caller's scope
|
|
225
|
+
caller_vars = {**frame.f_globals, **frame.f_locals}
|
|
226
|
+
id_to_name = {}
|
|
227
|
+
for name, val in caller_vars.items():
|
|
228
|
+
if isinstance(val, pd.DataFrame) and id(val) not in id_to_name:
|
|
229
|
+
id_to_name[id(val)] = name
|
|
230
|
+
|
|
231
|
+
dataframes = {}
|
|
232
|
+
used_names = set()
|
|
233
|
+
for i, val in enumerate(items):
|
|
234
|
+
name = id_to_name.get(id(val))
|
|
235
|
+
if not name or name in used_names:
|
|
236
|
+
base = name if (name and name not in used_names) else f"t{i + 1}"
|
|
237
|
+
name = base
|
|
238
|
+
suffix = 1
|
|
239
|
+
while name in used_names:
|
|
240
|
+
suffix += 1
|
|
241
|
+
name = f"{base}_{suffix}"
|
|
242
|
+
used_names.add(name)
|
|
243
|
+
dataframes[name] = val
|
|
244
|
+
|
|
245
|
+
return dataframes
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def insights(df, api_key=None, provider=None, model=None, base_url=None):
|
|
249
|
+
"""
|
|
250
|
+
Generates a short, plain-English summary of a DataFrame using an LLM —
|
|
251
|
+
notable trends, ranges, and data-quality issues like missing values.
|
|
252
|
+
|
|
253
|
+
Provider auto-detected from your API key's format, same as ask():
|
|
254
|
+
c = load("data.csv")
|
|
255
|
+
insights(c, api_key="AIza...")
|
|
256
|
+
|
|
257
|
+
API keys can also be set via environment variables: ANTHROPIC_API_KEY,
|
|
258
|
+
OPENAI_API_KEY, GEMINI_API_KEY.
|
|
259
|
+
"""
|
|
260
|
+
provider, api_key, base_url = _resolve_provider_and_key(api_key, provider, base_url)
|
|
261
|
+
|
|
262
|
+
stats_description = _describe_dataframe_stats(df)
|
|
263
|
+
prompt = (
|
|
264
|
+
"You are a data analyst. Given the dataset summary below, write a short, "
|
|
265
|
+
"plain-English overview (3-5 sentences) highlighting notable trends, ranges, "
|
|
266
|
+
"and any data-quality issues like missing values. Interpret the numbers — "
|
|
267
|
+
"don't just repeat them verbatim.\n\n"
|
|
268
|
+
f"{stats_description}\n\nSummary:"
|
|
269
|
+
)
|
|
270
|
+
|
|
271
|
+
summary = _call_llm(prompt, provider, api_key, model, base_url)
|
|
272
|
+
print(summary)
|
|
273
|
+
return summary
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def ask(data, question, api_key=None, provider=None, model=None, base_url=None):
|
|
277
|
+
"""
|
|
278
|
+
Converts a natural-language question into a SQL query using an LLM,
|
|
279
|
+
runs it against your data, and previews the result — same as sql(),
|
|
280
|
+
but you ask in plain English.
|
|
281
|
+
|
|
282
|
+
Pass one DataFrame:
|
|
283
|
+
c = load("sales.csv")
|
|
284
|
+
ask(c, "what were total sales by region last quarter?", api_key="...")
|
|
285
|
+
|
|
286
|
+
Pass two or more (for a JOIN, comparison, etc.) — scales to any count:
|
|
287
|
+
ask((c, d), "who are our top customers by total sales?", api_key="...")
|
|
288
|
+
ask((c, d, e), "...", api_key="...")
|
|
289
|
+
|
|
290
|
+
Each DataFrame's original variable name (c, d, e, ...) is automatically
|
|
291
|
+
recovered from your script and used as its table name in the generated
|
|
292
|
+
SQL. If a name can't be recovered (e.g. an inline, unassigned
|
|
293
|
+
DataFrame), it falls back to t1, t2, etc.
|
|
294
|
+
|
|
295
|
+
The provider is auto-detected from your API key's format, so you
|
|
296
|
+
usually don't need to specify it:
|
|
297
|
+
ask(c, "...", api_key="sk-ant-...") # detected as Anthropic
|
|
298
|
+
ask(c, "...", api_key="AIza...") # detected as Gemini
|
|
299
|
+
ask(c, "...", api_key="gsk_...") # detected as Groq (openai_compatible)
|
|
300
|
+
ask(c, "...", api_key="sk-...") # detected as OpenAI
|
|
301
|
+
|
|
302
|
+
You can still override or specify explicitly, which is required for
|
|
303
|
+
openai_compatible providers other than Groq:
|
|
304
|
+
ask(c, "...", provider="openai_compatible", api_key="...",
|
|
305
|
+
base_url="https://api.mistral.ai/v1", model="mistral-small-latest")
|
|
306
|
+
|
|
307
|
+
API keys can also be set via environment variables instead of passing
|
|
308
|
+
api_key= directly: ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY.
|
|
309
|
+
"""
|
|
310
|
+
frame = inspect.currentframe().f_back
|
|
311
|
+
dataframes = _resolve_dataframes(data, frame)
|
|
312
|
+
|
|
313
|
+
provider, api_key, base_url = _resolve_provider_and_key(api_key, provider, base_url)
|
|
314
|
+
|
|
315
|
+
schema_description = _describe_schema(dataframes)
|
|
316
|
+
generated_sql = _generate_sql(question, schema_description, provider, api_key, model, base_url)
|
|
317
|
+
|
|
318
|
+
print(f"Generated SQL: {generated_sql}\n")
|
|
319
|
+
|
|
320
|
+
con = duckdb.connect()
|
|
321
|
+
for name, val in dataframes.items():
|
|
322
|
+
con.register(name, val)
|
|
323
|
+
|
|
324
|
+
result = con.execute(generated_sql).df()
|
|
325
|
+
print(f"Query returned {len(result)} rows, {len(result.columns)} columns\n")
|
|
326
|
+
print(result.head())
|
|
327
|
+
return result
|