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

opendatasci/tools/mcp.py

@@ -0,0 +1,179 @@
+"""MCP (Model Context Protocol) adapter: fetch tools from MCP servers and wrap as LangChain Tools."""
+
+import json
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import httpx
+from langchain_core.tools import BaseTool, StructuredTool
+from pydantic import Field, create_model
+
+OPENDATASCI_DIRNAME = ".opendatasci"
+_MCP_CONFIG_FILE = "mcp.json"
+
+_MCP_TIMEOUT = 30.0
+_JSONRPC_VERSION = "2.0"
+
+_JSON_SCHEMA_TYPE_MAP: dict[str, type] = {
+    "string": str,
+    "integer": int,
+    "number": float,
+    "boolean": bool,
+    "array": list,
+    "object": dict,
+}
+
+
+def _jsonrpc(method: str, params: dict[str, Any] | None = None, req_id: int = 1) -> dict[str, Any]:
+    payload: dict[str, Any] = {"jsonrpc": _JSONRPC_VERSION, "id": req_id, "method": method}
+    if params is not None:
+        payload["params"] = params
+    return payload
+
+
+def _initialize(url: str) -> None:
+    """Perform the MCP initialization handshake with the server."""
+    payload = _jsonrpc(
+        "initialize",
+        {
+            "protocolVersion": "2024-11-05",
+            "capabilities": {},
+            "clientInfo": {"name": "opendatasci", "version": "1.0"},
+        },
+    )
+    with httpx.Client() as client:
+        client.post(url, json=payload, timeout=_MCP_TIMEOUT).raise_for_status()
+        # Best-effort initialized notification (servers may not require it)
+        notif = {"jsonrpc": _JSONRPC_VERSION, "method": "notifications/initialized"}
+        try:
+            client.post(url, json=notif, timeout=_MCP_TIMEOUT)
+        except Exception:
+            pass
+
+
+def _list_tools(url: str) -> list[dict[str, Any]]:
+    """Fetch the tool manifest from an MCP server."""
+    payload = _jsonrpc("tools/list", req_id=2)
+    with httpx.Client() as client:
+        resp = client.post(url, json=payload, timeout=_MCP_TIMEOUT)
+        resp.raise_for_status()
+    data: dict[str, Any] = resp.json()
+    return data.get("result", {}).get("tools", [])  # type: ignore[no-any-return]
+
+
+def _build_args_model(tool_name: str, input_schema: dict[str, Any]) -> type:
+    """Convert a JSON Schema object into a Pydantic model for StructuredTool."""
+    properties: dict[str, Any] = input_schema.get("properties", {})
+    required: set[str] = set(input_schema.get("required", []))
+    fields: dict[str, Any] = {}
+
+    for prop_name, prop in properties.items():
+        py_type = _JSON_SCHEMA_TYPE_MAP.get(prop.get("type", "string"), str)
+        description = prop.get("description", "")
+        if prop_name in required:
+            fields[prop_name] = (py_type, Field(description=description))
+        else:
+            fields[prop_name] = (Optional[py_type], Field(default=None, description=description))
+
+    return create_model(f"_{tool_name}_args", **fields)  # type: ignore[no-any-return]
+
+
+def _make_mcp_tool(server_url: str, tool_def: dict[str, Any]) -> BaseTool:
+    """Wrap a single MCP tool definition as an async LangChain StructuredTool."""
+    name: str = tool_def["name"]
+    description: str = tool_def.get("description", "")
+    input_schema: dict[str, Any] = tool_def.get("inputSchema", {"type": "object", "properties": {}})
+    args_model = _build_args_model(name, input_schema)
+
+    async def _call(**kwargs: Any) -> str:
+        payload = _jsonrpc("tools/call", {"name": name, "arguments": kwargs}, req_id=3)
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.post(server_url, json=payload, timeout=_MCP_TIMEOUT)
+                resp.raise_for_status()
+                data = resp.json()
+        except Exception as exc:
+            return f"Error calling MCP tool '{name}': {type(exc).__name__}: {exc}"
+
+        if data.get("error"):
+            err = data["error"]
+            return f"MCP error {err.get('code', '')}: {err.get('message', str(err))}"
+
+        content = data.get("result", {}).get("content", [])
+        parts = [item.get("text", "") for item in content if item.get("type") == "text"]
+        if not parts:
+            return json.dumps(data.get("result", {}))
+        return "\n".join(parts)
+
+    return StructuredTool.from_function(
+        coroutine=_call,
+        name=name,
+        description=description,
+        args_schema=args_model,
+    )
+
+
+def load_mcp_servers(workspace_path: Path) -> list[str]:
+    """Read MCP server URLs from ``<workspace>/.opendatasci/mcp.json``.
+
+    The file format mirrors the Cursor ``mcp.json`` convention::
+
+        {
+            "mcpServers": {
+                "my-server": { "url": "http://localhost:8080" },
+                "another":   { "url": "http://localhost:9000" }
+            }
+        }
+
+    Returns an empty list when the file is absent, empty, or malformed
+    (a warning is printed to stderr in the latter case).
+    """
+    config_path = workspace_path / OPENDATASCI_DIRNAME / _MCP_CONFIG_FILE
+    if not config_path.exists():
+        return []
+
+    try:
+        data = json.loads(config_path.read_text(encoding="utf-8"))
+        servers: dict[str, dict[str, str]] = data.get("mcpServers", {})
+        return [entry["url"] for entry in servers.values() if "url" in entry]
+    except Exception as exc:
+        print(
+            f"Warning: Failed to parse {config_path}: {type(exc).__name__}: {exc}",
+            file=sys.stderr,
+        )
+        return []
+
+
+def create_mcp_tools(server_urls: list[str]) -> list[BaseTool]:
+    """Query each MCP server URL, fetch its tool manifest, and return wrapped LangChain tools.
+
+    Each MCP server's tools self-register — no ``ToolName`` enum entry is required.
+    Servers that fail to connect or respond are skipped with a warning to stderr.
+
+    Args:
+        server_urls: List of MCP server base URLs (e.g. ``["http://localhost:8080"]``).
+
+    Returns:
+        Flat list of LangChain ``BaseTool`` instances ready to inject into any agent.
+    """
+    tools: list[BaseTool] = []
+    for url in server_urls:
+        try:
+            _initialize(url)
+            tool_defs = _list_tools(url)
+        except Exception as exc:
+            print(
+                f"Warning: Failed to connect to MCP server {url!r}: {type(exc).__name__}: {exc}",
+                file=sys.stderr,
+            )
+            continue
+        for td in tool_defs:
+            try:
+                tools.append(_make_mcp_tool(url, td))
+            except Exception as exc:
+                print(
+                    f"Warning: Failed to wrap MCP tool {td.get('name')!r} from {url!r}: {exc}",
+                    file=sys.stderr,
+                )
+    return tools

