dslighting 1.3.9__py3-none-any.whl

dsat/operators/aflow_ops.py

@@ -0,0 +1,90 @@
+from typing import List, Dict, Any
+from pydantic import BaseModel, Field
+
+from dsat.operators.base import Operator
+from dsat.services.llm import LLMService
+
+# --- Pydantic Models for Structured I/O ---
+
+class ScEnsembleResponse(BaseModel):
+    """Structured response for the Self-Consistency Ensemble operator."""
+    thought: str = Field(description="The step-by-step thinking process to determine the most consistent solution.")
+    solution_letter: str = Field(description="The single letter (A, B, C, etc.) of the most consistent solution.")
+
+class ReviewResponse(BaseModel):
+    """Structured response for the Review operator."""
+    is_correct: bool = Field(description="True if the solution is very likely correct, False otherwise.")
+    feedback: str = Field(description="If incorrect, detailed feedback for revision. If correct, a brief justification.")
+
+class ReviseResponse(BaseModel):
+    """Structured response for the Revise operator."""
+    solution: str = Field(description="The complete, revised solution based on the provided feedback.")
+
+
+# --- Operator Implementations ---
+
+class ScEnsembleOperator(Operator):
+    """
+    Performs a self-consistency check by asking the LLM to vote on the most
+    frequent or logical answer from a list of candidate solutions.
+    """
+    async def __call__(self, solutions: List[str], problem: str) -> str:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+        
+        solution_text = ""
+        solution_map = {}
+        for i, solution in enumerate(solutions):
+            letter = chr(65 + i)
+            solution_map[letter] = solution
+            solution_text += f"{letter}: \n{solution}\n\n"
+
+        prompt = (
+            f"Given the problem: '{problem}'\n\n"
+            f"Several solutions have been generated:\n{solution_text}\n"
+            "Carefully evaluate these solutions and identify the answer that appears most frequently or is most logical. "
+            "Respond with a JSON object containing your thought process and the letter of the most consistent solution."
+        )
+        
+        response_model = await self.llm_service.call_with_json(prompt, output_model=ScEnsembleResponse)
+        
+        chosen_letter = response_model.solution_letter.strip().upper()
+        return solution_map.get(chosen_letter, solutions[0]) # Default to first solution on failure
+
+class ReviewOperator(Operator):
+    """
+    Critically reviews a solution for correctness and provides structured feedback.
+    """
+    async def __call__(self, problem: str, solution: str) -> ReviewResponse:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+
+        prompt = (
+            "You are a meticulous reviewer. Given a problem and a solution, your task is to critically evaluate the solution's correctness. "
+            "If you are more than 95% confident the solution is incorrect, provide feedback for fixing it. Otherwise, confirm its correctness.\n\n"
+            f"# PROBLEM\n{problem}\n\n"
+            f"# SOLUTION\n{solution}\n\n"
+            "Respond with a JSON object containing your evaluation."
+        )
+        
+        return await self.llm_service.call_with_json(prompt, output_model=ReviewResponse)
+
+class ReviseOperator(Operator):
+    """
+    Revises a solution based on feedback from the Review operator.
+    """
+    async def __call__(self, problem: str, solution: str, feedback: str) -> str:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+
+        prompt = (
+            "You are an expert programmer. A previous solution was found to be incorrect. "
+            "Your task is to revise the solution based on the provided feedback.\n\n"
+            f"# PROBLEM\n{problem}\n\n"
+            f"# INCORRECT SOLUTION\n{solution}\n\n"
+            f"# FEEDBACK\n{feedback}\n\n"
+            "Provide a JSON object containing the complete, revised solution."
+        )
+        
+        response_model = await self.llm_service.call_with_json(prompt, output_model=ReviseResponse)
+        return response_model.solution

dsat/operators/autokaggle_ops.py

