File size: 3,545 Bytes
d21188b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import logging
import pandas as pd
from typing import Dict
from src.agents.agents import dataset_description_agent, data_context_gen
from src.utils.model_registry import mid_lm
from src.utils.logger import Logger
import dspy
# Initialize logger
logger = Logger("dataset_description_generator", see_time=False, console_log=False)
def generate_dataset_description(datasets: Dict[str, pd.DataFrame], existing_description: str = "", dataset_names: list = None) -> str:
"""
Generate AI-powered description for datasets
Args:
datasets: Dictionary of dataset names to DataFrames
existing_description: Existing description to improve upon (optional)
dataset_names: List of dataset names to use in the description format (optional)
Returns:
Generated description string with proper exact_python_name formatting
"""
try:
if not datasets or len(datasets) == 0:
return existing_description
# Build dataset view for description generation
dataset_view = ""
count = 0
for table_name, table_df in datasets.items():
head_data = table_df.head(3)
columns = [{col: str(head_data[col].dtype)} for col in head_data.columns]
dataset_view += f"exact_table_name={table_name}\n:columns:{str(columns)}\n{head_data.to_markdown()}\n"
count += 1
# Generate description using AI
with dspy.context(lm=mid_lm):
if count == 1:
data_context = dspy.Predict(dataset_description_agent)(
existing_description=existing_description,
dataset=dataset_view
)
generated_desc = data_context.description
elif count > 1:
data_context = dspy.Predict(data_context_gen)(
user_description=existing_description,
dataset_view=dataset_view
)
generated_desc = data_context.data_context
else:
generated_desc = existing_description
# Format the description with exact_python_name for all datasets
if dataset_names and len(dataset_names) > 0:
if len(dataset_names) == 1:
# Single dataset format
formatted_desc = f" exact_python_name: `{dataset_names[0]}` Dataset: {generated_desc}"
else:
# Multiple datasets format - list all dataset names
names_list = ", ".join([f"`{name}`" for name in dataset_names])
formatted_desc = f" exact_python_name: {names_list} Dataset: {generated_desc}"
else:
# Fallback to original format if no dataset names provided
dataset_keys = list(datasets.keys())
if len(dataset_keys) == 1:
formatted_desc = f" exact_python_name: `{dataset_keys[0]}` Dataset: {generated_desc}"
else:
names_list = ", ".join([f"`{name}`" for name in dataset_keys])
formatted_desc = f" exact_python_name: {names_list} Dataset: {generated_desc}"
logger.log_message(f"Successfully generated dataset description for {count} dataset(s)", level=logging.INFO)
return formatted_desc
except Exception as e:
logger.log_message(f"Failed to generate dataset description: {str(e)}", level=logging.WARNING)
# Return existing description if generation fails
return existing_description
|