dslighting 1.3.9__py3-none-any.whl

dsat/services/data_analyzer.py

@@ -0,0 +1,387 @@
+# dsat/services/data_analyzer.py
+
+import logging
+import traceback
+from pathlib import Path
+import pandas as pd
+from typing import Optional, List
+from itertools import islice
+
+from dsat.models.task import TaskType
+
+logger = logging.getLogger(__name__)
+
+
+def generate_file_tree(
+    start_path: Path,
+    max_depth: int = 5,
+    max_files: int = 100,
+    max_items_per_dir: int = 20,
+    display_root_name: Optional[str] = None
+) -> str:
+    """
+    Generates a textual representation of the file tree with intelligent truncation.
+
+    This version prevents a single large directory from consuming the entire file limit,
+    and it filters out common noise files.
+
+    Args:
+        start_path: The root directory to start the tree from.
+        max_depth: The maximum depth to traverse into directories.
+        max_files: The global limit for the total number of files to display.
+        max_items_per_dir: The maximum number of items (files and dirs) to show per directory.
+        display_root_name: An optional name to display for the root directory.
+    """
+    tree = []
+    start_path = Path(start_path)
+    if not start_path.exists():
+        return f"Directory not found: {start_path}"
+
+
+    base_name = display_root_name if display_root_name is not None else start_path.name
+    file_count = 0
+    global_limit_reached = False
+
+    def _walk(path: Path, prefix: str, depth: int):
+        nonlocal file_count, global_limit_reached
+        
+        if depth > max_depth or global_limit_reached:
+            return
+
+        try:
+            # Avoid listing huge directories fully (e.g., image folders).
+            # We only sample a small number of entries for display.
+            sampled = list(islice(path.iterdir(), max_items_per_dir + 1))
+        except OSError as e:
+            logger.warning(f"Error reading directory {path}: {e}")
+            tree.append(f"{prefix}└── [Error reading directory]")
+            return
+
+        truncated_in_dir = len(sampled) > max_items_per_dir
+        if truncated_in_dir:
+            display_items = sampled[: max(1, max_items_per_dir // 2)]
+        else:
+            display_items = sampled
+        display_items = sorted(display_items, key=lambda p: p.name)
+
+        pointers = ['├── '] * (len(display_items) - 1) + ['└── ']
+        # If we truncated this directory, the last visible item is not the true last item
+        if truncated_in_dir:
+            pointers[-1] = '├── '
+
+        for pointer, sub_path in zip(pointers, display_items):
+            if global_limit_reached:
+                return
+
+            if not sub_path.is_dir():
+                # Check global file limit *before* adding the next file
+                if file_count >= max_files:
+                    global_limit_reached = True
+                    return
+                file_count += 1
+
+            display_name = sub_path.name + ('/' if sub_path.is_dir() else '')
+            tree.append(f"{prefix}{pointer}{display_name}")
+
+            if sub_path.is_dir():
+                extension = '│   ' if pointer == '├── ' else '    '
+                _walk(sub_path, prefix=prefix + extension, depth=depth + 1)
+
+        if truncated_in_dir:
+            tree.append(f"{prefix}└── [... more items truncated ...]")
+
+    tree.append(f"{base_name}/")
+    _walk(start_path, prefix="", depth=1)
+
+    if global_limit_reached:
+        tree.append(f"\n[... Truncated. Total file limit ({max_files}) reached ...]")
+
+    return "\n".join(tree)
+
+
+class DataAnalyzer:
+    """
+    A centralized service for analyzing input data directories and generating
+    a comprehensive textual overview for the Agent.
+    """
+    def analyze(self, data_dir: Path, output_filename: str, task_type: Optional[TaskType] = None, optimization_context: bool = False) -> str:
+        """
+        Analyzes the data directory and returns a formatted overview string,
+        including critical I/O instructions.
+        """
+        report = self.analyze_data(data_dir, task_type)
+        report += self.generate_io_instructions(output_filename, optimization_context)
+        return report
+
+    def analyze_data(self, data_dir: Path, task_type: Optional[TaskType] = None) -> str:
+        """
+        Analyzes the data directory and returns ONLY the data report (structure, format, etc.).
+        Does NOT include I/O instructions.
+        """
+        if not data_dir or not data_dir.exists() or not data_dir.is_dir():
+            logger.error(f"Data directory issue during analysis: {data_dir}")
+            return "Error: Input data directory not found, not provided, or is not a directory."
+
+        report = "\n\n--- COMPREHENSIVE DATA REPORT ---\n\n"
+
+        # 1. Analyze File Structure (Universal)
+        report += self._analyze_structure(data_dir)
+
+        report += self._analyze_data_schema(data_dir)
+
+        # 3. Task-Specific Analysis
+        if task_type == "kaggle":
+            submission_analysis = self._analyze_kaggle_submission_format(data_dir)
+            if submission_analysis:
+                report += f"## Submission Format Requirements\n{submission_analysis}\n\n"
+        
+        # Remove the call to _generate_io_instructions from here
+        return report
+
+    def generate_io_instructions(self, output_filename: str, optimization_context: bool = False) -> str:
+        """
+        Generate standardized I/O instructions reflecting that CWD is the input directory.
+        """
+        output_suffix = Path(output_filename).suffix.lower()
+
+        input_instructions = (
+            "1. **INPUT DATA:**\n"
+            "   - All input files are located in the **current working directory** (./).\n"
+            "   - Example: Use `pd.read_csv('train.csv')`."
+        )
+
+        if optimization_context:
+            example_write = ""
+            if output_suffix == ".csv":
+                example_write = "   - **Example Write (Conceptual):** `final_df.to_csv(output_path.name, index=False)`"
+            elif output_suffix == ".npy":
+                example_write = "   - **Example Write (Conceptual):** `np.save(output_path.name, preds)`"
+
+            output_instructions = (
+                "2. **OUTPUT FILE (Dynamic Workflow Context):**\n"
+                "   - Your workflow's `solve` method receives an `output_path` argument.\n"
+                "   - You MUST save your final submission file using the filename derived from this argument (e.g., `output_path.name`).\n"
+                "   - The file must be saved in the current working directory (./).\n"
+                + (f"\n{example_write}" if example_write else "")
+            )
+        else:
+            example_write = ""
+            if output_suffix == ".csv":
+                example_write = f"   - **Correct Example:** `submission_df.to_csv('{output_filename}', index=False)`"
+            elif output_suffix == ".npy":
+                example_write = f"   - **Correct Example:** `np.save('{output_filename}', preds)`"
+
+            output_instructions = (
+                f"2. **OUTPUT FILE:**\n"
+                f"   - You MUST save your final submission file to the **current working directory** (./).\n"
+                f"   - The required output filename is: `{output_filename}`\n"
+                + (f"{example_write}\n" if example_write else "")
+            )
+
+        return f"""
+--- CRITICAL I/O REQUIREMENTS ---
+
+You MUST follow these file system rules precisely. Failure to do so will cause a fatal error.
+
+{input_instructions}
+
+{output_instructions}
+
+**IMPORTANT:** These path requirements are non-negotiable and must be followed exactly.
+"""
+
+    def _analyze_structure(self, data_dir: Path) -> str:
+        """Generates the file tree representation."""
+        try:
+            tree_output = generate_file_tree(data_dir, display_root_name=".")
+            return f"## Directory Structure (Current Working Directory)\n```text\n{tree_output}\n```\n\n"
+        except Exception as e:
+            logger.error(f"Failed to generate file tree for {data_dir}: {traceback.format_exc()}")
+            return f"## Directory Structure\nError analyzing structure: {traceback.format_exc()}\n\n"
+
+    def _analyze_kaggle_submission_format(self, data_dir: Path) -> str:
+        """
+        Analyzes the sample submission file for Kaggle tasks and extracts a
+        prescriptive schema.
+        """
+        sample_submission_file = self._find_sample_submission(data_dir)
+
+        if not sample_submission_file:
+            return ""
+
+        try:
+            # Read the sample submission file to get format details
+            sample_df = pd.read_csv(sample_submission_file)
+
+            # Get the first few rows and data types
+            head_info = sample_df.head().to_string(index=False)
+            dtypes_info = sample_df.dtypes.to_string()
+            
+            required_columns = sample_df.columns.tolist()
+            columns_instruction = f"""
+**Required Submission Columns:**
+Your submission file MUST contain the following columns in this exact order:
+```
+{required_columns}
+```
+This is a strict requirement for the submission to be graded correctly. The grading system uses the non-prediction columns (like 'Comment' or an 'id') to match your predictions against the ground truth.
+"""
+
+            return f"""
+**CRITICAL:** Your final submission file MUST EXACTLY match the format of the sample submission file provided (`{sample_submission_file.name}`).
+This includes the column names, column order, and data types. Failure to adhere to this format will result in a score of zero.
+
+{columns_instruction}
+
+**Format Details:**
+*First 5 rows:*
+```text
+{head_info}
+```
+
+*Data types:*
+```text
+{dtypes_info}
+```
+
+"""
+        except Exception as e:
+            # Fallback if pandas fails
+            logger.warning(f"Could not read sample submission file '{sample_submission_file}' for detailed analysis: {traceback.format_exc()}")
+            return f"""
+**CRITICAL:** Your final submission file MUST match the format of the sample submission file: `{sample_submission_file.name}`.
+(Note: Automatic format analysis failed, please inspect the file manually).
+"""
+
+    def _find_sample_submission(self, data_dir: Path) -> Optional[Path]:
+        """Helper to locate the sample submission file."""
+        try:
+            for file in data_dir.iterdir():
+                file_name_lower = file.name.lower()
+                if "sample" in file_name_lower and "submission" in file_name_lower and file_name_lower.endswith(".csv"):
+                    return file
+        except Exception as e:
+            logger.warning(f"Could not scan data directory '{data_dir}': {traceback.format_exc()}")
+        return None
+
+    def _analyze_data_schema(self, data_dir: Path) -> str:
+        """
+        Analyzes the schema of potential training and testing files to provide a
+        structured overview of columns, data types, missing values, and cardinality.
+        This helps the agent make better decisions about preprocessing.
+        """
+        report_parts = []
+        # Define supported extensions and keywords for more robust discovery
+        SUPPORTED_EXTENSIONS = ('.csv', '.tsv', '.parquet', '.xlsx')
+        KEYWORDS = ('train', 'test', 'val', 'eval', 'sample', 'submission', 'sub', 'data')
+
+        # Fast path: most prepared competitions keep train/test-like tables at the root.
+        files_to_analyze: List[Path] = []
+        try:
+            root_files = [p for p in data_dir.iterdir() if p.is_file()]
+        except OSError:
+            root_files = []
+
+        for p in root_files:
+            if p.suffix.lower() not in SUPPORTED_EXTENSIONS:
+                continue
+            if any(keyword in p.stem.lower() for keyword in KEYWORDS):
+                files_to_analyze.append(p)
+
+        # Fallback: if still nothing, just take the first few supported files at root
+        if not files_to_analyze and root_files:
+            files_to_analyze = [p for p in root_files if p.suffix.lower() in SUPPORTED_EXTENSIONS][:3]
+
+        # Deep discoveryFallback: bounded recursive search (avoid walking huge image folders).
+        if not files_to_analyze:
+            max_depth = 3
+            max_dirs = 200
+            max_files = 8
+            per_dir_limit = 200
+
+            queue: List[tuple[Path, int]] = [(data_dir, 0)]
+            visited = 0
+            while queue and visited < max_dirs and len(files_to_analyze) < max_files:
+                current, depth = queue.pop(0)
+                visited += 1
+                if depth > max_depth:
+                    continue
+                try:
+                    sampled = list(islice(current.iterdir(), per_dir_limit + 1))
+                except OSError:
+                    continue
+
+                truncated = len(sampled) > per_dir_limit
+                entries = sampled[:per_dir_limit] if truncated else sampled
+
+                for entry in entries:
+                    if entry.is_dir():
+                        # If this directory is huge, avoid descending further.
+                        if truncated and depth >= 1:
+                            continue
+                        queue.append((entry, depth + 1))
+                        continue
+                    if entry.suffix.lower() not in SUPPORTED_EXTENSIONS:
+                        continue
+                    if not any(keyword in entry.stem.lower() for keyword in KEYWORDS):
+                        continue
+                    files_to_analyze.append(entry)
+                    if len(files_to_analyze) >= max_files:
+                        break
+
+        files_to_analyze = sorted(set(files_to_analyze))
+
+        if not files_to_analyze:
+            return ""
+
+        max_rows = 5000
+        for file_path in files_to_analyze:
+            try:
+                # Dynamically choose the reader based on file extension
+                ext = file_path.suffix.lower()
+                if ext in ['.csv', '.tsv']:
+                    sep = "\t" if ext == ".tsv" else ","
+                    # Try multiple encodings
+                    df = None
+                    for enc in ['utf-8', 'gbk', 'latin1', 'utf-8-sig']:
+                        try:
+                            df = pd.read_csv(file_path, sep=sep, nrows=max_rows, encoding=enc)
+                            break
+                        except (UnicodeDecodeError, Exception):
+                            continue
+                    
+                    if df is None:
+                        raise Exception(f"Failed to read CSV with multiple encodings for {file_path.name}")
+                elif ext == '.parquet':
+                    # Note: This requires 'pyarrow' or 'fastparquet' to be installed
+                    df = pd.read_parquet(file_path).head(max_rows)
+                else:
+                    # Skip unsupported but matched files
+                    continue
+
+                report_parts.append(f"### Analysis of `{file_path.relative_to(data_dir)}`")
+
+                # Create a summary DataFrame
+                summary = pd.DataFrame({
+                    'Data Type': df.dtypes,
+                    'Missing (%)': (df.isnull().sum() * 100 / len(df)).round(2),
+                    'Cardinality': df.nunique(),
+                })
+                
+                sample_values = [col.dropna().head(2).tolist() for _, col in df.items()]
+                summary['Sample Values'] = sample_values
+
+                # Truncate sample values for readability
+                summary['Sample Values'] = summary['Sample Values'].apply(
+                    lambda x: str(x) if len(str(x)) < 40 else str(x)[:37] + '...'
+                )
+
+                report_parts.append(f"```text\n{summary.to_string()}\n```")
+            except Exception as e:
+                logger.warning(f"Could not analyze schema for {file_path.name}: {traceback.format_exc()}")
+                report_parts.append(f"### Analysis of `{file_path.relative_to(data_dir)}`\nCould not be analyzed due to error: {e}")
+
+        if not report_parts:
+            return ""
+
+        return "## Data Schema Analysis\n" + "\n\n".join(report_parts) + "\n\n"