@@ -0,0 +1,170 @@
+import json
+import logging
+import re
+from pathlib import Path
+from typing import Dict, Any, List
+
+from pydantic import BaseModel, Field
+
+from dsat.operators.base import Operator
+from dsat.services.llm import LLMService
+from dsat.services.sandbox import SandboxService
+from dsat.services.states.autokaggle_state import AutoKaggleState, TaskContract, PhaseMemory
+from dsat.prompts.autokaggle_prompt import (
+    get_deconstructor_prompt,
+    get_phase_planner_prompt,
+    get_step_planner_prompt,
+    get_developer_prompt,
+    get_validator_prompt,
+    get_reviewer_prompt,
+    get_summarizer_prompt,
+)
+from dsat.models.formats import StepPlan
+
+logger = logging.getLogger(__name__)
+
+
+class PhasePlanningResponse(BaseModel):
+    """Response model for phase planning."""
+    phases: List[str]
+
+
+class ValidationResponse(BaseModel):
+    """Response model for file validation."""
+    passed: bool
+    reason: str = ""
+
+
+class ReviewResponse(BaseModel):
+    """Response model for code review."""
+    score: int
+    suggestion: str = Field(default="", description="Constructive feedback or a suggestion for improvement.")
+
+
+class TaskDeconstructionOperator(Operator):
+    """Parses the natural language description into a structured TaskContract."""
+    
+    async def __call__(self, description: str) -> TaskContract:
+        logger.info("Deconstructing task description into a structured contract...")
+        prompt = get_deconstructor_prompt(description, TaskContract.model_json_schema())
+        contract = await self.llm_service.call_with_json(prompt, output_model=TaskContract)
+        logger.info(f"Task deconstructed. Goal: {contract.task_goal}")
+        return contract
+
+
+class AutoKagglePlannerOperator(Operator):
+    """Handles high-level phase planning and low-level step planning."""
+    
+    async def __call__(self, *args, **kwargs) -> Any:
+        """
+        Main entry point for the planner operator.
+        Can be called with different arguments for different planning tasks.
+        """
+        if len(args) == 1 and isinstance(args[0], TaskContract):
+            # Called for phase planning
+            return await self.plan_phases(args[0])
+        elif len(args) == 2 and isinstance(args[0], AutoKaggleState) and isinstance(args[1], str):
+            # Called for step planning
+            return await self.plan_step_details(args[0], args[1])
+        else:
+            raise ValueError(f"AutoKagglePlannerOperator called with unexpected arguments: {args}, {kwargs}")
+    
+    async def plan_phases(self, contract: TaskContract) -> List[str]:
+        logger.info("Planning dynamic phases for the workflow...")
+        prompt = get_phase_planner_prompt(contract)
+        response = await self.llm_service.call_with_json(prompt, output_model=PhasePlanningResponse)
+        phases = response.phases
+        logger.info(f"Dynamic phases planned: {phases}")
+        return phases
+
+    async def plan_step_details(self, state: AutoKaggleState, phase_goal: str) -> StepPlan:
+        logger.info(f"Planning detailed steps for phase: '{phase_goal}'...")
+        prompt = get_step_planner_prompt(state, phase_goal)
+        step_plan = await self.llm_service.call_with_json(prompt, output_model=StepPlan)
+        return step_plan
+
+
+class DynamicValidationOperator(Operator):
+    """Dynamically validates generated files against the TaskContract."""
+    
+    async def __call__(self, contract: TaskContract, workspace_dir: Path) -> Dict[str, Any]:
+        logger.info("Performing dynamic validation of output files...")
+        results = {}
+        for output_file in contract.output_files:
+            file_path = workspace_dir / output_file.filename
+            if not file_path.exists():
+                results[output_file.filename] = {"passed": False, "reason": "File was not generated."}
+                continue
+            
+            try:
+                content_snippet = "\\n".join(file_path.read_text().splitlines()[:20])
+            except Exception as e:
+                results[output_file.filename] = {"passed": False, "reason": f"Could not read file: {str(e)}"}
+                continue
+                
+            prompt = get_validator_prompt(contract, output_file.filename, content_snippet)
+            validation = await self.llm_service.call_with_json(prompt, output_model=ValidationResponse)
+            results[output_file.filename] = {"passed": validation.passed, "reason": validation.reason}
+        logger.info(f"Validation results: {results}")
+        return results
+
+
+class AutoKaggleDeveloperOperator(Operator):
+    """Writes, executes, and validates code."""
+    
+    def __init__(self, llm_service: LLMService, sandbox_service: SandboxService, validator: DynamicValidationOperator):
+        super().__init__(llm_service, name="AutoKaggleDeveloper")
+        self.sandbox = sandbox_service
+        self.validator = validator
+
+    async def __call__(self, state: AutoKaggleState, phase_goal: str, plan: str, attempt_history: List) -> Dict:
+        logger.info(f"Developer starting work for phase: '{phase_goal}'")
+        prompt = get_developer_prompt(state, phase_goal, plan, attempt_history)
+        
+        raw_reply = await self.llm_service.call(prompt)
+        match = re.search(r"```(?:python|py)?\s*([\s\S]*?)\s*```", raw_reply, re.DOTALL)
+        code = match.group(1).strip() if match else ""
+
+        if not code:
+            return {"code": "", "status": False, "output": "", "error": "No code was generated.", "validation_result": {}}
+
+        exec_result = self.sandbox.run_script(code)
+
+        validation_result = {}
+        if exec_result.success:
+            # Note: This still validates against the *final* contract outputs. The reviewer logic handles this.
+            validation_result = await self.validator(state.contract, self.sandbox.workspace.run_dir)
+        
+        return {
+            "code": code,
+            "status": exec_result.success,
+            "output": exec_result.stdout,
+            "error": exec_result.stderr,
+            "validation_result": validation_result
+        }
+
+
+class AutoKaggleReviewerOperator(Operator):
+    """Reviews the developer's work and provides a score and suggestions."""
+    
+    async def __call__(self, state: AutoKaggleState, phase_goal: str, dev_result: Dict, plan: str = "") -> Dict:
+        logger.info("Reviewer assessing developer's work...")
+        prompt = get_reviewer_prompt(phase_goal, dev_result, plan)
+        review = await self.llm_service.call_with_json(prompt, output_model=ReviewResponse)
+        review_dict = {
+            "score": review.score,
+            "suggestion": review.suggestion
+        }
+        logger.info(f"Review complete. Score: {review.score}")
+        return review_dict
+
+
+class AutoKaggleSummarizerOperator(Operator):
+    """Summarizes a successful phase into a report."""
+    
+    async def __call__(self, state: AutoKaggleState, phase_memory: PhaseMemory) -> str:
+        logger.info(f"Summarizer creating report for phase: '{phase_memory.phase_goal}'")
+        prompt = get_summarizer_prompt(state, phase_memory)
+        report = await self.llm_service.call(prompt)
+        logger.info("Report created.")
+        return report