opendatasci/tools/planning.py

@@ -0,0 +1,88 @@
+"""Plan-mode tools: enter_plan_mode and exit_plan_mode."""
+
+from typing import Annotated, Callable
+
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool, tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+
+from opendatasci.agents.states import AgentState
+
+
+def create_planning_tools(
+    save_plan: Callable[[str], None],
+) -> list[BaseTool]:
+    """Return ``enter_plan_mode`` and ``exit_plan_mode``.
+
+    Args:
+        save_plan: Callback that persists the final plan via ``BaseContextStore``.
+    """
+
+    @tool
+    def enter_plan_mode(
+        communication: str,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+    ) -> Command[AgentState]:
+        """Enter Plan Mode to decompose a complex task before executing it.
+
+        In Plan Mode you can think through the full problem and produce an ordered,
+        actionable plan. Call ``exit_plan_mode`` with the completed plan to return to execution.
+
+        # When to use this tool
+        - For tasks with more than two or three interdependent steps — e.g. building a
+          full ML pipeline, multi-stage analysis, or anything where step ordering matters.
+
+        # When NOT to use this tool
+        - For simple tasks — the overhead is wasteful.
+
+        Args:
+            communication: Brief message to the user about what you're doing
+                           (e.g. "This task has several interdependent steps — let me plan it first.").
+        """
+        return Command(
+            update={
+                "is_plan_mode": True,
+                "messages": [
+                    ToolMessage(
+                        content=(
+                            "Plan Mode active. Think through the full task carefully and produce a "
+                            "detailed, ordered plan. Call exit_plan_mode once your plan is complete."
+                        ),
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )
+
+    @tool
+    def exit_plan_mode(
+        final_plan: str,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+    ) -> Command[AgentState]:
+        """Exit Plan Mode and record the completed plan.
+
+        The plan is persisted and available as context throughout execution. Write each
+        step as a concise, single-action description; sequence steps so each one's output
+        feeds naturally into the next.
+
+        Args:
+            final_plan: The complete, ordered plan.
+        """
+        save_plan(final_plan)
+        return Command(
+            update={
+                "is_plan_mode": False,
+                "messages": [
+                    ToolMessage(
+                        content=(
+                            "Plan recorded and saved. You are back in execution mode. "
+                            "The plan is now part of your context — work through it step by step."
+                        ),
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )
+
+    return [enter_plan_mode, exit_plan_mode]

opendatasci/tools/skills.py

@@ -0,0 +1,90 @@
+"""Skills tool: load a specialised skill profile into the agent's system prompt."""
+
+from typing import Annotated
+
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool, tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+
+from opendatasci.agents.states import AgentState
+from opendatasci.skills import SKILL_LABELS, Skill
+from opendatasci.skills.base import BaseSkillStore
+
+
+def _label_for(skill: Skill) -> str:
+    return SKILL_LABELS.get(skill.name) or skill.name.replace("_", " ").title()
+
+
+def create_skill_tools(store: BaseSkillStore) -> list[BaseTool]:
+    """Return the skill tools bound to *store*."""
+    return [_create_skill_tool(store)]
+
+
+def _create_skill_tool(store: BaseSkillStore) -> BaseTool:
+    @tool
+    def load_skill(
+        skill: str,
+        summary: str,
+        communication: str,
+        state: Annotated[AgentState, InjectedState],
+        tool_call_id: Annotated[str, InjectedToolCallId],
+    ) -> Command[AgentState]:
+        """Load a specialised skill profile to sharpen domain-specific guidance.
+
+        Only one skill is active at a time; loading a new one replaces the previous.
+
+        # When to use this tool
+        - At the start of a domain-specific task to get targeted guidance and best practices.
+        - When switching task domains mid-session (e.g. from data wrangling to model training).
+
+        # Available skills
+        ``data_science``, ``competitive_data_science``,
+        ``competitive_data_science_v2``, ``machine_learning``,
+        ``deep_learning``, ``quantitative_analysis``, ``data_science_education``.
+
+        Args:
+            skill:         Profile name to load.
+            summary:       3-4 word status label (e.g. "Loading data science skill").
+            communication: Brief message to the user about what you're doing
+                           (e.g. "Let me load the data science skill for this task.").
+        """
+        active = state.active_skills
+        if active and active[0].name == skill:
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="This skill is already loaded.", tool_call_id=tool_call_id
+                        )
+                    ]
+                }
+            )
+
+        loaded = store.load(skill)
+        if loaded is None:
+            available = ", ".join(sorted(store.list()))
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=f"Unknown skill '{skill}'. Available: {available}",
+                            tool_call_id=tool_call_id,
+                        )
+                    ]
+                }
+            )
+
+        return Command(
+            update={
+                "active_skills": [loaded],
+                "messages": [
+                    ToolMessage(
+                        content=f"{_label_for(loaded)} skill loaded.", tool_call_id=tool_call_id
+                    )
+                ],
+            }
+        )
+
+    return load_skill

