diff --git a/.gitignore b/.gitignore index 9ac7fef..699e469 100644 --- a/.gitignore +++ b/.gitignore @@ -209,6 +209,7 @@ notes.txt testing_base /models +/notebooks/models/ models_bak settings.json diff --git a/docsite/docs/core-concepts.md b/docsite/docs/core-concepts.md deleted file mode 100644 index a56f55f..0000000 --- a/docsite/docs/core-concepts.md +++ /dev/null @@ -1,17 +0,0 @@ ---- -sidebar_position: 3 ---- - -# Core Concepts - -## Knowledge Builder - -The `KnowledgeBuilder` is the main entry point for building the semantic layer. It takes a dictionary of datasets as input and orchestrates the entire process of profiling, link prediction, and business glossary generation. - -## Data Product Builder - -The `DataProductBuilder` is used to generate data products from the semantic layer. It takes an config dictionary as input and generates a unified data product that can be used for analysis and exploration. - -## Semantic Search - -The semantic search feature allows you to search for columns in your datasets using natural language. It uses a hybrid search approach that combines dense and sparse vectors for more accurate results. diff --git a/docsite/docs/core-concepts/_category_.json b/docsite/docs/core-concepts/_category_.json new file mode 100644 index 0000000..49dd5c7 --- /dev/null +++ b/docsite/docs/core-concepts/_category_.json @@ -0,0 +1,8 @@ +{ + "label": "Core Concepts", + "position": 3, + "link": { + "type": "generated-index", + "description": "Learn about the core concepts and architecture of the Intugle Data Tools." + } +} diff --git a/docsite/docs/core-concepts/dataset.md b/docsite/docs/core-concepts/dataset.md new file mode 100644 index 0000000..5eefaa7 --- /dev/null +++ b/docsite/docs/core-concepts/dataset.md @@ -0,0 +1,133 @@ +--- +sidebar_position: 3 +title: DataSet +--- + +# DataSet + +The `DataSet` class is the heart of the analysis pipeline. It acts as a powerful, in-memory container for a single data source, holding not just the raw data but also all the rich metadata that is generated about it during the `KnowledgeBuilder`'s workflow. + +## Overview + +Think of a `DataSet` as a "unit of work" that gets progressively enriched as it moves through the analysis stages. Its key responsibilities are: + +1. **Data Abstraction**: It provides a consistent way to interact with data, regardless of its source. +2. **Metadata Storage**: It serves as the central object for storing all analysis results. +3. **Intelligent Caching**: It automatically persists and loads its own state, making the entire process efficient and resilient. + +## Key Features + +### Data Abstraction + +The `DataSet` uses a system of **Adapters** under the hood to connect to different data backends. When you initialize a `DataSet` with a file-based source, it uses the appropriate adapter (e.g., `DuckdbAdapter` for CSVs, Parquet, etc.) to handle the specific implementation details for profiling and querying the data. This makes the system easily extensible to support new data sources in the future. + +### Centralized Metadata + +All analysis results for a data source are stored within the `dataset.source_table_model` attribute. This attribute is a structured Pydantic model that makes accessing metadata predictable and easy. + +#### Metadata Structure + +The metadata is organized using the following Pydantic models: + +- `SourceTables`: The root object, containing table-level information. + - `name: str` + - `description: str` + - `key: Optional[str]` + - `columns: List[Column]` + - `profiling_metrics: Optional[ModelProfilingMetrics]` +- `Column`: Contains metadata for a single column. + - `name: str` + - `description: Optional[str]` + - `type: Optional[str]` (e.g., 'integer', 'date') + - `category: Literal["dimension", "measure"]` + - `tags: Optional[List[str]]` + - `profiling_metrics: Optional[ColumnProfilingMetrics]` +- `ColumnProfilingMetrics`: Detailed statistics for a column. + - `count: Optional[int]` + - `null_count: Optional[int]` + - `distinct_count: Optional[int]` + +#### Accessing Metadata + +You can access this rich metadata directly from the `DataSet` object. + +```python +# Assuming 'kb' is a built KnowledgeBuilder instance +customers_dataset = kb.datasets['customers'] + +# Access table-level metadata +print(f"Table Name: {customers_dataset.source_table_model.name}") +print(f"Table Description: {customers_dataset.source_table_model.description}") +print(f"Primary Key: {customers_dataset.source_table_model.key}") + +# Access column-level metadata for the first column +first_column = customers_dataset.source_table_model.columns[0] +print(f"Column Name: {first_column.name}") +print(f"Column Description: {first_column.description}") +print(f"Column Type: {first_column.type}") + +# Access profiling metrics for that column +metrics = first_column.profiling_metrics +if metrics: + print(f"Distinct Count: {metrics.distinct_count}") +``` + +### Automatic Caching + +The `DataSet` object is designed to be efficient by avoiding redundant work. When you initialize a `DataSet`, it automatically checks for a corresponding `.yml` file. If it finds one, it validates that the source data file hasn't changed since the last run. If the data is fresh, it loads the saved metadata, saving significant processing time. + +### Analysis Stage Functions + +You can run the analysis pipeline step-by-step for more granular control. Each of these methods includes a `save=True` option to persist the results of that specific stage. + +```python +from intugle import DataSet + +# Initialize the dataset +data_source = {"path": "path/to/my_data.csv", "type": "csv"} +dataset = DataSet(data_source, name="my_data") + +# Run each stage individually and save progress +print("Step 1: Profiling...") +dataset.profile(save=True) + +print("Step 2: Identifying Datatypes...") +dataset.identify_datatypes(save=True) + +print("Step 3: Identifying Keys...") +dataset.identify_keys(save=True) + +print("Step 4: Generating Glossary...") +dataset.generate_glossary(domain="my_domain", save=True) +``` + +### Other Useful Methods and Properties + +#### `save_yaml()` + +Manually trigger a save of the dataset's current state to its `.yml` file. + +```python +dataset.save_yaml() +``` + +#### `reload_from_yaml()` + +Force a reload from the YAML file, bypassing the staleness check. This can be useful if you manually edit the YAML and want to load your changes. + +```python +file_path = "path/to/my_data.yml" +dataset.reload_from_yaml(file_path) +``` + +#### `profiling_df` + +Access a Pandas DataFrame containing the complete profiling results for all columns in the dataset. This is very useful for exploration and validation in a notebook environment. + +```python +# Get a comprehensive DataFrame of all column profiles +profiles = dataset.profiling_df + +# Display the first 5 rows +print(profiles.head()) +``` \ No newline at end of file diff --git a/docsite/docs/core-concepts/index.md b/docsite/docs/core-concepts/index.md new file mode 100644 index 0000000..29d1215 --- /dev/null +++ b/docsite/docs/core-concepts/index.md @@ -0,0 +1,17 @@ +--- +sidebar_position: 1 +--- + +# Core Concepts + +Welcome to the core concepts section. Here, we'll dive deep into the fundamental components and architectural ideas that power the Intugle Data Tools library. + +Understanding these concepts is key to using the library effectively and leveraging its full potential for automating data intelligence. + +### Key Components: + +* [**Knowledge Builder**](./knowledge-builder.md): The main orchestrator that manages the end-to-end analysis pipeline. +* [**DataSet**](./dataset.md): The central in-memory container for a data source and all its rich metadata. +* [**Link Prediction**](./link-prediction.md): The process of automatically discovering relationships between your datasets. +* **Data Product Builder**: The tool used to generate unified data products from the enriched semantic layer. +* **Semantic Search**: The feature that enables natural language queries against your connected data assets. diff --git a/docsite/docs/core-concepts/knowledge-builder.md b/docsite/docs/core-concepts/knowledge-builder.md new file mode 100644 index 0000000..1db1ad9 --- /dev/null +++ b/docsite/docs/core-concepts/knowledge-builder.md @@ -0,0 +1,135 @@ +--- +sidebar_position: 2 +title: The Knowledge Builder +--- + +# The Knowledge Builder + +The `KnowledgeBuilder` is the primary orchestrator of the data intelligence pipeline. It is the main user-facing class designed to manage multiple data sources and run the end-to-end process of transforming them from raw, disconnected tables into a fully enriched and interconnected semantic layer. + +## Overview + +At a high level, the `KnowledgeBuilder` is responsible for: + +1. **Initializing and Managing Datasets**: It takes your raw data sources (e.g., file paths) and wraps each one in a `DataSet` object. +2. **Executing the Analysis Pipeline**: It runs a series of analysis stages in a specific, logical order to build up a rich understanding of your data. +3. **Ensuring Resilience**: The pipeline is designed to be modular and resilient. It automatically saves its progress after each major stage, allowing you to resume an interrupted run without losing completed work. + +## Initialization + +You can initialize the `KnowledgeBuilder` in two ways: + +1. **With a Dictionary of File-Based Sources**: This is the most common method. You provide a dictionary where keys are the desired names for your datasets and values are dictionary configs pointing to your data. + + ```python + from intugle import KnowledgeBuilder + + data_sources = { + "customers": {"path": "path/to/customers.csv", "type": "csv"}, + "orders": {"path": "path/to/orders.csv", "type": "csv"}, + } + + kb = KnowledgeBuilder(data_input=data_sources, domain="e-commerce") + ``` + +2. **With a List of `DataSet` Objects**: If you have already created `DataSet` objects, you can pass a list of them directly. + + ```python + from intugle import KnowledgeBuilder, DataSet + + # Create DataSet objects from file-based sources + customers_data = {"path": "path/to/customers.csv", "type": "csv"} + orders_data = {"path": "path/to/orders.csv", "type": "csv"} + + dataset_one = DataSet(customers_data, name="customers") + dataset_two = DataSet(orders_data, name="orders") + + datasets = [dataset_one, dataset_two] + + kb = KnowledgeBuilder(data_input=datasets, domain="e-commerce") + ``` + +The `domain` parameter is an optional but highly recommended string that provides context to the underlying AI models, helping them generate more accurate and relevant business glossary terms. + +## The Analysis Pipeline + +The `KnowledgeBuilder` executes its workflow in distinct, modular stages. This allows for greater control and makes the process resilient to interruptions. + +### 1. `profile()` + +This is the first and most foundational stage. It performs a deep analysis of each dataset to understand its structure and content, covering profiling, datatype identification, and key identification. + +```python +# Run only the profiling and key identification stage +kb.profile() +``` + +Progress from this stage is automatically saved to a `.yml` file for each dataset. + +### 2. `predict_links()` + +Once the datasets are profiled, this stage uses the `LinkPredictor` to analyze the metadata from all datasets and discover potential relationships between them. You can learn more about the [Link Prediction](./link-prediction.md) process in its dedicated section. + +```python +# Run the link prediction stage +# This assumes profile() has already been run +kb.predict_links() + +# Access the links via the `links` attribute, which is a shortcut +discovered_links = kb.links +print(discovered_links) + +# You can also access the full LinkPredictor instance for more options +# See the section below for more details. +``` + +The discovered relationships are saved to a central `__relationships__.yml` file. + +### 3. `generate_glossary()` + +In the final stage, the `KnowledgeBuilder` uses a Large Language Model (LLM) to generate business-friendly context for your data. + +```python +# Run the glossary generation stage +# This assumes profile() has already been run +kb.generate_glossary() +``` + +This information is saved back into each dataset's `.yml` file. + +### The `build()` Method + +For convenience, the `build()` method runs all three stages (`profile`, `predict_links`, `generate_glossary`) in the correct sequence. + +```python +# Run the full pipeline from start to finish +kb.build() + +# You can also force it to re-run everything, ignoring any cached results +kb.build(force_recreate=True) +``` + +This modular design means that if your process is interrupted during the `generate_glossary` stage, you can simply re-run `kb.build()`, and it will quickly skip the already-completed stages, picking up right where it left off. + +## Accessing Processed Datasets and Predictor + +After running any stage of the pipeline, you can easily access the enriched `DataSet` objects and the `LinkPredictor` instance to explore the results programmatically. + +```python +# Run the full build +kb.build() + +# Access the 'customers' dataset +customers_dataset = kb.datasets['customers'] + +# Access the LinkPredictor instance +link_predictor = kb.link_predictor + +# Now you can explore rich metadata or results +print(f"Primary Key for customers: {customers_dataset.source_table_model.description}") +print("Discovered Links:") +print(link_predictor.get_links_df()) + +``` +Learn more about what you can do with these objects. See the [DataSet](./dataset.md) and [Link Prediction](./link-prediction.md) documentation. + diff --git a/docsite/docs/core-concepts/link-prediction.md b/docsite/docs/core-concepts/link-prediction.md new file mode 100644 index 0000000..de7592b --- /dev/null +++ b/docsite/docs/core-concepts/link-prediction.md @@ -0,0 +1,98 @@ +--- +sidebar_position: 4 +title: Link Prediction +--- + +# Link Prediction + +Link Prediction is one of the most powerful features of the Intugle Data Tools library. It is the process of automatically discovering meaningful relationships and potential join keys between different, isolated datasets. This turns a collection of separate tables into a connected semantic graph, which is the foundation for building unified data products. + +## The `LinkPredictor` Class + +The core component responsible for this process is the `LinkPredictor`. While the `KnowledgeBuilder` manages this process for you, you can also use the `LinkPredictor` directly for more granular control. + +### Accessing the `LinkPredictor` + +After running the `predict_links()` or `build()` method on a `KnowledgeBuilder` instance, you can access the underlying `LinkPredictor` instance via the `link_predictor` attribute. + +```python +# After running the pipeline... +predictor_instance = kb.link_predictor + +# Now you can use all the methods of the LinkPredictor +links_list = predictor_instance.links +``` + +### Manual Usage + +To use the `LinkPredictor` manually, you must provide it with a list of fully profiled `DataSet` objects. + +```python +from intugle import DataSet, LinkPredictor + +# 1. Initialize and fully profile your DataSet objects first +customers_data = {"path": "path/to/customers.csv", "type": "csv"} +orders_data = {"path": "path/to/orders.csv", "type": "csv"} + +customers_dataset = DataSet(customers_data, name="customers") +customers_dataset.profile().identify_datatypes().identify_keys() + +orders_dataset = DataSet(orders_data, name="orders") +orders_dataset.profile().identify_datatypes().identify_keys() + +# 2. Initialize the LinkPredictor with the processed datasets +predictor = LinkPredictor([customers_dataset, orders_dataset]) + +# 3. Run the prediction +predictor.predict(save=True) + +# 4. Access the results +# The discovered links are stored as a list of PredictedLink objects in the `links` attribute +links_list = predictor.links +for link in links_list: + print(f"Found link from {link.from_dataset}.{link.from_column} to {link.to_dataset}.{link.to_column}") +``` + +### Caching Mechanism + +The `predict()` method is designed to be efficient. It saves its results to a `__relationships__.yml` file and only re-runs the analysis if it detects that any of the underlying dataset analyses have changed since the last run. + +### Useful Methods and Attributes + +#### `links` + +The primary way to access the results. This attribute holds a list of `PredictedLink` Pydantic objects, giving you structured access to the discovered relationships. + +#### `get_links_df()` + +A utility function that converts the `links` list into a Pandas DataFrame. This is useful for quick exploration, analysis, or display in a notebook environment. + +```python +# Get the results as a DataFrame for easy viewing +links_df = predictor.get_links_df() + +# Display the DataFrame +# columns: from_dataset, from_column, to_dataset, to_column +print(links_df) +``` + +#### `save_yaml()` and `load_from_yaml()` + +You can manually save the state of the predictor or load results from a specific file. + +```python +# Save the discovered links to a custom file +predictor.save_yaml("my_custom_links.yml") + +# Load links from a file +predictor.load_from_yaml("my_custom_links.yml") +``` + +#### `show_graph()` + +After running the prediction, you can easily visualize the discovered relationships as a graph. This is an excellent way to understand the overall structure of your connected data. + +```python +# This will render a graph of the relationships +predictor.show_graph() +``` \ No newline at end of file diff --git a/notebooks/quickstart_healthcare.ipynb b/notebooks/quickstart_healthcare.ipynb index d264b9a..1e0f796 100644 --- a/notebooks/quickstart_healthcare.ipynb +++ b/notebooks/quickstart_healthcare.ipynb @@ -3706,7 +3706,8 @@ "id": "9c8aa361", "metadata": {}, "source": [ - "You can also visualize these relationships as a graph:\n" + "You can also visualize these relationships as a graph. In case you run into an error, make sure you install/upgrade your ipykernel package:\n", + "> %pip install --upgrade ipykernel" ] }, { @@ -3763,7 +3764,7 @@ } ], "source": [ - "kb.visualize()\n" + "kb.visualize() # To visualize the relationships as a graph\n" ] }, { @@ -4417,7 +4418,7 @@ ], "metadata": { "kernelspec": { - "display_name": "env", + "display_name": "intugle", "language": "python", "name": "python3" }, diff --git a/notebooks/quickstart_tech_company.ipynb b/notebooks/quickstart_tech_company.ipynb index aa2e550..db647d9 100644 --- a/notebooks/quickstart_tech_company.ipynb +++ b/notebooks/quickstart_tech_company.ipynb @@ -3986,12 +3986,13 @@ "id": "9c8aa361", "metadata": {}, "source": [ - "You can also visualize these relationships as a graph:\n" + "You can also visualize these relationships as a graph. In case you run into an error, make sure you install/upgrade your ipykernel package:\n", + "> %pip install --upgrade ipykernel" ] }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "a59704ba", "metadata": {}, "outputs": [ @@ -4007,7 +4008,7 @@ } ], "source": [ - "kb.visualize()\n" + "kb.visualize() # To visualize the relationships as a graph\n" ] }, { diff --git a/pyproject.toml b/pyproject.toml index b90c813..9d7c850 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "intugle" -version = "0.1.5" +version = "0.1.6" authors = [ { name="Intugle", email="hello@intugle.ai" }, ] @@ -53,7 +53,7 @@ dependencies = [ "scikit-learn<=1.7.1", "langchain[anthropic,google-genai,openai]>=0.3.27", "qdrant-client>=1.15.1", - "ipykernel>=6.30.1", + "rich>=14.1.0", ] [project.urls] diff --git a/src/intugle/adapters/types/duckdb/duckdb.py b/src/intugle/adapters/types/duckdb/duckdb.py index eeea2bd..1670a5a 100644 --- a/src/intugle/adapters/types/duckdb/duckdb.py +++ b/src/intugle/adapters/types/duckdb/duckdb.py @@ -62,7 +62,7 @@ def __format_dtype__(dtype: Any) -> str: # Fetching total count of table query = f""" - SELECT count(*) as count FROM {table_name} + SELECT count(*) as count FROM "{table_name}" """ data = duckdb.execute(query).fetchall() @@ -203,12 +203,12 @@ def _get_load_func(data: DuckdbConfig): return f"{ld_func}('{data.path}')" def load_view(self, data: DuckdbConfig, table_name: str): - query = f"""CREATE OR REPLACE VIEW {table_name} AS {data.path}""" + query = f"""CREATE OR REPLACE VIEW "{table_name}" AS {data.path}""" duckdb.execute(query) def load_file(self, data: DuckdbConfig, table_name: str): ld_func = self._get_load_func(data) - query = f"""CREATE TABLE IF NOT EXISTS {table_name} AS SELECT * FROM {ld_func};""" + query = f"""CREATE VIEW IF NOT EXISTS "{table_name}" AS SELECT * FROM {ld_func};""" duckdb.execute(query) @@ -224,7 +224,7 @@ def execute_df(self, query): return df def to_df(self, _: DuckdbConfig, table_name: str): - query = f"SELECT * from {table_name}" + query = f'''SELECT * from "{table_name}"''' df = self.execute_df(query) return df diff --git a/src/intugle/analysis/models.py b/src/intugle/analysis/models.py index dc713de..8564ad4 100644 --- a/src/intugle/analysis/models.py +++ b/src/intugle/analysis/models.py @@ -1,35 +1,33 @@ import json +import logging import os -import time import uuid -from typing import Any, Dict, Optional, Self +from typing import Dict, Optional, Self import pandas as pd import yaml from intugle.adapters.factory import AdapterFactory from intugle.adapters.models import ( - BusinessGlossaryOutput, - ColumnGlossary, - ColumnProfile, DataSetData, DataTypeIdentificationL1Output, DataTypeIdentificationL2Input, DataTypeIdentificationL2Output, KeyIdentificationOutput, - ProfilingOutput, ) -from intugle.common.exception import errors from intugle.core import settings +from intugle.core.console import console, warning_style from intugle.core.pipeline.business_glossary.bg import BusinessGlossary from intugle.core.pipeline.datatype_identification.l2_model import L2Model from intugle.core.pipeline.datatype_identification.pipeline import DataTypeIdentificationPipeline from intugle.core.pipeline.key_identification.ki import KeyIdentificationLLM from intugle.core.utilities.processing import string_standardization -from intugle.models.resources.model import Column, ColumnProfilingMetrics +from intugle.models.resources.model import Column, ColumnProfilingMetrics, ModelProfilingMetrics from intugle.models.resources.source import Source, SourceTables +log = logging.getLogger(__name__) + class DataSet: """ @@ -47,29 +45,74 @@ def __init__(self, data: DataSetData, name: str): self.adapter = AdapterFactory().create(data) # A dictionary to store the results of each analysis step - self.results: Dict[str, Any] = {} + self.source_table_model: SourceTables = SourceTables(name=name, description="") + self._columns_map: Dict[str, Column] = {} # A convenience map for quick column lookup + + # Check if a YAML file exists and load it + file_path = os.path.join(settings.PROJECT_BASE, f"{self.name}.yml") + if os.path.exists(file_path): + print(f"Found existing YAML for '{self.name}'. Checking for staleness.") + self.load_from_yaml(file_path) self.load() + def _is_yaml_stale(self, yaml_data: dict) -> bool: + """Check if the YAML data is stale by comparing source modification times.""" + if not isinstance(self.data, dict) or "path" not in self.data or not os.path.exists(self.data["path"]): + # Not a file-based source, so we cannot check for staleness. + return False + + try: + source = yaml_data.get("sources", [])[0] + table = source.get("table", {}) + source_last_modified = table.get("source_last_modified") + + if source_last_modified: + current_mtime = os.path.getmtime(self.data["path"]) + if current_mtime > source_last_modified: + console.print( + f"Warning: Source file for '{self.name}' has been modified since the last analysis.", + style=warning_style, + ) + return True + return False + except (IndexError, KeyError, TypeError): + # If YAML is malformed, treat it as stale. + console.print(f"Warning: Could not parse existing YAML for '{self.name}'. Treating as stale.", style=warning_style) + return True + + def _populate_from_yaml(self, yaml_data: dict): + """Populate the DataSet object from YAML data.""" + source = yaml_data.get("sources", [])[0] + table = source.get("table", {}) + self.source_table_model = SourceTables.model_validate(table) + self._columns_map = {col.name: col for col in self.source_table_model.columns} + @property def sql_query(self): - if 'type' in self.data and self.data['type'] == 'query': - return self.data['path'] + if "type" in self.data and self.data["type"] == "query": + return self.data["path"] return None - + def load(self): try: self.adapter.load(self.data, self.name) print(f"{self.name} loaded") except Exception as e: - print("eee", e) + log.error(e) ... def profile_table(self) -> Self: """ Profiles the table and stores the result in the 'results' dictionary. """ - self.results["table_profile"] = self.adapter.profile(self.data, self.name) + table_profile = self.adapter.profile(self.data, self.name) + if self.source_table_model.profiling_metrics is None: + self.source_table_model.profiling_metrics = ModelProfilingMetrics() + self.source_table_model.profiling_metrics.count = table_profile.count + + self.source_table_model.columns = [Column(name=col_name) for col_name in table_profile.columns] + self._columns_map = {col.name: col for col in self.source_table_model.columns} return self def profile_columns(self) -> Self: @@ -77,16 +120,24 @@ def profile_columns(self) -> Self: Profiles each column in the dataset and stores the results in the 'results' dictionary. This method relies on the 'table_profile' result to get the list of columns. """ - if "table_profile" not in self.results: + if not self.source_table_model.columns: raise RuntimeError("TableProfiler must be run before profiling columns.") - table_profile: ProfilingOutput = ProfilingOutput.model_validate(self.results["table_profile"]) - self.results["column_profiles"] = { - col_name: self.adapter.column_profile( - self.data, self.name, col_name, table_profile.count, settings.UPSTREAM_SAMPLE_LIMIT + count = self.source_table_model.profiling_metrics.count + + for column in self.source_table_model.columns: + column_profile = self.adapter.column_profile( + self.data, self.name, column.name, count, settings.UPSTREAM_SAMPLE_LIMIT ) - for col_name in table_profile.columns - } + if column_profile: + if column.profiling_metrics is None: + column.profiling_metrics = ColumnProfilingMetrics() + + column.profiling_metrics.count = column_profile.count + column.profiling_metrics.null_count = column_profile.null_count + column.profiling_metrics.distinct_count = column_profile.distinct_count + column.profiling_metrics.sample_data = column_profile.sample_data + column.profiling_metrics.dtype_sample = column_profile.dtype_sample return self def identify_datatypes_l1(self) -> "DataSet": @@ -94,13 +145,16 @@ def identify_datatypes_l1(self) -> "DataSet": Identifies the data types at Level 1 for each column based on the column profiles. This method relies on the 'column_profiles' result. """ - if "column_profiles" not in self.results: + if not self.source_table_model.columns or any( + c.profiling_metrics is None for c in self.source_table_model.columns + ): raise RuntimeError("TableProfiler and ColumnProfiler must be run before data type identification.") - column_profiles: dict[str, ColumnProfile] = self.results["column_profiles"] records = [] - for column_name, stats in column_profiles.items(): - records.append({"table_name": self.name, "column_name": column_name, "values": stats.dtype_sample}) + for column in self.source_table_model.columns: + records.append( + {"table_name": self.name, "column_name": column.name, "values": column.profiling_metrics.dtype_sample} + ) l1_df = pd.DataFrame(records) di_pipeline = DataTypeIdentificationPipeline() @@ -108,10 +162,8 @@ def identify_datatypes_l1(self) -> "DataSet": column_datatypes_l1 = [DataTypeIdentificationL1Output(**row) for row in l1_result.to_dict(orient="records")] - for column in column_datatypes_l1: - column_profiles[column.column_name].datatype_l1 = column.datatype_l1 - - self.results["column_datatypes_l1"] = column_datatypes_l1 + for col_l1 in column_datatypes_l1: + self._columns_map[col_l1.column_name].type = col_l1.datatype_l1 return self def identify_datatypes_l2(self) -> "DataSet": @@ -119,93 +171,132 @@ def identify_datatypes_l2(self) -> "DataSet": Identifies the data types at Level 2 for each column based on the column profiles. This method relies on the 'column_profiles' result. """ - if "column_profiles" not in self.results: + if not self.source_table_model.columns or any(c.type is None for c in self.source_table_model.columns): raise RuntimeError("TableProfiler and ColumnProfiler must be run before data type identification.") - column_profiles: dict[str, ColumnProfile] = self.results["column_profiles"] - columns_with_samples = [DataTypeIdentificationL2Input(**col.model_dump()) for col in column_profiles.values()] + columns_with_samples = [] + for column in self.source_table_model.columns: + columns_with_samples.append( + DataTypeIdentificationL2Input( + column_name=column.name, + table_name=self.name, + sample_data=column.profiling_metrics.sample_data, + datatype_l1=column.type, + ) + ) column_values_df = pd.DataFrame([item.model_dump() for item in columns_with_samples]) l2_model = L2Model() l2_result = l2_model(l1_pred=column_values_df) column_datatypes_l2 = [DataTypeIdentificationL2Output(**row) for row in l2_result.to_dict(orient="records")] - for column in column_datatypes_l2: - column_profiles[column.column_name].datatype_l2 = column.datatype_l2 - - self.results["column_datatypes_l2"] = column_datatypes_l2 + for col_l2 in column_datatypes_l2: + self._columns_map[col_l2.column_name].category = col_l2.datatype_l2 return self - def identify_keys(self) -> Self: + def identify_keys(self, save: bool = False) -> Self: """ Identifies potential primary keys in the dataset based on column profiles. This method relies on the 'column_profiles' result. """ - if "column_datatypes_l1" not in self.results or "column_datatypes_l2" not in self.results: + if not self.source_table_model.columns or any( + c.type is None or c.category is None for c in self.source_table_model.columns + ): raise RuntimeError("DataTypeIdentifierL1 and L2 must be run before KeyIdentifier.") - column_profiles: dict[str, ColumnProfile] = self.results["column_profiles"] - column_profiles_df = pd.DataFrame([col.model_dump() for col in column_profiles.values()]) + column_profiles_data = [] + for column in self.source_table_model.columns: + metrics = column.profiling_metrics + count = metrics.count if metrics.count is not None else 0 + null_count = metrics.null_count if metrics.null_count is not None else 0 + distinct_count = metrics.distinct_count if metrics.distinct_count is not None else 0 + column_profiles_data.append( + { + "column_name": column.name, + "table_name": self.name, + "datatype_l1": column.type, + "datatype_l2": column.category, + "count": count, + "null_count": null_count, + "distinct_count": distinct_count, + "uniqueness": distinct_count / count if count > 0 else 0.0, + "completeness": (count - null_count) / count if count > 0 else 0.0, + "sample_data": metrics.sample_data, + } + ) + column_profiles_df = pd.DataFrame(column_profiles_data) ki_model = KeyIdentificationLLM(profiling_data=column_profiles_df) ki_result = ki_model() output = KeyIdentificationOutput(**ki_result) - key = output.column_name + self.source_table_model.key = output.column_name or "" - if key is not None: - self.results["key"] = key - else: - self.results['key'] = None + if save: + self.save_yaml() return self - def profile(self) -> Self: + def profile(self, save: bool = False) -> Self: """ Profiles the dataset including table and columns and stores the result in the 'results' dictionary. This is a convenience method to run profiling on the raw dataframe. """ self.profile_table().profile_columns() + if save: + self.save_yaml() return self - def identify_datatypes(self) -> Self: + def identify_datatypes(self, save: bool = False) -> Self: """ Identifies the data types for the dataset and stores the result in the 'results' dictionary. This is a convenience method to run data type identification on the raw dataframe. """ self.identify_datatypes_l1().identify_datatypes_l2() + if save: + self.save_yaml() return self - def generate_glossary(self, domain: str = "") -> Self: + def generate_glossary(self, domain: str = "", save: bool = False) -> Self: """ Generates a business glossary for the dataset and stores the result in the 'results' dictionary. This method relies on the 'column_datatypes_l1' results. """ - if "column_datatypes_l1" not in self.results: + if not self.source_table_model.columns or any(c.type is None for c in self.source_table_model.columns): raise RuntimeError("DataTypeIdentifierL1 must be run before Business Glossary Generation.") - column_profiles: dict[str, ColumnProfile] = self.results["column_profiles"] - column_profiles_df = pd.DataFrame([col.model_dump() for col in column_profiles.values()]) + column_profiles_data = [] + for column in self.source_table_model.columns: + metrics = column.profiling_metrics + count = metrics.count if metrics.count is not None else 0 + null_count = metrics.null_count if metrics.null_count is not None else 0 + distinct_count = metrics.distinct_count if metrics.distinct_count is not None else 0 + column_profiles_data.append( + { + "column_name": column.name, + "table_name": self.name, + "datatype_l1": column.type, + "datatype_l2": column.category, + "count": count, + "null_count": null_count, + "distinct_count": distinct_count, + "uniqueness": distinct_count / count if count > 0 else 0.0, + "completeness": (count - null_count) / count if count > 0 else 0.0, + "sample_data": metrics.sample_data, + } + ) + column_profiles_df = pd.DataFrame(column_profiles_data) bg_model = BusinessGlossary(profiling_data=column_profiles_df) table_glossary, glossary_df = bg_model(table_name=self.name, domain=domain) - columns_glossary = [] - for _, row in glossary_df.iterrows(): - columns_glossary.append( - ColumnGlossary( - column_name=row["column_name"], - business_glossary=row.get("business_glossary", ""), - business_tags=row.get("business_tags", []), - ) - ) - glossary_output = BusinessGlossaryOutput( - table_name=self.name, table_glossary=table_glossary, columns=columns_glossary - ) - for column in glossary_output.columns: - column_profiles[column.column_name].business_glossary = column.business_glossary - column_profiles[column.column_name].business_tags = column.business_tags + self.source_table_model.description = table_glossary - self.results["business_glossary_and_tags"] = glossary_output - self.results["table_glossary"] = glossary_output.table_glossary + for _, row in glossary_df.iterrows(): + column = self._columns_map[row["column_name"]] + column.description = row.get("business_glossary", "") + column.tags = row.get("business_tags", []) + + if save: + self.save_yaml() return self def run(self, domain: str, save: bool = True) -> Self: @@ -218,117 +309,79 @@ def run(self, domain: str, save: bool = True) -> Self: return self - # FIXME - this is a temporary solution to save the results of the analysis - # need to use model while executing the pipeline def save_yaml(self, file_path: Optional[str] = None) -> None: if file_path is None: file_path = f"{self.name}.yml" file_path = os.path.join(settings.PROJECT_BASE, file_path) - column_profiles = self.results.get("column_profiles") - key = self.results.get("key") - - table_description = self.results.get("table_glossary") - table_tags = self.results.get("business_glossary_and_tags") - - if column_profiles is None or table_description is None or table_tags is None: - raise errors.NotFoundError( - "Column profiles not found in the dataset results. Ensure profiling steps were executed." - ) - - columns: list[Column] = [] - - for column_profile in column_profiles.values(): - column_profile = ColumnProfile.model_validate(column_profile) - column = Column( - name=column_profile.column_name, - description=column_profile.business_glossary, - type=column_profile.datatype_l1, - category=column_profile.datatype_l2, - tags=column_profile.business_tags, - profiling_metrics=ColumnProfilingMetrics( - count=column_profile.count, - null_count=column_profile.null_count, - distinct_count=column_profile.distinct_count, - sample_data=column_profile.sample_data, - ), - ) - columns.append(column) - details = self.adapter.get_details(self.data) - - table = SourceTables(name=self.name, description=table_description, columns=columns, details=details, key=key) - - source = Source(name="healthcare", description=table_description, schema="public", database="", table=table) + self.source_table_model.details = details + + # Store the source's last modification time + if isinstance(self.data, dict) and "path" in self.data and os.path.exists(self.data["path"]): + self.source_table_model.source_last_modified = os.path.getmtime(self.data["path"]) + + source = Source( + name="healthcare", + description=self.source_table_model.description, + schema="public", + database="", + table=self.source_table_model, + ) sources = {"sources": [json.loads(source.model_dump_json())]} # Save the YAML representation of the sources with open(file_path, "w") as file: yaml.dump(sources, file, sort_keys=False, default_flow_style=False) - + def to_df(self): return self.adapter.to_df(self.data, self.name) def load_from_yaml(self, file_path: str) -> None: - with open(file_path, "r") as file: - data = yaml.safe_load(file) - - source = data.get("sources", [])[0] - table = source.get("table", {}) - - self.results["table_glossary"] = table.get("description") - - columns = table.get("columns", []) - column_profiles = {} - for col in columns: - profiling_metrics = col.get("profiling_metrics", {}) - count = profiling_metrics.get("count", 0) - null_count = profiling_metrics.get("null_count", 0) - distinct_count = profiling_metrics.get("distinct_count", 0) - - column_profiles[col["name"]] = ColumnProfile( - column_name=col["name"], - business_name=string_standardization(col["name"]), - table_name=self.name, - business_glossary=col.get("description"), - datatype_l1=col.get("type"), - datatype_l2=col.get("category"), - business_tags=col.get("tags"), - count=count, - null_count=null_count, - distinct_count=distinct_count, - uniqueness=distinct_count / count if count > 0 else 0.0, - completeness=(count - null_count) / count if count > 0 else 0.0, - sample_data=profiling_metrics.get("sample_data"), - ts=time.time(), - ) - self.results["column_profiles"] = column_profiles - - self.results["business_glossary_and_tags"] = BusinessGlossaryOutput( - table_name=self.name, - table_glossary=self.results["table_glossary"], - columns=[ - ColumnGlossary( - column_name=col.column_name, - business_glossary=col.business_glossary, - business_tags=col.business_tags, - ) - for col in column_profiles.values() - ], - ) - - self.results["key"] = table.get('key') - - def to_df(self): - return self.adapter.to_df(self.data, self.name) + """Loads the dataset from a YAML file, checking for staleness.""" + with open(file_path, "r") as f: + yaml_data = yaml.safe_load(f) + if not self._is_yaml_stale(yaml_data): + self._populate_from_yaml(yaml_data) + + def reload_from_yaml(self, file_path: str) -> None: + """Forces a reload from a YAML file, bypassing staleness checks.""" + with open(file_path, "r") as f: + yaml_data = yaml.safe_load(f) + self._populate_from_yaml(yaml_data) @property def profiling_df(self): - column_profiles = self.results.get("column_profiles") - if column_profiles is None: + if not self.source_table_model.columns: return "