dsat/operators/automind_ops.py

@@ -0,0 +1,38 @@
+import logging
+from typing import Dict
+
+from dsat.models.formats import ComplexityScore, DecomposedPlan
+from dsat.operators.base import Operator
+from dsat.services.llm import LLMService
+
+logger = logging.getLogger(__name__)
+
+class ComplexityScorerOperator(Operator):
+    """Scores the complexity of a natural language plan."""
+    async def __call__(self, plan: str, task_goal: str) -> ComplexityScore:
+        logger.info("Scoring plan complexity...")
+        prompt = (
+            "You are an expert project manager. On a scale of 1 to 5, where 1 is trivial "
+            "(e.g., load a file and report basic properties) and 5 is highly complex (e.g., multi-stage pipeline involving custom algorithms, extensive preprocessing, or complex simulations), how complex is the following plan?\n\n"
+            f"## Task Goal:\n{task_goal}\n\n"
+            f"## Proposed Plan:\n{plan}\n\n"
+            "Respond with a JSON object containing your score and a brief justification."
+        )
+        score = await self.llm_service.call_with_json(prompt, output_model=ComplexityScore)
+        logger.info(f"Plan complexity scored at {score.complexity}/5.")
+        return score
+
+class PlanDecomposerOperator(Operator):
+    """Decomposes a complex plan into a structured list of sequential steps."""
+    async def __call__(self, plan: str, task_goal: str) -> DecomposedPlan:
+        logger.info("Decomposing complex plan into steps...")
+        prompt = (
+            "You are an expert data scientist. Decompose the following high-level plan into a sequence of "
+            "small, logical, and executable steps. Each step should represent a single cell in a data science notebook.\n\n"
+            f"## Task Goal:\n{task_goal}\n\n"
+            f"## High-Level Plan:\n{plan}\n\n"
+            "Respond with a JSON object containing a list of tasks. Each task must have a unique `task_id` and an `instruction`."
+        )
+        decomposed_plan = await self.llm_service.call_with_json(prompt, output_model=DecomposedPlan)
+        logger.info(f"Plan decomposed into {len(decomposed_plan.tasks)} steps.")
+        return decomposed_plan

