data-atlas 0.0.1__py3-none-any.whl
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.
- data_atlas-0.0.1.dist-info/METADATA +139 -0
- data_atlas-0.0.1.dist-info/RECORD +15 -0
- data_atlas-0.0.1.dist-info/WHEEL +5 -0
- data_atlas-0.0.1.dist-info/entry_points.txt +2 -0
- data_atlas-0.0.1.dist-info/licenses/LICENSE +21 -0
- data_atlas-0.0.1.dist-info/top_level.txt +1 -0
- data_dictionary_generator/__init__.py +0 -0
- data_dictionary_generator/column_stats.py +56 -0
- data_dictionary_generator/data_type.py +66 -0
- data_dictionary_generator/description_generator.py +331 -0
- data_dictionary_generator/generate_dictionary.py +157 -0
- data_dictionary_generator/load_dataset.py +34 -0
- data_dictionary_generator/metadata_extraction.py +206 -0
- data_dictionary_generator/pattern_extraction.py +120 -0
- data_dictionary_generator/relationships.py +284 -0
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: data-atlas
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Generate clinical data dictionaries using language models
|
|
5
|
+
Author-email: Raffaele Giancotti <raffaele.giancotti@unical.it>
|
|
6
|
+
License: MIT
|
|
7
|
+
Requires-Python: >=3.8
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: click>=8.0.0
|
|
11
|
+
Requires-Dist: pandas>=2.0.0
|
|
12
|
+
Requires-Dist: numpy>=1.20.0
|
|
13
|
+
Requires-Dist: requests>=2.25.0
|
|
14
|
+
Requires-Dist: ollama>=0.1.0
|
|
15
|
+
Requires-Dist: torch>=2.0.0
|
|
16
|
+
Requires-Dist: transformers>=4.20.0
|
|
17
|
+
Requires-Dist: sentence-transformers>=2.0.0
|
|
18
|
+
Requires-Dist: scikit-learn>=1.0.0
|
|
19
|
+
Requires-Dist: reportlab>=3.6.0
|
|
20
|
+
Requires-Dist: tqdm>=4.60.0
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: pytest>=7.0.0; extra == "dev"
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# Data Dictionary Generator
|
|
26
|
+
|
|
27
|
+
A Python package to automatically generate data dictionaries for clinical datasets using a large language model (LLM) via Ollama. This tool takes in dataset files (CSV format), processes them, and generates descriptions for each column in the dataset, as well as other metadata like data types, sample data, table descriptions and some quality information such as missing values, outliers and redundant values. It is also able to find relationships between tables and columns.
|
|
28
|
+
|
|
29
|
+
## Features
|
|
30
|
+
### Metadata Generation
|
|
31
|
+
- **Column Descriptions**: AI-generated clinical context for each field
|
|
32
|
+
- **Table Summaries**: Dataset-level documentation
|
|
33
|
+
- **Smart Sampling**: Automatic data type detection with sample values
|
|
34
|
+
|
|
35
|
+
### Quality Analysis
|
|
36
|
+
- Missing value statistics
|
|
37
|
+
- Outlier detection (numeric fields)
|
|
38
|
+
- Duplicate row/column identification
|
|
39
|
+
|
|
40
|
+
### Advanced Capabilities
|
|
41
|
+
- **Semantic Relationship Detection**: Finds connected columns across tables
|
|
42
|
+
- **Multi-Format Outputs**: JSON, Markdown
|
|
43
|
+
- **Custom LLM Integration**: Supports any Ollama, OpenAI and Google Gemini models
|
|
44
|
+
|
|
45
|
+
## Requirements
|
|
46
|
+
|
|
47
|
+
Make sure you have Python 3.8+ installed. All dependencies are automatically installed when you install the package.
|
|
48
|
+
|
|
49
|
+
## Installation
|
|
50
|
+
|
|
51
|
+
Clone the repository:
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
git clone https://github.com/rafgia/data-dictionary-generator.git
|
|
55
|
+
cd code
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
Install the package:
|
|
59
|
+
|
|
60
|
+
```bash
|
|
61
|
+
pip install -e .
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## Usage
|
|
65
|
+
|
|
66
|
+
### Command-line Interface
|
|
67
|
+
|
|
68
|
+
Once the package is installed, you can use the command line to generate metadata for your dataset(s).
|
|
69
|
+
|
|
70
|
+
To run the tool, use the following command:
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
python generate_dictionary.py <folder_path> --output-dir <output_dir> --model llama3.1
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
#### Parameters:
|
|
77
|
+
- `<folder_path>`: The path to the folder containing your CSV files.
|
|
78
|
+
- `<output-dir>`: The name of the path where the metadata will be saved.
|
|
79
|
+
- `--model`: Specify the Ollama or OpenAI model to use for generating metadata.
|
|
80
|
+
|
|
81
|
+
### Example
|
|
82
|
+
|
|
83
|
+
1. **Prepare your data files**:
|
|
84
|
+
Place all your CSV files (representing tables in your dataset) in a folder, e.g., `data/MIMIC`.
|
|
85
|
+
|
|
86
|
+
2. **Run the generator**:
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
python -m generate_dictionary.py data/MIMIC --output-dir output --model gpt-4o-mini
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
This will generate metadata for each column in the dataset and save it to a JSON file and a Markdown summary for the dictionary, and a .png plot and a markdown summary for the detected relationships.
|
|
93
|
+
|
|
94
|
+
### Sample Output
|
|
95
|
+
|
|
96
|
+
For each **table** in your dataset, the following metadata will be generated:
|
|
97
|
+
- **Table Name**: The name of the table (CSV filename).
|
|
98
|
+
- **Number of Rows**: The total number of rows in the table.
|
|
99
|
+
- **Number of Columns**: The total number of columns in the table.
|
|
100
|
+
- **Table Description**: A generated description of what the table contains.
|
|
101
|
+
|
|
102
|
+
For each **column** in your dataset:
|
|
103
|
+
- **Column Name**: The name of the column.
|
|
104
|
+
- **Sample Data**: A sample of 5 data points from the column.
|
|
105
|
+
- **Data Type**: The inferred data type (e.g., integer, float, string).
|
|
106
|
+
- **Column Description**: A description generated by the model for the column.
|
|
107
|
+
|
|
108
|
+
### Example of generated metadata:
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
## Troubleshooting
|
|
112
|
+
|
|
113
|
+
### 1. If you encounter an error such as `model not found`, make sure you have set up Ollama correctly and the model is available.
|
|
114
|
+
- Ensure that you can manually run the model using `ollama run llama3.1` from the command line before using it in the Python script.
|
|
115
|
+
|
|
116
|
+
### 2. If the dependencies are not installing, make sure you're using the correct Python version and have all required libraries listed in `pyproject.toml`.
|
|
117
|
+
|
|
118
|
+
### 3. If the dataset is very large, consider breaking it down into smaller CSV files for more efficient processing.
|
|
119
|
+
|
|
120
|
+
## Contributing
|
|
121
|
+
|
|
122
|
+
If you would like to contribute to this project, feel free to fork the repository and submit a pull request. Make sure to add tests and document any new features.
|
|
123
|
+
|
|
124
|
+
### To contribute:
|
|
125
|
+
1. Fork the repository.
|
|
126
|
+
2. Create a new branch (`git checkout -b feature-name`).
|
|
127
|
+
3. Make your changes and commit them (`git commit -am 'Add new feature'`).
|
|
128
|
+
4. Push to the branch (`git push origin feature-name`).
|
|
129
|
+
5. Open a pull request.
|
|
130
|
+
|
|
131
|
+
## License
|
|
132
|
+
|
|
133
|
+
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
|
|
134
|
+
|
|
135
|
+
---
|
|
136
|
+
|
|
137
|
+
## Author
|
|
138
|
+
|
|
139
|
+
Raffaele Giancotti
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
data_atlas-0.0.1.dist-info/licenses/LICENSE,sha256=Fs1PK6wJ87X7i9kIGoRcHOm_HrAFazHyD-yCYZrtvnk,1096
|
|
2
|
+
data_dictionary_generator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
3
|
+
data_dictionary_generator/column_stats.py,sha256=o3BA1R6o9-LJQKvDbnlvL8HmcdmwaDr0QAGm2a2klZA,1736
|
|
4
|
+
data_dictionary_generator/data_type.py,sha256=i1_WaIkSw_oGnNYSJ8DQ7TGansJS9P9exTbJ1Bzg3-k,2057
|
|
5
|
+
data_dictionary_generator/description_generator.py,sha256=W6ZEoNAf9lludseReSZyFFDJcCIIa5pQBd5M7eoSb58,13487
|
|
6
|
+
data_dictionary_generator/generate_dictionary.py,sha256=sXkpSkEaeMv28O9T2WfUwmfFduBhuH1oHt3TBRa1U0Y,5506
|
|
7
|
+
data_dictionary_generator/load_dataset.py,sha256=33zbTfXve83FfWw1sQgu_D2I-FrFz5ZUm8HFDaSPRWI,1165
|
|
8
|
+
data_dictionary_generator/metadata_extraction.py,sha256=q1M22qAxWERAmtFDvBi3FSP2f5GFlOAW07olUkyPy3s,7650
|
|
9
|
+
data_dictionary_generator/pattern_extraction.py,sha256=yCCSzbWbYH4w4nAgMJfGhxG9xg-UV1C6M5duyucCpkc,4196
|
|
10
|
+
data_dictionary_generator/relationships.py,sha256=-lu4m-yM1_wRdGaYNmDAIJbGGDBa28yscCbj7pmTEM4,11436
|
|
11
|
+
data_atlas-0.0.1.dist-info/METADATA,sha256=6e61nz0lzDevqqstDa_fvIo8Yb-lwaLwDPJL6pNNr0w,5092
|
|
12
|
+
data_atlas-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
|
|
13
|
+
data_atlas-0.0.1.dist-info/entry_points.txt,sha256=0bveK9ODg6B1uqanTxwVSCppNgdCETz8JCter0r0cLg,91
|
|
14
|
+
data_atlas-0.0.1.dist-info/top_level.txt,sha256=toz52AXzeut2aQUPFNKxYmD8aZf-2-BBY4ZHLGLrgAc,26
|
|
15
|
+
data_atlas-0.0.1.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Raffaele Giancotti
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
data_dictionary_generator
|
|
File without changes
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
import pandas as pd
|
|
2
|
+
import numpy as np
|
|
3
|
+
from collections import Counter
|
|
4
|
+
from data_type import infer_data_type
|
|
5
|
+
|
|
6
|
+
def compute_column_stats(series: pd.Series, sample_size: int = 5000) -> dict:
|
|
7
|
+
"""
|
|
8
|
+
Compute statistics for a column, based on inferred data type.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
if series.empty:
|
|
12
|
+
return None
|
|
13
|
+
|
|
14
|
+
non_null_series = series.dropna()
|
|
15
|
+
if non_null_series.empty:
|
|
16
|
+
return None
|
|
17
|
+
|
|
18
|
+
if len(non_null_series) > sample_size:
|
|
19
|
+
non_null_series = non_null_series.sample(sample_size, random_state=42)
|
|
20
|
+
|
|
21
|
+
inferred_types = non_null_series.apply(infer_data_type)
|
|
22
|
+
col_type = inferred_types.value_counts().idxmax()
|
|
23
|
+
|
|
24
|
+
stats = {}
|
|
25
|
+
|
|
26
|
+
if col_type in ["integer", "float"]:
|
|
27
|
+
numeric = pd.to_numeric(non_null_series, errors="coerce").dropna()
|
|
28
|
+
if not numeric.empty:
|
|
29
|
+
stats = {
|
|
30
|
+
"min": float(numeric.min()),
|
|
31
|
+
"max": float(numeric.max()),
|
|
32
|
+
"mean": float(numeric.mean()),
|
|
33
|
+
"median": float(numeric.median()),
|
|
34
|
+
"std": float(numeric.std())
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
elif col_type == "datetime":
|
|
38
|
+
dates = pd.to_datetime(non_null_series, errors="coerce", format="mixed").dropna()
|
|
39
|
+
if not dates.empty:
|
|
40
|
+
stats = {
|
|
41
|
+
"min_date": dates.min().isoformat(),
|
|
42
|
+
"max_date": dates.max().isoformat()
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
elif col_type in ["string", "boolean"]:
|
|
46
|
+
counts = Counter(non_null_series.astype(str))
|
|
47
|
+
top_values = counts.most_common(10)
|
|
48
|
+
stats = {
|
|
49
|
+
"top_values": top_values,
|
|
50
|
+
"unique_count": len(counts)
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
else:
|
|
54
|
+
return None
|
|
55
|
+
|
|
56
|
+
return stats if stats else None
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
from dateutil.parser import parse as dateparse
|
|
2
|
+
import pandas as pd
|
|
3
|
+
import re
|
|
4
|
+
|
|
5
|
+
_date_pattern = re.compile(
|
|
6
|
+
r"(\d{4}[-/]\d{2}[-/]\d{2})" # 2024-01-31
|
|
7
|
+
r"|(\d{2}[-/]\d{2}[-/]\d{4})" # 31-01-2024
|
|
8
|
+
r"|(\d{4}[-/]\d{2}[-/]\d{2}T\d{2}:\d{2})" # 2024-01-31T12:00
|
|
9
|
+
)
|
|
10
|
+
|
|
11
|
+
def infer_data_type(value) -> str:
|
|
12
|
+
if pd.isna(value):
|
|
13
|
+
return "null"
|
|
14
|
+
if isinstance(value, bool):
|
|
15
|
+
return "boolean"
|
|
16
|
+
if isinstance(value, int):
|
|
17
|
+
return "integer"
|
|
18
|
+
if isinstance(value, float):
|
|
19
|
+
return "float"
|
|
20
|
+
if not isinstance(value, str):
|
|
21
|
+
return "string"
|
|
22
|
+
v = value.strip()
|
|
23
|
+
if v == "":
|
|
24
|
+
return "null"
|
|
25
|
+
lower = v.lower()
|
|
26
|
+
if lower in ("true", "false", "yes", "no", "y", "n", "t", "f"):
|
|
27
|
+
return "boolean"
|
|
28
|
+
temp_v = v.replace(",", "")
|
|
29
|
+
has_sign = temp_v.startswith(("+", "-"))
|
|
30
|
+
numeric_part = temp_v[1:] if has_sign else temp_v
|
|
31
|
+
if "+" in numeric_part or "-" in numeric_part:
|
|
32
|
+
return "string"
|
|
33
|
+
if numeric_part.isdigit():
|
|
34
|
+
return "integer"
|
|
35
|
+
try:
|
|
36
|
+
if "." in numeric_part:
|
|
37
|
+
if numeric_part.replace(".", "", 1).isdigit():
|
|
38
|
+
return "float"
|
|
39
|
+
float(temp_v)
|
|
40
|
+
return "float"
|
|
41
|
+
except ValueError:
|
|
42
|
+
pass
|
|
43
|
+
|
|
44
|
+
if _date_pattern.search(v):
|
|
45
|
+
try:
|
|
46
|
+
dateparse(v, fuzzy=False)
|
|
47
|
+
return "datetime"
|
|
48
|
+
except (ValueError, OverflowError):
|
|
49
|
+
pass
|
|
50
|
+
|
|
51
|
+
return "string"
|
|
52
|
+
|
|
53
|
+
def infer_column_data_type(series: pd.Series, sample_size: int = 100) -> str:
|
|
54
|
+
"""
|
|
55
|
+
Infer the most likely data type of a column by sampling non-null values
|
|
56
|
+
and selecting the most frequent inferred type.
|
|
57
|
+
"""
|
|
58
|
+
non_null_series = series.dropna()
|
|
59
|
+
if non_null_series.empty:
|
|
60
|
+
return "null"
|
|
61
|
+
sample = non_null_series.sample(
|
|
62
|
+
min(sample_size, len(non_null_series)), random_state=42
|
|
63
|
+
)
|
|
64
|
+
inferred_types = sample.apply(infer_data_type)
|
|
65
|
+
most_common = inferred_types.value_counts().idxmax()
|
|
66
|
+
return most_common
|
|
@@ -0,0 +1,331 @@
|
|
|
1
|
+
import re
|
|
2
|
+
import os
|
|
3
|
+
import requests
|
|
4
|
+
from typing import Optional
|
|
5
|
+
import logging
|
|
6
|
+
from pydantic import BaseModel, Field
|
|
7
|
+
from metadata_extraction import DatasetMeta, TableMeta, ColumnMeta
|
|
8
|
+
|
|
9
|
+
logging.basicConfig(level=logging.INFO)
|
|
10
|
+
logger = logging.getLogger(__name__)
|
|
11
|
+
|
|
12
|
+
OPENAI_MODELS = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
|
|
13
|
+
OLLAMA_URL = "http://localhost:11434"
|
|
14
|
+
|
|
15
|
+
DATASET_PROMPT_TEMPLATE = """
|
|
16
|
+
You are an expert in dataset documentation.
|
|
17
|
+
Generate a short, precise, and factual description of the dataset.
|
|
18
|
+
Do NOT write introductions, explanations, examples, or subjective opinions.
|
|
19
|
+
Do NOT repeat the dataset name or make assumptions.
|
|
20
|
+
|
|
21
|
+
Dataset overview:
|
|
22
|
+
- Domain: {{domain}}
|
|
23
|
+
- Tables: {{table_list}}
|
|
24
|
+
- Row counts per table: {{table_row_counts}}
|
|
25
|
+
- Time coverage: {{time_range}}
|
|
26
|
+
|
|
27
|
+
Output:
|
|
28
|
+
- A single sentence (max 15 words) describing the dataset’s content, purpose, and domain.
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
TABLE_PROMPT_TEMPLATE = """
|
|
32
|
+
You are an expert in data documentation.
|
|
33
|
+
Generate a concise and precise description of the table.
|
|
34
|
+
Do NOT give examples, assumptions, or repeat the table name.
|
|
35
|
+
Use only the provided metadata.
|
|
36
|
+
|
|
37
|
+
Dataset context:
|
|
38
|
+
- Name: {{dataset_name}}
|
|
39
|
+
- Domain: {{dataset_domain}}
|
|
40
|
+
- Description: {{dataset_description}}
|
|
41
|
+
|
|
42
|
+
Table metadata:
|
|
43
|
+
- Name: {{table_name}}
|
|
44
|
+
- Columns (name and type): {{column_list_with_types}}
|
|
45
|
+
- Number of rows: {{n_rows}}
|
|
46
|
+
- Sample rows: {{sample_rows}}
|
|
47
|
+
|
|
48
|
+
Output:
|
|
49
|
+
- One short sentence (max 15 words) describing what the table represents and its content.
|
|
50
|
+
- Include factual details only.
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
COLUMN_PROMPT_TEMPLATE = """
|
|
54
|
+
You are an expert in data documentation and metadata analysis.
|
|
55
|
+
Generate a precise, factual, single-sentence description of the column.
|
|
56
|
+
|
|
57
|
+
Constraints:
|
|
58
|
+
1. Do NOT speculate, give examples, or repeat column/table names.
|
|
59
|
+
2. Do NOT repeat null percentage, distinct counts, or stats summary in the text.
|
|
60
|
+
3. Do NOT start descriptions with "Represents", "Indicates", "Contains", or "Column". Start directly with the subject (e.g., "Patient age..." instead of "Indicates patient age...").
|
|
61
|
+
|
|
62
|
+
Goals:
|
|
63
|
+
- Capture the unit of measurements (e.g., hours, kg, ml) if implied by sample values.
|
|
64
|
+
- Use the stats only to infer if the data is categorical or continuous, then write the description accordingly.
|
|
65
|
+
|
|
66
|
+
Dataset context:
|
|
67
|
+
- Name: {{dataset_name}}
|
|
68
|
+
- Domain: {{dataset_domain}}
|
|
69
|
+
- Description: {{dataset_description}}
|
|
70
|
+
|
|
71
|
+
Table context:
|
|
72
|
+
- Name: {{table_name}}
|
|
73
|
+
- Description: {{table_description}}
|
|
74
|
+
|
|
75
|
+
Column metadata:
|
|
76
|
+
- Name: {{column_name}}
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
Output:
|
|
80
|
+
- One factual sentence (max 15 words).
|
|
81
|
+
"""
|
|
82
|
+
|
|
83
|
+
def _calculate_null_percent(col_meta: ColumnMeta) -> str:
|
|
84
|
+
if col_meta.total_rows == 0:
|
|
85
|
+
return "N/A"
|
|
86
|
+
percent = (col_meta.null_count / col_meta.total_rows) * 100
|
|
87
|
+
return f"{percent:.1f}%"
|
|
88
|
+
|
|
89
|
+
def _format_column_stats(col_meta: ColumnMeta) -> str:
|
|
90
|
+
stats = col_meta.stats
|
|
91
|
+
if not stats or stats.get("type") is None:
|
|
92
|
+
return f"Unique values: {stats.get('unique_count', 'N/A')}" if stats else "N/A"
|
|
93
|
+
|
|
94
|
+
if stats["type"] == "numeric":
|
|
95
|
+
return (f"numerical: [min: {stats.get('min', 'N/A')}, max: {stats.get('max', 'N/A')}, "
|
|
96
|
+
f"mean: {stats.get('mean', 'N/A')}, median: {stats.get('median', 'N/A')}]")
|
|
97
|
+
elif stats["type"] == "datetime":
|
|
98
|
+
return f"datetime range: [{stats.get('min_date', 'N/A')} to {stats.get('max_date', 'N/A')}]"
|
|
99
|
+
elif stats["type"] == "categorical":
|
|
100
|
+
top_vals = stats.get('top_values', [])
|
|
101
|
+
if top_vals:
|
|
102
|
+
top_5 = ", ".join([str(v[0]) for v in top_vals[:5]])
|
|
103
|
+
return f"categorical: top_5_values={{{top_5}}}"
|
|
104
|
+
return f"categorical: unique_count={stats.get('unique_count', 'N/A')}"
|
|
105
|
+
return "N/A"
|
|
106
|
+
|
|
107
|
+
def _format_sample_rows(table_meta: TableMeta) -> str:
|
|
108
|
+
if not table_meta.sample_rows:
|
|
109
|
+
return "N/A"
|
|
110
|
+
return "\n".join([str(row) for row in table_meta.sample_rows])
|
|
111
|
+
|
|
112
|
+
def generate_dataset_description(dataset_meta: DatasetMeta, model: str, llm_function) -> str:
|
|
113
|
+
table_counts = {t_name: t_meta.row_count for t_name, t_meta in dataset_meta.tables.items()}
|
|
114
|
+
time_range = (f"{dataset_meta.time_period_covered.get('start')} to "
|
|
115
|
+
f"{dataset_meta.time_period_covered.get('end')}") if dataset_meta.time_period_covered else "N/A"
|
|
116
|
+
|
|
117
|
+
prompt = DATASET_PROMPT_TEMPLATE.replace("{{domain}}", dataset_meta.domain_covered or "Unspecified")\
|
|
118
|
+
.replace("{{table_list}}", ", ".join(dataset_meta.list_of_tables))\
|
|
119
|
+
.replace("{{table_row_counts}}", str(table_counts))\
|
|
120
|
+
.replace("{{time_range}}", time_range)
|
|
121
|
+
|
|
122
|
+
description = llm_function(prompt, model)
|
|
123
|
+
return description or f"Dataset containing {len(dataset_meta.list_of_tables)} tables covering {dataset_meta.domain_covered or 'unspecified domain'}."
|
|
124
|
+
|
|
125
|
+
def generate_table_description(table_meta: TableMeta, dataset_name: str, dataset_description: str, dataset_domain: str, model: str, llm_function) -> str:
|
|
126
|
+
col_list_types = ", ".join([f"{col_name} ({meta.data_type})" for col_name, meta in table_meta.columns.items()])
|
|
127
|
+
sample_rows = _format_sample_rows(table_meta)
|
|
128
|
+
|
|
129
|
+
prompt = TABLE_PROMPT_TEMPLATE.replace("{{dataset_name}}", dataset_name)\
|
|
130
|
+
.replace("{{dataset_domain}}", dataset_domain)\
|
|
131
|
+
.replace("{{dataset_description}}", dataset_description)\
|
|
132
|
+
.replace("{{table_name}}", table_meta.table_name)\
|
|
133
|
+
.replace("{{column_list_with_types}}", col_list_types)\
|
|
134
|
+
.replace("{{n_rows}}", str(table_meta.row_count))\
|
|
135
|
+
.replace("{{sample_rows}}", sample_rows)
|
|
136
|
+
|
|
137
|
+
description = llm_function(prompt, model)
|
|
138
|
+
return description or f"Table {table_meta.table_name} with {table_meta.number_of_columns} columns and {table_meta.row_count} rows."
|
|
139
|
+
|
|
140
|
+
def generate_single_column_description(column_meta: ColumnMeta, table_name: str, table_description: str, dataset_name: str, dataset_description: str, dataset_domain: str, model: str, llm_function) -> str:
|
|
141
|
+
null_percent = _calculate_null_percent(column_meta)
|
|
142
|
+
stats_summary = _format_column_stats(column_meta)
|
|
143
|
+
samples_str = ", ".join([str(s) for s in column_meta.sample_values]) if column_meta.sample_values else "N/A"
|
|
144
|
+
|
|
145
|
+
type_str = column_meta.semantic_type if column_meta.semantic_type else column_meta.data_type
|
|
146
|
+
if column_meta.is_primary_key:
|
|
147
|
+
type_str += " (Primary Key)"
|
|
148
|
+
if column_meta.foreign_key_to:
|
|
149
|
+
type_str += f" (Foreign Key to {column_meta.foreign_key_to})"
|
|
150
|
+
|
|
151
|
+
prompt = COLUMN_PROMPT_TEMPLATE.replace("{{dataset_name}}", dataset_name)\
|
|
152
|
+
.replace("{{dataset_domain}}", dataset_domain)\
|
|
153
|
+
.replace("{{dataset_description}}", dataset_description)\
|
|
154
|
+
.replace("{{table_name}}", table_name)\
|
|
155
|
+
.replace("{{table_description}}", table_description)\
|
|
156
|
+
.replace("{{column_name}}", column_meta.column_name)\
|
|
157
|
+
.replace("{{type}}", type_str)\
|
|
158
|
+
.replace("{{samples}}", samples_str)\
|
|
159
|
+
.replace("{{unique}}", str(column_meta.unique_count))\
|
|
160
|
+
.replace("{{null_percent}}", null_percent)\
|
|
161
|
+
.replace("{{stats_summary}}", stats_summary)
|
|
162
|
+
|
|
163
|
+
response = llm_function(prompt, model)
|
|
164
|
+
return response if response else f"Column of type {type_str}."
|
|
165
|
+
|
|
166
|
+
def _post_process_response(response_text: str) -> str:
|
|
167
|
+
"""Cleans LLM response and strictly enforces word count limit."""
|
|
168
|
+
if not response_text:
|
|
169
|
+
return ""
|
|
170
|
+
|
|
171
|
+
clean_response = re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL)
|
|
172
|
+
clean_response = re.sub(r"```.*?```", "", clean_response, flags=re.DOTALL)
|
|
173
|
+
clean_response = re.sub(r"[*_#>`\"']", "", clean_response)
|
|
174
|
+
clean_response = re.sub(r"\s+", " ", clean_response).strip()
|
|
175
|
+
words = clean_response.split()
|
|
176
|
+
if len(words) > 15:
|
|
177
|
+
clean_response = " ".join(words[:15])
|
|
178
|
+
|
|
179
|
+
return clean_response
|
|
180
|
+
|
|
181
|
+
COST_TRACKER = {
|
|
182
|
+
"total_cost": 0.0,
|
|
183
|
+
"models": {}
|
|
184
|
+
}
|
|
185
|
+
|
|
186
|
+
def update_cost(model: str, cost: float):
|
|
187
|
+
COST_TRACKER["total_cost"] += cost
|
|
188
|
+
if model not in COST_TRACKER["models"]:
|
|
189
|
+
COST_TRACKER["models"][model] = {"calls": 0, "cost": 0.0}
|
|
190
|
+
COST_TRACKER["models"][model]["calls"] += 1
|
|
191
|
+
COST_TRACKER["models"][model]["cost"] += cost
|
|
192
|
+
|
|
193
|
+
def get_cost_summary() -> dict:
|
|
194
|
+
"""
|
|
195
|
+
Returns the entire cost summary.
|
|
196
|
+
"""
|
|
197
|
+
return COST_TRACKER
|
|
198
|
+
|
|
199
|
+
def _run_openai_model(prompt: str, model: str) -> Optional[str]:
|
|
200
|
+
"""
|
|
201
|
+
Call OpenAI for description generation with incremental cost tracking.
|
|
202
|
+
"""
|
|
203
|
+
try:
|
|
204
|
+
api_key = os.environ.get("OPENAI_API_KEY")
|
|
205
|
+
if not api_key:
|
|
206
|
+
raise ValueError("OPENAI_API_KEY environment variable not set")
|
|
207
|
+
|
|
208
|
+
from openai import OpenAI
|
|
209
|
+
client = OpenAI(api_key=api_key)
|
|
210
|
+
response = client.chat.completions.create(
|
|
211
|
+
model=model,
|
|
212
|
+
messages=[
|
|
213
|
+
{
|
|
214
|
+
"role": "system",
|
|
215
|
+
"content": (
|
|
216
|
+
"You are a precise metadata expert. "
|
|
217
|
+
"You write extremely concise, factual descriptions based ONLY on provided context. "
|
|
218
|
+
"Do not use filler words. Strictly adhere to word limits (max 15 words)."
|
|
219
|
+
)
|
|
220
|
+
},
|
|
221
|
+
{"role": "user", "content": prompt}
|
|
222
|
+
],
|
|
223
|
+
temperature=0.1,
|
|
224
|
+
max_tokens=50,
|
|
225
|
+
)
|
|
226
|
+
|
|
227
|
+
response_text = response.choices[0].message.content.strip()
|
|
228
|
+
|
|
229
|
+
# Token usage
|
|
230
|
+
usage = getattr(response, "usage", None)
|
|
231
|
+
if usage:
|
|
232
|
+
prompt_tokens = usage.prompt_tokens
|
|
233
|
+
completion_tokens = usage.completion_tokens
|
|
234
|
+
else:
|
|
235
|
+
prompt_tokens = completion_tokens = 0
|
|
236
|
+
|
|
237
|
+
# Pricing table (per 1k tokens)
|
|
238
|
+
model_costs = {
|
|
239
|
+
"gpt-4": (0.03, 0.06),
|
|
240
|
+
"gpt-4-turbo": (0.01, 0.03),
|
|
241
|
+
"gpt-4o": (0.005, 0.015),
|
|
242
|
+
"gpt-4o-mini": (0.00015, 0.0006),
|
|
243
|
+
"gpt-3.5-turbo": (0.0005, 0.0015),
|
|
244
|
+
}
|
|
245
|
+
prompt_cost_per_1k, completion_cost_per_1k = model_costs.get(model, (0.005, 0.015))
|
|
246
|
+
|
|
247
|
+
# Cost of this single call
|
|
248
|
+
cost = (
|
|
249
|
+
(prompt_tokens / 1000) * prompt_cost_per_1k +
|
|
250
|
+
(completion_tokens / 1000) * completion_cost_per_1k
|
|
251
|
+
)
|
|
252
|
+
|
|
253
|
+
# Incremental cost update
|
|
254
|
+
update_cost(model, cost)
|
|
255
|
+
|
|
256
|
+
logger.info(
|
|
257
|
+
f"[COST] Model={model} | Prompt={prompt_tokens} | Completion={completion_tokens} "
|
|
258
|
+
f"| This call=${cost:.6f} | Total=${COST_TRACKER['total_cost']:.6f}"
|
|
259
|
+
)
|
|
260
|
+
|
|
261
|
+
return _post_process_response(response_text)
|
|
262
|
+
|
|
263
|
+
except Exception as e:
|
|
264
|
+
logger.error(f"OpenAI API error for model {model}: {e}")
|
|
265
|
+
return None
|
|
266
|
+
|
|
267
|
+
def _run_gemini_model(prompt: str, model: str) -> Optional[str]:
|
|
268
|
+
"""
|
|
269
|
+
Runs Google Gemini models (via google-generativeai).
|
|
270
|
+
"""
|
|
271
|
+
try:
|
|
272
|
+
api_key = os.environ.get("GEMINI_API_COST_KEY")
|
|
273
|
+
if not api_key:
|
|
274
|
+
raise ValueError("GEMINI_API_KEY environment variable not set")
|
|
275
|
+
|
|
276
|
+
import google.generativeai as genai
|
|
277
|
+
genai.configure(api_key=api_key)
|
|
278
|
+
gem_model = genai.GenerativeModel(model)
|
|
279
|
+
response = gem_model.generate_content(prompt)
|
|
280
|
+
response_text = response.text.strip() if hasattr(response, "text") else ""
|
|
281
|
+
return _post_process_response(response_text)
|
|
282
|
+
except Exception as e:
|
|
283
|
+
logger.error(f"Gemini API error for model {model}: {e}")
|
|
284
|
+
return None
|
|
285
|
+
|
|
286
|
+
def run_ollama_model(prompt: str, model: str) -> Optional[str]:
|
|
287
|
+
"""
|
|
288
|
+
Call local Ollama API for concise descriptions.
|
|
289
|
+
"""
|
|
290
|
+
API_ENDPOINT = f"{OLLAMA_URL}/api/generate"
|
|
291
|
+
try:
|
|
292
|
+
num_predict = 50
|
|
293
|
+
|
|
294
|
+
payload = {
|
|
295
|
+
"model": model,
|
|
296
|
+
"prompt": prompt,
|
|
297
|
+
"stream": False,
|
|
298
|
+
"options": {
|
|
299
|
+
"temperature": 0.1,
|
|
300
|
+
"num_predict": num_predict
|
|
301
|
+
}
|
|
302
|
+
}
|
|
303
|
+
|
|
304
|
+
response = requests.post(API_ENDPOINT, json=payload)
|
|
305
|
+
response.raise_for_status()
|
|
306
|
+
|
|
307
|
+
response_data = response.json()
|
|
308
|
+
raw_text = response_data.get("response", "").strip()
|
|
309
|
+
return _post_process_response(raw_text)
|
|
310
|
+
|
|
311
|
+
except requests.exceptions.RequestException as e:
|
|
312
|
+
logger.error(f"Ollama API call failed for model {model}: {e}")
|
|
313
|
+
return None
|
|
314
|
+
|
|
315
|
+
|
|
316
|
+
def run_llm_dispatcher(prompt: str, model: str) -> Optional[str]:
|
|
317
|
+
"""
|
|
318
|
+
Routes the prompt to the correct LLM implementation (OpenAI, Gemini or Ollama).
|
|
319
|
+
"""
|
|
320
|
+
model_lower = model.lower()
|
|
321
|
+
|
|
322
|
+
# OpenAI Models
|
|
323
|
+
if model_lower in [m.lower() for m in OPENAI_MODELS] or model_lower.startswith("gpt-"):
|
|
324
|
+
return _run_openai_model(prompt, model)
|
|
325
|
+
|
|
326
|
+
# Google Gemini Models
|
|
327
|
+
if model_lower.startswith("gemini"):
|
|
328
|
+
return _run_gemini_model(prompt, model)
|
|
329
|
+
|
|
330
|
+
# Local Ollama
|
|
331
|
+
return run_ollama_model(prompt, model)
|