dslighting 1.3.9__py3-none-any.whl

dsat/services/vdb.py

@@ -0,0 +1,76 @@
+"""
+Service for managing an in-memory vector database for case-based reasoning.
+"""
+import logging
+from pathlib import Path
+from typing import List
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+class VDBService:
+    """
+    Manages embedding and retrieving text documents (cases) for retrieval-augmented generation.
+    """
+    def __init__(self, case_dir: str, model_name: str = "BAAI/llm-embedder"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name).to(self.device).eval()
+        self.case_files: List[Path] = []
+        self.embedding_bank: torch.Tensor = None
+        self._build_index(Path(case_dir))
+
+    def _build_index(self, case_dir: Path):
+        """Loads cases from a directory and builds the vector index."""
+        logger.info(f"Building vector index from cases in: {case_dir}")
+        if not case_dir.exists():
+            logger.warning(f"Case directory not found: {case_dir}. Creating empty index.")
+            return
+        
+        self.case_files = sorted(list(case_dir.glob("*.py")))
+        case_texts = []
+        for file_path in self.case_files:
+            with open(file_path, "r", encoding="utf-8") as f:
+                case_texts.append(f.read())
+        
+        if not case_texts:
+            logger.warning("No case files found to build index.")
+            return
+
+        with torch.no_grad():
+            inputs = self.tokenizer(case_texts, padding=True, truncation=True, return_tensors='pt', max_length=512).to(self.device)
+            outputs = self.model(**inputs)
+            # Use CLS pooling
+            embeddings = outputs.last_hidden_state[:, 0]
+            self.embedding_bank = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        logger.info(f"Successfully built index with {len(self.case_files)} cases.")
+
+    def retrieve(self, query: str, top_k: int) -> List[str]:
+        """Retrieves the top_k most similar case texts for a given query."""
+        if self.embedding_bank is None:
+            return []
+            
+        with torch.no_grad():
+            inputs = self.tokenizer([query], padding=True, truncation=True, return_tensors='pt', max_length=512).to(self.device)
+            query_embedding = self.model(**inputs).last_hidden_state[:, 0]
+            query_embedding = torch.nn.functional.normalize(query_embedding, p=2, dim=1)
+
+        similarity = (query_embedding @ self.embedding_bank.T).squeeze()
+        _, indices = torch.topk(similarity, min(top_k, len(self.case_files)))
+
+        retrieved_cases = []
+        for idx in indices.tolist():
+            with open(self.case_files[idx], "r", encoding="utf-8") as f:
+                retrieved_cases.append(f.read())
+        return retrieved_cases
+
+    async def store_documents(self, documents: list):
+        """Store documents in vector database."""
+        # This method is kept for compatibility but not used in DS-Agent workflow
+        pass
+    
+    async def search(self, query: str, top_k: int = 5):
+        """Search for similar documents."""
+        return self.retrieve(query, top_k)

dsat/services/workspace.py