dsat/operators/base.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Any
+
+from dsat.services.llm import LLMService
+
+class Operator(ABC):
+    """
+    Abstract base class for a self-contained capability.
+
+    Operators are the "verbs" of the agent framework, representing discrete
+    actions like generating code, executing it, or reviewing results.
+    """
+    def __init__(self, llm_service: Optional[LLMService] = None, name: Optional[str] = None):
+        self.name = name or self.__class__.__name__
+        self.llm_service = llm_service
+
+    @abstractmethod
+    async def __call__(self, *args, **kwargs) -> Any:
+        """
+        Executes the operator's logic. All operator calls are asynchronous.
+        """
+        raise NotImplementedError

dsat/operators/code.py

@@ -0,0 +1,45 @@
+import logging
+from typing import Any
+
+from dsat.common.typing import ExecutionResult
+from dsat.operators.base import Operator
+from dsat.services.sandbox import SandboxService, ProcessIsolatedNotebookExecutor
+
+logger = logging.getLogger(__name__)
+
+class ExecuteAndTestOperator(Operator):
+    """
+    An operator that acts as a clean interface to the SandboxService,
+    handling both script and notebook execution modes.
+    """
+    def __init__(self, sandbox_service: SandboxService):
+        super().__init__(name="ExecuteAndTest")
+        self.sandbox = sandbox_service
+
+    async def __call__(self, code: str, mode: str = "script", executor_context: Any = None) -> ExecutionResult:
+        """
+        Executes code using the configured sandbox.
+
+        Args:
+            code (str): The Python code to execute.
+            mode (str): The execution mode, either 'script' or 'notebook'.
+            executor_context (Any): For notebook mode, this must be the active ProcessIsolatedNotebookExecutor instance.
+
+        Returns:
+            ExecutionResult: The outcome of the execution.
+        """
+        if mode == "script":
+            logger.info("Executing code as a script...")
+            # run_script is synchronous in the sandbox service, but we call it from an async operator
+            # A fully async sandbox would use asyncio.to_thread
+            return self.sandbox.run_script(code)
+        
+        elif mode == "notebook":
+            if not isinstance(executor_context, ProcessIsolatedNotebookExecutor):
+                raise TypeError("Notebook mode requires a valid ProcessIsolatedNotebookExecutor instance passed via executor_context.")
+            
+            logger.info("Executing notebook cell...")
+            return await executor_context.execute_cell(code)
+            
+        else:
+            raise ValueError(f"Unknown execution mode: {mode}")

dsat/operators/dsagent_ops.py

