open-data-sci 0.1.0__py3-none-any.whl

opendatasci/tools/dataset_info.py

@@ -0,0 +1,391 @@
+"""Dataset context tools and profiling code generation.
+
+Exposes:
+- ``build_profile_code(path)``    – sandbox-ready profiling snippet generator.
+- ``get_load_dataset_info_tool(session)``  – load persistent dataset info.
+- ``get_profile_dataset_tool(session)``    – profile a dataset and cache the card.
+- ``get_data_context_tools(session)``      – all three tools as a list.
+"""
+
+from langchain_core.tools import BaseTool, tool
+
+from opendatasci.context.base import BaseContextStore
+from opendatasci.sandbox.base import BaseSandbox
+
+# ---------------------------------------------------------------------------
+# Profiling code generation
+# ---------------------------------------------------------------------------
+# All curly-braces that belong to *Python* code (f-strings, dict literals,
+# format specs) are doubled so they survive the .format(path_repr=...) call.
+# Only ``{path_repr}`` is a real format placeholder.
+# ---------------------------------------------------------------------------
+
+_PROFILE_CODE_TEMPLATE = """\
+import datetime
+import pathlib
+
+import pandas as pd
+
+_path = pathlib.Path({path_repr})
+_ext = _path.suffix.lower()
+
+_LOADERS = {{
+    ".csv":    lambda p: pd.read_csv(p, low_memory=False),
+    ".tsv":    lambda p: pd.read_csv(p, sep="\\t", low_memory=False),
+    ".txt":    lambda p: pd.read_csv(p, low_memory=False),
+    ".parquet": lambda p: pd.read_parquet(p),
+    ".pq":     lambda p: pd.read_parquet(p),
+    ".json":   lambda p: pd.read_json(p),
+    ".jsonl":  lambda p: pd.read_json(p, lines=True),
+    ".ndjson": lambda p: pd.read_json(p, lines=True),
+    ".xlsx":   lambda p: pd.read_excel(p),
+    ".xls":    lambda p: pd.read_excel(p),
+    ".xlsm":   lambda p: pd.read_excel(p),
+    ".feather": lambda p: pd.read_feather(p),
+    ".ftr":    lambda p: pd.read_feather(p),
+    ".pkl":    lambda p: pd.read_pickle(p),
+}}
+
+_loader = _LOADERS.get(_ext)
+try:
+    if _loader is not None:
+        _df = _loader(_path)
+    else:
+        _df = pd.read_csv(_path, low_memory=False)
+except Exception as _load_err:
+    result = f"__PROFILE_SKIP__Cannot load {{_path.name}} ({{type(_load_err).__name__}}: {{_load_err}})"
+else:
+    _rows, _cols = _df.shape
+    _mem_mb = _df.memory_usage(deep=True).sum() / 1_048_576
+    _dupes = int(_df.duplicated().sum())
+    _ts = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
+
+    # --- Per-column unique counts (computed once, reused below) ---
+    _unique_counts = {{col: int(_df[col].nunique(dropna=True)) for col in _df.columns}}
+
+    # --- Pandera schema inference ---
+    # Used for (a) nullable-integer detection and (b) the Inferred Schema section.
+    _pa_schema = None
+    _pa_nullable = {{}}   # col -> bool
+    try:
+        import pandera as pa
+        _pa_schema = pa.infer_schema(_df)
+        _pa_nullable = {{n: c.nullable for n, c in _pa_schema.columns.items()}}
+    except Exception:
+        pass
+
+    # --- Data quality flags ---
+    _dq_flags = []
+    for _col in _df.columns:
+        _nuniq   = _unique_counts[_col]
+        _nonnull = int(_df[_col].notna().sum())
+        _null_pct = ((_rows - _nonnull) / _rows * 100) if _rows > 0 else 0.0
+        _dtype_s  = str(_df[_col].dtype)
+
+        if _nuniq == 1:
+            _dq_flags.append((_col, "Constant", "Only 1 unique value across all rows"))
+        if _rows > 0 and _nuniq == _rows and _null_pct == 0:
+            _dq_flags.append((_col, "Likely ID", "100% unique, no nulls"))
+        elif _rows > 0 and _nuniq / _rows > 0.95 and _null_pct < 1 and _dtype_s == "object":
+            _dq_flags.append((_col, "High cardinality", f"{{_nuniq / _rows * 100:.1f}}% unique values"))
+        if _null_pct > 50:
+            _dq_flags.append((_col, "High nulls", f"{{_null_pct:.1f}}% missing"))
+        # Pandera marks these nullable=True; cross-check that all non-null values are
+        # whole numbers to confirm the column is really an integer column with nulls.
+        if _dtype_s.startswith("float") and _pa_nullable.get(_col, False):
+            _sample = _df[_col].dropna()
+            if len(_sample) > 0 and (_sample % 1 == 0).all():
+                _dq_flags.append((
+                    _col, "Nullable int",
+                    f"Stored as {{_dtype_s}}; all non-null values are integers — "
+                    "consider casting to a nullable integer dtype",
+                ))
+
+    _lines = [
+        f"# Dataset Profile: {{_path.name}}",
+        "",
+        f"**Profiled:** {{_ts}}  ",
+        f"**Path:** `{{_path}}`",
+        "",
+        "## Overview",
+        "",
+        "| Metric | Value |",
+        "|--------|-------|",
+        f"| Rows | {{_rows:,}} |",
+        f"| Columns | {{_cols:,}} |",
+        f"| Memory | {{_mem_mb:.2f}} MB |",
+        f"| Duplicate rows | {{_dupes:,}} |",
+        "",
+        "## Columns",
+        "",
+        "| Column | Type | Non-null | Null % | Unique | Unique % |",
+        "|--------|------|----------|--------|--------|----------|",
+    ]
+
+    for _col in _df.columns:
+        _dtype   = str(_df[_col].dtype)
+        _nonnull = int(_df[_col].notna().sum())
+        _null_pct = ((_rows - _nonnull) / _rows * 100) if _rows > 0 else 0.0
+        _nuniq   = _unique_counts[_col]
+        _uniq_pct = (_nuniq / _rows * 100) if _rows > 0 else 0.0
+        _lines.append(
+            f"| `{{_col}}` | {{_dtype}} | {{_nonnull:,}} | {{_null_pct:.1f}}% "
+            f"| {{_nuniq:,}} | {{_uniq_pct:.1f}}% |"
+        )
+
+    # --- Data quality flags ---
+    if _dq_flags:
+        _lines += [
+            "",
+            "## Data Quality Flags",
+            "",
+            "| Column | Flag | Detail |",
+            "|--------|------|--------|",
+        ]
+        for _dq_col, _dq_flag, _dq_detail in _dq_flags:
+            _lines.append(f"| `{{_dq_col}}` | {{_dq_flag}} | {{_dq_detail}} |")
+
+    # --- Numeric summary ---
+    _num_cols = _df.select_dtypes(include="number").columns.tolist()
+    if _num_cols:
+        _lines += [
+            "",
+            "## Numeric Summary",
+            "",
+            "| Column | Mean | Std | Min | p25 | p50 | p75 | Max |",
+            "|--------|------|-----|-----|-----|-----|-----|-----|",
+        ]
+        _desc = _df[_num_cols].describe(percentiles=[0.25, 0.5, 0.75])
+
+        def _fmt(v):
+            if pd.isna(v):
+                return "—"
+            return f"{{v:,.2f}}" if abs(v) < 1_000_000 else f"{{v:.3e}}"
+
+        for _col in _num_cols:
+            _lines.append(
+                f"| `{{_col}}` "
+                f"| {{_fmt(_desc.loc['mean', _col])}} "
+                f"| {{_fmt(_desc.loc['std', _col])}} "
+                f"| {{_fmt(_desc.loc['min', _col])}} "
+                f"| {{_fmt(_desc.loc['25%', _col])}} "
+                f"| {{_fmt(_desc.loc['50%', _col])}} "
+                f"| {{_fmt(_desc.loc['75%', _col])}} "
+                f"| {{_fmt(_desc.loc['max', _col])}} |"
+            )
+
+    # --- Top categoricals ---
+    try:
+        _cat_dtypes = ["object", "category", "string"]
+        _cat_cols = _df.select_dtypes(include=_cat_dtypes).columns.tolist()
+    except Exception:
+        _cat_cols = _df.select_dtypes(include=["object", "category"]).columns.tolist()
+
+    if _cat_cols:
+        _lines.append("")
+        _lines.append("## Top Categoricals")
+        for _col in _cat_cols:
+            _vc = _df[_col].value_counts().head(5)
+            if _vc.empty:
+                continue
+            _lines += [
+                "",
+                f"### `{{_col}}`",
+                "",
+                "| Value | Count | % |",
+                "|-------|-------|---|",
+            ]
+            for _val, _cnt in _vc.items():
+                _pct = (_cnt / _rows * 100) if _rows > 0 else 0.0
+                _display = str(_val)
+                if len(_display) > 40:
+                    _display = _display[:37] + "..."
+                _lines.append(f"| {{_display!r}} | {{_cnt:,}} | {{_pct:.1f}}% |")
+
+    # --- Pandera inferred schema (dtype + nullable per column) ---
+    if _pa_schema is not None:
+        _lines += [
+            "",
+            "## Inferred Schema (pandera)",
+            "",
+            "| Column | Inferred dtype | Nullable |",
+            "|--------|----------------|----------|",
+        ]
+        for _pa_col_name, _pa_col in _pa_schema.columns.items():
+            _pa_col_nullable = "yes" if _pa_col.nullable else "no"
+            _lines.append(
+                f"| `{{_pa_col_name}}` | {{_pa_col.dtype}} | {{_pa_col_nullable}} |"
+            )
+
+    result = "\\n".join(_lines) + "\\n"
+"""
+
+
+def build_profile_code(path: str) -> str:
+    """Return a sandbox-ready Python snippet that profiles the dataset at *path*.
+
+    The snippet sets ``result`` to the Markdown profile card on success, or to
+    a ``__PROFILE_SKIP__<reason>`` sentinel string when the file cannot be
+    loaded (so the caller can surface a helpful message without saving garbage).
+
+    Args:
+        path: Absolute filesystem path to the dataset file.  Must be absolute
+              so the snippet runs correctly regardless of the sandbox CWD.
+    """
+    return _PROFILE_CODE_TEMPLATE.format(path_repr=repr(path))
+
+
+# ---------------------------------------------------------------------------
+# LangChain tools
+# ---------------------------------------------------------------------------
+
+
+def create_read_dataset_info_tools(context: BaseContextStore | None) -> list[BaseTool]:
+    """Return the ``read_dataset_info`` tool bound to *context*."""
+
+    @tool
+    async def read_dataset_info(path: str, summary: str, communication: str) -> str:
+        """Load all accumulated knowledge about a dataset — profile card and agent notes.
+
+        Returns a Markdown string with two sections:
+        - ``# DATASET PROFILING`` — auto-generated profile card (shape, dtypes, null rates,
+          numeric summary, top categoricals).
+        - ``# DATASET NOTES`` — cumulative agent notes from all prior sessions. Primary source
+          of institutional knowledge; read carefully before exploring.
+
+        # When to use this tool
+        - Always, before writing any code that reads or processes a dataset.
+        - Even if you profiled the dataset earlier in the session — notes may have been
+          updated by ``update_dataset_info`` since then.
+
+        Args:
+            path:          Absolute or relative path to the dataset file or directory.
+            summary:       3-4 word status label (e.g. "Reading train.csv info").
+            communication: Brief message to the user about what you're doing
+                           (e.g. "Let me read existing notes about the dataset.").
+        """
+        if context is None:
+            return "Error: No workspace path available."
+        try:
+            return await context.read_dataset_info(path)
+        except FileNotFoundError as exc:
+            return f"Error: {exc}"
+        except Exception as exc:
+            return f"Error loading dataset info: {type(exc).__name__}: {exc}"
+
+    return [read_dataset_info]
+
+
+def create_profile_dataset_tools(
+    context: BaseContextStore | None, sandbox: BaseSandbox, persist: bool = True
+) -> list[BaseTool]:
+    """Return the ``profile_dataset`` tool bound to *context* and *sandbox*."""
+
+    @tool
+    async def profile_dataset(path: str, summary: str, communication: str) -> str:
+        """Auto-profile a dataset and return its card (shape, dtypes, null rates, distributions).
+
+        Profiles are cached — subsequent calls return the existing card without re-scanning.
+        The card covers: shape, dtypes, null rates, numeric summary, top categoricals,
+        data-quality flags, and inferred schema.
+
+        # When to use this tool
+        - Once per new dataset, before exploring it.
+        - When ``read_dataset_info`` returns an empty or stale profile for a dataset.
+
+        # When NOT to use this tool
+        - When a profile was already returned by ``read_dataset_info`` this session — it is current.
+
+        Args:
+            path:          Absolute or relative path to the dataset file or directory.
+            summary:       3-4 word status label (e.g. "Profiling train.csv").
+            communication: Brief message to the user about what you're doing
+                           (e.g. "Let me profile the dataset to understand its structure.").
+        """
+        if context is None:
+            return "Error: No workspace path available."
+        try:
+            wc = context
+            resolved, hash_hex, existing = await wc.get_profile_info(path)
+
+            if existing is not None:
+                return existing
+
+            code = build_profile_code(resolved)
+            exec_result = await sandbox.execute(code)
+
+            if not exec_result.success:
+                return f"Profiling failed: {exec_result.error}"
+
+            if exec_result.output is None:
+                return "Profiling produced no output — the dataset may be empty."
+
+            content = str(exec_result.output)
+            if content.startswith("__PROFILE_SKIP__"):
+                return content[len("__PROFILE_SKIP__") :]
+
+            if persist:
+                wc.save_dataset_profile(hash_hex, content)
+            return content
+
+        except FileNotFoundError as exc:
+            return f"Error: {exc}"
+        except Exception as exc:
+            return f"Error profiling dataset: {type(exc).__name__}: {exc}"
+
+    return [profile_dataset]
+
+
+def create_update_dataset_info_tools(context: BaseContextStore | None) -> list[BaseTool]:
+    """Return the ``update_dataset_info`` tool bound to *context*."""
+
+    @tool
+    async def update_dataset_info(path: str, update: str, merge: bool = False) -> str:
+        """Persist findings, observations, and decisions for a dataset across sessions.
+
+        This is how knowledge carries forward — everything written here is surfaced
+        automatically next session via ``read_dataset_info``. Skipping means the
+        next session starts blind to the new knowledge and context about you've gathered
+        about the data and tasks you've tackled around it.
+
+        # When to use this tool
+        - After every turn that touches a dataset: findings, quality issues, surprises,
+            failed approaches, user decisions, and hypotheses to revisit.
+        - Err on the side of capturing more — even minor or confirmatory findings.
+
+        Args:
+            path:   Absolute or relative path to the dataset file or directory.
+            update: Markdown content to write or append. Make it structured.
+            merge:  ``True`` to append to existing notes (default);
+                    ``False`` to overwrite all prior notes.
+        """
+        if context is None:
+            return "Error: No workspace path available."
+        try:
+            return await context.update_dataset_info(path, update, merge=merge)
+        except FileNotFoundError as exc:
+            return f"Error: {exc}"
+        except Exception as exc:
+            return f"Error updating dataset info: {type(exc).__name__}: {exc}"
+
+    return [update_dataset_info]
+
+
+def create_data_context_tools(
+    context: BaseContextStore | None, sandbox: BaseSandbox, persist: bool = True
+) -> list[BaseTool]:
+    """Return dataset-context tools bound to *context* and *sandbox*.
+
+    Args:
+        context: I/O boundary for dataset notes and profiles.
+        sandbox:      Execution sandbox used by ``profile_dataset``.
+        persist:      When ``False``, ``update_dataset_info`` is excluded and
+                      ``profile_dataset`` will not write profiles to disk.
+    """
+    tools = [
+        *create_read_dataset_info_tools(context),
+        *create_profile_dataset_tools(context, sandbox, persist=persist),
+    ]
+    if persist:
+        tools.extend(create_update_dataset_info_tools(context))
+    return tools