@@ -0,0 +1,178 @@
+"""
+Service for managing the run-specific workspace, including data, logs, and artifacts.
+"""
+import logging
+import shutil
+import uuid
+import os
+import contextlib
+from pathlib import Path
+from typing import Dict, Optional
+
+from dsat.common.constants import DEFAULT_WORKSPACE_DIR
+
+logger = logging.getLogger(__name__)
+
+class WorkspaceService:
+    """
+    Manages the file system for a single, isolated agent run.
+    It creates a directory based on the provided unique run name and provides structured access to it.
+    The responsibility for generating unique run names is delegated to the DSATRunner.
+    """
+    def __init__(self, run_name: str, base_dir: str = None):
+        """
+        Initializes the workspace for a new run.
+
+        Args:
+            run_name (str): A descriptive and unique name for the run, provided by the runner.
+            base_dir (str, optional): The base directory where all run folders will be stored.
+                                    If None, uses DEFAULT_WORKSPACE_DIR from constants.
+        """
+        
+        # Use the constant if base_dir is not provided
+        if base_dir is None:
+            base_dir = DEFAULT_WORKSPACE_DIR
+            
+        if not Path(base_dir).is_absolute():
+            # Use Path.cwd() as the base for relative paths
+            base_dir_path = (Path.cwd() / base_dir).resolve()
+        else:
+            base_dir_path = Path(base_dir).resolve()
+
+        base_dir_path.mkdir(parents=True, exist_ok=True)
+            
+        self.run_dir = base_dir_path / run_name
+        self.sandbox_workdir = self.run_dir / "sandbox"
+
+        self.paths: Dict[str, Path] = {
+            "run_dir": self.run_dir,
+            "sandbox_workdir": self.sandbox_workdir,
+            "config": self.run_dir / "config.yaml",
+            "workflow": self.run_dir / "workflow.py",
+            "logs": self.run_dir / "logs",
+            "state": self.run_dir / "state",
+            "candidates": self.run_dir / "candidates",
+            "artifacts": self.run_dir / "artifacts", 
+            "results": self.run_dir / "results.json",
+        }
+
+        self._create_directories()
+        logger.info(f"Workspace initialized at: {self.run_dir.resolve()}. Sandbox Workdir: {self.sandbox_workdir.resolve()}")
+
+    def _create_directories(self):
+        """Creates the full directory structure for the run."""
+        for path in self.paths.values():
+            if not path.suffix:  # Check if it's a directory
+                path.mkdir(parents=True, exist_ok=True)
+
+    def get_path(self, name: str) -> Path:
+        """
+        Retrieves a managed path from the workspace.
+
+        Args:
+            name (str): The key of the path to retrieve (e.g., 'logs', 'artifacts').
+
+        Returns:
+            Path: The absolute Path object for the requested resource.
+
+        Raises:
+            KeyError: If the requested path name is not defined.
+        """
+        if name == 'sandbox_cwd':
+             logger.warning("Accessing deprecated 'sandbox_cwd'. Use 'sandbox_workdir' instead.")
+             name = 'sandbox_workdir'
+
+        if name not in self.paths:
+            raise KeyError(f"Path '{name}' is not a defined workspace path.")
+        return self.paths[name]
+
+    def write_file(self, content: str, path_name: str, sub_path: str = None):
+        """
+        Writes content to a file within a managed directory.
+
+        Args:
+            content (str): The string content to write.
+            path_name (str): The key of the managed directory (e.g., 'logs').
+            sub_path (str, optional): A filename or relative path within the managed directory.
+        """
+        target_dir = self.get_path(path_name)
+        file_path = target_dir / sub_path if sub_path else target_dir
+        
+        # Ensure parent directory of the file exists if sub_path contains folders
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+        logger.debug(f"Wrote {len(content)} bytes to {file_path}")
+
+    def link_data_to_workspace(self, source_data_dir: Path):
+        """
+        Links or copies the CONTENTS of a source data directory into the run's sandbox_workdir.
+        This ensures the agent runs in an isolated environment containing all inputs.
+        """
+        # Use sandbox_workdir as the destination
+        destination_dir = self.get_path("sandbox_workdir")
+
+        src = source_data_dir.resolve()
+        if not src.exists() or not src.is_dir():
+            raise FileNotFoundError(f"Source data directory not found: {src}")
+
+        for item in src.iterdir():
+            source_item = item
+            destination_item = destination_dir / item.name
+
+            # If the destination item already exists, skip it (idempotent behavior)
+            if destination_item.exists() or destination_item.is_symlink():
+                continue
+
+            # Try to create a symlink for the item
+            try:
+                # Determine if the target is a directory for Windows compatibility
+                target_is_directory = source_item.is_dir()
+                os.symlink(source_item, destination_item, target_is_directory=target_is_directory)
+                logger.debug(f"Linked {source_item.name} into sandbox.")
+            
+            except (OSError, NotImplementedError) as e:
+                # Symlink not permitted. Fallback to copy.
+                warning_message = (
+                    f"Symlink creation failed for {item.name} ({e}). Falling back to copying. "
+                )
+                if os.name == 'nt':
+                     warning_message += " On Windows, enable 'Developer Mode' or run as administrator for symlinks."
+                
+                logger.warning(warning_message)
+                
+                try:
+                    if source_item.resolve() == destination_item.resolve():
+                        logger.warning(
+                            f"Skipping copy of {source_item.name} because source and destination "
+                            "resolve to the same file. This may indicate a workspace configuration issue."
+                        )
+                        continue
+
+                    if source_item.is_dir():
+                        shutil.copytree(source_item, destination_item)
+                    else:
+                        shutil.copy2(source_item, destination_item)
+                    logger.debug(f"Copied {source_item.name} into sandbox.")
+                except Exception as copy_e:
+                    logger.error(f"Failed to copy item {item.name}: {copy_e}", exc_info=True)
+                    raise
+
+        logger.info(f"Data from {src} successfully populated into {destination_dir}")
+
+    def cleanup(self, keep_workspace: bool = False):
+        """
+        Removes the entire run directory unless explicitly told to keep it.
+        """
+        if keep_workspace:
+            logger.info(f"Workspace preserved as requested: {self.run_dir.resolve()}")
+            return
+
+        logger.info(f"Cleaning up workspace: {self.run_dir.resolve()}")
+        try:
+            if self.run_dir.exists():
+                shutil.rmtree(self.run_dir)
+        except Exception as e:
+            logger.error(f"Failed to clean up workspace {self.run_dir}: {e}")
+