@@ -0,0 +1,123 @@
+import logging
+from typing import Optional, Tuple
+
+from dsat.operators.base import Operator
+from dsat.prompts.dsagent_prompt import (
+    PLAN_PROMPT_TEMPLATE, PROGRAMMER_PROMPT_TEMPLATE, 
+    DEBUGGER_PROMPT_TEMPLATE, LOGGER_PROMPT_TEMPLATE
+)
+from dsat.services.llm import LLMService
+from dsat.services.sandbox import SandboxService
+from dsat.services.vdb import VDBService
+from dsat.common.typing import ExecutionResult
+from dsat.utils.parsing import parse_plan_and_code
+from dsat.utils.context import MAX_HISTORY_CHARS, MAX_OUTPUT_CHARS, truncate_output
+
+# Define how much recent context to keep verbatim during summarization
+RECENT_CONTEXT_VERBATIM = 8000  # Increased from 2000
+
+logger = logging.getLogger(__name__)
+
+class DevelopPlanOperator(Operator):
+    """Retrieves cases and generates an experiment plan."""
+    def __init__(self, llm_service: LLMService, vdb_service: Optional[VDBService] = None):
+        super().__init__(llm_service, name="DevelopPlan")
+        self.vdb = vdb_service
+
+    async def __call__(self, research_problem: str, io_instructions: str, running_log: str) -> str:
+        safe_running_log = truncate_output(running_log, MAX_HISTORY_CHARS)
+        query = f"{research_problem}\n{safe_running_log}"
+        if self.vdb:
+            retrieved_cases = self.vdb.retrieve(query, top_k=1)
+            case = retrieved_cases[0] if retrieved_cases else "No relevant cases found."
+        else:
+            case = "No relevant cases found."
+        
+        prompt = PLAN_PROMPT_TEMPLATE.format(
+            research_problem=research_problem,
+            io_instructions=io_instructions,
+            running_log=safe_running_log,
+            case=case
+        )
+        plan = await self.llm_service.call(prompt)
+        return plan.strip()
+
+class ExecutePlanOperator(Operator):
+    """Manages the programmer-debugger loop to implement a plan."""
+    def __init__(self, llm_service: LLMService, sandbox_service: SandboxService, max_retries: int = 10):
+        super().__init__(llm_service, name="ExecutePlan")
+        self.sandbox = sandbox_service
+        self.max_retries = max_retries
+
+    async def __call__(self, initial_code: str, plan: str, research_problem: str, io_instructions: str, running_log: str = "") -> Tuple[ExecutionResult, str]:
+        safe_running_log = truncate_output(running_log, MAX_HISTORY_CHARS)
+        current_code = initial_code
+        for attempt in range(self.max_retries):
+            if attempt == 0:
+                prompt = PROGRAMMER_PROMPT_TEMPLATE.format(
+                    code=current_code, plan=plan, research_problem=research_problem, 
+                    io_instructions=io_instructions, running_log=safe_running_log
+                )
+            else:
+                safe_error_log = truncate_output(exec_result.stderr, MAX_OUTPUT_CHARS)
+                prompt = DEBUGGER_PROMPT_TEMPLATE.format(
+                    plan=plan, code=current_code, error_log=safe_error_log, 
+                    research_problem=research_problem, io_instructions=io_instructions, running_log=safe_running_log
+                )
+
+            response = await self.llm_service.call(prompt)
+            _, current_code = parse_plan_and_code(response)
+            
+            exec_result = self.sandbox.run_script(current_code)
+            if exec_result.success:
+                return exec_result, current_code
+        
+        logger.error(f"Execution failed after {self.max_retries} attempts.")
+        return exec_result, current_code # Return last failed result
+
+class ReviseLogOperator(Operator):
+    """Summarizes the results of an experiment step."""
+    async def __call__(self, running_log: str, plan: str, exec_result: ExecutionResult, diff: str) -> str:
+        safe_running_log = truncate_output(running_log, MAX_HISTORY_CHARS)
+        execution_log = exec_result.stdout or exec_result.stderr
+        safe_execution_log = truncate_output(execution_log, MAX_OUTPUT_CHARS)
+        safe_diff = truncate_output(diff, MAX_OUTPUT_CHARS)
+        prompt = LOGGER_PROMPT_TEMPLATE.format(
+            plan=plan,
+            execution_log=safe_execution_log,
+            diff=safe_diff,
+            running_log=safe_running_log
+        )
+        new_summary = await self.llm_service.call(prompt)
+        
+        updated_log = running_log + "\n\n---\n" + new_summary.strip()
+        
+        if len(updated_log) > MAX_HISTORY_CHARS:
+            logger.warning("Running log exceeds character limit; summarizing older history...")
+
+            # Split the log: recent history (verbatim) and older history (to be summarized)
+            split_point = max(0, len(updated_log) - RECENT_CONTEXT_VERBATIM)
+            older_history = updated_log[:split_point]
+            recent_history = updated_log[split_point:]
+
+            if not older_history:
+                return updated_log
+
+            # Summarize only the older part
+            summarize_prompt = (
+                "The following is the older history of a data science project. "
+                "Summarize this history concisely. You MUST preserve the following elements if present:\n"
+                "1. Key findings and model performance metrics.\n"
+                "2. Any explicit data format constraints (e.g., required CSV columns).\n"
+                "3. Any explicit I/O details (e.g., filenames used, successful data loading patterns).\n"
+                "4. Specific error messages or tracebacks from failed attempts, as they are crucial context for avoiding repeated mistakes.\n"
+                f"\n\n# OLDER HISTORY\n{older_history}"
+            )
+            summarized_older_history = await self.llm_service.call(summarize_prompt)
+            
+            # Combine summarized older history with verbatim recent history
+            updated_log = f"--- SUMMARIZED HISTORY ---\n{summarized_older_history}\n\n--- RECENT HISTORY (VERBATIM) ---\n{recent_history}"
+
+            logger.info("Running log has been updated with summarized history.")
+            
+        return updated_log

