zoo_mcp 0.9.2__py3-none-any.whl

zoo_mcp/__init__.py

@@ -0,0 +1,77 @@
+"""Zoo Model Context Protocol (MCP) Server.
+
+A lightweight service that enables AI assistants to execute Zoo commands through the Model Context Protocol (MCP).
+"""
+
+import asyncio
+import logging
+import ssl
+import sys
+from importlib.metadata import PackageNotFoundError, version
+
+import truststore
+from kittycad import KittyCAD
+
+FORMAT = "%(asctime)s | %(levelname)-7s | %(filename)s:%(lineno)d | %(funcName)s | %(message)s"
+
+logging.basicConfig(
+    level=logging.INFO, format=FORMAT, handlers=[logging.StreamHandler(sys.stderr)]
+)
+logger = logging.getLogger("zoo_mcp")
+
+
+try:
+    __version__ = version("zoo_mcp")
+except PackageNotFoundError:
+    # package is not installed
+    logger.error("zoo-mcp package is not installed.")
+
+
+class ZooMCPException(Exception):
+    """Custom exception for Zoo MCP Server."""
+
+
+ctx = truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+kittycad_client = KittyCAD(verify_ssl=ctx)
+# set the websocket receive timeout to 5 minutes
+kittycad_client.websocket_recv_timeout = 300
+
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.WARNING)
+
+
+def _initialize_kcl_docs() -> None:
+    """Initialize KCL documentation cache at module load time."""
+    from zoo_mcp.kcl_docs import initialize_docs_cache
+
+    try:
+        try:
+            loop = asyncio.get_running_loop()
+            # Already in an event loop, schedule the task
+            loop.create_task(initialize_docs_cache())
+        except RuntimeError:
+            # No running loop, use asyncio.run()
+            asyncio.run(initialize_docs_cache())
+    except Exception as e:
+        logger.warning(f"Failed to initialize KCL docs cache: {e}")
+
+
+def _initialize_kcl_samples() -> None:
+    """Initialize KCL samples cache at module load time."""
+    from zoo_mcp.kcl_samples import initialize_samples_cache
+
+    try:
+        try:
+            loop = asyncio.get_running_loop()
+            # Already in an event loop, schedule the task
+            loop.create_task(initialize_samples_cache())
+        except RuntimeError:
+            # No running loop, use asyncio.run()
+            asyncio.run(initialize_samples_cache())
+    except Exception as e:
+        logger.warning(f"Failed to initialize KCL samples cache: {e}")
+
+
+# Initialize caches when module is imported
+_initialize_kcl_docs()
+_initialize_kcl_samples()

zoo_mcp/__main__.py

@@ -0,0 +1,15 @@
+import sys
+
+from zoo_mcp import logger
+from zoo_mcp.server import mcp
+
+if __name__ == "__main__":
+    try:
+        logger.info("Starting MCP server...")
+        mcp.run(transport="stdio")
+
+    except KeyboardInterrupt:
+        logger.info("Shutting down MCP server...")
+    except Exception as e:
+        logger.exception("Server encountered an error: %s", e)
+        sys.exit(1)

zoo_mcp/ai_tools.py