opendatasci/tools/factory.py

@@ -0,0 +1,172 @@
+"""Tool factories: assemble the right tool sets for main and worker agents."""
+
+from collections.abc import Callable
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from langchain_core.tools import BaseTool
+
+from opendatasci.context.base import BaseContextStore
+from opendatasci.sandbox.base import BaseSandbox, BaseSandboxFactory
+from opendatasci.skills import BaseSkillStore
+from opendatasci.skills.local import LocalSkillStore
+from opendatasci.tools.coding import (
+    create_cli_tools,
+    create_code_verification_tools,
+    create_coding_tools,
+)
+from opendatasci.tools.critic import create_critic_tools
+from opendatasci.tools.dataset_info import create_data_context_tools
+from opendatasci.tools.mcp import create_mcp_tools
+from opendatasci.tools.planning import create_planning_tools
+from opendatasci.tools.skills import create_skill_tools
+from opendatasci.tools.user_interaction import create_user_interaction_tools
+from opendatasci.tools.web import create_web_tools
+from opendatasci.tools.workers import create_worker_tools
+from opendatasci.tools.workspace import create_workspace_tools
+from opendatasci.workspace.base import BaseWorkspace
+from opendatasci.workspace.local import LocalWorkspace
+
+if TYPE_CHECKING:
+    from opendatasci.configs import OpenDataSciConfig
+
+
+class ToolName(str, Enum):
+    """Canonical names for all agent tools."""
+
+    EXECUTE_PYTHON_CODE = "execute_python_code"
+    EXECUTE_CLI = "execute_cli_command"
+    LIST_PYTHON_LIBS = "list_python_libs"
+    LOAD_SKILL = "load_skill"
+    ENTER_PLAN_MODE = "enter_plan_mode"
+    EXIT_PLAN_MODE = "exit_plan_mode"
+    ENTER_SELF_REVIEW_MODE = "enter_self_review_mode"
+    EXIT_SELF_REVIEW_MODE = "exit_self_review_mode"
+    SPAWN_WORKERS = "spawn_workers"
+    READ_DATASET_INFO = "read_dataset_info"
+    UPDATE_DATASET_INFO = "update_dataset_info"
+    PROFILE_DATASET = "profile_dataset"
+    LIST_WORKSPACE_FILES = "list_workspace_files"
+    WEB_SEARCH = "web_search"
+    FETCH_URL = "fetch_url"
+    ASK_USER_MCQ = "ask_user_mcq"
+    VERIFY_PYTHON_CODE = "verify_python_code"
+
+
+def _base_tools(
+    workspace: BaseWorkspace,
+    sandbox: BaseSandbox,
+    context: "BaseContextStore | None",
+    store: BaseSkillStore,
+    persist: bool = True,
+) -> list[BaseTool]:
+    """Return the tools shared by both main and worker agents.
+
+    Args:
+        workspace:    Workspace container.
+        sandbox:      Code execution sandbox.
+        context: I/O boundary for dataset notes and profiles.
+        store:        Skill store used by the ``load_skill`` tool.
+        persist:      When ``False``, write-side tools (``update_dataset_info``)
+                      are excluded and ``profile_dataset`` will not write profiles
+                      to disk.
+    """
+    tools: list[BaseTool] = [
+        *create_coding_tools(sandbox),
+        *create_cli_tools(sandbox),
+        *create_data_context_tools(context, sandbox, persist=persist),
+        *create_skill_tools(store),
+    ]
+    if isinstance(workspace, LocalWorkspace):
+        tools.extend(create_workspace_tools(Path(workspace.get_reference())))
+    return tools
+
+
+def create_worker_agent_tools(
+    workspace: BaseWorkspace,
+    context: "BaseContextStore | None",
+    sandbox: BaseSandbox | None = None,
+    store: BaseSkillStore | None = None,
+) -> list[BaseTool]:
+    """Return the tool list for a worker agent.
+
+    Workers share the same core tools as the main agent but cannot spawn
+    further workers, plan, or access the web.
+
+    Args:
+        workspace:    Workspace container.
+        context: I/O boundary for dataset notes and profiles.
+        sandbox:      Code execution sandbox.  A new :class:`~opendatasci.sandbox.srt.SRTSandbox`
+                      is created when ``None``.
+        store:        Skill store injected from the caller.  Defaults to a
+                      :class:`~opendatasci.skills.local.LocalSkillStore` rooted
+                      at ``<context.root>/skills``.
+    """
+    if sandbox is None:
+        from opendatasci.sandbox.srt import SRTSandbox
+
+        sandbox = SRTSandbox(workspace_path=Path(workspace.get_reference()))
+    if store is None:
+        user_skills_dir = Path(context.root) / "skills" if context is not None else None
+        store = LocalSkillStore([user_skills_dir] if user_skills_dir is not None else None)
+    return _base_tools(workspace, sandbox, context, store, persist=False)
+
+
+def create_agent_tools(
+    workspace: BaseWorkspace,
+    sandbox: BaseSandbox,
+    context: "BaseContextStore | None",
+    sandbox_factory: BaseSandboxFactory,
+    store: BaseSkillStore | None = None,
+    datasci_config: "OpenDataSciConfig | None" = None,
+    save_plan: "Callable[[str], None] | None" = None,
+) -> list[BaseTool]:
+    """Return the tool list for the main agent.
+
+    Extends the worker tool set with planning, worker spawning, web access,
+    and user interaction.
+
+    Args:
+        workspace:       Workspace container.
+        sandbox:         Code execution sandbox.
+        context:         I/O boundary for dataset notes and profiles.
+        sandbox_factory: Factory used by spawned workers to create their own
+                         isolated sandboxes.
+        store:           Skill store injected from the caller.  Defaults to a
+                         :class:`~opendatasci.skills.local.LocalSkillStore` rooted
+                         at ``<context.root>/skills``.
+        datasci_config:  LLM configuration forwarded to spawned workers.
+        save_plan:       Callback that persists the final plan via
+                         ``BaseContextStore``.  When provided,
+                         ``enter_plan_mode`` and ``exit_plan_mode`` are added
+                         to the tool list.
+    """
+    if store is None:
+        user_skills_dir = Path(context.root) / "skills" if context is not None else None
+        store = LocalSkillStore([user_skills_dir] if user_skills_dir is not None else None)
+    tools = _base_tools(workspace, sandbox, context, store)
+    if datasci_config is not None:
+        tools.extend(create_code_verification_tools(datasci_config))
+    if save_plan is not None:
+        tools.extend(create_planning_tools(save_plan))
+    tools.extend(create_critic_tools(store))
+    tools.extend(
+        create_worker_tools(
+            workspace,
+            context,
+            datasci_config,
+            store=store,
+            sandbox_factory=sandbox_factory,
+        )
+    )
+    tools.extend(
+        create_web_tools(
+            datasci_config.extra_web_domains if datasci_config else (),
+            datasci_config.override_web_domains if datasci_config else None,
+        )
+    )
+    tools.extend(create_user_interaction_tools())
+    if datasci_config is not None and datasci_config.mcp_servers:
+        tools.extend(create_mcp_tools(datasci_config.mcp_servers))
+    return tools