dsat/operators/llm_basic.py

@@ -0,0 +1,84 @@
+import logging
+from typing import Dict
+
+from dsat.models.formats import Plan, ReviewResult, Task
+from dsat.operators.base import Operator
+from dsat.services.llm import LLMService
+from dsat.utils.parsing import parse_plan_and_code
+from dsat.utils.context import summarize_repetitive_logs
+
+logger = logging.getLogger(__name__)
+
+class GenerateCodeAndPlanOperator(Operator):
+    """Generates a plan and corresponding code based on a prompt."""
+    async def __call__(self, system_prompt: str, user_prompt: str = "") -> tuple[str, str]:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+
+        logger.info("Generating new code and plan...")
+        full_prompt = f"{system_prompt}\n\n{user_prompt}"
+        
+        # Use the new standard call method
+        response = await self.llm_service.call(full_prompt)
+        plan, code = parse_plan_and_code(response)
+
+        if "# ERROR" in code:
+            logger.warning("Failed to parse a valid code block from the LLM response.")
+        else:
+            logger.info("Successfully generated code and plan.")
+        
+        return plan, code
+
+class PlanOperator(Operator):
+    """Creates a structured, multi-step plan based on a user request."""
+    async def __call__(self, user_request: str) -> Plan:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+        
+        logger.info(f"Generating a plan for request: '{user_request[:100]}...'")
+        
+        prompt = f"Create a structured JSON plan for this user request: {user_request}"
+        # No more placeholder! This is a real structured call.
+        try:
+            plan_model = await self.llm_service.call_with_json(prompt, output_model=Plan)
+        except Exception as e:
+            logger.warning(f"Structured plan failed ({e}); falling back to text plan.")
+            text = await self.llm_service.call(prompt)
+            plan_model = Plan(tasks=[Task(task_id="1", instruction=text.strip(), dependent_task_ids=[])])
+        logger.info(f"Successfully generated a plan with {len(plan_model.tasks)} tasks.")
+        return plan_model
+
+class ReviewOperator(Operator):
+    """Reviews code execution output and provides a structured score and analysis."""
+    async def __call__(self, prompt_context: Dict) -> ReviewResult:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+
+        logger.info("Reviewing execution output...")
+        
+        raw_output = prompt_context.get('output', '# N/A')
+        processed_output = summarize_repetitive_logs(raw_output)
+
+        prompt = (
+            "You are a data science judge. Review the following code and its output.\n\n"
+            f"# TASK\n{prompt_context.get('task', 'N/A')}\n\n"
+            f"# CODE\n```python\n{prompt_context.get('code', '# N/A')}\n```\n\n"
+            f"# OUTPUT\n```\n{processed_output}\n```\n\n"
+            "Respond with a JSON object containing your evaluation."
+        )
+
+        # No more simulation! This is a real structured call.
+        review_model = await self.llm_service.call_with_json(prompt, output_model=ReviewResult)
+        return review_model
+
+class SummarizeOperator(Operator):
+    """Generates a concise summary of a completed phase or task."""
+    async def __call__(self, context: str) -> str:
+        if not self.llm_service:
+            raise ValueError("LLMService is required for this operator.")
+
+        logger.info("Generating summary...")
+        prompt = f"Please provide a concise summary of the following events:\n\n{context}"
+        summary = await self.llm_service.call(prompt)
+        logger.info("Summary generated successfully.")
+        return summary