@@ -0,0 +1,256 @@
+import asyncio
+import time
+from pathlib import Path
+
+from kittycad.models import (
+    ApiCallStatus,
+    FileExportFormat,
+    TextToCadCreateBody,
+    TextToCadMultiFileIterationBody,
+)
+from kittycad.models.ml_copilot_server_message import EndOfStream, Reasoning, ToolOutput
+from kittycad.models.reasoning_message import (
+    OptionCreatedKclFile,
+    OptionDeletedKclFile,
+    OptionDesignPlan,
+    OptionFeatureTreeOutline,
+    OptionGeneratedKclCode,
+    OptionKclCodeError,
+    OptionKclCodeExamples,
+    OptionKclDocs,
+    OptionMarkdown,
+    OptionText,
+    OptionUpdatedKclFile,
+)
+from kittycad.models.text_to_cad_response import (
+    OptionTextToCad,
+    OptionTextToCadMultiFileIteration,
+)
+from websockets.exceptions import ConnectionClosedError
+
+from zoo_mcp import ZooMCPException, kittycad_client, logger
+
+
+def log_websocket_message(conn_id: str) -> bool:
+    logger.info("Connecting to Text-To-CAD websocket...")
+    with kittycad_client.ml.ml_reasoning_ws(id=conn_id) as ws:
+        logger.info(
+            "Successfully connected to Text-To-CAD websocket with id %s", conn_id
+        )
+        while True:
+            try:
+                message = ws.recv()
+                if isinstance(message.root, Reasoning):
+                    message_option = message.root.reasoning.root
+                    match message_option:
+                        case OptionCreatedKclFile():
+                            logger.info(
+                                "Created %s: %s"
+                                % (message_option.file_name, message_option.content),
+                            )
+                        case OptionDeletedKclFile():
+                            logger.info("Deleted %s", message_option.file_name)
+                        case OptionDesignPlan():
+                            design_steps = " ".join(
+                                [
+                                    f"Editing: {step.filepath_to_edit} with these instruction {step.edit_instructions}"
+                                    for step in message_option.steps
+                                ]
+                            )
+                            logger.info("Design Plan: %s", design_steps)
+                        case OptionFeatureTreeOutline():
+                            logger.info(
+                                "Feature Tree Outline: %s", message_option.content
+                            )
+                        case OptionGeneratedKclCode():
+                            logger.info("Generated KCL code: %s", message_option.code)
+                        case OptionKclCodeError():
+                            logger.info("KCL Code Error: %s", message_option.error)
+                        case OptionKclDocs():
+                            logger.info("KCL Docs: %s", message_option.content)
+                        case OptionKclCodeExamples():
+                            logger.info("KCL Code Examples: %s", message_option.content)
+                        case OptionMarkdown():
+                            logger.info(message_option.content)
+                        case OptionText():
+                            logger.info(message_option.content)
+                        case OptionUpdatedKclFile():
+                            logger.info(
+                                "Updated %s: %s"
+                                % (message_option.file_name, message_option.content),
+                            )
+                        case _:
+                            logger.info(
+                                "Received unhandled reasoning message: %s",
+                                type(message_option),
+                            )
+                if isinstance(message.root, ToolOutput):
+                    tool_result = message.root.result.root
+                    if tool_result.error:
+                        logger.info(
+                            "Tool: %s, Error: %s"
+                            % (tool_result.type, tool_result.error)
+                        )
+                    else:
+                        logger.info(
+                            "Tool: %s, Output: %s"
+                            % (tool_result.type, tool_result.outputs)
+                        )
+                if isinstance(message.root, EndOfStream):
+                    logger.info("Text-To-CAD reasoning complete.")
+                    return True
+
+            except ConnectionClosedError as e:
+                logger.info(
+                    "Text To CAD could still be running but the websocket connection closed with error: %s",
+                    e,
+                )
+                return False
+
+            except Exception as e:
+                logger.info(
+                    "Text To CAD could still be running but an unexpected error occurred: %s",
+                    e,
+                )
+                return False
+
+
+async def text_to_cad(prompt: str) -> str:
+    """Send a prompt to Zoo's Text-To-CAD create endpoint
+
+    Args:
+        prompt (str): a description of the CAD model to be created
+
+    Returns:
+        A string containing the complete KCL code of the CAD model if Text-To-CAD was successful, otherwise an error
+        message from Text-To-CAD
+    """
+
+    logger.info("Sending prompt to Text-To-CAD")
+
+    # send prompt via the kittycad client
+    t2c = kittycad_client.ml.create_text_to_cad(
+        output_format=FileExportFormat.STEP,
+        kcl=True,
+        body=TextToCadCreateBody(
+            prompt=prompt,
+        ),
+    )
+
+    # get the response based on the request id
+    result = kittycad_client.ml.get_text_to_cad_part_for_user(id=t2c.id)
+
+    # check if the request has either completed or failed, otherwise sleep and try again
+    time_start = time.time()
+    ws_complete = False
+    while result.root.status not in [ApiCallStatus.COMPLETED, ApiCallStatus.FAILED]:
+        if (
+            result.root.status == ApiCallStatus.QUEUED
+            and (time.time() - time_start) % 5 == 0
+        ):
+            logger.info("Text-To-CAD queued...")
+        if result.root.status == ApiCallStatus.IN_PROGRESS:
+            logger.info("Text-To-CAD in progress...")
+            if not ws_complete:
+                ws_complete = log_websocket_message(t2c.id)
+        logger.info(
+            "Waiting for Text-To-CAD to complete... status %s", result.root.status
+        )
+        result = kittycad_client.ml.get_text_to_cad_part_for_user(id=t2c.id)
+        await asyncio.sleep(1)
+
+    logger.info("Received response from Text-To-CAD")
+
+    # get the data object (root) of the response
+    response = result.root
+
+    # check the data type of the response
+    if not isinstance(response, OptionTextToCad):
+        return "Error: Text-to-CAD response is not of type OptionTextToCad."
+
+    # if Text To CAD was successful return the KCL code, otherwise return the error
+    if response.status == ApiCallStatus.COMPLETED:
+        if response.code is None:
+            return "Error: Text-to-CAD response is null."
+        return response.code
+    else:
+        if response.error is None:
+            return "Error: Text-to-CAD response is null."
+        return response.error
+
+
+async def edit_kcl_project(
+    prompt: str,
+    proj_path: Path | str,
+) -> dict | str:
+    """Send a prompt and a KCL project to Zoo's Text-To-CAD edit KCL project endpoint. The proj_path will upload all contained files to the endpoint. There must be a main.kcl file in the root of the project.
+
+    Args:
+        prompt (str): A description of the changes to be made to the KCL project associated with the provided KCL files.
+        proj_path (Path | str): A path to a directory containing a main.kcl file. All contained files (found recursively) will be sent to the endpoint.
+
+    Returns:
+        dict | str: A dictionary containing the complete KCL code of the CAD model if Text-To-CAD edit KCL project was successful.
+                    Each key in the dict refers to a KCL file path relative to the project path, and each value is the complete KCL code for that file.
+                    If unsuccessful, returns an error message from Text-To-CAD.
+    """
+    logger.info("Sending KCL code prompt to Text-To-CAD edit kcl project")
+
+    logger.info("Finding all files in project path")
+    proj_path = Path(proj_path)
+    file_paths = list(proj_path.rglob("*"))
+    file_paths = [fp for fp in file_paths if fp.is_file()]
+    logger.info("Found %s files in project path", len(file_paths))
+
+    if not file_paths:
+        logger.error("No files paths provided or found in project path")
+        raise ZooMCPException("No file paths provided or found in project path")
+
+    if ".kcl" not in [fp.suffix for fp in file_paths]:
+        logger.error("No .kcl files found in the provided project path")
+        raise ZooMCPException("No .kcl files found in the provided project path")
+
+    if not (proj_path / "main.kcl").is_file():
+        logger.error("No main.kcl file found in the root of the provided project path")
+        raise ZooMCPException(
+            "No main.kcl file found in the root of the provided project path"
+        )
+
+    file_attachments = {
+        str(fp.relative_to(proj_path)): str(fp.resolve()) for fp in file_paths
+    }
+
+    t2cmfi = kittycad_client.ml.create_text_to_cad_multi_file_iteration(
+        body=TextToCadMultiFileIterationBody(
+            source_ranges=[],
+            prompt=prompt,
+        ),
+        file_attachments=file_attachments,
+    )
+
+    log_websocket_message(t2cmfi.id)
+
+    # get the response based on the request id
+    result = kittycad_client.ml.get_text_to_cad_part_for_user(id=t2cmfi.id)
+
+    # check if the request has either completed or failed, otherwise sleep and try again
+    while result.root.status not in [ApiCallStatus.COMPLETED, ApiCallStatus.FAILED]:
+        result = kittycad_client.ml.get_text_to_cad_part_for_user(id=t2cmfi.id)
+        await asyncio.sleep(1)
+
+    # get the data object (root) of the response
+    response = result.root
+
+    # check the data type of the response
+    if not isinstance(response, OptionTextToCadMultiFileIteration):
+        return "Error: Text-to-CAD response is not of type OptionTextToCadMultiFileIteration."
+
+    # if Text To CAD iteration was successful return the KCL code, otherwise return the error
+    if response.status == ApiCallStatus.COMPLETED:
+        if response.outputs is None:
+            return "Error: Text-to-CAD edit kcl project response is null."
+        return response.outputs
+    else:
+        if response.error is None:
+            return "Error: Text-to-CAD edit kcl project response is null."
+        return response.error