opendatasci/tools/user_interaction.py

@@ -0,0 +1,54 @@
+"""User interaction tools: ask_user_mcq."""
+
+from langchain_core.tools import BaseTool, tool
+from langgraph.types import interrupt
+
+
+def create_user_interaction_tools() -> list[BaseTool]:
+    """Return user interaction tools that pause the graph to ask the user a question.
+
+    Uses LangGraph's ``interrupt()`` mechanism: the graph is paused and its state
+    is persisted to the checkpointer until the caller resumes it via
+    ``Command(resume=answer)``.
+
+    Identical questions are deduplicated: the first answer is cached per tool
+    instance so the agent never asks the user the same MCQ twice.
+    """
+    _cache: dict[tuple[str, str, str, str], str] = {}
+
+    @tool
+    def ask_user_mcq(
+        question: str,
+        choice_a: str,
+        choice_b: str,
+        choice_c: str,
+    ) -> str:
+        """Ask the user a multiple-choice question when the task cannot proceed without their input.
+
+        Presents three predefined choices (A, B, C). The user may also type a free-form answer —
+        treat any response that doesn't match a choice as a custom answer.
+
+        # When to use this tool
+        - When the problem is genuinely underspecified and the right approach depends on
+          an unstated user goal.
+        - When you need the user's input to make an assumption — ask only when correctness cannot be verified by available means.
+
+        # When NOT to use this tool
+        - For technical decisions you can make yourself — do not delegate judgment.
+        - When a reasonable assumption would unblock the task — ask only if truly blocked.
+
+        Args:
+            question: The question to ask.
+            choice_a: Text for option A.
+            choice_b: Text for option B.
+            choice_c: Text for option C.
+        """
+        key = (question, choice_a, choice_b, choice_c)
+        if key in _cache:
+            return _cache[key]
+
+        answer: str = interrupt({"question": question, "choices": [choice_a, choice_b, choice_c]})
+        _cache[key] = answer
+        return answer
+
+    return [ask_user_mcq]

