# pytorch-mcp

`pytorch-mcp` is a full-featured MCP server for PyTorch documentation workflows. It indexes the repository's local `docs/` tree and exposes search, page retrieval, symbol lookup, code-example extraction, troubleshooting, and question-answering tools that an LLM can use to help developers work with PyTorch.

## Features

- Uses `docs/` as the source of truth for documentation-aware tools.
- Indexes local Markdown and reStructuredText PyTorch docs.
- Searches by workflow, concept, topic, or exact symbol such as `torch.compile`.
- Returns grounded snippets, headings, and page metadata for follow-up exploration.
- Extracts code examples from relevant docs pages.
- Understands declared Torch ecosystem libraries from `pyproject.toml` and inspects runtime availability.
- Recommends reading paths for tasks like training, compilation, data loading, profiling, and troubleshooting.
- Exposes MCP tools, resources, prompts, plus HTTP health/readiness routes.

## Server Instructions

The MCP server is intended to behave like a PyTorch development copilot:

- Use local `docs/` content as the authoritative source for documentation-aware answers.
- Inspect declared and installed Torch libraries before making environment-specific recommendations.
- Use planning, template-generation, code-inspection, and runtime-validation tools to help developers build models faster.
- Keep debugging and optimization advice grounded in retrieved docs, parsed traces, profiler data, and runtime checks when available.

## Tools

- `list_doc_topics`
- `search_docs`
- `get_doc_page`
- `get_symbol_reference`
- `extract_code_examples`
- `answer_pytorch_question`
- `recommend_docs`
- `troubleshoot_pytorch`
- `plan_model_build`
- `assemble_training_stack`
- `generate_training_loop_template`
- `generate_task_specific_template`
- `generate_training_project_template`
- `review_training_code`
- `suggest_model_architecture`
- `choose_loss_and_optimizer`
- `optimize_data_pipeline`
- `diagnose_training_issue`
- `inspect_pytorch_code`
- `inspect_runtime_environment`
- `execute_pytorch_snippet`
- `run_forward_pass_check`
- `benchmark_compile_candidate`
- `validate_training_setup`
- `list_torch_libraries`
- `inspect_torch_library`
- `recommend_torch_libraries`
- `audit_torch_stack`
- `parse_stack_trace`
- `analyze_stack_trace`
- `analyze_shape_mismatch`
- `parse_profiler_export`
- `analyze_profiler_summary`

## Resources

- `pytorch://server/capabilities`
- `pytorch://project/settings`
- `pytorch://docs/index`
- `pytorch://docs/categories`
- `pytorch://docs/page/{doc_path}`
- `pytorch://docs/category/{category}`
- `pytorch://docs/search/{query}?limit=5`
- `pytorch://reference/overview`

## Prompts

- `explain pytorch topic`
- `plan pytorch implementation`
- `debug pytorch issue`
- `compare pytorch approaches`
- `build pytorch model`
- `review pytorch training code`
- `choose pytorch training objective`
- `diagnose pytorch training issue`
- `inspect pytorch code`
- `analyze pytorch stack trace`
- `recommend torch libraries`

## Run

Install dependencies:

```bash
uv sync
```

Run over stdio:

```bash
uv run python mcp_server.py --transport stdio
```

Run over HTTP:

```bash
uv run python mcp_server.py --transport http --host 127.0.0.1 --port 8000
```

Health endpoints:

- `GET /healthz`
- `GET /readyz`

## Configuration

Important environment variables:

- `PYTORCH_MCP_DOCS_ROOT`
- `PYTORCH_MCP_MAX_SEARCH_RESULTS`
- `PYTORCH_MCP_MAX_PAGE_CHARACTERS`
- `PYTORCH_MCP_MAX_CODE_EXAMPLES`
- `PYTORCH_MCP_TRANSPORT`
- `PYTORCH_MCP_HOST`
- `PYTORCH_MCP_PORT`

Example:

```bash
PYTORCH_MCP_DOCS_ROOT=/path/to/pytorch/docs \
uv run python mcp_server.py --transport stdio
```

By default the server reads from this repository's `docs/` directory. If you package or deploy the server elsewhere, point `PYTORCH_MCP_DOCS_ROOT` at a local PyTorch documentation checkout.

## Testing

```bash
just test
```