dataface 0.1.2__py3-none-any.whl

dataface/ai/external_mcp.py

@@ -0,0 +1,610 @@
+"""External MCP client: discover and connect to user-configured MCP servers.
+
+On `dft chat` start, reads the user's existing MCP client configs (same files
+that `dft init mcp` writes), filters out the dft self-entry, spawns the remaining
+servers over stdio using the `mcp` Python client, and merges their tools into the
+agent's tool list namespaced as `<server>__<tool>` (double underscore, same convention
+as Claude Code's MCP client).
+
+Lifecycle: `ExternalMCPManager` is created once per `dft chat` session, opened
+before the chat loop starts, and closed in the CLI's top-level finally block.
+Per-session is mandatory because dbt MCP cold-starts in 5–10s; per-message
+respawn would make the agent unusable.
+
+Async/sync bridge: a single asyncio event loop lives on a background thread.
+`call_tool` bridges to it via `asyncio.run_coroutine_threadsafe(...).result()`.
+Never use `asyncio.run()` here — it tears down the loop and terminates the long-
+lived stdio subprocess sessions.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+import sys
+import threading
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+import mcp
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from rich.console import Console
+
+_console = Console(stderr=True)
+
+
+# ---------------------------------------------------------------------------
+# Data types
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ServerConfig:
+    """Parsed MCP server entry from a client config file."""
+
+    name: str
+    command: str
+    args: list[str] = field(default_factory=list)
+    env: dict[str, str] | None = None
+
+
+# ---------------------------------------------------------------------------
+# Self-MCP detection
+# ---------------------------------------------------------------------------
+
+_SELF_BASENAMES = {"dft", "dataface"}
+_UV_BASENAMES = {"uvx", "uv"}
+# Match python, python3, python3.10, python3.13, etc. — but not `pythonista` or similar
+_PYTHON_RE = re.compile(r"python\d*(\.\d+)?$")
+
+
+def is_self_mcp(cfg: ServerConfig) -> bool:
+    """Return True if cfg points at dft's own MCP server.
+
+    Patterns detected:
+    1. command basename in {dft, dataface} AND args[:2] == ["mcp", "serve"]
+    2. command basename in {uvx, uv} AND args contains a dataface/dft token
+       followed eventually by "mcp" "serve"
+    3. command is a Python interpreter AND args invoke -m dataface/-m dft
+       followed by "mcp" "serve"
+    """
+    basename = Path(cfg.command).name
+    args = cfg.args
+
+    # Pattern 1: dft/dataface directly
+    if basename in _SELF_BASENAMES:
+        return _has_mcp_serve(args)
+
+    # Pattern 2: uvx / uv run dataface mcp serve
+    # Require the three tokens to be immediately adjacent: <dft|dataface> mcp serve
+    if basename in _UV_BASENAMES:
+        for i, token in enumerate(args):
+            if token in _SELF_BASENAMES:
+                return (
+                    i + 2 < len(args)
+                    and args[i + 1] == "mcp"
+                    and args[i + 2] == "serve"
+                )
+        return False
+
+    # Pattern 3: python / python3 / python3.13 -m dataface mcp serve
+    # Require the three tokens immediately: <dft|dataface> mcp serve (no flags between)
+    if _PYTHON_RE.match(basename):
+        try:
+            m_idx = args.index("-m")
+        except ValueError:
+            return False
+        module_args = args[m_idx + 1 :]
+        if not module_args:
+            return False
+        if module_args[0] in _SELF_BASENAMES:
+            return (
+                len(module_args) >= 3
+                and module_args[1] == "mcp"
+                and module_args[2] == "serve"
+            )
+        return False
+
+    return False
+
+
+def _has_mcp_serve(args: list[str]) -> bool:
+    """True if args contains 'mcp' immediately followed by 'serve' anywhere."""
+    for i in range(len(args) - 1):
+        if args[i] == "mcp" and args[i + 1] == "serve":
+            return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Config discovery
+# ---------------------------------------------------------------------------
+
+# (relative_path, servers_key) for each JSON-based client config
+_JSON_CLIENT_CONFIGS: list[tuple[Path, str]] = [
+    (Path(".cursor/mcp.json"), "mcpServers"),
+    (Path(".mcp.json"), "mcpServers"),  # claude-code
+    (Path(".vscode/mcp.json"), "servers"),
+    (Path(".github/copilot/mcp.json"), "servers"),
+]
+
+# (relative_path, servers_key) for each TOML-based client config
+_TOML_CLIENT_CONFIGS: list[tuple[Path, str]] = [
+    (Path(".codex/config.toml"), "mcp_servers"),  # OpenAI Codex CLI
+]
+
+# Absolute paths (global configs, not project-relative).
+# These mirror the paths that `dft init mcp` writes, plus vendor-specific locations.
+#
+# ~/.cursor/mcp.json             — Cursor global MCP config
+# ~/.config/claude/config.json   — written by `dft init mcp claude` (Claude Code/CLI)
+# ~/Library/Application Support/Claude/claude_desktop_config.json  — macOS Claude Desktop
+_GLOBAL_JSON_CONFIGS: list[tuple[Path, str]] = [
+    (Path.home() / ".cursor" / "mcp.json", "mcpServers"),
+    (Path.home() / ".config" / "claude" / "config.json", "mcpServers"),
+]
+if sys.platform == "darwin":
+    _GLOBAL_JSON_CONFIGS.append(
+        (
+            Path.home()
+            / "Library"
+            / "Application Support"
+            / "Claude"
+            / "claude_desktop_config.json",
+            "mcpServers",
+        )
+    )
+
+
+def discover_servers(
+    config_dirs: list[Path] | None = None,
+    *,
+    global_configs: list[tuple[Path, str]] | None = None,
+) -> list[ServerConfig]:
+    """Read user's MCP configs, return external server configs (self filtered out).
+
+    Args:
+        config_dirs: directories to search for relative configs (default: [Path.cwd()])
+        global_configs: absolute-path configs to search (default: _GLOBAL_JSON_CONFIGS).
+            Pass [] in tests to prevent reading the developer's home directory.
+    """
+    search_dirs = config_dirs if config_dirs is not None else [Path.cwd()]
+    effective_globals = (
+        _GLOBAL_JSON_CONFIGS if global_configs is None else global_configs
+    )
+    seen_names: set[str] = set()
+    results: list[ServerConfig] = []
+
+    for search_dir in search_dirs:
+        for rel_path, servers_key in _JSON_CLIENT_CONFIGS:
+            cfg_path = search_dir / rel_path
+            _parse_json_config(cfg_path, servers_key, seen_names, results)
+        for rel_path, servers_key in _TOML_CLIENT_CONFIGS:
+            cfg_path = search_dir / rel_path
+            _parse_toml_config(cfg_path, servers_key, seen_names, results)
+
+    for cfg_path, servers_key in effective_globals:
+        _parse_json_config(cfg_path, servers_key, seen_names, results)
+
+    return results
+
+
+def _accept_server(
+    name: str,
+    cfg: ServerConfig,
+    seen_names: set[str],
+    results: list[ServerConfig],
+) -> None:
+    """Validate and register a ServerConfig if it passes all guards."""
+    if "__" in name:
+        _console.print(
+            f'[yellow]MCP: skipping "{name}" — server names must not contain __ '
+            f"(would break tool-name routing)[/yellow]"
+        )
+        return
+    if name in seen_names:
+        return  # first config wins
+    if is_self_mcp(cfg):
+        # Claim the name so a later global config with the same name cannot sneak in
+        # a non-dft server that happens to share the self-entry's name.
+        seen_names.add(name)
+        _console.print(
+            f'[dim]MCP: skipping self-entry "{name}" (using in-process tools)[/dim]'
+        )
+        return
+    seen_names.add(name)
+    results.append(cfg)
+
+
+def _parse_json_config(
+    path: Path,
+    servers_key: str,
+    seen_names: set[str],
+    results: list[ServerConfig],
+) -> None:
+    if not path.exists():
+        return
+    try:
+        data = json.loads(path.read_text())
+    except json.JSONDecodeError:
+        _console.print(f"[yellow]MCP: {path} has invalid JSON; skipping[/yellow]")
+        return
+    except OSError as exc:
+        _console.print(f"[yellow]MCP: could not read {path}: {exc}; skipping[/yellow]")
+        return
+    if not isinstance(data, dict):
+        return
+    servers = data.get(servers_key, {})
+    if not isinstance(servers, dict):
+        return
+    for name, entry in servers.items():
+        if not isinstance(entry, dict):
+            continue
+        command = entry.get("command", "")
+        if not command:
+            continue
+        cfg = ServerConfig(
+            name=name,
+            command=command,
+            args=entry.get("args") or [],
+            env=entry.get("env"),
+        )
+        _accept_server(name, cfg, seen_names, results)
+
+
+def _parse_toml_config(
+    path: Path,
+    servers_key: str,
+    seen_names: set[str],
+    results: list[ServerConfig],
+) -> None:
+    """Parse a TOML-format MCP config (Codex CLI's config.toml)."""
+    if not path.exists():
+        return
+    try:
+        data = tomllib.loads(path.read_text())
+    except (tomllib.TOMLDecodeError, OSError) as exc:
+        _console.print(
+            f"[yellow]MCP: {path} has invalid TOML: {exc}; skipping[/yellow]"
+        )
+        return
+    servers = data.get(servers_key, {})
+    if not isinstance(servers, dict):
+        return
+    for name, entry in servers.items():
+        if not isinstance(entry, dict):
+            continue
+        command = entry.get("command", "")
+        if not command:
+            continue
+        cfg = ServerConfig(
+            name=name,
+            command=command,
+            args=entry.get("args") or [],
+            env=entry.get("env"),
+        )
+        _accept_server(name, cfg, seen_names, results)
+
+
+# ---------------------------------------------------------------------------
+# Tool namespacing helpers
+# ---------------------------------------------------------------------------
+
+
+def _namespace_tool(server_name: str, tool: dict[str, Any]) -> dict[str, Any]:
+    """Rewrite tool name to <server>__<tool>.
+
+    Uses `__` as the separator — matches Claude Code's MCP convention and avoids
+    `:`, which both Anthropic and OpenAI APIs reject in tool names. Built-in
+    Dataface tool names do not contain `__`, so no collision is possible.
+    """
+    return {**tool, "name": f"{server_name}__{tool['name']}"}
+
+
+def _tools_for_config(server_name: str, mcp_tools: list[Any]) -> list[dict[str, Any]]:
+    """Convert mcp.types.Tool list to canonical tool dicts with namespaced names."""
+    result = []
+    for t in mcp_tools:
+        tool_dict = {
+            "name": t.name,
+            "description": t.description or "",
+            "input_schema": t.inputSchema,
+        }
+        result.append(_namespace_tool(server_name, tool_dict))
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Per-server holder: persistent async task that keeps context managers alive
+# ---------------------------------------------------------------------------
+
+
+class _ServerHolder:
+    """Runs a long-lived MCP session inside a persistent asyncio task.
+
+    The MCP stdio_client and ClientSession use anyio context managers that must
+    be entered AND exited within the same task. This holder runs a coroutine on
+    the background loop that keeps the `async with` blocks open for the session's
+    lifetime and processes call_tool requests via a shared Queue.
+    """
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.tools: list[dict[str, Any]] = []
+        # Set once the session is ready (or failed)
+        self._ready_event = threading.Event()
+        self._error: Exception | None = None
+        # Async queue: each item is (tool_name, args, threading.Event, result_box)
+        # None is the shutdown sentinel.
+        self._call_queue: (
+            asyncio.Queue[tuple[str, dict[str, Any], threading.Event, list[Any]] | None]
+            | None
+        ) = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    def error(self) -> Exception | None:
+        """Return the startup error, if any. Thread-safe after _ready_event is set."""
+        return self._error
+
+    def wait_ready(self, timeout: float) -> bool:
+        """Block until ready or timed out. Returns True if connected with no error."""
+        self._ready_event.wait(timeout=timeout)
+        return self.is_connected()
+
+    async def run(self, cfg: ServerConfig, startup_timeout: float) -> None:
+        """Persistent coroutine: open the MCP session, signal ready, service calls.
+
+        The startup_timeout applies only to initialize + tools/list. The service
+        loop has no timeout — it runs until the shutdown sentinel is received.
+        """
+        self._loop = asyncio.get_running_loop()
+        self._call_queue = asyncio.Queue()
+
+        params = StdioServerParameters(
+            command=cfg.command,
+            args=cfg.args,
+            env=cfg.env,
+        )
+        try:
+            async with stdio_client(params) as (read, write):  # noqa: SIM117
+                async with ClientSession(read, write) as session:
+                    # Apply timeout only to the startup phase — cannot combine these
+                    # three async with blocks: stdio_client + ClientSession must stay
+                    # alive for the whole session; the timeout wraps only startup.
+                    await asyncio.wait_for(
+                        session.initialize(), timeout=startup_timeout
+                    )
+                    tools_result = await asyncio.wait_for(
+                        session.list_tools(), timeout=startup_timeout
+                    )
+                    self.tools = _tools_for_config(cfg.name, tools_result.tools)
+                    self._ready_event.set()
+                    # Service call requests until shutdown sentinel (no timeout)
+                    while True:
+                        item = await self._call_queue.get()
+                        if item is None:
+                            break
+                        tool_name, arguments, done_event, result_box = item
+                        try:
+                            result = await session.call_tool(tool_name, arguments)
+                            result_box.append(result)
+                        except Exception as exc:  # noqa: BLE001 — keep loop alive
+                            result_box.append(exc)
+                        done_event.set()
+        except (mcp.McpError, asyncio.TimeoutError, OSError, EOFError) as exc:
+            # Startup/connection failure — unblock waiters with the error recorded.
+            self._error = exc
+            self._ready_event.set()
+
+    def is_connected(self) -> bool:
+        """True if session started successfully (no error, ready event set)."""
+        return self._error is None and self._ready_event.is_set()
+
+    def call_tool_sync(
+        self, tool_name: str, arguments: dict[str, Any], timeout: float | None
+    ) -> Any:
+        """Enqueue a call_tool request and wait for the result (sync)."""
+        if self._call_queue is None or self._loop is None:
+            raise RuntimeError("Session not running")
+        done_event = threading.Event()
+        result_box: list[Any] = []
+        item = (tool_name, arguments, done_event, result_box)
+        self._loop.call_soon_threadsafe(self._call_queue.put_nowait, item)
+        if not done_event.wait(timeout=timeout):
+            raise TimeoutError(f"call_tool({tool_name!r}) timed out after {timeout}s")
+        result = result_box[0]
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def shutdown(self) -> None:
+        if self._call_queue is not None and self._loop is not None:
+            self._loop.call_soon_threadsafe(self._call_queue.put_nowait, None)
+
+
+# ---------------------------------------------------------------------------
+# ExternalMCPManager — lifecycle + dispatch
+# ---------------------------------------------------------------------------
+
+
+class ExternalMCPManager:
+    """Manages external MCP server sessions for the lifetime of a dft chat session.
+
+    Usage::
+
+        manager = ExternalMCPManager(discover_servers())
+        manager.start(timeout=10.0)   # parallel startup; prints status
+        # ...run agent with manager.all_tools() merged in...
+        manager.close()               # in a finally block
+
+    Tool calls route via::
+
+        manager.call_tool("dbt__list_metrics", {"filter": "rev"})
+    """
+
+    def __init__(self, configs: list[ServerConfig]) -> None:
+        self._configs = configs
+        self._holders: dict[str, _ServerHolder] = {}
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._loop_thread: threading.Thread | None = None
+
+    # ------------------------------------------------------------------
+    # Startup
+    # ------------------------------------------------------------------
+
+    def start(self, timeout: float = 10.0) -> None:
+        """Spawn all servers in parallel and collect their tools.
+
+        Per-server timeout: `timeout` seconds for initialize + tools/list.
+        On failure, log a yellow warning and continue without that server.
+        """
+        if not self._configs:
+            return
+
+        self._loop = asyncio.new_event_loop()
+        self._loop_thread = threading.Thread(
+            target=self._loop.run_forever, daemon=True, name="dft-mcp-loop"
+        )
+        self._loop_thread.start()
+
+        names = ", ".join(c.name for c in self._configs)
+        _console.print(f"[dim]MCP: starting {names}...[/dim]")
+
+        # Launch one persistent task per server on the background loop.
+        for cfg in self._configs:
+            holder = _ServerHolder(cfg.name)
+            self._holders[cfg.name] = holder
+            asyncio.run_coroutine_threadsafe(
+                holder.run(cfg, startup_timeout=timeout), self._loop
+            )
+
+        # Wait for all servers in parallel — each holder has its own threading.Event.
+        # Use a shared deadline so servers that start fast don't wait for slow ones.
+        deadline = time.monotonic() + timeout + 1.0
+        for name, holder in self._holders.items():
+            remaining = max(0.0, deadline - time.monotonic())
+            holder.wait_ready(remaining)
+            if not holder.is_connected():
+                err = holder.error()
+                _console.print(
+                    f"[yellow]MCP: {name} failed to start "
+                    f"({err}); continuing without it[/yellow]"
+                )
+
+        loaded = [
+            f"{n} ({len(h.tools)} tools)"
+            for n, h in self._holders.items()
+            if h.is_connected()
+        ]
+        failed = [n for n, h in self._holders.items() if not h.is_connected()]
+        parts = loaded + ([f"{', '.join(failed)} failed"] if failed else [])
+        _console.print(
+            f"[dim]MCP: {'; '.join(parts) if parts else 'no external tools'}[/dim]"
+        )
+
+    # ------------------------------------------------------------------
+    # Tool access
+    # ------------------------------------------------------------------
+
+    def all_tools(self) -> list[dict[str, Any]]:
+        """Return all namespaced tool dicts from all loaded servers."""
+        result = []
+        for h in self._holders.values():
+            if h.is_connected():
+                result.extend(h.tools)
+        return result
+
+    # Default per-call timeout. dbt semantic-layer queries can legitimately take
+    # 60–120s on large warehouses; 120s gives a wide margin without hanging forever.
+    _DEFAULT_TOOL_TIMEOUT: float = 120.0
+
+    def call_tool(
+        self,
+        namespaced_name: str,
+        arguments: dict[str, Any],
+        timeout: float | None = None,
+    ) -> dict[str, Any]:
+        """Call an external tool by 'server__tool' name. Bridges async → sync.
+
+        Args:
+            timeout: per-call timeout in seconds (default: 120s). Pass None only
+                     if you are sure the upstream tool is bounded by its own logic.
+        """
+        if timeout is None:
+            timeout = self._DEFAULT_TOOL_TIMEOUT
+        if "__" not in namespaced_name:
+            raise ValueError(
+                f"Expected 'server__tool' format, got: {namespaced_name!r}"
+            )
+        server_name, tool_name = namespaced_name.split("__", 1)
+        holder = self._holders.get(server_name)
+        if holder is None or not holder.is_connected():
+            return {"error": f"MCP server '{server_name}' is not connected"}
+
+        try:
+            result = holder.call_tool_sync(tool_name, arguments, timeout=timeout)
+        except TimeoutError:
+            return {
+                "error": f"Tool call '{namespaced_name}' timed out after {timeout}s"
+            }
+        except Exception as exc:  # noqa: BLE001 — surface any unexpected failure
+            return {"error": f"Tool call '{namespaced_name}' failed: {exc}"}
+
+        return _flatten_call_result(result)
+
+    # ------------------------------------------------------------------
+    # Shutdown
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Signal all sessions to shut down and stop the event loop."""
+        for holder in self._holders.values():
+            holder.shutdown()
+        if self._loop is None:
+            return
+        # Brief sleep lets the shutdown sentinel drain the queue and the
+        # stdio_client/ClientSession context managers exit cleanly before we
+        # stop the loop. This is best-effort CLI teardown — precision is not
+        # required here; the daemon thread and subprocess will be killed on
+        # process exit regardless.
+        time.sleep(0.3)
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        if self._loop_thread:
+            self._loop_thread.join(timeout=3.0)
+        self._holders.clear()
+
+
+# ---------------------------------------------------------------------------
+# Flatten CallToolResult → plain dict
+# ---------------------------------------------------------------------------
+
+
+def _flatten_call_result(result: Any) -> dict[str, Any]:
+    """Convert mcp.types.CallToolResult to a plain dict the agent can use."""
+    if result.isError:
+        texts = [c.text for c in result.content if hasattr(c, "text")]
+        return {"error": "\n".join(texts) or "Tool returned an error"}
+
+    if result.structuredContent:
+        return result.structuredContent
+
+    texts = [c.text for c in result.content if hasattr(c, "text")]
+    if len(texts) == 1:
+        # Try to parse as JSON for structured results; only treat dicts as structured.
+        try:
+            parsed = json.loads(texts[0])
+            if isinstance(parsed, dict):
+                return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+        return {"content": texts[0]}
+    return {"content": "\n".join(texts)}

dataface/ai/generate_sql.py

@@ -0,0 +1,96 @@
+"""Shared text-to-SQL generation function.
+
+Single code path for SQL generation used by:
+- Eval runner (standalone question → SQL)
+- Cloud AIService (thin async wrapper)
+- Playground (SQL guidance for system prompts)
+
+Uses the rich schema context from get_schema_context() and the same
+prompt quality across all consumers.
+"""
+
+from __future__ import annotations
+
+import json
+
+from dataface.ai.llm import OpenAIClient
+
+_SQL_SYSTEM_PROMPT = """You are a SQL expert. Generate a single SQL query that answers the user's question.
+
+{schema_context}
+
+Rules:
+- Return ONLY valid SQL for the database dialect shown above.
+- Use clear column aliases so results are self-describing.
+- Include appropriate aggregations, groupings, and ordering.
+- Prefer explicit column names over SELECT *.
+- Do not wrap the SQL in markdown code fences.
+
+Respond with JSON: {{"sql": "<your SQL query>"}}"""
+
+_SQL_GUIDANCE = """## Shared SQL Generation Rules
+
+- Write SQL that matches the database dialect in the schema context.
+- Use real table and column names from the provided schema context.
+- Use clear column aliases so results are self-describing.
+- Include appropriate aggregations, groupings, and ordering.
+- Prefer explicit column names over `SELECT *`.
+- Return executable SQL, not pseudocode or placeholders.
+"""
+
+
+def get_sql_generation_guidance() -> str:
+    """Return the shared SQL rules used across text-to-SQL consumers."""
+    return _SQL_GUIDANCE
+
+
+def build_sql_system_prompt(schema_context: str) -> str:
+    """Build the shared system prompt for standalone SQL generation."""
+    return _SQL_SYSTEM_PROMPT.format(schema_context=schema_context)
+
+
+def generate_sql(
+    question: str,
+    schema_context: str,
+    *,
+    client: OpenAIClient,
+) -> str:
+    """Generate SQL from a natural-language question using the LLM.
+
+    Args:
+        question: Natural-language question about the data.
+        schema_context: Rich schema context string from get_schema_context()
+            or format_schema_context(). Includes table names, column types,
+            roles, semantic types, and value distributions.
+        client: An ``OpenAIClient`` instance.
+
+    Returns:
+        The generated SQL query string.
+    """
+    system_prompt = build_sql_system_prompt(schema_context)
+
+    response = client.create(
+        model=client.model,
+        input=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": question},
+        ],
+        text={
+            "format": {
+                "type": "json_schema",
+                "name": "SQLResponse",
+                "schema": {
+                    "type": "object",
+                    "properties": {
+                        "sql": {"type": "string"},
+                    },
+                    "required": ["sql"],
+                    "additionalProperties": False,
+                },
+                "strict": True,
+            }
+        },
+    )
+
+    parsed = json.loads(response.output_text)
+    return parsed["sql"]