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.
@@ -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,5 @@
1
+ Wheel-Version: 1.0
2
+ Generator: setuptools (82.0.1)
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
5
+
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ generate-dictionary = data_dictionary_generator.generate_dictionary:main
@@ -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)