zoo_mcp/kcl_docs.py

@@ -0,0 +1,268 @@
+"""KCL Documentation fetching and search.
+
+This module fetches KCL documentation from the modeling-app GitHub repository
+at server startup and provides search functionality for LLMs.
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import ClassVar
+
+import httpx
+
+from zoo_mcp import logger
+
+GITHUB_TREE_URL = (
+    "https://api.github.com/repos/KittyCAD/modeling-app/git/trees/main?recursive=1"
+)
+RAW_CONTENT_BASE = "https://raw.githubusercontent.com/KittyCAD/modeling-app/main/"
+
+
+@dataclass
+class KCLDocs:
+    """Container for documentation data."""
+
+    docs: dict[str, str] = field(default_factory=dict)
+    index: dict[str, list[str]] = field(
+        default_factory=lambda: {
+            "kcl-lang": [],
+            "kcl-std-functions": [],
+            "kcl-std-types": [],
+            "kcl-std-consts": [],
+            "kcl-std-modules": [],
+        }
+    )
+
+    _instance: ClassVar["KCLDocs | None"] = None
+
+    @classmethod
+    def get(cls) -> "KCLDocs":
+        """Get the cached docs instance, or empty cache if not initialized."""
+        return cls._instance if cls._instance is not None else cls()
+
+    @classmethod
+    async def initialize(cls) -> None:
+        """Initialize the docs cache from GitHub."""
+        if cls._instance is None:
+            cls._instance = await _fetch_docs_from_github()
+
+
+def _categorize_doc_path(path: str) -> str | None:
+    """Categorize a doc path into one of the index categories."""
+    if path.startswith("docs/kcl-lang/"):
+        return "kcl-lang"
+    elif path.startswith("docs/kcl-std/functions/"):
+        return "kcl-std-functions"
+    elif path.startswith("docs/kcl-std/types/"):
+        return "kcl-std-types"
+    elif path.startswith("docs/kcl-std/consts/"):
+        return "kcl-std-consts"
+    elif path.startswith("docs/kcl-std/modules/"):
+        return "kcl-std-modules"
+    elif path.startswith("docs/kcl-std/"):
+        # Other kcl-std files (index.md, README.md)
+        return None
+    return None
+
+
+def _extract_title(content: str) -> str:
+    """Extract the title from Markdown content (first # heading)."""
+    for line in content.split("\n"):
+        line = line.strip()
+        if line.startswith("# "):
+            return line[2:].strip()
+    return ""
+
+
+def _extract_excerpt(content: str, query: str, context_chars: int = 200) -> str:
+    """Extract an excerpt around the first match of query in content."""
+    query_lower = query.lower()
+    content_lower = content.lower()
+
+    pos = content_lower.find(query_lower)
+    if pos == -1:
+        # Return first context_chars of content as fallback
+        return content[:context_chars].strip() + "..."
+
+    # Find start and end positions for excerpt
+    start = max(0, pos - context_chars // 2)
+    end = min(len(content), pos + len(query) + context_chars // 2)
+
+    # Adjust to word boundaries
+    if start > 0:
+        # Find the start of the word
+        while start > 0 and content[start - 1] not in " \n\t":
+            start -= 1
+
+    if end < len(content):
+        # Find the end of the word
+        while end < len(content) and content[end] not in " \n\t":
+            end += 1
+
+    excerpt = content[start:end].strip()
+
+    # Add ellipsis
+    prefix = "..." if start > 0 else ""
+    suffix = "..." if end < len(content) else ""
+
+    return f"{prefix}{excerpt}{suffix}"
+
+
+async def _fetch_doc_content(
+    client: httpx.AsyncClient, path: str
+) -> tuple[str, str | None]:
+    """Fetch a single doc file's content."""
+    url = f"{RAW_CONTENT_BASE}{path}"
+    try:
+        response = await client.get(url)
+        response.raise_for_status()
+        return path, response.text
+    except httpx.HTTPError as e:
+        logger.warning(f"Failed to fetch {path}: {e}")
+        return path, None
+
+
+async def _fetch_docs_from_github() -> KCLDocs:
+    """Fetch all docs from GitHub and return a KCLDocs."""
+    docs = KCLDocs()
+
+    logger.info("Fetching KCL documentation from GitHub...")
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        # 1. Get file tree from GitHub API
+        try:
+            response = await client.get(GITHUB_TREE_URL)
+            response.raise_for_status()
+            tree_data = response.json()
+        except httpx.HTTPError as e:
+            logger.warning(f"Failed to fetch GitHub tree: {e}")
+            return docs
+
+        # 2. Filter for docs/*.md files
+        doc_paths: list[str] = []
+        for item in tree_data.get("tree", []):
+            path = item.get("path", "")
+            if (
+                path.startswith("docs/")
+                and path.endswith(".md")
+                and item.get("type") == "blob"
+            ):
+                doc_paths.append(path)
+
+        logger.info(f"Found {len(doc_paths)} documentation files")
+
+        # 3. Fetch raw content in parallel
+        tasks = [_fetch_doc_content(client, path) for path in doc_paths]
+        results = await asyncio.gather(*tasks)
+
+        # 4. Populate cache and index
+        for path, content in results:
+            if content is not None:
+                docs.docs[path] = content
+
+                # Categorize the doc
+                category = _categorize_doc_path(path)
+                if category and category in docs.index:
+                    docs.index[category].append(path)
+
+    # Sort the index lists
+    for category in docs.index:
+        docs.index[category].sort()
+
+    logger.info(f"KCL documentation cache initialized with {len(docs.docs)} files")
+    return docs
+
+
+async def initialize_docs_cache() -> None:
+    """Initialize the docs cache from GitHub."""
+    await KCLDocs.initialize()
+
+
+def list_available_docs() -> dict[str, list[str]]:
+    """Return categorized list of available documentation.
+
+    Returns a dictionary with the following categories:
+    - kcl-lang: KCL language documentation (syntax, types, functions, etc.)
+    - kcl-std-functions: Standard library function documentation
+    - kcl-std-types: Standard library type documentation
+    - kcl-std-consts: Standard library constants documentation
+    - kcl-std-modules: Standard library module documentation
+
+    Each category contains a list of documentation file paths that can be
+    retrieved using get_kcl_doc().
+
+    Returns:
+        dict: Categories mapped to lists of available documentation paths.
+    """
+    return KCLDocs.get().index
+
+
+def search_docs(query: str, max_results: int = 5) -> list[dict]:
+    """Search docs by keyword.
+
+    Searches across all KCL language and standard library documentation
+    for the given query. Returns relevant excerpts with surrounding context.
+
+    Args:
+        query (str): The search query (case-insensitive).
+        max_results (int): Maximum number of results to return (default: 5).
+
+    Returns:
+        list[dict]: List of search results, each containing:
+            - path: The documentation file path
+            - title: The document title (from first heading)
+            - excerpt: A relevant excerpt with the match highlighted in context
+            - match_count: Number of times the query appears in the document
+    """
+
+    if not query or not query.strip():
+        return [{"error": "Empty search query"}]
+
+    query = query.strip()
+    query_lower = query.lower()
+    results: list[dict] = []
+
+    for path, content in KCLDocs.get().docs.items():
+        content_lower = content.lower()
+
+        # Count matches
+        match_count = content_lower.count(query_lower)
+        if match_count > 0:
+            title = _extract_title(content)
+            excerpt = _extract_excerpt(content, query)
+
+            results.append(
+                {
+                    "path": path,
+                    "title": title,
+                    "excerpt": excerpt,
+                    "match_count": match_count,
+                }
+            )
+
+    # Sort by match count (descending)
+    results.sort(key=lambda x: x["match_count"], reverse=True)
+
+    return results[:max_results]
+
+
+def get_doc_content(doc_path: str) -> str | None:
+    """Get the full content of a specific KCL documentation file.
+
+    Use list_kcl_docs() to see available documentation paths, or
+    search_kcl_docs() to find relevant documentation by keyword.
+
+    Args:
+        doc_path (str): The path to the documentation file
+            (e.g., "docs/kcl-lang/functions.md" or "docs/kcl-std/functions/extrude.md")
+
+    Returns:
+        str: The full Markdown content of the documentation file,
+            or an error message if not found.
+    """
+
+    # Basic validation to prevent path traversal
+    if ".." in doc_path or not doc_path.startswith("docs/"):
+        return None
+
+    return KCLDocs.get().docs.get(doc_path)