satif-ai 0.1.1__py3-none-any.whl

satif_ai/adapters/tidy.py

@@ -0,0 +1,462 @@
+import inspect
+import json
+import logging
+import re
+import shutil
+import sqlite3
+import tempfile
+from pathlib import Path
+from typing import Optional
+
+# MCP and Agent imports
+from agents import Agent, Runner, function_tool
+from agents.mcp.server import MCPServerStdio
+from libs.core.satif_core.types import Datasource
+from mcp import ClientSession
+
+# SATIF imports
+from satif_core.adapters.base import Adapter
+from satif_sdk import SDIFDatabase
+from satif_sdk.adapters.code import AdapterError, CodeAdapter
+
+logger = logging.getLogger(__name__)
+
+
+# --- Tidy Transformation Prompt ---
+TIDY_TRANSFORMATION_PROMPT = """
+You are an expert Data Tidying Agent for SDIF databases.
+Your task is to write Python code to transform tables within a given SDIF database into a 'tidy' format, modifying the database *in place*.
+
+**Tidy Data Principles:**
+1. Each variable forms a column.
+2. Each observation forms a row.
+3. Each type of observational unit forms a table (you might need to create new tables).
+
+**Input SDIF Context:**
+You will be provided with:
+- The schema of the input SDIF database (`input_schema`).
+- A sample analysis of the input SDIF database (`input_sample`).
+
+<input_schema>
+{input_schema}
+</input_schema>
+
+<input_sample>
+{input_sample}
+</input_sample>
+
+**Available `SDIFDatabase` Methods:**
+Here are the public methods you can call on the `db` object passed to your `adapt_sdif` function:
+```python
+{sdif_database_methods}
+```
+
+**Your Goal:**
+Generate Python code for an adaptation function named `adapt_sdif`. This function MUST:
+- Accept an `SDIFDatabase` object (`db`) which represents the database to be modified.
+- Perform tidying operations directly on this `db` instance using its available methods (see list above) and potentially pandas for intermediate processing.
+- Examples of operations:
+    - Read a messy table: `df = db.read_table('messy_table')`
+    - Tidy the DataFrame using pandas (melt, split, etc.).
+    - Write the tidy DataFrame back: `db.write_dataframe(tidy_df, 'tidy_table_name', source_id=1, if_exists='replace')` (obtain a source_id if needed or use a default).
+    - Drop the original messy table if desired: `db.drop_table('messy_table')`
+- The function should modify the `db` object **in place** and MUST return `None`.
+
+**Tools Available:**
+- `execute_sql(query: str)`: Execute a read-only SQL query against the **input** SDIF database to inspect data further before generating adaptation code.
+- `execute_tidy_adaptation(code: str)`: Execute the Python code string containing your `adapt_sdif` function.
+    - This tool will run your code against a **copy** of the input SDIF.
+    - It will return a **sample analysis** of the **modified SDIF file**.
+
+**Workflow:**
+1. **Analyze:** Examine the `input_schema` and `input_sample`. Use `execute_sql` if needed to understand the data structure and identify tables needing tidying.
+2. **Code:** Write the `adapt` Python code function using `SDIFDatabase` methods and pandas.
+3. **Execute & Verify:** Use `execute_tidy_adaptation` with your code.
+4. **Review:** Examine the returned `output_sample_analysis`. Check if the tables in the modified SDIF meet the tidy data principles.
+5. **Refine:** If the output is not tidy, modify your Python code and repeat step 3.
+6. **Finalize:** Once the `output_sample_analysis` confirms the data is tidy, respond **only** with the final, validated Python code string for the `adapt_sdif` function enclosed in triple backticks (```python ... ```). Do not include any other text before or after the code block.
+
+**Example `adapt_sdif` function:**
+```python
+import pandas as pd
+from satif.adapters.code import AdapterError # Import for raising errors if needed
+from typing import Dict, Any
+
+# Assume input db has table 'wide_sales' with columns: 'Region', 'Q1_Sales', 'Q2_Sales'
+
+def adapt_sdif(db: SDIFDatabase) -> None:
+    try:
+        # Get a default source_id or create one if needed
+        sources = db.list_sources()
+        source_id = sources[0]['source_id'] if sources else db.add_source('tidy_adapter', 'script')
+
+        # Read the table to tidy
+        df_sales = db.read_table('wide_sales')
+
+        # Tidy using pandas
+        tidy_sales = pd.melt(df_sales,
+                             id_vars=['Region'],
+                             value_vars=['Q1_Sales', 'Q2_Sales'],
+                             var_name='Quarter',
+                             value_name='Sales')
+        tidy_sales['Quarter'] = tidy_sales['Quarter'].str.replace('_Sales', '')
+
+        # Write the tidy table back, replacing the original if desired,
+        # or writing to a new table.
+        # Here, we write to a new table and drop the old one.
+        db.write_dataframe(tidy_sales,
+                           'tidy_sales_data',
+                           source_id=source_id,
+                           if_exists='replace', # Replace if 'tidy_sales_data' already exists
+                           description="Tidied sales data")
+
+        # Optionally drop the original table
+        db.drop_table('wide_sales')
+
+    except Exception as e:
+        print(f"Error during tidying: {{e}}") # Log errors
+        # Re-raise the exception to signal failure to the execution framework
+        raise AdapterError(f"Error in adapt_sdif: {{e}}") from e
+
+    # IMPORTANT: Function must return None (can be implicit)
+    return None
+```
+
+**Important:**
+- Your Python code string MUST define the `adapt_sdif(db: SDIFDatabase)` function.
+- The function MUST return `None`.
+- Use `db.write_dataframe` with `if_exists='replace'` or `if_exists='append'` or write to new tables and potentially use `db.drop_table` for the old ones.
+- Handle potential errors during data reading or processing within your function and raise an `AdapterError` or similar to indicate failure.
+- Ensure pandas and other necessary libraries (like `typing`, `AdapterError`) are imported within the code string if you use them.
+"""
+
+# --- Global context for tools ---
+# These will be set within the TidyAdapter instance when adapt is called
+TOOL_CONTEXT = {
+    "copied_input_sdif_path": None,
+    "temp_dir": None,
+    "current_output_sdif_path": None,  # Path generated by the tool
+}
+
+
+@function_tool
+async def execute_tidy_adaptation(code: str) -> str:
+    """
+    Tool implementation for the agent to execute the tidying adaptation code.
+    Runs the code against a *copy* of the input SDIF, creating a *new* output SDIF,
+    and returns a sample analysis of the modified output.
+    """
+    copied_input_path = TOOL_CONTEXT.get("copied_input_sdif_path")
+    temp_dir = TOOL_CONTEXT.get("temp_dir")
+
+    if not copied_input_path or not copied_input_path.exists():
+        return (
+            "Error: Input SDIF copy not found for transformation. Tool context issue."
+        )
+    if not temp_dir:
+        return "Error: Temporary directory not set up. Tool context issue."
+
+    # Define path for the *output* SDIF generated by *this tool execution*
+    # This path is temporary just for this tool's run
+    tool_output_sdif_path = temp_dir / "tidy_adaptation_output.sdif"
+    # Update context for potential internal use, though CodeAdapter calculates its own path
+    TOOL_CONTEXT["current_output_sdif_path"] = tool_output_sdif_path
+
+    logger.info(
+        f"Executing adaptation code via tool. Output will be: {tool_output_sdif_path}"
+    )
+
+    try:
+        # 1. Instantiate CodeAdapter with the provided code
+        #    It will operate on a *copy* specified by copied_input_path
+        #    and write to a *new* file (_adapted suffix by default).
+        adapter = CodeAdapter(
+            function=code,
+            function_name="adapt_sdif",  # As specified in prompt
+            output_suffix="_adapted_tool_run",  # Give tool runs a distinct suffix
+        )
+        # Run the adaptation. It copies `copied_input_path` and modifies the copy.
+        # The returned path is the newly created, adapted file.
+        adapted_sdif_path = adapter.adapt(copied_input_path)
+
+        # 2. Get sample analysis of the *adapted* SDIF file
+        with SDIFDatabase(adapted_sdif_path, read_only=True) as adapted_db_read:
+            analysis = adapted_db_read.get_sample_analysis(
+                num_sample_rows=5, top_n_common_values=5
+            )
+            # Store the path generated by the tool for reference if needed elsewhere
+            TOOL_CONTEXT["current_output_sdif_path"] = adapted_sdif_path
+            return json.dumps({"output_sample_analysis": analysis})
+
+    except AdapterError as e:
+        logger.error(f"Error during adaptation code execution via tool: {e}")
+        # Clean up the failed output file if it exists
+        if (
+            TOOL_CONTEXT["current_output_sdif_path"]
+            and TOOL_CONTEXT["current_output_sdif_path"].exists()
+        ):
+            try:
+                TOOL_CONTEXT["current_output_sdif_path"].unlink()
+            except OSError:
+                pass
+        return f"Error: Adaptation code failed: {e}"
+    except (sqlite3.Error, ValueError, TypeError, FileNotFoundError) as e:
+        logger.error(f"Error analyzing adapted SDIF or file issue: {e}")
+        if (
+            TOOL_CONTEXT["current_output_sdif_path"]
+            and TOOL_CONTEXT["current_output_sdif_path"].exists()
+        ):
+            try:
+                TOOL_CONTEXT["current_output_sdif_path"].unlink()
+            except OSError:
+                pass
+        return f"Error: Failed to analyze adapted SDIF: {e}"
+    except Exception as e:
+        logger.exception("Unexpected error in execute_tidy_adaptation tool")
+        if (
+            TOOL_CONTEXT["current_output_sdif_path"]
+            and TOOL_CONTEXT["current_output_sdif_path"].exists()
+        ):
+            try:
+                TOOL_CONTEXT["current_output_sdif_path"].unlink()
+            except OSError:
+                pass
+        return f"Unexpected Error: {e}"
+
+
+class TidyAdapter(Adapter):
+    """
+    Uses an AI agent (via agents library) to generate and execute Python code
+    that transforms tables in an SDIF file into a tidy format using CodeAdapter.
+    """
+
+    def __init__(
+        self,
+        mcp_server: MCPServerStdio,  # Use the server instance
+        mcp_session: ClientSession,  # Use the client session
+        llm_model: str = "o4-mini",  # Specify the LLM model
+        max_iterations: int = 5,
+    ):
+        """
+        Initialize the TidyAdapter.
+
+        Args:
+            mcp_server: An instance of MCPServerStdio for agent communication.
+            mcp_session: An instance of ClientSession for resource/prompt fetching.
+            llm_model: Name of the language model to use for the agent.
+            max_iterations: Maximum number of attempts the agent gets to refine the code.
+        """
+        self.mcp_server = mcp_server
+        self.mcp_session = mcp_session
+        self.llm_model = llm_model
+        self.max_iterations = max_iterations  # Not directly used by Runner, but good for context/potential future use
+
+        # Temporary environment management (could be context managed)
+        self._temp_dir: Optional[Path] = None
+        self._copied_input_sdif_path: Optional[Path] = None
+
+    def _get_sdif_methods(self) -> str:
+        """Introspects SDIFDatabase and returns formatted public method signatures."""
+        signatures = []
+        # Exclude known internal/base methods explicitly
+        exclude_methods = {
+            "__init__",
+            "__enter__",
+            "__exit__",
+            "__del__",
+            "_validate_connection",
+            "_create_metadata_tables",
+        }
+        try:
+            # Iterate through members of the class
+            for name, member in inspect.getmembers(SDIFDatabase):
+                # Check if it's a function/method and not excluded/private
+                if (
+                    inspect.isfunction(member)
+                    and not name.startswith("_")
+                    and name not in exclude_methods
+                ):
+                    try:
+                        sig = inspect.signature(member)
+                        # Format the signature string
+                        sig_str = f"db.{name}{sig}"
+                        signatures.append(sig_str)
+                    except (ValueError, TypeError) as e:
+                        # Handle methods that might not have clear signatures (e.g., built-ins if any slip through)
+                        logger.debug(f"Could not get signature for method {name}: {e}")
+                        signatures.append(f"db.{name}(...) # Signature unavailable")
+
+            return "\n".join(sorted(signatures))
+        except Exception as e:
+            logger.error(f"Failed to introspect SDIFDatabase methods: {e}")
+            return "# Failed to retrieve method signatures."
+
+    def _setup_temp_env(self, input_sdif_path: Path) -> Path:
+        """Creates a temporary directory and copies the input SDIF."""
+        self._temp_dir = Path(tempfile.mkdtemp(prefix="satif_tidy_adapter_"))
+        self._copied_input_sdif_path = (
+            self._temp_dir / f"input_copy_{input_sdif_path.name}"
+        )
+        shutil.copy(input_sdif_path, self._copied_input_sdif_path)
+        logger.info(
+            f"Copied input SDIF to temporary location: {self._copied_input_sdif_path}"
+        )
+        # Set global tool context
+        TOOL_CONTEXT["copied_input_sdif_path"] = self._copied_input_sdif_path
+        TOOL_CONTEXT["temp_dir"] = self._temp_dir
+        TOOL_CONTEXT["current_output_sdif_path"] = None  # Reset output path context
+        return self._copied_input_sdif_path
+
+    def _cleanup_temp_env(self):
+        """Removes the temporary directory."""
+        if self._temp_dir and self._temp_dir.exists():
+            try:
+                shutil.rmtree(self._temp_dir)
+                logger.info(f"Cleaned up temporary directory: {self._temp_dir}")
+            except Exception as e:
+                logger.error(
+                    f"Error cleaning up temporary directory {self._temp_dir}: {e}"
+                )
+        # Clear global tool context
+        TOOL_CONTEXT["copied_input_sdif_path"] = None
+        TOOL_CONTEXT["temp_dir"] = None
+        TOOL_CONTEXT["current_output_sdif_path"] = None
+        self._temp_dir = None
+        self._copied_input_sdif_path = None
+
+    def parse_code(self, code_text: str) -> Optional[str]:
+        """Extracts Python code from markdown code blocks."""
+        match = re.search(r"```(?:python)?(.*?)```", code_text, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        else:
+            # If no markdown block, assume the whole text might be code (less reliable)
+            # Check for keywords common in the expected code
+            if "def adapt_sdif(" in code_text and "SDIFDatabase" in code_text:
+                logger.warning(
+                    "No markdown code block found, attempting to use entire response as code."
+                )
+                return code_text.strip()
+            return None  # Indicate no valid code found
+
+    async def adapt(self, sdif_database: SDIFDatabase) -> Datasource:
+        """
+        Transforms the data in the input SDIF to be tidy using an AI agent.
+
+        Args:
+            sdif_database: The input SDIF database instance. Connection will be closed.
+
+        Returns:
+            Path to the new SDIF file containing the tidied data.
+
+        Raises:
+            FileNotFoundError: If the input SDIF path doesn't exist.
+            RuntimeError: If the agent fails to produce valid tidy code.
+            Exception: For unexpected errors during the process.
+        """
+        input_path = Path(sdif_database.path)
+        if not input_path.exists():
+            raise FileNotFoundError(f"Input SDIF file not found: {input_path}")
+
+        # Ensure the input DB connection is closed before copying
+        try:
+            sdif_database.close()
+        except Exception:
+            pass
+
+        copied_input_path = self._setup_temp_env(input_path)
+
+        try:
+            # Get Initial Context using SDIFDatabase methods directly
+            with SDIFDatabase(copied_input_path, read_only=True) as db:
+                input_schema_dict = db.get_schema()
+                input_sample_dict = db.get_sample_analysis()
+
+            # Get SDIFDatabase method signatures
+            sdif_methods_str = self._get_sdif_methods()
+
+            # Prepare context for the prompt
+            initial_context = {
+                "input_schema": json.dumps(input_schema_dict, indent=2),
+                "input_sample": json.dumps(input_sample_dict, indent=2),
+                "sdif_database_methods": sdif_methods_str,
+            }
+
+            # Instantiate the Agent
+            agent = Agent(
+                name="Tidy SDIF Adapter Agent",
+                mcp_servers=[self.mcp_server],
+                tools=[execute_tidy_adaptation],  # Use the decorated tools
+                model=self.llm_model,
+            )
+
+            # Run the agent using the Runner
+            # Pass the prompt and initial context
+            logger.info(f"Running Tidy Agent with model {self.llm_model}...")
+            result = await Runner.run(
+                agent,
+                input=TIDY_TRANSFORMATION_PROMPT.format(
+                    input_schema=initial_context["input_schema"],
+                    input_sample=initial_context["input_sample"],
+                    sdif_database_methods=initial_context["sdif_database_methods"],
+                ),
+            )
+
+            if not result or not result.final_output:
+                raise RuntimeError("Agent execution failed or returned no output.")
+
+            logger.info(
+                f"Agent finished. Final output message:\n{result.final_output[:500]}..."
+            )
+
+            # Parse the final code from the agent's response
+            final_code = self.parse_code(result.final_output)
+
+            if not final_code:
+                raise RuntimeError(
+                    f"Agent failed to produce valid final Python code in its response."
+                    f" Full response:\n{result.final_output}"
+                )
+
+            logger.info(
+                "Successfully parsed final adaptation code from agent response."
+            )
+            # print(f"--- Final Code ---\n{final_code}\n------------------") # Debugging
+
+            # Execute the *final* code using CodeAdapter directly to create the definitive output
+            logger.info("Executing final adaptation code...")
+            final_adapter = CodeAdapter(
+                function=final_code,
+                function_name="adapt_sdif",
+                output_suffix="_tidy_final",  # Use a distinct suffix for the final output
+            )
+            # Adapt the *original* copied input path
+            final_adapted_path = final_adapter.adapt(copied_input_path)
+
+            # Move the final successful output SDIF to a persistent location
+            # Example: place it next to the original input file
+            persistent_output_path = (
+                input_path.parent / final_adapted_path.name
+            ).resolve()
+            if persistent_output_path.exists():
+                logger.warning(
+                    f"Overwriting existing file at final destination: {persistent_output_path}"
+                )
+                persistent_output_path.unlink()
+
+            shutil.move(
+                str(final_adapted_path), persistent_output_path
+            )  # Move needs strings sometimes
+            logger.info(
+                f"Successfully generated final tidy SDIF: {persistent_output_path}"
+            )
+
+            return persistent_output_path
+
+        except Exception as e:
+            logger.exception(f"Error during TidyAdapter adapt process: {e}")
+            # Re-raise or handle as appropriate
+            raise
+        finally:
+            # Always clean up temporary files
+            self._cleanup_temp_env()

satif_ai/code_builders/adaptation.py

@@ -0,0 +1,9 @@
+from satif_core.code_builders.base import AsyncCodeBuilder, CodeBuilder
+
+
+class AdaptationCodeBuilder(CodeBuilder):
+    pass
+
+
+class AdaptationAsyncCodeBuilder(AsyncCodeBuilder):
+    pass

satif_ai/code_builders/transformation.py

@@ -0,0 +1,152 @@
+import os
+import re
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+from agents import Agent, Runner, function_tool
+from agents.mcp.server import MCPServerStdio
+from mcp import ClientSession
+from satif_core import AsyncCodeBuilder, CodeBuilder, SDIFDatabase
+from satif_sdk.comparators import get_comparator
+from satif_sdk.representers import get_representer
+from satif_sdk.transformers import CodeTransformer
+
+# Global variables for transformation
+INPUT_SDIF_PATH: Optional[Path] = None
+OUTPUT_TARGET_FILES: Optional[Dict[Union[str, Path], str]] = None
+
+
+@function_tool
+async def execute_transformation(code: str) -> str:
+    """Executes the transformation code on the input and returns the
+    comparison difference between the transformed output and the target output example.
+
+    Args:
+        code: The code to execute on the input.
+    """
+    if INPUT_SDIF_PATH is None or OUTPUT_TARGET_FILES is None:
+        return "Error: Transformation context not initialized"
+
+    code_transformer = CodeTransformer(function=code)
+    generated_output_path = code_transformer.export(INPUT_SDIF_PATH)
+
+    comparisons = []
+
+    if os.path.isdir(generated_output_path):
+        # If it's a directory, compare each file with its corresponding target
+        generated_files = os.listdir(generated_output_path)
+
+        for (
+            output_base_file,
+            output_target_file_name,
+        ) in OUTPUT_TARGET_FILES.items():
+            if output_target_file_name in generated_files:
+                generated_file_path = os.path.join(
+                    generated_output_path, output_target_file_name
+                )
+                comparator = get_comparator(output_target_file_name.split(".")[-1])
+                comparison = comparator.compare(generated_file_path, output_base_file)
+                comparisons.append(
+                    f"Comparison for {generated_file_path} [SOURCE] with {output_target_file_name} [TARGET]: {comparison}"
+                )
+            else:
+                comparisons.append(
+                    f"Error: {output_target_file_name} not found in the generated output"
+                )
+    else:
+        # If it's a single file, ensure there's only one target and compare
+        if len(OUTPUT_TARGET_FILES) == 1:
+            output_file = list(OUTPUT_TARGET_FILES.keys())[0]
+            output_target_file_name = list(OUTPUT_TARGET_FILES.values())[0]
+            comparator = get_comparator(output_file.split(".")[-1])
+            comparison = comparator.compare(generated_output_path, output_file)
+            comparisons.append(
+                f"Comparison for {generated_output_path} [SOURCE] with {output_target_file_name} [TARGET]: {comparison}"
+            )
+        else:
+            comparisons.append(
+                "Error: Single output file generated but multiple target files expected"
+            )
+
+    return "\n".join(comparisons)
+
+
+class TransformationCodeBuilder(CodeBuilder):
+    def __init__(self, output_example: Path | List[Path] | Dict[str, Path]):
+        self.output_example = output_example
+
+    def build(
+        self,
+        sdif: Path | SDIFDatabase,
+        instructions: Optional[str] = None,
+    ) -> str:
+        pass
+
+
+class TransformationAsyncCodeBuilder(AsyncCodeBuilder):
+    """This class is used to build a transformation code that will be used to transform a SDIF database into a set of files following the format of the given output files."""
+
+    def __init__(
+        self,
+        mcp_server: MCPServerStdio,
+        mcp_session: ClientSession,
+        llm_model: str = "o3-mini",
+    ):
+        self.mcp_server = mcp_server
+        self.mcp_session = mcp_session
+        self.llm_model = llm_model
+
+    async def build(
+        self,
+        sdif: Path,
+        output_target_files: Dict[Union[str, Path], str] | List[Path],
+        output_sdif: Optional[Path] = None,
+        instructions: Optional[str] = None,
+    ) -> str:
+        global INPUT_SDIF_PATH, OUTPUT_TARGET_FILES
+        INPUT_SDIF_PATH = Path(sdif)
+
+        if isinstance(output_target_files, list):
+            OUTPUT_TARGET_FILES = {file: file.name for file in output_target_files}
+        else:
+            OUTPUT_TARGET_FILES = output_target_files
+
+        input_schema = await self.mcp_session.read_resource(f"schema://{sdif}")
+        input_sample = await self.mcp_session.read_resource(f"sample://{sdif}")
+
+        output_schema = await self.mcp_session.read_resource(f"schema://{output_sdif}")
+        output_sample = await self.mcp_session.read_resource(f"sample://{output_sdif}")
+        output_representation = {
+            file: get_representer(file).represent(file)
+            for file in list(OUTPUT_TARGET_FILES.keys())
+        }
+
+        prompt = await self.mcp_session.get_prompt(
+            "create_transformation",
+            arguments={
+                "input_file": INPUT_SDIF_PATH.name,
+                "input_schema": input_schema.contents[0].text,
+                "input_sample": input_sample.contents[0].text,
+                "output_files": str(list(OUTPUT_TARGET_FILES.values())),
+                "output_schema": output_schema.contents[0].text,
+                "output_sample": output_sample.contents[0].text,
+                "output_representation": str(output_representation),
+            },
+        )
+        agent = Agent(
+            name="Transformation Builder",
+            mcp_servers=[self.mcp_server],
+            tools=[execute_transformation],
+            model=self.llm_model,
+        )
+        result = await Runner.run(agent, prompt.messages[0].content.text)
+        transformation_code = self.parse_code(result.final_output)
+        return transformation_code
+
+    def parse_code(self, code) -> str:
+        match = re.search(r"```(?:python)?(.*?)```", code, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        else:
+            # Handle case where no code block is found
+            return code.strip()

satif_ai/plot_builders/agent.py

@@ -0,0 +1,204 @@
+import logging
+import re
+from pathlib import Path
+from typing import Optional, Union
+
+from agents import Agent, Runner
+from agents.mcp import MCPServerStdio
+from mcp import ClientSession
+
+from satif_ai.plot_builders.prompt import PLOTTING_AGENT_PROMPT
+from satif_ai.plot_builders.tool import PLOTTING_TOOL_CONTEXT, execute_plotting_code
+
+logger = logging.getLogger(__name__)
+
+
+class PlottingAgent:
+    """Agent that generates Plotly plots from SDIF data based on user instructions."""
+
+    def __init__(
+        self,
+        mcp_server: MCPServerStdio,
+        mcp_session: ClientSession,
+        llm_model: str = "o4-mini",
+    ):
+        self.mcp_server = mcp_server
+        self.mcp_session = mcp_session
+        self.llm_model = llm_model
+
+    def _parse_final_path(self, final_text: str) -> Optional[Path]:
+        """Extracts the path from the success message."""
+        # Regex to find the path after "Success: Plot saved to "
+        match = re.search(r"Success: Plot saved to (.*)", final_text)
+        if match:
+            path_str = match.group(1).strip()
+            try:
+                p = Path(path_str)
+                # Check if it seems plausible (e.g., ends with .html and absolute)
+                # Check for existence here is important
+                if p.is_absolute() and p.name.endswith(".html") and p.exists():
+                    return p
+                elif (
+                    p.exists()
+                ):  # Accept relative path if it exists (less ideal but maybe happens)
+                    logger.warning(
+                        f"Parsed path {p} is not absolute but exists. Accepting."
+                    )
+                    return p.resolve()  # Return resolved absolute path
+            except Exception as e:
+                logger.warning(f"Error validating parsed path '{path_str}': {e}")
+                pass
+        # Fallback checks remain the same
+        if "plot.html" in final_text:
+            potential_path_str = final_text.strip()
+            # Try to extract if it's just the path
+            if Path(potential_path_str).name == "plot.html":
+                try:
+                    potential_path = Path(
+                        potential_path_str
+                    ).resolve()  # Resolve relative paths
+                    if potential_path.exists():
+                        logger.warning(
+                            "Agent returned path directly instead of success message."
+                        )
+                        return potential_path
+                except Exception:
+                    pass
+
+        return None
+
+    async def generate_plot(
+        self, sdif_path: Union[str, Path], instructions: str
+    ) -> Optional[Path]:
+        """
+        Generates a Plotly plot HTML file based on instructions and SDIF data.
+
+        Args:
+            sdif_path: Path to the input SDIF database file.
+            instructions: Natural language instructions for the plot.
+
+        Returns:
+            Path to the generated HTML plot file, or None if generation failed.
+
+        Raises:
+            FileNotFoundError: If the input SDIF file does not exist.
+            RuntimeError: If agent execution fails or context cannot be fetched or plot fails.
+            Exception: For other unexpected errors.
+        """
+        input_path = sdif_path
+        # Set tool context
+        PLOTTING_TOOL_CONTEXT["input_sdif_path"] = input_path
+        PLOTTING_TOOL_CONTEXT["user_instructions"] = instructions
+        PLOTTING_TOOL_CONTEXT["output_plot_path"] = None
+
+        agent_final_output_text = (
+            "Agent did not produce final output."  # Default message
+        )
+
+        try:
+            # Get Initial Context from MCP Resources
+            logger.info(
+                f"Fetching schema and sample for {input_path}..."
+            )  # Changed level to INFO
+            input_schema_str = "Error: Could not get schema."
+            input_sample_str = "Error: Could not get sample."
+            try:
+                input_path_str = str(input_path)
+                schema_uri = f"schema://{input_path_str}"
+                sample_uri = f"sample://{input_path_str}"
+                logger.debug(f"Requesting schema URI: {schema_uri}")
+                logger.debug(f"Requesting sample URI: {sample_uri}")
+
+                input_schema_resource = await self.mcp_session.read_resource(schema_uri)
+                input_sample_resource = await self.mcp_session.read_resource(sample_uri)
+
+                input_schema_str = (
+                    input_schema_resource.contents[0].text
+                    if input_schema_resource.contents
+                    else "Error: Could not get schema (empty response)."
+                )
+                input_sample_str = (
+                    input_sample_resource.contents[0].text
+                    if input_sample_resource.contents
+                    else "Error: Could not get sample (empty response)."
+                )
+
+            except Exception as mcp_err:
+                logger.error(f"Failed to get schema/sample via MCP: {mcp_err}")
+                raise RuntimeError(
+                    f"Failed to get required context via MCP: {mcp_err}"
+                ) from mcp_err
+
+            # Format the prompt
+            formatted_prompt = PLOTTING_AGENT_PROMPT.format(
+                input_sdif_path=str(input_path),
+                input_schema=input_schema_str,
+                input_sample=input_sample_str,
+                user_instructions=instructions,
+            )
+
+            # Instantiate the Agent
+            agent = Agent(
+                name="Plotting Agent",
+                mcp_servers=[self.mcp_server],
+                tools=[execute_plotting_code],
+                model=self.llm_model,
+            )
+
+            # Run the agent
+            logger.info(f"Running Plotting Agent with model {self.llm_model}...")
+            result = await Runner.run(
+                agent,
+                input=formatted_prompt,
+            )
+
+            if not result or not result.final_output:
+                raise RuntimeError(
+                    "Plotting agent execution failed or returned no output."
+                )
+
+            agent_final_output_text = (
+                result.final_output
+            )  # Store for potential error message
+            logger.info(
+                f"Plotting Agent finished. Final output:\n{agent_final_output_text}"
+            )
+
+            # Attempt to parse the path from the agent's final confirmation
+            final_plot_path = self._parse_final_path(agent_final_output_text)
+
+            if final_plot_path:  # Path found and exists
+                logger.info(
+                    f"Successfully confirmed plot generation at: {final_plot_path}"
+                )
+                return final_plot_path
+            else:
+                final_plot_path_from_context = PLOTTING_TOOL_CONTEXT.get(
+                    "output_plot_path"
+                )
+                if (
+                    final_plot_path_from_context
+                    and final_plot_path_from_context.exists()
+                ):
+                    logger.warning(
+                        "Parsed path from final output failed, but tool context has valid path."
+                    )
+                    return final_plot_path_from_context
+                else:
+                    logger.error(
+                        "Agent finished, but could not confirm successful plot generation or find output file."
+                    )
+                    # Include agent output in error for debugging
+                    raise RuntimeError(
+                        f"Agent finished, but plot generation failed or output path couldn't be determined. Agent final output: '{agent_final_output_text}'"
+                    )  # Modified Error
+
+        except Exception as e:
+            logger.exception(f"Error during PlottingAgent generate_plot: {e}")
+            raise  # Re-raise other exceptions
+        finally:
+            # Robust context cleanup using pop
+            PLOTTING_TOOL_CONTEXT.pop("input_sdif_path", None)
+            PLOTTING_TOOL_CONTEXT.pop("user_instructions", None)
+            PLOTTING_TOOL_CONTEXT.pop("output_plot_path", None)
+            logger.debug("Cleared plotting tool context.")

satif_ai/plot_builders/prompt.py

@@ -0,0 +1,92 @@
+# satif/plot_builders/prompt.py
+
+PLOTTING_AGENT_PROMPT = """
+You are an expert Data Visualization Agent specialized in creating insightful and interactive plots using Plotly from data stored in SDIF (SQLite) databases. You are autonomous and **must not ask clarifying questions**.
+
+**Goal:** Generate Python **script code** to create a Plotly visualization based on user instructions and data within the provided SDIF file. **Critically analyze the data (schema, sample) and instructions to infer the user's likely analytical goal. Prepare and transform the data as needed (e.g., cleaning types, handling missing values appropriately for the plot, calculating new fields), choose the most appropriate chart type (e.g., line for trends, bar for comparisons, scatter for correlations, histogram for distributions) and apply necessary data transformations (grouping, aggregation, pivoting) to best represent the data and answer the implied question in the instructions.** Use standard visualization best practices. Your objective is to produce an effective plot, not engage in conversation.
+
+**Execution Context:**
+Your code will be executed in an environment where the following variables are **already defined**:
+- `db`: An instance of `SDIFDatabase`, connected in read-only mode to the input SDIF file (`{input_sdif_path}`).
+- `instructions`: A string containing the user's request (`{user_instructions}`).
+
+**Input SDIF Context:**
+You have access to the following information about the input SDIF database (accessible via the `db` object):
+
+<input_schema>
+{input_schema}
+</input_schema>
+
+<input_sample>
+{input_sample}
+</input_sample>
+
+**Available Tools:**
+1.  `execute_sql(query: str) -> str`: Execute a read-only SQL query against the **input** SDIF database (using the available `db` object, e.g., `db.query(...)`) to inspect data further *before* writing your main plotting code. Use this only if absolutely necessary to confirm data characteristics crucial for choosing the **correct** plot type or transformation (e.g., checking cardinality for grouping, range for binning).
+2.  `execute_plotting_code(code: str) -> str`: Executes the Python **script code** you generate. Your script **MUST** use the pre-defined `db` and `instructions` variables, generate a Plotly figure, and **save it to an HTML file** named `plot.html` in the current directory (e.g., `fig.write_html('plot.html')`). This tool will return the absolute path to the generated 'plot.html' on success, or an error message on failure.
+
+**Workflow:**
+1.  **Analyze & Infer & Select:** Carefully review the user instructions, input schema, and sample data. **Infer the analytical goal. Based on the data types, cardinality, and instructions, determine the necessary data preparation steps (cleaning, type conversion, handling missing values suitable for the plot), select the *most appropriate* Plotly chart type, and identify required data aggregations (e.g., sum, mean, count) or transformations (e.g., grouping, calculating percentages, date extraction) needed to create an insightful visualization.** Do not ask for clarification.
+2.  **Explore (Minimal Use):** Only use `execute_sql` if essential for confirming data properties needed for your chosen preparation/chart/transformation strategy.
+3.  **Code Generation:** Write Python **script code** (NOT a function definition) that:
+    *   Imports necessary libraries (`pandas as pd`, `plotly.express as px` or `plotly.graph_objects as go`).
+    *   Uses the pre-defined `db` object to read the relevant data.
+    *   Uses the `instructions` string variable if helpful for parameterizing the plot (e.g., titles).
+    *   **Performs the necessary data preparation (cleaning, type conversion, handling NaNs/nulls appropriately) and transformations/aggregations identified in step 1 using pandas.**
+    *   Creates the Plotly figure using the **chosen appropriate chart type** and the prepared/transformed/aggregated data. Make axes labels clear and add an informative title.
+    *   **Crucially:** Saves the figure using `fig.write_html('plot.html')`.
+4.  **Execute:** Call the `execute_plotting_code` tool with your generated Python script code string. **You must call this tool.**
+5.  **Finalize:**
+    *   **If `execute_plotting_code` returns a success message:** Respond **only** with the success message provided by the tool (e.g., "Success: Plot saved to /path/to/plot.html").
+    *   **If `execute_plotting_code` returns an error message:** Respond **only** with the error message provided by the tool.
+
+**Example Script Code (Illustrating Transformation & Chart Choice):**
+```python
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go # Import go if needed
+
+# Assume 'db' and 'instructions' are pre-defined
+# Assume instructions = "Show average monthly revenue trend"
+try:
+    # Infer table and columns (e.g., 'transactions' with 'date', 'revenue')
+    df = db.read_table('transactions')
+
+    # --- Data Preparation ---
+    # Ensure date is datetime type
+    df['date'] = pd.to_datetime(df['date'], errors='coerce')
+    # Ensure revenue is numeric, handle errors (e.g., fill with 0 or drop)
+    df['revenue'] = pd.to_numeric(df['revenue'], errors='coerce').fillna(0)
+    # Drop rows where date conversion failed if necessary for plot
+    df = df.dropna(subset=['date'])
+
+    # --- Transformation for Plot ---
+    # Infer appropriate transformation: Group by month and calculate mean revenue
+    df['month'] = df['date'].dt.to_period('M').astype(str)
+    df_agg = df.groupby('month')['revenue'].mean().reset_index()
+
+    # --- Plotting ---
+    # Infer appropriate chart type: Line chart for trend
+    title = f"Average Monthly Revenue Trend (based on: {{instructions[:30]}}...)"
+    fig = px.line(df_agg, x='month', y='revenue', title=title, markers=True,
+                  labels={{'revenue':'Average Revenue', 'month':'Month'}}) # Clear labels
+
+    # Save plot - THIS IS REQUIRED
+    output_path = 'plot.html'
+    fig.write_html(output_path)
+    print(f"Plot successfully saved to {{output_path}}") # Optional print
+
+except Exception as e:
+    print(f"Error during plotting script execution: {{e}}")
+    raise # Re-raise exception
+```
+
+**CRITICAL INSTRUCTIONS:**
+- **DO NOT ask clarifying questions.** Analyze the data and instructions to infer the best approach.
+- **Prepare and transform the data as needed before plotting (handle types, NaNs, aggregate, etc.).**
+- **Choose the MOST APPROPRIATE chart type.**
+- **You MUST generate Python script code, NOT a function definition.**
+- **Your script code MUST use the pre-defined `db` and `instructions` variables.**
+- **You MUST call the `execute_plotting_code` tool with your generated script code.**
+- **Your final response MUST be ONLY the exact success or error message returned by the `execute_plotting_code` tool.** No extra explanations or conversation.
+"""

satif_ai/plot_builders/tool.py

@@ -0,0 +1,146 @@
+import logging
+import sqlite3  # Ensure sqlite3 is imported if needed
+from pathlib import Path
+
+import pandas as pd
+from agents import function_tool
+from satif_sdk import SDIFDatabase
+from satif_sdk.code_executors import CodeExecutionError, LocalCodeExecutor
+
+logger = logging.getLogger(__name__)
+
+# Global or context for the tool (similar to TidyAdapter)
+PLOTTING_TOOL_CONTEXT = {
+    "input_sdif_path": None,
+    "user_instructions": None,
+    "output_plot_path": None,
+}
+
+
+@function_tool
+async def execute_plotting_code(code: str) -> str:
+    """
+    Executes the provided Python plotting script code.
+    The code should use the pre-defined 'db' (SDIFDatabase instance)
+    and 'instructions' (string) variables.
+    The code MUST save the generated Plotly plot to 'plot.html'.
+    Returns the path to the saved plot on success or an error message.
+    """
+    input_sdif_path = PLOTTING_TOOL_CONTEXT.get("input_sdif_path")
+    user_instructions = PLOTTING_TOOL_CONTEXT.get("user_instructions")
+
+    if not input_sdif_path:
+        return "Error: Input SDIF path not found in tool context."
+    # Check existence *before* trying to open
+    if not Path(input_sdif_path).exists():
+        return f"Error: Input SDIF file not found at {input_sdif_path}."
+    if user_instructions is None:
+        # Allow empty instructions, pass empty string
+        user_instructions = ""
+        logger.warning(
+            "User instructions not found in tool context, using empty string."
+        )
+        # return "Error: User instructions not found in tool context."
+
+    # Use LocalCodeExecutor - WARNING: Insecure for untrusted code
+    executor = LocalCodeExecutor()
+
+    expected_output_filename = "plot.html"
+    # Resolve path relative to the current working directory
+    expected_output_path = Path(expected_output_filename).resolve()
+
+    # Clear previous plot if it exists
+    if expected_output_path.exists():
+        try:
+            expected_output_path.unlink()
+        except OSError as e:
+            logger.error(
+                f"Could not remove existing plot file {expected_output_path}: {e}"
+            )
+
+    # Prepare the extra context for the executor
+    db_instance = None
+    try:
+        # Instantiate the DB instance to be passed to the code
+        # Use read-only mode as the code should only read for plotting
+        db_instance = SDIFDatabase(input_sdif_path, read_only=True)
+
+        # Define the context that will be available to the executed code
+        code_context = {
+            "db": db_instance,
+            "instructions": user_instructions,
+            # Add any other context variables if needed by the executor/code
+        }
+
+        # The code provided by the agent is now expected to be a script
+        script_code = f"""
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+from pathlib import Path
+from satif.sdif_database.database import SDIFDatabase # Make class available for type hints etc.
+import sqlite3 # For potential use by pandas/db interaction
+
+# Pre-defined variables available:
+# db: SDIFDatabase instance connected to {input_sdif_path}
+# instructions: str = User's instructions
+
+# --- User's Script Code Start ---
+{code}
+# --- User's Script Code End ---
+"""
+        logger.debug(f"Executing plotting script code:\n---\n{code[:500]}...\n---")
+
+        # LocalCodeExecutor.execute expects 'db' and 'datasource' in its signature
+        # We pass our *actual* db instance as part of extra_context which overrides
+        # the dummy 'db' argument passed for signature compliance.
+        executor.execute(
+            code=script_code,
+            db=db_instance,  # Pass instance for signature, but it's also in extra_context
+            datasource=None,  # Not needed
+            extra_context=code_context,  # Pass db and instructions here
+        )
+
+        # Check if the expected output file was created
+        if expected_output_path.exists():
+            logger.info(
+                f"Plotting code executed successfully. Output: {expected_output_path}"
+            )
+            PLOTTING_TOOL_CONTEXT["output_plot_path"] = expected_output_path
+            return f"Success: Plot saved to {expected_output_path}"
+        else:
+            logger.error(
+                "Plotting code executed but output file 'plot.html' not found."
+            )
+            # Check if the error might be within the script's own error handling
+            # (This requires parsing the execution output, which LocalExecutor doesn't provide easily)
+            return "Error: Code executed (possibly with internal errors), but the expected output file 'plot.html' was not created."
+
+    except (
+        CodeExecutionError,
+        sqlite3.Error,
+        FileNotFoundError,
+        ValueError,
+        TypeError,
+    ) as e:
+        logger.error(f"Error executing plotting code or accessing DB: {e}")
+        # Attempt to provide more specific feedback if possible
+        error_message = f"Error executing plotting code: {e}"
+        # Look for common issues like table not found
+        if isinstance(e, pd.io.sql.DatabaseError) and "no such table" in str(e).lower():
+            error_message = f"Error executing plotting code: Table not found. {e}"
+        elif isinstance(e, KeyError):  # Pandas KeyError on column access
+            error_message = f"Error executing plotting code: Column not found or data processing error. {e}"
+
+        return error_message  # Return the formatted error
+    except Exception as e:
+        logger.exception("Unexpected error during plotting code execution via tool.")
+        return f"Unexpected Error during execution: {e}"
+    finally:
+        # Ensure the db instance created here is closed
+        if db_instance:
+            try:
+                db_instance.close()
+                logger.debug("Closed DB instance in plotting tool.")
+            except Exception as close_err:
+                logger.error(f"Error closing DB instance in plotting tool: {close_err}")

satif_ai-0.1.1.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 SyncPulse, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

satif_ai-0.1.1.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.3
+Name: satif-ai
+Version: 0.1.1
+Summary: AI Agents for Satif
+License: MIT
+Author: Bryan Djafer
+Author-email: bryan.djafer@syncpulse.fr
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: openai-agents (>=0.0.9,<0.0.10)
+Requires-Dist: satif-sdk (>=0.1.0,<1.0.0)
+Requires-Dist: sdif-mcp (>=0.1.0,<1.0.0)
+Description-Content-Type: text/markdown
+
+# SATIF AI
+

satif_ai-0.1.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+satif_ai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+satif_ai/adapters/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+satif_ai/adapters/tidy.py,sha256=cvMQvXi2D4mWI_eBgXNWwvNiVZ5Lgff6CgrxPKnNA6Y,19894
+satif_ai/code_builders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+satif_ai/code_builders/adaptation.py,sha256=E29YM0S6pMtAfB0uzSUexoeWKwXfF8iJVyYUCKWQz5k,188
+satif_ai/code_builders/transformation.py,sha256=fdegXe2jNYxbPWOoDKhpKws8ltO5we-aJQQAhBBYVWQ,5970
+satif_ai/plot_builders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+satif_ai/plot_builders/agent.py,sha256=Ncw7SL9qkpRN0hw76ezSo1K8vVQK6gcXFp8x8VFwqUI,8291
+satif_ai/plot_builders/prompt.py,sha256=m0W1SsxnB9_FhIYRumkthImJbK-7KUm4dygN3kjAXGk,6877
+satif_ai/plot_builders/tool.py,sha256=MeLnG_wFoITSVWZcNFsQLCi157O4L3wItQgolBa4fAw,5994
+satif_ai-0.1.1.dist-info/LICENSE,sha256=kS8EN6yAaGZd7V5z6GKSn_x3ozcZltrfRky4vMPRCw8,1072
+satif_ai-0.1.1.dist-info/METADATA,sha256=WyMhb5M6Qgj2Hp897-L0x26Wgx-nNhXAg4anRzoR5S4,670
+satif_ai-0.1.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+satif_ai-0.1.1.dist-info/entry_points.txt,sha256=_8RP7sSos2cwqYblhwHOK3hhRkLbP6XkbAju3vHWft8,40
+satif_ai-0.1.1.dist-info/RECORD,,

satif_ai-0.1.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.3
+Root-Is-Purelib: true
+Tag: py3-none-any

satif_ai-0.1.1.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+satif=satif.cli:main
+