opendatasci/tools/web.py

@@ -0,0 +1,236 @@
+"""Web tools: web_search and fetch_url."""
+
+import re
+from collections.abc import Iterable
+from functools import lru_cache
+from typing import Annotated
+from urllib.parse import urlparse
+
+from annotated_types import Ge
+from langchain_core.tools import BaseTool, tool
+
+# Domains permitted for fetch_url.  A URL is allowed when its hostname equals
+# one of these entries *or* is a subdomain of one (e.g. "en.wikipedia.org"
+# matches "wikipedia.org").
+_SEARCH_limit: int = 10
+_SEARCH_SNIPPET_MAX_CHARS: int = 300
+
+_FETCH_ALLOWED_DOMAINS: frozenset[str] = frozenset(
+    {
+        # Code & documentation
+        "raw.githubusercontent.com",
+        "github.com",
+        "docs.python.org",
+        "pandas.pydata.org",
+        "numpy.org",
+        "scikit-learn.org",
+        "matplotlib.org",
+        "scipy.org",
+        # ONNX ecosystem
+        "onnx.ai",
+        "onnxruntime.ai",
+        # Research
+        "arxiv.org",
+        # Finance & open APIs
+        "finance.yahoo.com",
+        # Competitive data science
+        "kaggle.com",
+    }
+)
+
+
+_EMPTY_DOMAINS: frozenset[str] = frozenset()
+
+
+def _is_domain_allowed(
+    url: str,
+    extra: frozenset[str] = _EMPTY_DOMAINS,
+    override: frozenset[str] | None = None,
+) -> bool:
+    """Return True when *url*'s hostname is in or is a subdomain of the allowed set.
+
+    Args:
+        url:      The URL to check.
+        extra:    Additional domains unioned with the base set.
+        override: When provided, replaces ``_FETCH_ALLOWED_DOMAINS`` entirely
+                  before ``extra`` is applied.
+    """
+    try:
+        host = (urlparse(url).hostname or "").lower()
+    except Exception:
+        return False
+    base = override if override is not None else _FETCH_ALLOWED_DOMAINS
+    allowed = base | extra
+    return any(host == d or host.endswith("." + d) for d in allowed)
+
+
+def _clean_html(content: str) -> str:
+    """Return clean plain text extracted from *content* (HTML)."""
+    try:
+        from lxml import html as lxml_html
+
+        doc = lxml_html.fromstring(content)
+        for el in doc.xpath(
+            "//script | //style | //nav | //header | //footer | //aside | //noscript"
+        ):
+            parent = el.getparent()
+            if parent is not None:
+                parent.remove(el)
+        text: str = doc.text_content()
+    except Exception:
+        text = re.sub(
+            r"<\s*script\b[^>]*>.*?<\s*/\s*script\b[^>]*>",
+            "",
+            content,
+            flags=re.DOTALL | re.IGNORECASE,
+        )
+        text = re.sub(
+            r"<\s*style\b[^>]*>.*?<\s*/\s*style\b[^>]*>", "", text, flags=re.DOTALL | re.IGNORECASE
+        )
+        text = re.sub(r"<[^>]+>", "", text)
+
+    lines = [ln.strip() for ln in text.splitlines()]
+    text = "\n".join(ln for ln in lines if ln)
+    text = re.sub(r"\n{3,}", "\n\n", text)
+    return text
+
+
+@lru_cache(maxsize=16)
+async def _web_search_impl(query: str, limit: int) -> str:
+    try:
+        from duckduckgo_search import DDGS
+    except ImportError:
+        return "Error: duckduckgo-search is not installed. Run: pip install duckduckgo-search"
+
+    n = max(1, min(int(limit), _SEARCH_limit))
+    try:
+        results = [r async for r in DDGS().atext(query, limit=n)]  # type: ignore[attr-defined]
+    except Exception as exc:
+        return f"Error performing web search: {type(exc).__name__}: {exc}"
+
+    if not results:
+        return "No results found."
+
+    lines = []
+    for i, r in enumerate(results, 1):
+        line = f"{i}. {r['title']} — {r['href']}"
+        body = (r.get("body") or "").strip()[:_SEARCH_SNIPPET_MAX_CHARS]
+        if body:
+            line += f"\n   {body}"
+        lines.append(line)
+    return "\n".join(lines)
+
+
+@tool
+async def web_search(
+    query: str, summary: str, communication: str, limit: Annotated[int, Ge(1)] = 10
+) -> str:
+    """Search the web for resources, documentation, data sources, or reference pages.
+
+    Returns titles, URLs, and short snippets. Follow up with ``fetch_url`` to retrieve full content.
+
+    # When to use this tool
+    - To discover data sources, APIs, documentation, or research papers.
+    - When you don't know the exact URL of the resource you need.
+
+    # How to use this tool
+    - Keep queries specific: include key terms rather than full sentences.
+    - Follow up with ``fetch_url`` on the most relevant result to get full content.
+
+    Args:
+        query:         Search query (natural language or keywords).
+        summary:       3-4 word status label (e.g. "Searching BLS data").
+        communication: Brief message to the user about what you're doing
+                       (e.g. "Let me search for data sources that could be useful for this task.").
+        limit:         Number of results to return.
+    """
+    return await _web_search_impl(query, limit)
+
+
+def _make_fetch_url_tool(
+    extra: frozenset[str],
+    override: frozenset[str] | None = None,
+) -> BaseTool:
+    """Return a fetch_url tool bound to the given domain sets."""
+
+    if override is not None:
+        allowed_domains = override
+    else:
+        allowed_domains = _FETCH_ALLOWED_DOMAINS | extra
+
+    @lru_cache(maxsize=16)
+    async def _fetch_url_impl(url: str) -> str:
+        try:
+            import httpx
+        except ImportError:
+            return "Error: httpx is not installed. Run: pip install httpx"
+
+        if not _is_domain_allowed(url, extra, override):
+            host = urlparse(url).hostname or url
+            return (
+                f"Error: Domain '{host}' is not in the fetch allowlist. "
+                "Use web_search to find content from an allowed domain, then fetch that URL."
+            )
+
+        try:
+            async with httpx.AsyncClient(
+                follow_redirects=True,
+                timeout=20,
+                headers={"User-Agent": "Mozilla/5.0 (compatible; OpenDataSci/1.0)"},
+            ) as client:
+                response = await client.get(url)
+            response.raise_for_status()
+        except httpx.TimeoutException:
+            return "Error: Request timed out after 20 seconds."
+        except httpx.HTTPStatusError as exc:
+            return f"Error: HTTP {exc.response.status_code} fetching {url}"
+        except Exception as exc:
+            return f"Error fetching URL: {type(exc).__name__}: {exc}"
+
+        content_type = response.headers.get("content-type", "").lower()
+        if "html" in content_type:
+            return _clean_html(response.text)
+        return response.text
+
+    async def fetch_url(url: str, summary: str, communication: str) -> str:
+        return await _fetch_url_impl(url)
+
+    sorted_domains = ", ".join(sorted(allowed_domains))
+    fetch_url.__doc__ = (
+        f"Fetch the full plain-text content of a URL from an allowed domain.\n\n"
+        f"Allowed domains: {sorted_domains}\n\n"
+        f"# When to use this tool\n"
+        f"- When you have a specific URL from an allowed domain to retrieve.\n"
+        f"- To read documentation, papers, or data from a page found via ``web_search``.\n\n"
+        f"# When NOT to use this tool\n"
+        f"- When the target domain is not in the allowlist — use ``web_search`` instead\n"
+        f"  to find useful links that resolve to an allowed domain.\n\n"
+        f"Args:\n"
+        f"    url:           Full URL to fetch (must be from an allowed domain).\n"
+        f'    summary:       3-4 word status label (e.g. "Fetching BLS report").\n'
+        f"    communication: Brief message to the user about what you're doing\n"
+        f'                   (e.g. "Let me fetch this research paper.").\n'
+    )
+
+    return tool(fetch_url)
+
+
+def create_web_tools(
+    extra_web_domains: Iterable[str] = (),
+    override_web_domains: Iterable[str] | None = None,
+) -> list[BaseTool]:
+    """Return the web_search and fetch_url tools (main agent only).
+
+    Args:
+        extra_web_domains:   Additional hostnames (or apex domains) to permit
+            in ``fetch_url``, beyond the base allowlist.
+        override_web_domains: When provided, replaces the built-in allowlist
+            entirely.  ``extra_web_domains`` is still applied on top.
+    """
+    extra = frozenset(d.lower().strip() for d in extra_web_domains)
+    override = (
+        frozenset(d.lower().strip() for d in override_web_domains)
+        if override_web_domains is not None
+        else None
+    )
+    return [web_search, _make_fetch_url_tool(extra, override)]