dsat/tasks/__init__.py

@@ -0,0 +1,3 @@
+# dsat/tasks/__init__.py
+
+# This file makes the 'tasks' directory a Python package.

dsat/tasks/handlers.py

@@ -0,0 +1,376 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+import tempfile
+import logging
+from typing import Tuple, Any
+from dsat.models.task import TaskDefinition
+from dsat.services.data_analyzer import DataAnalyzer
+
+logger = logging.getLogger(__name__)
+
+
+class TaskHandler(ABC):
+    """
+    Base class for handlers that translate between logical TaskDefinition and physical file interfaces required by DSATWorkflow.
+
+    Each handler encapsulates preparation and parsing logic for specific task types (e.g., Kaggle, QA),
+    allowing the workflow itself to remain task-agnostic.
+    """
+    def __init__(self):
+        """
+        Initialize the handler and create a temporary, self-managed directory
+        for storing physical files generated for the task.
+        """
+        try:
+            self.temp_dir = tempfile.TemporaryDirectory()
+        except Exception as e:
+            logger.error(f"Failed to create temporary directory for TaskHandler: {e}")
+            self.temp_dir = None
+        
+        self.analyzer = DataAnalyzer()
+
+    @abstractmethod
+    def prepare_input(self, task: TaskDefinition) -> Tuple[str, str, Path, Path]:
+        """
+        Prepare the physical file input required by the workflow.
+
+        This method converts logical tasks into physical parameters needed by DSATWorkflow.solve().
+
+        Args:
+            task: Logical task definition.
+
+        Returns:
+            A tuple (description, io_instructions, data_dir, output_path) to pass to workflow.solve().
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def parse_output(self, output_path: Path) -> Any:
+        """
+        Parse the workflow's output file into structured results required by benchmarking.
+
+        This method converts physical output files back into logical answers.
+
+        Args:
+            output_path: Path where the workflow saved its output.
+
+        Returns:
+            Final answer in the format expected by benchmarking (e.g., string for QA, Path object for Kaggle).
+        """
+        raise NotImplementedError
+
+    def cleanup(self):
+        """
+        Explicitly clean up the temporary directory.
+        """
+        if self.temp_dir:
+            try:
+                self.temp_dir.cleanup()
+                logger.debug(f"Successfully cleaned up temporary directory for {self.__class__.__name__}.")
+            except Exception as e:
+                logger.error(f"Error cleaning up temporary directory for {self.__class__.__name__}: {e}")
+
+    def __del__(self):
+        """Ensure cleanup is called when the object is garbage collected."""
+        self.cleanup()
+
+
+class KaggleTaskHandler(TaskHandler):
+    """
+    Handler for Kaggle-style file input/file output tasks.
+    This is a "pass-through" implementation since tasks are already file-based.
+    """
+    def prepare_input(self, task: TaskDefinition) -> Tuple[str, str, Path, Path]:
+        """Extract paths, analyze data, and separate description from I/O instructions."""
+        if task.task_type != "kaggle":
+            raise ValueError("KaggleTaskHandler can only handle tasks of type 'kaggle'.")
+
+        description = task.payload.get("description")
+        data_dir = Path(task.payload.get("public_data_dir"))
+        output_path = Path(task.payload.get("output_submission_path"))
+
+        if not all([description, data_dir, output_path]):
+            raise ValueError("Kaggle task payload is missing required keys: 'description', 'public_data_dir', 'output_submission_path'.")
+        if not data_dir.exists() or not data_dir.is_dir():
+            raise FileNotFoundError(f"Kaggle public_data_dir not found: {data_dir}")
+
+        logger.info(f"Analyzing input data for task '{task.task_id}'...")
+        
+        data_report = self.analyzer.analyze_data(data_dir, task_type="kaggle")
+        io_instructions = self.analyzer.generate_io_instructions(output_path.name, optimization_context=False)
+
+        augmented_description = f"{description}\n{data_report}"
+
+        logger.debug(f"Preparing Kaggle task '{task.task_id}': data_dir='{data_dir}', output_path='{output_path}'")
+        return augmented_description, io_instructions, data_dir, output_path
+
+    def parse_output(self, output_path: Path) -> Path:
+        """
+        For Kaggle tasks, the result is the output file itself.
+        This just validates that the file was created.
+        """
+        if not output_path.exists():
+            # In actual evaluation, this will be caught and reported as a failure.
+            logger.warning(f"Agent did not produce the required submission file at: {output_path}")
+            # Return the path even if it doesn't exist, let the caller (e.g., benchmark) handle the file not found case.
+        return output_path
+
+
+class QATaskHandler(TaskHandler):
+    """
+    Handler for simple question-answer (QA) tasks.
+    This is a "translation" implementation that converts string questions to files and expects answers as files.
+    """
+    def prepare_input(self, task: TaskDefinition) -> Tuple[str, str, Path, Path]:
+        """Convert QA question to physical file input."""
+        if task.task_type != "qa":
+            raise ValueError("QATaskHandler can only handle tasks of type 'qa'.")
+        if not self.temp_dir:
+            raise RuntimeError("Temporary directory not available for QATaskHandler.")
+
+        question = task.payload.get("question")
+        if not question:
+            raise ValueError("QA task payload is missing required key: 'question'.")
+
+        data_dir = Path(self.temp_dir.name)
+        
+        # Create physical task representation
+        problem_file = data_dir / "problem.txt"
+        problem_file.write_text(question, encoding='utf-8')
+        
+        # Define output contract
+        output_path = data_dir / "answer.txt"
+        
+        # This core instruction is now simpler
+        core_instruction = (
+            "Your task is to answer the question found in `problem.txt`. "
+            "Write ONLY the final answer into the required output file."
+        )
+
+        data_report = self.analyzer.analyze_data(data_dir, task_type="qa")
+        io_instructions = self.analyzer.generate_io_instructions(output_path.name, optimization_context=False)
+        
+        description = f"{core_instruction}\n{data_report}"
+        
+        logger.debug(f"Preparing QA task '{task.task_id}': input file='{problem_file}', expected output='{output_path}'")
+        return description, io_instructions, data_dir, output_path
+
+    def parse_output(self, output_path: Path) -> str:
+        """Read and return the final answer string from the output file."""
+        if not output_path.exists() or not output_path.is_file():
+            logger.warning(f"Agent did not produce the answer file for QA task at: {output_path}")
+            return "[ERROR] Agent did not produce an answer file."
+        
+        try:
+            answer = output_path.read_text(encoding='utf-8').strip()
+            logger.debug(f"Parsed QA answer from '{output_path}': '{answer[:50]}...'")
+            return answer
+        except Exception as e:
+            logger.error(f"Failed to read or parse QA answer file '{output_path}': {e}")
+            return f"[ERROR] Failed to parse answer file: {e}"
+
+
+class DataSciTaskHandler(TaskHandler):
+    """
+    Handler for DataSciBench tasks.
+    These are multi-step data science tasks with prompts and optional input files.
+    """
+    def prepare_input(self, task: TaskDefinition) -> Tuple[str, str, Path, Path]:
+        """Prepare DataSciBench task input."""
+        if task.task_type != "datasci":
+            raise ValueError("DataSciTaskHandler can only handle tasks of type 'datasci'.")
+
+        prompt = task.payload.get("prompt", "")
+        input_dir = task.payload.get("input_dir", "")
+        output_dir = task.payload.get("output_dir", "")
+
+        if not prompt:
+            raise ValueError("DataSci task payload is missing required key: 'prompt'.")
+
+        # Use input_dir as data_dir, or temp_dir if no input files
+        if input_dir and Path(input_dir).exists():
+            data_dir = Path(input_dir)
+        elif self.temp_dir:
+            data_dir = Path(self.temp_dir.name)
+        else:
+            raise RuntimeError("No data directory available for DataSciTaskHandler.")
+
+        # Output directory
+        if output_dir:
+            output_path = Path(output_dir) / "output.csv"
+        else:
+            output_path = data_dir / "output.csv"
+
+        # Build description with the prompt
+        description = prompt
+
+        # Analyze data if available
+        try:
+            data_report = self.analyzer.analyze_data(data_dir, task_type="datasci")
+            description = f"{prompt}\n\n{data_report}"
+        except Exception as e:
+            logger.debug(f"Data analysis skipped: {e}")
+
+        # Generate I/O instructions
+        io_instructions = (
+            f"All input data files are in the current working directory.\n"
+            f"Save all output files to the current working directory.\n"
+            f"Follow the task instructions carefully and generate the required output files."
+        )
+
+        logger.debug(f"Preparing DataSci task '{task.task_id}': data_dir='{data_dir}', output_dir='{output_dir}'")
+        return description, io_instructions, data_dir, output_path
+
+    def parse_output(self, output_path: Path) -> Path:
+        """
+        For DataSci tasks, return the output directory path.
+        The actual evaluation is done by the benchmark class using metric.yaml.
+        """
+        if output_path.parent.exists():
+            return output_path.parent
+        return output_path
+
+
+class OpenEndedTaskHandler(TaskHandler):
+    """
+    Handler for open-ended tasks (mathematical modeling, simulations, strategy tasks).
+    These tasks don't have ground truth answers and are evaluated via LLM judges.
+    """
+    def prepare_input(self, task: TaskDefinition) -> Tuple[str, str, Path, Path]:
+        """Prepare open-ended task input."""
+        if task.task_type != "open_ended":
+            raise ValueError("OpenEndedTaskHandler can only handle tasks of type 'open_ended'.")
+        if not self.temp_dir:
+            raise RuntimeError("Temporary directory not available for OpenEndedTaskHandler.")
+
+        # Get task paths from payload
+        raw_dir_str = task.payload.get("raw_data_dir", "")
+        description_file = task.payload.get("description_file", "")
+        rubric_file = task.payload.get("rubric_file", "")
+
+        # Use temp directory as working directory
+        data_dir = Path(self.temp_dir.name)
+
+        # Copy ONLY data files (CSV, JSON, etc.) - exclude description and rubric files
+        if raw_dir_str:
+            raw_dir = Path(raw_dir_str)
+            if raw_dir.exists():
+                import shutil
+                for file in raw_dir.iterdir():
+                    if file.is_file() and file.suffix in ['.csv', '.json', '.txt', '.xlsx', '.parquet']:
+                        # Exclude description.md and rubric.md from being treated as data files
+                        if file.name not in ['description.md', 'rubric.md']:
+                            shutil.copy2(file, data_dir / file.name)
+                            logger.debug(f"Copied data file: {file.name}")
+
+        # Read task description and rubric from files if provided
+        description = task.payload.get("description", "")
+        rubric = task.payload.get("rubric", "")
+
+        # Read from files if specified
+        if description_file and Path(description_file).exists():
+            try:
+                description = Path(description_file).read_text(encoding='utf-8')
+                logger.debug(f"Read description from {description_file} ({len(description)} chars)")
+            except Exception as e:
+                logger.warning(f"Failed to read description file {description_file}: {e}")
+
+        if rubric_file and Path(rubric_file).exists():
+            try:
+                rubric = Path(rubric_file).read_text(encoding='utf-8')
+                logger.debug(f"Read rubric from {rubric_file} ({len(rubric)} chars)")
+            except Exception as e:
+                logger.warning(f"Failed to read rubric file {rubric_file}: {e}")
+
+        if not description:
+            raise ValueError("Open-ended task payload is missing required key: 'description'.")
+
+        # Output path - agent should create an artifacts directory or report
+        output_path = data_dir / "artifacts"
+
+        # Build the FULL task description directly in the prompt
+        # Include description and evaluation criteria
+        # IMPORTANT: For open-ended tasks, explicitly require artifacts directory creation
+
+        task_description_section = f"""## Task Description
+
+{description}
+"""
+
+        if rubric:
+            task_description_section += f"""
+
+## Evaluation Criteria
+
+{rubric}
+"""
+
+        # Analyze available data files to provide schema information (excluding task files)
+        data_report = self.analyzer.analyze_data(data_dir, task_type="datasci")
+
+        # Combine everything into the full description
+        # Note: data_report already contains "--- COMPREHENSIVE DATA REPORT ---" header
+        full_description = f"""{task_description_section}
+
+{data_report}
+
+## CRITICAL OUTPUT INSTRUCTIONS
+
+**YOU MUST CREATE AN `artifacts/` DIRECTORY AND SAVE ALL OUTPUTS THERE:**
+
+```python
+import os
+artifact_dir = 'artifacts'
+os.makedirs(artifact_dir, exist_ok=True)
+
+# Save all your work to the artifacts directory:
+# - Analysis code: artifacts/analysis.py
+# - Visualizations: artifacts/plot_*.png
+# - Data files: artifacts/results.csv
+# - Models, notebooks, etc.
+```
+
+## Task Goals
+- Your goal is to complete this task to the best of your ability
+- Create appropriate output files (code, analysis, visualizations, etc.) in the `artifacts/` subdirectory
+- The evaluation will be based on the quality and completeness of your work according to the evaluation criteria
+"""
+
+        # Generate I/O instructions - VERY EXPLICIT for open-ended tasks
+        io_instructions = f"""**OUTPUT DIRECTORY STRUCTURE (MANDATORY):**
+
+```python
+# At the START of your code, create the artifacts directory:
+import os
+artifact_dir = 'artifacts'
+os.makedirs(artifact_dir, exist_ok=True)
+
+# Save ALL outputs to this directory:
+# - Code: f"{{artifact_dir}}/solution.py"
+# - Plots: f"{{artifact_dir}}/visualization_{{i}}.png"
+# - Data: f"{{artifact_dir}}/results.csv"
+```
+
+**REQUIREMENTS:**
+1. Create the `artifacts/` directory at the beginning of your code
+2. Save ALL generated files (plots, models, data, code) to this directory
+3. Do NOT save files to the current directory - use the artifacts/ subdirectory
+4. Focus on quality, completeness, and following the evaluation criteria
+"""
+
+        logger.debug(f"Preparing open-ended task '{task.task_id}': output_path='{output_path}', description_len={len(description)}")
+        return full_description, io_instructions, data_dir, output_path
+
+    def parse_output(self, output_path: Path) -> Path:
+        """
+        For open-ended tasks, return the artifacts directory path.
+        The actual evaluation is done by LLM judges, not CSV grading.
+        """
+        if not output_path.exists():
+            # If artifacts directory doesn't exist, return the parent temp dir
+            # This allows evaluation to proceed even if no artifacts were created
+            logger.warning(f"Open-ended task did not create artifacts directory at: {output_path}")
+            return output_path.parent
+
+        logger.debug(f"Parsed open-ended task artifacts from: {output_path}")
+        return output_path