dsat/prompts/aflow_prompt.py

@@ -0,0 +1,76 @@
+from pydantic import BaseModel, Field
+
+# 1. Pydantic model for the optimizer's structured output
+class GraphOptimize(BaseModel):
+    modification: str = Field(description="A brief, one-sentence summary of the change made to the workflow code.")
+    graph: str = Field(description="The complete, runnable Python code for the new 'Workflow' class.")
+
+# 2. Prompt templates
+# --- FIX START: MAKE PROMPT MORE GENERIC AND STRATEGY-FOCUSED ---
+WORKFLOW_OPTIMIZE_PROMPT = """
+You are an expert AI workflow engineer. Your task is to iteratively optimize a Python-based AI workflow to improve its problem-solving score.
+
+You will be given the code of a parent workflow, its performance score, and a history of modifications.
+Your goal is to propose a single, small, logical modification to the workflow code. The new code must be a complete and runnable Python class named `Workflow`.
+
+RULES:
+1.  **Focus on Logic**: Your modifications should improve the problem-solving STRATEGY. Examples: add a data cleaning step, try a different model, change a prompt, add a self-correction loop, etc.
+2.  **Adhere to the Standard Interface**: The workflow you generate MUST correctly implement the `solve` method: `async def solve(self, description: str, io_instructions: str, data_dir: Path, output_path: Path)`.
+3.  **Use the Arguments Correctly**:
+    - `description`: Contains the task goal and a COMPLETE data analysis report.
+    - `io_instructions`: Contains CRITICAL I/O requirements.
+    Your code's logic MUST use both arguments effectively.
+4.  **DO NOT Hardcode Filenames**: Your generated workflow code should extract the required output filename from the `output_path.name` attribute provided to the `solve` method, or ensure that any internal prompts correctly instruct the final code-generation step to do so.
+5.  **Make only ONE logical change.** (e.g., add one new operator call, change a prompt, add a loop).
+6.  **Analyze the experience log.** Avoid modifications that have failed in the past. Learn from successful ones.
+7.  **Ensure the `graph` output is the complete, final Python code**, including all necessary imports and the full class definition.
+8.  **Inherit from DSATWorkflow**: The generated class MUST be `class Workflow(DSATWorkflow):`.
+
+Your response MUST be a JSON object that adheres to the provided schema. Do not include any other text, markdown, or explanations.
+"""
+# --- FIX END ---
+
+WORKFLOW_INPUT_TEMPLATE = """
+# PARENT WORKFLOW CONTEXT
+
+## Experience Log
+{experience}
+
+## Parent Score
+{score:.4f}
+
+## Parent Code
+```python
+{graph_code}
+```
+
+## Available Operators
+{operator_description}
+"""
+
+# 3. Helper functions to assemble the final prompt
+def get_operator_description() -> str:
+    """Returns a formatted string of available operators for the prompt."""
+    return """
+    You can call operators from the `self.operators` dictionary inside the `solve` method. For example: `final_answer = await self.operators['ScEnsemble'](...)`
+
+    - `ScEnsemble`: Performs self-consistency voting to find the best solution from a list.
+      - `__call__(self, solutions: List[str], problem: str) -> str`
+    - `Review`: Critically reviews a solution and returns structured feedback.
+      - `__call__(self, problem: str, solution: str) -> ReviewResponse(is_correct: bool, feedback: str)`
+    - `Revise`: Revises a solution based on feedback.
+      - `__call__(self, problem: str, solution: str, feedback: str) -> str`
+    """
+
+# --- FIX: Update function signature ---
+def get_graph_optimize_prompt(experience: str, score: float, graph_code: str) -> str:
+    """Assembles the full prompt for the optimizer LLM."""
+    main_prompt = WORKFLOW_OPTIMIZE_PROMPT # Use the updated prompt directly
+    
+    inputs = WORKFLOW_INPUT_TEMPLATE.format(
+        experience=experience,
+        score=score,
+        graph_code=graph_code,
+        operator_description=get_operator_description()
+    )
+    return main_prompt + "\n" + inputs