No column profiles available.

" - df = pd.DataFrame([col.model_dump() for col in column_profiles.values()]) + + column_profiles_data = [] + for column in self.source_table_model.columns: + metrics = column.profiling_metrics + if metrics: + count = metrics.count if metrics.count is not None else 0 + null_count = metrics.null_count if metrics.null_count is not None else 0 + distinct_count = metrics.distinct_count if metrics.distinct_count is not None else 0 + + column_profiles_data.append( + { + "column_name": column.name, + "table_name": self.name, + "business_name": string_standardization(column.name), + "datatype_l1": column.type, + "datatype_l2": column.category, + "business_glossary": column.description, + "business_tags": column.tags, + "count": count, + "null_count": null_count, + "distinct_count": distinct_count, + "uniqueness": distinct_count / count if count > 0 else 0.0, + "completeness": (count - null_count) / count if count > 0 else 0.0, + "sample_data": metrics.sample_data, + } + ) + df = pd.DataFrame(column_profiles_data) return df def _repr_html_(self): diff --git a/src/intugle/core/__init__.py b/src/intugle/core/__init__.py index a2cc6d6..207fc01 100644 --- a/src/intugle/core/__init__.py +++ b/src/intugle/core/__init__.py @@ -1 +1,2 @@ -from .settings import settings as settings \ No newline at end of file +from .settings import settings as settings +from .console import console \ No newline at end of file diff --git a/src/intugle/core/console.py b/src/intugle/core/console.py new file mode 100644 index 0000000..f8621da --- /dev/null +++ b/src/intugle/core/console.py @@ -0,0 +1,8 @@ +from rich.console import Console +from rich.style import Style + +console = Console() + +warning_style = Style(color="yellow", bold=True) +success_style = Style(color="green", bold=True) +danger_style = Style(color="red", bold=True) \ No newline at end of file diff --git a/src/intugle/knowledge_builder.py b/src/intugle/knowledge_builder.py index f5fc361..ba0b630 100644 --- a/src/intugle/knowledge_builder.py +++ b/src/intugle/knowledge_builder.py @@ -1,13 +1,11 @@ -import logging - -from intugle.analysis.models import DataSet -from intugle.link_predictor.predictor import LinkPredictor import asyncio +import logging import threading from typing import TYPE_CHECKING, Any, Awaitable, Dict, List, TypeVar from intugle.analysis.models import DataSet +from intugle.core.console import console, success_style from intugle.link_predictor.predictor import LinkPredictor from intugle.semantic_search import SemanticSearch @@ -63,7 +61,7 @@ def __init__(self, data_input: Dict[str, Any] | List[DataSet], domain: str = "") self._initialize_from_list(data_input) else: raise TypeError("Input must be a dictionary of named dataframes or a list of DataSet objects.") - + def _initialize_from_dict(self, data_dict: Dict[str, Any]): """Creates and processes DataSet objects from a dictionary of raw dataframes.""" for name, df in data_dict.items(): @@ -77,26 +75,47 @@ def _initialize_from_list(self, data_list: List[DataSet]): raise ValueError("DataSet objects provided in a list must have a 'name' attribute.") self.datasets[dataset.name] = dataset - def build(self, force_recreate: bool = False): - import os - - from intugle.core import settings - - # run analysis on all datasets + def profile(self, force_recreate: bool = False): + """Run profiling, datatype identification, and key identification for all datasets.""" + console.print("Starting profiling and key identification stage...", style="yellow") for dataset in self.datasets.values(): - file_path = os.path.join(settings.PROJECT_BASE, f"{dataset.name}.yml") - if os.path.exists(file_path) and not force_recreate: - print(f"Dataset {dataset.name} already processed. Loading from file.") - dataset.load_from_yaml(file_path) + # Check if this stage is already complete + if dataset.source_table_model.key is not None and not force_recreate: + print(f"Dataset '{dataset.name}' already profiled. Skipping.") continue - dataset.run(domain=self.domain, save=True) - # Initialize the predictor - self.link_predictor = LinkPredictor(list(self.datasets.values())) + console.print(f"Processing dataset: {dataset.name}", style="orange1") + dataset.profile(save=True) + dataset.identify_datatypes(save=True) + dataset.identify_keys(save=True) + console.print("Profiling and key identification complete.", style="bold green") - # Run the prediction + def predict_links(self): + """Run link prediction across all datasets.""" + console.print("Starting link prediction stage...", style="yellow") + self.link_predictor = LinkPredictor(list(self.datasets.values())) self.link_predictor.predict(save=True) self.links: list[PredictedLink] = self.link_predictor.links + console.print("Link prediction complete.", style="bold green") + + def generate_glossary(self, force_recreate: bool = False): + """Generate business glossary for all datasets.""" + console.print("Starting business glossary generation stage...", style="yellow") + for dataset in self.datasets.values(): + # Check if this stage is already complete + if dataset.source_table_model.description and not force_recreate: + console.print(f"Glossary for '{dataset.name}' already exists. Skipping.") + continue + + console.print(f"Generating glossary for dataset: {dataset.name}", style=success_style) + dataset.generate_glossary(domain=self.domain, save=True) + console.print("Business glossary generation complete.", style="bold green") + + def build(self, force_recreate: bool = False): + """Run the full end-to-end knowledge building pipeline.""" + self.profile(force_recreate=force_recreate) + self.predict_links() + self.generate_glossary(force_recreate=force_recreate) # Initialize semantic search try: diff --git a/src/intugle/link_predictor/predictor.py b/src/intugle/link_predictor/predictor.py index 057a2a7..a54ae83 100644 --- a/src/intugle/link_predictor/predictor.py +++ b/src/intugle/link_predictor/predictor.py @@ -18,6 +18,7 @@ TableProfiler, ) from intugle.core import settings +from intugle.core.console import console, warning_style from intugle.core.pipeline.link_prediction.lp import LinkPredictionAgentic from intugle.libs.smart_query_generator.utils.join import Join from intugle.models.resources.relationship import ( @@ -31,6 +32,11 @@ log = logging.getLogger(__name__) +class NoLinksFoundError(Exception): + """Custom exception raised when no links are found to save.""" + pass + + class LinkPredictor: """ Analyzes a collection of datasets to predict column links between all @@ -89,7 +95,7 @@ def _initialize_from_list(self, data_list: List[DataSet]): for dataset in data_list: if not dataset.name: raise ValueError("DataSet objects provided in a list must have a 'name' attribute.") - if "key" not in dataset.results: + if dataset.source_table_model.key is None: print(f"Dataset '{dataset.name}' is missing key identification. Running prerequisite analysis...") self._run_prerequisites(dataset) else: @@ -102,7 +108,11 @@ def _create_table_combination_id(self, table_a: str, table_b: str) -> str: return f"{assets[0]}--{assets[1]}" def _predict_for_pair( - self, name_a: str, dataset_a: DataSet, name_b: str, dataset_b: DataSet + self, + name_a: str, + dataset_a: DataSet, + name_b: str, + dataset_b: DataSet, ) -> List[PredictedLink]: """ Contains the core logic for finding links between TWO dataframes. @@ -111,19 +121,19 @@ def _predict_for_pair( table_combination = self._create_table_combination_id(name_a, name_b) if table_combination in self.already_executed_combo: log.warning(f"[!] Skipping already executed combination: {table_combination}") - return + return [] - dataset_a_column_profiles = [col.model_dump() for col in dataset_a.results["column_profiles"].values()] - dataset_b_column_profiles = [col.model_dump() for col in dataset_b.results["column_profiles"].values()] + dataset_a_column_profiles = dataset_a.profiling_df + dataset_b_column_profiles = dataset_b.profiling_df profiling_data = pd.concat( - [pd.DataFrame(dataset_a_column_profiles), pd.DataFrame(dataset_b_column_profiles)], ignore_index=True + [dataset_a_column_profiles, dataset_b_column_profiles], ignore_index=True ) primary_keys = [] - if dataset_a.results.get("key"): - primary_keys.append((name_a, dataset_a.results["key"])) - if dataset_b.results.get("key"): - primary_keys.append((name_b, dataset_b.results["key"])) + if dataset_a.source_table_model.key: + primary_keys.append((name_a, dataset_a.source_table_model.key)) + if dataset_b.source_table_model.key: + primary_keys.append((name_b, dataset_b.source_table_model.key)) pipeline = LinkPredictionAgentic( profiling_data=profiling_data, @@ -143,11 +153,27 @@ def _predict_for_pair( ] return pair_links - def predict(self, filename='__relationships__.yml', save: bool = False) -> Self: + def predict(self, filename='__relationships__.yml', save: bool = False, force_recreate: bool = False) -> Self: """ Iterates through all unique pairs of datasets, predicts the links for each pair, and returns the aggregated results. """ + relationships_file = os.path.join(settings.PROJECT_BASE, filename) + + if not force_recreate and os.path.exists(relationships_file): + is_stale = False + relationships_mtime = os.path.getmtime(relationships_file) + for dataset in self.datasets.values(): + dataset_yml = os.path.join(settings.PROJECT_BASE, f"{dataset.name}.yml") + if os.path.exists(dataset_yml) and os.path.getmtime(dataset_yml) > relationships_mtime: + is_stale = True + break + + if not is_stale: + console.print("Link predictions are up-to-date. Loading from cache.", style="green") + self.load_from_yaml(relationships_file) + return self + all_links: List[PredictedLink] = [] dataset_names = list(self.datasets.keys()) @@ -165,13 +191,21 @@ def predict(self, filename='__relationships__.yml', save: bool = False) -> Self: print("No links found for this pair.") self.links = all_links + + if len(self.links) == 0: + console.print("No links found between any datasets.", style=warning_style) + return self + if save: self.save_yaml(file_path=filename) return self - def get_links_df(self): - ... + def get_links_df(self) -> pd.DataFrame: + """Returns the predicted links as a pandas DataFrame.""" + if not self.links: + return pd.DataFrame() + return pd.DataFrame([link.model_dump() for link in self.links]) def show_graph(self): links = [link.relationship.link for link in self.links] @@ -186,7 +220,7 @@ def save_yaml(self, file_path: str) -> None: file_path = os.path.join(settings.PROJECT_BASE, file_path) if len(self.links) == 0: - raise ValueError("No links found to save.") + raise NoLinksFoundError("No links found to save.") relationships = {"relationships": [json.loads(link.relationship.model_dump_json()) for link in self.links]} @@ -194,6 +228,23 @@ def save_yaml(self, file_path: str) -> None: with open(file_path, "w") as file: yaml.dump(relationships, file, sort_keys=False, default_flow_style=False) + def load_from_yaml(self, file_path: str) -> None: + """Loads link predictions from a YAML file.""" + with open(file_path, "r") as f: + data = yaml.safe_load(f) + + relationships = data.get("relationships", []) + loaded_links = [] + for rel in relationships: + link = PredictedLink( + from_dataset=rel["source"]["table"], + from_column=rel["source"]["column"], + to_dataset=rel["target"]["table"], + to_column=rel["target"]["column"], + ) + loaded_links.append(link) + self.links = loaded_links + class LinkPredictionSaver: @classmethod diff --git a/src/intugle/models/resources/model.py b/src/intugle/models/resources/model.py index 9525385..4cff1da 100644 --- a/src/intugle/models/resources/model.py +++ b/src/intugle/models/resources/model.py @@ -11,6 +11,7 @@ class ColumnProfilingMetrics(SchemaBase): null_count: Optional[int] = None distinct_count: Optional[int] = None sample_data: Optional[List[Any]] = Field(default_factory=list) + dtype_sample: Optional[List[Any]] = Field(default_factory=list, exclude=True) class Column(SchemaBase): diff --git a/src/intugle/models/resources/source.py b/src/intugle/models/resources/source.py index 2b302ad..7423698 100644 --- a/src/intugle/models/resources/source.py +++ b/src/intugle/models/resources/source.py @@ -11,10 +11,11 @@ class SourceTables(SchemaBase): name: str description: str tags: Optional[List[str]] = Field(default_factory=list) - details: Optional[dict] + details: Optional[dict] = None columns: List[Column] = Field(default_factory=list) profiling_metrics: Optional[ModelProfilingMetrics] = None key: Optional[str] = None + source_last_modified: Optional[float] = None class Source(BaseResource): diff --git a/src/intugle/parser/table_schema.py b/src/intugle/parser/table_schema.py index 82a2f87..9929a33 100644 --- a/src/intugle/parser/table_schema.py +++ b/src/intugle/parser/table_schema.py @@ -33,7 +33,7 @@ def generate_table_schema(self, table_name: str) -> str: # Iterate through the columns of the table and create the column definitions columns_statements = [ - f"{column.name} {column.type}, -- {column.description}" for column in table_detail.table.columns + f"\"{column.name}\" {column.type}, -- {column.description}" for column in table_detail.table.columns ] # join the column definitions into a single string diff --git a/tests/analysis/test_business_glossary.py b/tests/analysis/test_business_glossary.py index 244503a..6e36cda 100644 --- a/tests/analysis/test_business_glossary.py +++ b/tests/analysis/test_business_glossary.py @@ -1,6 +1,5 @@ import pandas as pd -from intugle.adapters.models import BusinessGlossaryOutput, ColumnGlossary from intugle.analysis.models import DataSet from intugle.analysis.pipeline import Pipeline from intugle.analysis.steps import ( @@ -34,34 +33,19 @@ def test_business_glossary_generator(): BusinessGlossaryGenerator(domain=domain), ]) - # 2. Initialize DataSet - dataset = DataSet(df, name=table_name) - - # 3. Run prerequisite steps (TableProfiler, ColumnProfiler) + # 2. Run the pipeline dataset = pipeline.run(df, table_name) - # 4. Assert the results - assert "business_glossary_and_tags" in dataset.results - assert "table_glossary" in dataset.results - - glossary_output = dataset.results["business_glossary_and_tags"] - table_glossary_str = dataset.results["table_glossary"] - - assert isinstance(glossary_output, BusinessGlossaryOutput) - assert glossary_output.table_name == table_name - assert isinstance(table_glossary_str, str) - assert len(table_glossary_str) > 0 - assert len(glossary_output.columns) == len(df.columns) + # 3. Assert the results + assert dataset.source_table_model.description is not None + assert len(dataset.source_table_model.description) > 0 + assert len(dataset.source_table_model.columns) == len(df.columns) # Check a specific column's glossary entry - product_id_glossary = next((col for col in glossary_output.columns if col.column_name == "product_id"), None) - assert product_id_glossary is not None - assert isinstance(product_id_glossary, ColumnGlossary) - assert isinstance(product_id_glossary.business_glossary, str) - assert len(product_id_glossary.business_glossary) > 0 - assert len(product_id_glossary.business_tags) > 0 - - # Verify that column profiles were updated - assert dataset.results["column_profiles"]["product_id"].business_glossary is not None - assert len(dataset.results["column_profiles"]["product_id"].business_tags) > 0 + product_id_column = dataset._columns_map.get("product_id") + assert product_id_column is not None + assert product_id_column.description is not None + assert len(product_id_column.description) > 0 + assert product_id_column.tags is not None + assert len(product_id_column.tags) > 0 dataset.save_yaml() diff --git a/tests/analysis/test_datatype_identification.py b/tests/analysis/test_datatype_identification.py index df3af76..dc86013 100644 --- a/tests/analysis/test_datatype_identification.py +++ b/tests/analysis/test_datatype_identification.py @@ -26,30 +26,10 @@ def test_datatype_identification_l1_end_to_end(): analysis_results = pipeline.run(COMPLEX_DF, DF_NAME) # Check the final output of the L1 step - datatype_l1_results = analysis_results.results.get("column_datatypes_l1") - assert datatype_l1_results is not None - assert len(datatype_l1_results) == 3 - - # Create a dictionary for easy lookup - results_map = {res.column_name: res for res in datatype_l1_results} - - # Assertions for 'user_id' - assert 'user_id' in results_map - assert results_map['user_id'].datatype_l1 == 'integer' - - # Assertions for 'product_name' - assert 'product_name' in results_map - assert results_map['product_name'].datatype_l1 == 'close_ended_text' - - # Assertions for 'price' - assert 'price' in results_map - assert results_map['price'].datatype_l1 == 'float' - - # Also check that the original column profiles were updated - column_profiles = analysis_results.results.get("column_profiles") - assert column_profiles['user_id'].datatype_l1 == 'integer' - assert column_profiles['product_name'].datatype_l1 == 'close_ended_text' - assert column_profiles['price'].datatype_l1 == 'float' + columns_map = analysis_results._columns_map + assert columns_map['user_id'].type == 'integer' + assert columns_map['product_name'].type == 'close_ended_text' + assert columns_map['price'].type == 'float' def test_datatype_identification_l2_end_to_end(): @@ -65,28 +45,8 @@ def test_datatype_identification_l2_end_to_end(): analysis_results = pipeline.run(COMPLEX_DF, DF_NAME) - # Check the final output of the L1 step - datatype_l2_results = analysis_results.results.get("column_datatypes_l2") - assert datatype_l2_results is not None - assert len(datatype_l2_results) == 3 - - # Create a dictionary for easy lookup - results_map = {res.column_name: res for res in datatype_l2_results} - - # Assertions for 'user_id' - assert 'user_id' in results_map - assert results_map['user_id'].datatype_l2 == 'dimension' - - # Assertions for 'product_name' - assert 'product_name' in results_map - assert results_map['product_name'].datatype_l2 == 'dimension' - - # Assertions for 'price' - assert 'price' in results_map - assert results_map['price'].datatype_l2 == 'measure' - - # Also check that the original column profiles were updated - column_profiles = analysis_results.results.get("column_profiles") - assert column_profiles['user_id'].datatype_l2 == 'dimension' - assert column_profiles['product_name'].datatype_l2 == 'dimension' - assert column_profiles['price'].datatype_l2 == 'measure' \ No newline at end of file + # Check the final output of the L2 step + columns_map = analysis_results._columns_map + assert columns_map['user_id'].category == 'dimension' + assert columns_map['product_name'].category == 'dimension' + assert columns_map['price'].category == 'measure' diff --git a/tests/analysis/test_high_level.py b/tests/analysis/test_high_level.py index 9564b3b..f1fcfb2 100644 --- a/tests/analysis/test_high_level.py +++ b/tests/analysis/test_high_level.py @@ -22,16 +22,13 @@ def test_profile(sample_dataframe): dataset = DataSet(sample_dataframe, name="test_table") dataset.profile() - assert "table_profile" in dataset.results - table_profile = dataset.results["table_profile"] - assert table_profile is not None - assert table_profile.count == 5 - assert set(table_profile.columns) == {"user_id", "product_name", "price", "purchase_date"} + table_model = dataset.source_table_model + assert table_model.profiling_metrics is not None + assert table_model.profiling_metrics.count == 5 + assert len(table_model.columns) == 4 + assert {col.name for col in table_model.columns} == {"user_id", "product_name", "price", "purchase_date"} - assert "column_profiles" in dataset.results - column_profiles = dataset.results["column_profiles"] - assert column_profiles is not None - assert len(column_profiles) == 4 + assert all(col.profiling_metrics is not None for col in table_model.columns) def test_identify_datatypes(sample_dataframe): @@ -40,15 +37,9 @@ def test_identify_datatypes(sample_dataframe): dataset.profile() dataset.identify_datatypes() - assert "column_datatypes_l1" in dataset.results - column_datatypes_l1 = dataset.results["column_datatypes_l1"] - assert column_datatypes_l1 is not None - assert len(column_datatypes_l1) == 4 - - assert "column_datatypes_l2" in dataset.results - column_datatypes_l2 = dataset.results["column_datatypes_l2"] - assert column_datatypes_l2 is not None - assert len(column_datatypes_l2) == 4 + table_model = dataset.source_table_model + assert all(col.type is not None for col in table_model.columns) + assert all(col.category is not None for col in table_model.columns) def test_identify_keys(sample_dataframe): @@ -58,9 +49,7 @@ def test_identify_keys(sample_dataframe): dataset.identify_datatypes() dataset.identify_keys() - assert "key" in dataset.results - key = dataset.results["key"] - assert key is not None + assert dataset.source_table_model.key is not None def test_generate_glossary(sample_dataframe): @@ -70,12 +59,9 @@ def test_generate_glossary(sample_dataframe): dataset.identify_datatypes() dataset.generate_glossary(domain="ecommerce") - assert "business_glossary_and_tags" in dataset.results - glossary = dataset.results["business_glossary_and_tags"] - assert glossary is not None - assert "table_glossary" in dataset.results - table_glossary = dataset.results["table_glossary"] - assert table_glossary is not None + table_model = dataset.source_table_model + assert table_model.description is not None + assert all(col.description is not None for col in table_model.columns) def test_save_yaml(sample_dataframe, tmp_path): diff --git a/tests/analysis/test_key_identification.py b/tests/analysis/test_key_identification.py index c740242..153bc52 100644 --- a/tests/analysis/test_key_identification.py +++ b/tests/analysis/test_key_identification.py @@ -40,12 +40,6 @@ def test_key_identification_end_to_end(): analysis_results = pipeline.run(KEY_TEST_DF, DF_NAME) # Check the final output of the KeyIdentifier step - identified_key = analysis_results.results.get("key") + identified_key = analysis_results.source_table_model.key assert identified_key is not None - - # The result should identify 'order_id' as the primary key. - # Based on the implementation, the output is a KeyIdentificationOutput object - # which contains the identified key information. - # We expect one identified key. - assert identified_key == "order_id" diff --git a/tests/analysis/test_pipeline.py b/tests/analysis/test_pipeline.py index 7be089c..c884738 100644 --- a/tests/analysis/test_pipeline.py +++ b/tests/analysis/test_pipeline.py @@ -3,6 +3,7 @@ from intugle.analysis.pipeline import Pipeline from intugle.analysis.steps import ColumnProfiler, TableProfiler +from intugle.core.utilities.processing import string_standardization # --- Test Data --- # A more complex and realistic DataFrame for testing. @@ -26,18 +27,12 @@ def test_pipeline_with_complex_data(): """ pipeline = Pipeline([TableProfiler()]) analysis_results = pipeline.run(COMPLEX_DF, DF_NAME) - profile = analysis_results.results.get("table_profile") + table_model = analysis_results.source_table_model - assert profile is not None - assert profile.count == 10 - assert profile.columns == ['user_id', 'product_name', 'price', 'purchase_date', 'is returned'] - assert profile.dtypes == { - 'user_id': 'float', # Floats because of NaN - 'product_name': 'string', - 'price': 'float', - 'purchase_date': 'date & time', - 'is returned': 'string' # Objects (mixed types) are treated as strings - } + assert table_model.profiling_metrics is not None + assert table_model.profiling_metrics.count == 10 + assert len(table_model.columns) == 5 + assert {col.name for col in table_model.columns} == {'user_id', 'product_name', 'price', 'purchase_date', 'is returned'} def test_column_profiling_with_complex_data(): @@ -49,44 +44,34 @@ def test_column_profiling_with_complex_data(): ColumnProfiler() ]) analysis_results = pipeline.run(COMPLEX_DF, DF_NAME) - column_profiles = analysis_results.results.get("column_profiles") - assert column_profiles is not None - assert len(column_profiles) == 5 + columns_map = analysis_results._columns_map + assert len(columns_map) == 5 + # --- Assertions for 'user_id' --- - user_id_profile = column_profiles.get('user_id') + user_id_profile = columns_map.get('user_id').profiling_metrics assert user_id_profile is not None - assert user_id_profile.column_name == 'user_id' - assert user_id_profile.table_name == DF_NAME assert user_id_profile.count == 10 assert user_id_profile.null_count == 1 assert user_id_profile.distinct_count == 7 # 101, 102, 103, 104, 105, 106, 107 -> 101 is repeated - assert user_id_profile.uniqueness == 0.7 # 7 distinct out of 10 total - assert user_id_profile.completeness == 0.9 # 9 non-null out of 10 total # --- Assertions for 'product_name' --- - product_profile = column_profiles.get('product_name') + product_profile = columns_map.get('product_name').profiling_metrics assert product_profile is not None assert product_profile.count == 10 assert product_profile.null_count == 1 assert product_profile.distinct_count == 6 # Laptop, Mouse, Keyboard, Monitor, Webcam, HDMI Cable - assert user_id_profile.uniqueness == 0.7 # 7 distinct out of 10 total - assert user_id_profile.completeness == 0.9 # 9 non-null out of 10 total # --- Assertions for 'price' --- - price_profile = column_profiles.get('price') + price_profile = columns_map.get('price').profiling_metrics assert price_profile is not None assert price_profile.count == 10 assert price_profile.null_count == 1 assert price_profile.distinct_count == 8 # 1200.50 is repeated - assert user_id_profile.uniqueness == 0.7 # 7 distinct out of 10 total - assert user_id_profile.completeness == 0.9 # 9 non-null out of 10 total # --- Assertions for 'is returned' --- - returned_profile = column_profiles.get('is returned') + returned_profile = columns_map.get('is returned').profiling_metrics assert returned_profile is not None assert returned_profile.count == 10 assert returned_profile.null_count == 1 assert returned_profile.distinct_count == 2 # True, False - assert returned_profile.business_name == "is_returned" - assert user_id_profile.uniqueness == 0.7 # 7 distinct out of 10 total - assert user_id_profile.completeness == 0.9 # 9 non-null out of 10 total \ No newline at end of file + assert string_standardization("is returned") == "is_returned" diff --git a/tests/link_predictor/test_predictor.py b/tests/link_predictor/test_predictor.py index 6589ab6..87442f0 100644 --- a/tests/link_predictor/test_predictor.py +++ b/tests/link_predictor/test_predictor.py @@ -49,11 +49,10 @@ def test_predictor_with_dict_input(mock_predict_for_pair): assert "orders" in predictor.datasets assert isinstance(predictor.datasets["customers"], DataSet) # Check that the prerequisite step was completed - assert "key" in predictor.datasets["customers"].results - assert predictor.datasets["customers"].results["key"] == "id" + assert predictor.datasets["customers"].source_table_model.key == "id" # 4. Run prediction - results = predictor.predict() + results = predictor.predict(force_recreate=True) # 5. Verify that the mocked prediction was called and results are correct assert mock_predict_for_pair.call_count == 1 @@ -70,7 +69,7 @@ def test_predictor_with_list_input(mock_predict_for_pair): customers_df = pd.DataFrame({"id": [1, 2, 3]}) processed_dataset = DataSet(customers_df, name="customers") # Manually add the key to simulate it being pre-analyzed - processed_dataset.results["key"] = "id" + processed_dataset.source_table_model.key = "id" # 2. Prepare a raw DataSet that needs analysis orders_df = pd.DataFrame({"order_id": [101, 102], "customer_id": [1, 3]}) @@ -82,12 +81,12 @@ def test_predictor_with_list_input(mock_predict_for_pair): # 4. Verify that both datasets are now fully processed assert "customers" in predictor.datasets assert "orders" in predictor.datasets - assert "key" in predictor.datasets["customers"].results - assert "key" in predictor.datasets["orders"].results - assert predictor.datasets["orders"].results["key"] == "order_id" + assert predictor.datasets["customers"].source_table_model.key is not None + assert predictor.datasets["orders"].source_table_model.key is not None + assert predictor.datasets["orders"].source_table_model.key == "order_id" # 5. Run prediction - results = predictor.predict() + results = predictor.predict(force_recreate=True) # 6. Verify results assert mock_predict_for_pair.call_count == 1 @@ -134,7 +133,7 @@ def test_predictor_end_to_end_complex(): predictor = LinkPredictor(datasets) # 3. Run the prediction - results = predictor.predict() + results = predictor.predict(force_recreate=True) # 4. Assert that the correct links were found assert len(results.links) >= 2, "Expected at least two links to be found" diff --git a/uv.lock b/uv.lock index bc1642e..20bb35b 100644 --- a/uv.lock +++ b/uv.lock @@ -1458,13 +1458,12 @@ wheels = [ [[package]] name = "intugle" -version = "0.1.5" +version = "0.1.6" source = { editable = "." } dependencies = [ { name = "asyncpg" }, { name = "duckdb" }, { name = "fastapi", extra = ["standard"] }, - { name = "ipykernel" }, { name = "langchain", extra = ["anthropic", "google-genai", "openai"] }, { name = "langchain-community" }, { name = "langchain-deepseek" }, @@ -1488,6 +1487,7 @@ dependencies = [ { name = "python-dotenv" }, { name = "pyyaml" }, { name = "qdrant-client" }, + { name = "rich" }, { name = "scikit-learn" }, { name = "symspellpy" }, { name = "trieregex" }, @@ -1517,7 +1517,6 @@ requires-dist = [ { name = "asyncpg", specifier = ">=0.30.0" }, { name = "duckdb", specifier = ">=1.3.2" }, { name = "fastapi", extras = ["standard"], specifier = ">=0.116.1" }, - { name = "ipykernel", specifier = ">=6.30.1" }, { name = "langchain", extras = ["anthropic", "google-genai", "openai"], specifier = ">=0.3.27" }, { name = "langchain-community", specifier = ">=0.3.21" }, { name = "langchain-deepseek", specifier = ">=0.1.4" }, @@ -1539,6 +1538,7 @@ requires-dist = [ { name = "python-dotenv", specifier = ">=1.1.1" }, { name = "pyyaml", specifier = ">=6.0.2" }, { name = "qdrant-client", specifier = ">=1.15.1" }, + { name = "rich", specifier = ">=14.1.0" }, { name = "scikit-learn", specifier = "<=1.7.1" }, { name = "symspellpy", specifier = ">=6.9.0" }, { name = "trieregex", specifier = ">=1.0.0" },