diff --git a/mcp_server/.env.example b/mcp_server/.env.example new file mode 100644 index 000000000..2fb994bbe --- /dev/null +++ b/mcp_server/.env.example @@ -0,0 +1,44 @@ +# Copy this file to .env and fill in your credentials. +# Obtain them from https://platform.mat3ra.com → Account Preferences → API Tokens + +MAT3RA_ACCOUNT_ID=your_account_id_here +MAT3RA_AUTH_TOKEN=your_auth_token_here + +# Optional – only needed when working with an organization account +# MAT3RA_ORGANIZATION_ID=your_org_id_here + +# API connection settings – defaults are correct for the hosted platform +# MAT3RA_API_HOST=platform.mat3ra.com +# MAT3RA_API_PORT=443 +# MAT3RA_API_VERSION=2018-10-01 +# MAT3RA_API_SECURE=true + +# ── Intent detection backend ────────────────────────────────────────────────── +# choices: rules (default) | openai | claude | gemini | hf +# INTENT_BACKEND=rules + +# ── HuggingFace local backend (hf) ─────────────────────────────────────────── +# Two runtimes – select with HF_RUNTIME: +# +# transformers (default) – runs the model in-process via HuggingFace Transformers +# pip install transformers accelerate torch +# HF_RUNTIME=transformers +# HF_MODEL=Qwen/Qwen2.5-1.5B-Instruct # ~3 GB, good balance +# HF_DEVICE=auto # auto | cpu | cuda | mps +# +# ollama – delegates to a locally running Ollama server +# Ollama serves HuggingFace-compatible (GGUF) models +# Install: https://ollama.com then: ollama serve +# HF_RUNTIME=ollama +# OLLAMA_HOST=http://localhost:11434 +# OLLAMA_MODEL=llama3.2 + +# ── Cloud LLM API keys ──────────────────────────────────────────────────────── +# OPENAI_API_KEY=sk-... +# OPENAI_MODEL=gpt-4o-mini + +# ANTHROPIC_API_KEY=sk-ant-... +# ANTHROPIC_MODEL=claude-3-5-haiku-latest + +# GEMINI_API_KEY=... +# GEMINI_MODEL=gemini-2.0-flash diff --git a/mcp_server/AGENTS.md b/mcp_server/AGENTS.md new file mode 100644 index 000000000..5049f3512 --- /dev/null +++ b/mcp_server/AGENTS.md @@ -0,0 +1,210 @@ +# AGENTS.md — Mat3ra MCP Server + +Reference document for AI agents, coding assistants, and future contributors +working on this codebase. Describes the architecture, extension points, and +conventions to follow. + +--- + +## Repository Layout + +``` +mcp_server/ +├── server.py # MCP stdio server (tools exposed to LLM clients) +├── code_server.py # HTTP playground (prompt → Python code, browser UI) +├── intent.py # Intent detection (rules + 6 LLM backends) +├── pyproject.toml # Package manifest + optional deps +├── .env.example # Credential template +├── README.md # User-facing documentation +└── AGENTS.md # This file +``` + +--- + +## Key Abstractions + +### 1. Tool Catalogue (`intent.py → TOOLS`) + +`TOOLS` is the single source of truth for all Mat3ra API operations. +Each entry drives **three** things simultaneously: + +- The **LLM system prompt** (tool names, args, descriptions) +- The **JSON schema** exposed by the MCP server (`server.py → list_tools`) +- The **Python code generation** (`intent.py → _build_call`) + +When adding a new Mat3ra API operation, add it to `TOOLS` first — everything else derives from it. + +```python +TOOLS: dict[str, dict] = { + "my_new_tool": { + "module": "materials", # mat3ra.api_client.endpoints. + "cls": "MaterialEndpoints", # endpoint class + "description": "...", # shown to the LLM + "args_schema": { # used in system prompt + MCP schema + "param_name": "description of param", + }, + }, +} +``` + +Then add a matching `elif tool == "my_new_tool":` branch in `_build_call()` and +a matching `Tool(...)` entry in `server.py → list_tools()` and handler branch in +`server.py → _dispatch()`. + +### 2. Intent Detection (`intent.py → detect_intent`) + +```python +def detect_intent(prompt: str, backend: str | None = None) -> dict: + ... +``` + +Returns: + +```python +{ + "tool": str, # e.g. "list_materials" + "args": dict, # e.g. {"formula": "Si"} + "module": str, # e.g. "materials" + "cls": str, # e.g. "MaterialEndpoints" + "call": str, # Python call snippet + "code": str, # Full runnable Python script + "backend": str, # Which backend was actually used +} +``` + +All LLM backends receive the same `SYSTEM_PROMPT` (auto-generated from `TOOLS`) +and return raw text that is parsed by `_parse_llm_json()`. Any backend that +raises an exception automatically falls back to `rules`. + +### 3. LLM Backends (`intent.py → _BACKENDS`) + +| Key | Function | Transport | +|---|---|---| +| `rules` | `_rules()` | regex, no network | +| `openai` | `_openai()` | HTTPS to api.openai.com | +| `claude` | `_claude()` | HTTPS to api.anthropic.com | +| `gemini` | `_gemini()` | HTTPS to generativelanguage.googleapis.com | +| `local` | `_local()` | HTTP to localhost:11434 (Ollama) | +| `hf` | `_huggingface()` | in-process, HuggingFace transformers | + +Each backend function signature: `(prompt: str) -> tuple[str, dict]` +i.e. `(tool_name, args_dict)`. + +To add a new backend: +1. Write `def _mybackend(prompt: str) -> tuple[str, dict]` in `intent.py` +2. Add `"mybackend": _mybackend` to `_BACKENDS` +3. Add `` to the ` +
+ LLM backend + +
+ + + + +
+ Try: + + + + + + +
+ +
Result
+
+ + + +
+ +
Generated Python code
+
+ + + + + + +""" + +# --------------------------------------------------------------------------- +# HTTP handler +# --------------------------------------------------------------------------- + + +class Handler(BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + print(f" {self.address_string()} {self.command} {self.path}" + f" → {fmt % args}") + + def _send(self, status: int, content_type: str, body: bytes) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(body) + + def do_OPTIONS(self): + self.send_response(204) + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Headers", "Content-Type") + self.send_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS") + self.end_headers() + + def do_GET(self): + if self.path in ("/", "/index.html"): + self._send(200, "text/html; charset=utf-8", HTML.encode()) + else: + self._send(404, "text/plain", b"Not found") + + def do_POST(self): + if self.path != "/ask": + self._send(404, "text/plain", b"Not found") + return + + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) + try: + payload = json.loads(body) + except json.JSONDecodeError: + self._send(400, "application/json", + json.dumps({"error": "invalid JSON"}).encode()) + return + + prompt = payload.get("prompt", "").strip() + backend = payload.get("backend", os.environ.get("INTENT_BACKEND", "rules")) + + if not prompt: + self._send(400, "application/json", + json.dumps({"error": "prompt is required"}).encode()) + return + + try: + intent = detect_intent(prompt, backend) + except Exception as exc: + self._send(200, "application/json; charset=utf-8", + json.dumps({"error": str(exc)}).encode()) + return + + self._send(200, "application/json; charset=utf-8", + json.dumps(intent, indent=2).encode()) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + HOST, PORT = "localhost", 8888 + httpd = HTTPServer((HOST, PORT), Handler) + backend = os.environ.get("INTENT_BACKEND", "rules") + print(f"\n Mat3ra Prompt-to-Code server") + print(f" Playground : http://{HOST}:{PORT}/") + print(f" API : POST http://{HOST}:{PORT}/ask") + print(f" Default LLM : {backend}") + print(f" Backends available: rules (always), openai, claude, gemini, hf") + print(f" hf runtime: HF_RUNTIME=transformers (default) | ollama") + print(f" Set INTENT_BACKEND=openai (etc.) or select in the UI.") + print(f"\n Press Ctrl+C to stop.\n") + try: + httpd.serve_forever() + except KeyboardInterrupt: + print("\n Server stopped.") diff --git a/mcp_server/intent.py b/mcp_server/intent.py new file mode 100644 index 000000000..8acda1542 --- /dev/null +++ b/mcp_server/intent.py @@ -0,0 +1,533 @@ +""" +intent.py – prompt → Mat3ra API intent detection +================================================== +Supports six backends, selected via the INTENT_BACKEND env var +or the `backend` argument passed at call time: + + rules no deps, keyword matching (default / fallback) + openai requires: pip install openai + OPENAI_API_KEY + claude requires: pip install anthropic + ANTHROPIC_API_KEY + gemini requires: pip install google-genai + GEMINI_API_KEY + hf local open-source model — two runtimes (set via HF_RUNTIME): + transformers run a model in-process via HuggingFace Transformers (default) + pip install transformers accelerate torch + HF_MODEL (default: Qwen/Qwen2.5-1.5B-Instruct) + HF_DEVICE (default: auto) + ollama delegate to a running Ollama server (serves HF/GGUF models) + requires: Ollama running on OLLAMA_HOST (default: localhost:11434) + OLLAMA_MODEL (default: llama3.2) + +All backends return the same dict: + { + "tool": str, # one of the 7 tool names + "args": dict, # tool arguments (minus owner._id which is added later) + "module": str, # python module name e.g. "materials" + "cls": str, # endpoint class name e.g. "MaterialEndpoints" + "call": str, # python code snippet for the API call + } +""" + +from __future__ import annotations + +import json +import os +import re +import textwrap +import urllib.error +import urllib.request +from typing import Any + +# --------------------------------------------------------------------------- +# Tool catalogue – drives both code generation and the LLM system prompt +# --------------------------------------------------------------------------- + +TOOLS: dict[str, dict[str, Any]] = { + "list_materials": { + "module": "materials", "cls": "MaterialEndpoints", + "description": "Search / list materials by formula or Mongo query.", + "args_schema": {"formula": "chemical formula e.g. Si, GaAs (optional)"}, + }, + "get_material": { + "module": "materials", "cls": "MaterialEndpoints", + "description": "Fetch a single material by its 24-char hex _id.", + "args_schema": {"material_id": "24-char hex string"}, + }, + "create_material": { + "module": "materials", "cls": "MaterialEndpoints", + "description": "Create a new material from a JSON config.", + "args_schema": {}, + }, + "list_workflows": { + "module": "workflows", "cls": "WorkflowEndpoints", + "description": "List workflows in the account.", + "args_schema": {}, + }, + "list_jobs": { + "module": "jobs", "cls": "JobEndpoints", + "description": "List calculation jobs in the account.", + "args_schema": {}, + }, + "create_and_submit_job": { + "module": "jobs", "cls": "JobEndpoints", + "description": "Create and immediately submit a calculation job.", + "args_schema": { + "material_id": "24-char hex _id of the material", + "workflow_id": "24-char hex _id of the workflow", + "job_name": "optional human-readable name", + }, + }, + "get_job": { + "module": "jobs", "cls": "JobEndpoints", + "description": "Fetch a job and its current status by _id.", + "args_schema": {"job_id": "24-char hex string"}, + }, +} + +# --------------------------------------------------------------------------- +# System prompt shared by all LLM backends +# --------------------------------------------------------------------------- + +_TOOL_DESCRIPTIONS = "\n".join( + f" - {name}({', '.join(k for k in meta['args_schema'])})" + f" # {meta['description']}" + for name, meta in TOOLS.items() +) + +SYSTEM_PROMPT = f"""\ +You are a Mat3ra materials-science API assistant. + +Given a natural-language request, identify which Mat3ra API tool to call \ +and what arguments to pass. + +Available tools: +{_TOOL_DESCRIPTIONS} + +Rules: +- Extract chemical formulas from the text (e.g. "silicon" → "Si", "GaAs" → "GaAs"). +- If an ID-like string appears (24 hex chars), use it as material_id / job_id. +- If the intent is unclear, default to list_materials with no formula. + +Respond ONLY with a single JSON object, no markdown fences, no explanation: +{{ + "tool": "", + "args": {{ ... }} +}} + +Examples: + "List my silicon materials" → {{"tool":"list_materials","args":{{"formula":"Si"}}}} + "Find GaAs compounds" → {{"tool":"list_materials","args":{{"formula":"GaAs"}}}} + "Create a copper FCC crystal" → {{"tool":"create_material","args":{{}}}} + "Show my workflows" → {{"tool":"list_workflows","args":{{}}}} + "Submit a DFT job" → {{"tool":"create_and_submit_job","args":{{"material_id":"","workflow_id":""}}}} + "Check job 507f1f77bcf86cd799439011" → {{"tool":"get_job","args":{{"job_id":"507f1f77bcf86cd799439011"}}}} +""" + +# --------------------------------------------------------------------------- +# Code generation from a resolved (tool, args) pair +# --------------------------------------------------------------------------- + +BOILERPLATE = """\ +import os +from mat3ra.api_client.endpoints.{module} import {cls} + +ACCOUNT_ID = os.environ["MAT3RA_ACCOUNT_ID"] +AUTH_TOKEN = os.environ["MAT3RA_AUTH_TOKEN"] +OWNER_ID = os.environ.get("MAT3RA_ORGANIZATION_ID") or ACCOUNT_ID + +endpoint = {cls}( + "platform.mat3ra.com", 443, + ACCOUNT_ID, AUTH_TOKEN, + "2018-10-01", secure=True +) +""" + + +def _build_call(tool: str, args: dict) -> str: + """Generate the Python call snippet for a given tool + args.""" + if tool == "list_materials": + query: dict = {"owner._id": "OWNER_ID"} + if f := args.get("formula"): + query["formula"] = f + q = json.dumps(query, indent=4).replace('"OWNER_ID"', "OWNER_ID") + return f"materials = endpoint.list(\n {q}\n)" + + if tool == "get_material": + mid = args.get("material_id", "") + return f'material = endpoint.get("{mid}")' + + if tool == "create_material": + return textwrap.dedent("""\ + config = { + "name": "My Material", + "basis": { + "elements": [{"id": 1, "value": "Si"}, {"id": 2, "value": "Si"}], + "coordinates": [{"id": 1, "value": [0, 0, 0]}, + {"id": 2, "value": [0.25, 0.25, 0.25]}], + "units": "crystal", + }, + "lattice": { + "type": "FCC", "a": 3.867, "b": 3.867, "c": 3.867, + "alpha": 60, "beta": 60, "gamma": 60, + "units": {"length": "angstrom", "angle": "degree"}, + }, + "tags": ["MCP"], + } + material = endpoint.create(config, OWNER_ID)""") + + if tool == "list_workflows": + return 'workflows = endpoint.list({"owner._id": OWNER_ID})' + + if tool == "list_jobs": + return 'jobs = endpoint.list({"owner._id": OWNER_ID})' + + if tool == "create_and_submit_job": + mid = args.get("material_id", "") + wid = args.get("workflow_id", "") + name = args.get("job_name", "MCP Job") + return textwrap.dedent(f"""\ + config = {{ + "owner": {{"_id": OWNER_ID}}, + "_material": {{"_id": "{mid}"}}, + "workflow": {{"_id": "{wid}"}}, + "name": "{name}", + }} + job = endpoint.create(config) + endpoint.submit(job["_id"]) + job = endpoint.get(job["_id"]) # re-fetch to see updated status""") + + if tool == "get_job": + jid = args.get("job_id", "") + return f'job = endpoint.get("{jid}")' + + return f"# Unknown tool: {tool}" + + +def build_intent(tool: str, args: dict, backend: str = "rules") -> dict: + """Assemble the full intent dict from an (tool, args) pair.""" + if tool not in TOOLS: + tool = "list_materials" + args = {} + meta = TOOLS[tool] + call = _build_call(tool, args) + code = BOILERPLATE.format(module=meta["module"], cls=meta["cls"]).rstrip() + code += "\n\n" + call + "\n" + return { + "tool": tool, + "args": {k: v for k, v in args.items() if k != "owner._id"}, + "module": meta["module"], + "cls": meta["cls"], + "call": call, + "code": code, + "backend": backend, + } + + +def _parse_llm_json(text: str) -> tuple[str, dict]: + """ + Extract tool + args from raw LLM output. + Strips markdown fences if present. + """ + # strip ```json ... ``` fences + text = re.sub(r"```(?:json)?\s*", "", text).strip().rstrip("`").strip() + # find first {...} + m = re.search(r"\{.*\}", text, re.DOTALL) + if not m: + raise ValueError(f"No JSON object found in LLM response: {text!r}") + data = json.loads(m.group()) + return data.get("tool", "list_materials"), data.get("args", {}) + + +# --------------------------------------------------------------------------- +# Rule-based backend (no deps) +# --------------------------------------------------------------------------- + +def _rules(prompt: str) -> tuple[str, dict]: + p = prompt.lower() + fm = re.search(r"\b([A-Z][a-z]?(?:\d*[A-Z][a-z]?)*\d*)\b", prompt) + formula = fm.group(1) if fm else None + junk = {"List", "Show", "Find", "Get", "Create", "Submit", "Check", "Run"} + + if re.search(r"\b(list|search|find|get|show)\b.*\bmaterial", p): + args: dict = {} + if formula and formula not in junk: + args["formula"] = formula + return "list_materials", args + + if re.search(r"\b(create|add|upload|make)\b.*\bmaterial", p): + return "create_material", {} + + if re.search(r"\bget\b.*\bmaterial", p): + id_m = re.search(r"\b([a-f0-9]{24})\b", prompt) + return "get_material", {"material_id": id_m.group(1) if id_m else ""} + + if re.search(r"\b(list|search|find|show)\b.*\bworkflow", p): + return "list_workflows", {} + + if re.search(r"\b(list|search|find|show)\b.*\bjob", p): + return "list_jobs", {} + + if re.search(r"\b(create|submit|run|launch)\b.*\bjob", p): + return "create_and_submit_job", { + "material_id": "", + "workflow_id": "", + } + + if re.search(r"\b(get|status|check)\b.*\bjob", p): + id_m = re.search(r"\b([a-f0-9]{24})\b", prompt) + return "get_job", {"job_id": id_m.group(1) if id_m else ""} + + # default + args2: dict = {} + if formula and formula not in junk: + args2["formula"] = formula + return "list_materials", args2 + + +# --------------------------------------------------------------------------- +# OpenAI backend +# --------------------------------------------------------------------------- + +def _openai(prompt: str) -> tuple[str, dict]: + try: + import openai # noqa: PLC0415 + except ImportError: + raise RuntimeError("pip install openai") + + client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"]) + model = os.environ.get("OPENAI_MODEL", "gpt-4o-mini") + resp = client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + temperature=0, + ) + return _parse_llm_json(resp.choices[0].message.content or "") + + +# --------------------------------------------------------------------------- +# Anthropic / Claude backend +# --------------------------------------------------------------------------- + +def _claude(prompt: str) -> tuple[str, dict]: + try: + import anthropic # noqa: PLC0415 + except ImportError: + raise RuntimeError("pip install anthropic") + + client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"]) + model = os.environ.get("ANTHROPIC_MODEL", "claude-3-5-haiku-latest") + msg = client.messages.create( + model=model, + max_tokens=256, + system=SYSTEM_PROMPT, + messages=[{"role": "user", "content": prompt}], + ) + return _parse_llm_json(msg.content[0].text) + + +# --------------------------------------------------------------------------- +# Google Gemini backend +# --------------------------------------------------------------------------- + +def _gemini(prompt: str) -> tuple[str, dict]: + try: + from google import genai # noqa: PLC0415 + except ImportError: + raise RuntimeError("pip install google-genai") + + api_key = os.environ["GEMINI_API_KEY"] + model = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash") + client = genai.Client(api_key=api_key) + resp = client.models.generate_content( + model=model, + contents=SYSTEM_PROMPT + "\n\nUser request: " + prompt, + ) + return _parse_llm_json(resp.text) + + +# --------------------------------------------------------------------------- +# Local / Ollama backend +# --------------------------------------------------------------------------- + +def _local(prompt: str) -> tuple[str, dict]: + host = os.environ.get("OLLAMA_HOST", "http://localhost:11434") + model = os.environ.get("OLLAMA_MODEL", "llama3.2") + full_prompt = SYSTEM_PROMPT + "\n\nUser request: " + prompt + + payload = json.dumps({ + "model": model, + "prompt": full_prompt, + "stream": False, + }).encode() + + req = urllib.request.Request( + f"{host}/api/generate", + data = payload, + headers = {"Content-Type": "application/json"}, + method = "POST", + ) + try: + with urllib.request.urlopen(req, timeout=60) as r: + data = json.loads(r.read()) + except urllib.error.URLError as e: + raise RuntimeError( + f"Cannot reach Ollama at {host}. " + f"Is it running? (ollama serve) Error: {e}" + ) + return _parse_llm_json(data.get("response", "")) + + +# --------------------------------------------------------------------------- +# HuggingFace backend — two runtimes selected by HF_RUNTIME env var +# --------------------------------------------------------------------------- +# transformers in-process HuggingFace Transformers pipeline (default) +# ollama HTTP call to a local Ollama server (serves HF/GGUF models) +# --------------------------------------------------------------------------- + +# Pipeline cache for the transformers runtime (loaded once per server process). +_hf_pipeline_cache: dict = {} + + +def _hf_ollama(prompt: str) -> tuple[str, dict]: + """HuggingFace via Ollama — Ollama is a local server that serves HF-compatible + (GGUF) models. No Python dependencies beyond stdlib urllib.""" + host = os.environ.get("OLLAMA_HOST", "http://localhost:11434") + model = os.environ.get("OLLAMA_MODEL", "llama3.2") + full_prompt = SYSTEM_PROMPT + "\n\nUser request: " + prompt + + payload = json.dumps({ + "model": model, + "prompt": full_prompt, + "stream": False, + }).encode() + req = urllib.request.Request( + f"{host}/api/generate", + data = payload, + headers = {"Content-Type": "application/json"}, + method = "POST", + ) + try: + with urllib.request.urlopen(req, timeout=60) as r: + data = json.loads(r.read()) + except urllib.error.URLError as e: + raise RuntimeError( + f"Cannot reach Ollama at {host}. Is it running? (ollama serve)\n" + f"Set OLLAMA_HOST if it is on a different address. Error: {e}" + ) + return _parse_llm_json(data.get("response", "")) + + +def _hf_transformers(prompt: str) -> tuple[str, dict]: + """HuggingFace via Transformers — runs the model in-process. + + Recommended small models (auto-downloaded on first use): + Qwen/Qwen2.5-0.5B-Instruct ~1 GB fastest, any CPU + Qwen/Qwen2.5-1.5B-Instruct ~3 GB default, good balance + HuggingFaceTB/SmolLM2-1.7B-Instruct ~3 GB very capable for size + microsoft/Phi-3-mini-4k-instruct ~8 GB high quality + + Env vars: + HF_MODEL model repo id (default: Qwen/Qwen2.5-1.5B-Instruct) + HF_DEVICE cpu | cuda | mps | auto (default: auto) + """ + try: + from transformers import pipeline as hf_pipeline # noqa: PLC0415 + except ImportError: + raise RuntimeError( + "pip install transformers accelerate\n" + "Mac (Apple Silicon / CPU): pip install torch\n" + "Linux (CUDA): pip install torch " + "--index-url https://download.pytorch.org/whl/cu121" + ) + + model_id = os.environ.get("HF_MODEL", "Qwen/Qwen2.5-1.5B-Instruct") + device = os.environ.get("HF_DEVICE", "auto") + cache_key = (model_id, device) + + if cache_key not in _hf_pipeline_cache: + print(f"[hf] Loading {model_id} on device={device} … (one-time download/load)") + _hf_pipeline_cache[cache_key] = hf_pipeline( + "text-generation", + model = model_id, + device_map = device, + torch_dtype = "auto", + ) + print(f"[hf] {model_id} ready.") + + pipe = _hf_pipeline_cache[cache_key] + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ] + try: + result = pipe(messages, max_new_tokens=256, do_sample=False, + return_full_text=False) + text = result[0]["generated_text"] + if isinstance(text, list): # some pipelines return message list + text = text[-1].get("content", "") + except Exception: + full = SYSTEM_PROMPT + "\n\nUser request: " + prompt + result = pipe(full, max_new_tokens=256, do_sample=False, + return_full_text=False) + text = result[0]["generated_text"] + return _parse_llm_json(text) + + +def _huggingface(prompt: str) -> tuple[str, dict]: + """Dispatch to the correct HF runtime based on HF_RUNTIME env var.""" + runtime = os.environ.get("HF_RUNTIME", "transformers").lower() + if runtime == "ollama": + return _hf_ollama(prompt) + if runtime == "transformers": + return _hf_transformers(prompt) + raise ValueError( + f"Unknown HF_RUNTIME={runtime!r}. Choose 'transformers' or 'ollama'." + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +_BACKENDS = { + "rules": _rules, + "openai": _openai, + "claude": _claude, + "gemini": _gemini, + "hf": _huggingface, +} + +DEFAULT_BACKEND = os.environ.get("INTENT_BACKEND", "rules") + + +def detect_intent(prompt: str, backend: str | None = None) -> dict: + """ + Detect the Mat3ra API intent from a natural-language prompt. + + Args: + prompt: user text + backend: "rules" | "openai" | "claude" | "gemini" | "hf" + Defaults to INTENT_BACKEND env var, then "rules". + + Returns: + Intent dict with keys: tool, args, module, cls, call, code, backend + """ + backend = (backend or DEFAULT_BACKEND).lower() + if backend not in _BACKENDS: + raise ValueError(f"Unknown backend {backend!r}. Choose from: {list(_BACKENDS)}") + + fn = _BACKENDS[backend] + try: + tool, args = fn(prompt) + except Exception as exc: + if backend != "rules": + # graceful fallback to rules on any LLM error + print(f"[intent] {backend} failed ({exc}), falling back to rules") + tool, args = _rules(prompt) + backend = f"rules (fallback from {backend})" + else: + raise + + return build_intent(tool, args, backend) diff --git a/mcp_server/pyproject.toml b/mcp_server/pyproject.toml new file mode 100644 index 000000000..8d4580173 --- /dev/null +++ b/mcp_server/pyproject.toml @@ -0,0 +1,32 @@ +[project] +name = "mat3ra-mcp-server" +version = "0.1.0" +description = "Proof-of-concept MCP server for the Mat3ra API" +requires-python = ">=3.10" +dependencies = [ + "mcp>=1.0", + "mat3ra-api-client", +] + +[project.optional-dependencies] +openai = ["openai>=1.0"] +claude = ["anthropic>=0.25"] +gemini = ["google-genai>=1.0"] +# hf backend: transformers runtime needs these; ollama runtime uses only stdlib urllib +hf = ["transformers>=4.40", "accelerate>=0.27"] +# torch must be installed separately (platform-specific wheel) +# Mac/CPU: pip install torch +# CUDA: pip install torch --index-url https://download.pytorch.org/whl/cu121 +all-llm = [ + "mat3ra-mcp-server[openai]", + "mat3ra-mcp-server[claude]", + "mat3ra-mcp-server[gemini]", + "mat3ra-mcp-server[hf]", +] + +[project.scripts] +mat3ra-mcp = "mcp_server.server:main" + +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/mcp_server/server.py b/mcp_server/server.py new file mode 100644 index 000000000..e5dafb302 --- /dev/null +++ b/mcp_server/server.py @@ -0,0 +1,334 @@ +""" +Mat3ra MCP Server – proof of concept +===================================== +Exposes a handful of Mat3ra API operations as MCP tools so that any +MCP-compatible client (Claude Desktop, Cursor, etc.) can call them. + +Tools +----- +* list_materials – search materials by formula or free-form query +* get_material – fetch a single material by ID +* create_material – create a new material from a JSON config +* list_workflows – search workflows +* list_jobs – search jobs +* create_and_submit_job – create + immediately submit a job +* get_job – fetch a single job by ID + +Authentication +-------------- +Set the following environment variables before starting the server +(or copy them from examples/config/settings.json): + + MAT3RA_ACCOUNT_ID= + MAT3RA_AUTH_TOKEN= + +Optional: + MAT3RA_ORGANIZATION_ID= # use org account instead of personal + MAT3RA_API_HOST=platform.mat3ra.com + MAT3RA_API_PORT=443 + MAT3RA_API_VERSION=2018-10-01 + MAT3RA_API_SECURE=true + +Running +------- + # install deps once (inside a venv) + pip install mcp mat3ra-api-client + + # start the server (stdio transport – default for MCP) + python mcp_server/server.py + + # or with the MCP CLI + mcp run mcp_server/server.py +""" + +from __future__ import annotations + +import json +import os +import sys + +# --------------------------------------------------------------------------- +# MCP SDK – https://github.com/modelcontextprotocol/python-sdk +# --------------------------------------------------------------------------- +try: + import mcp.server.stdio + from mcp.server import NotificationOptions, Server + from mcp.server.models import InitializationOptions + from mcp.types import TextContent, Tool +except ImportError: + sys.exit( + "The 'mcp' package is not installed.\n" + "Run: pip install mcp mat3ra-api-client\n" + ) + +# --------------------------------------------------------------------------- +# Mat3ra API client +# --------------------------------------------------------------------------- +try: + from mat3ra.api_client.endpoints.jobs import JobEndpoints + from mat3ra.api_client.endpoints.materials import MaterialEndpoints + from mat3ra.api_client.endpoints.workflows import WorkflowEndpoints +except ImportError: + sys.exit( + "The 'mat3ra-api-client' package is not installed.\n" + "Run: pip install mcp mat3ra-api-client\n" + ) + +# --------------------------------------------------------------------------- +# Read credentials from environment (fall back to empty string so the error +# surfaces at call time, not at server start-up). +# --------------------------------------------------------------------------- +ACCOUNT_ID = os.environ.get("MAT3RA_ACCOUNT_ID", "") +AUTH_TOKEN = os.environ.get("MAT3RA_AUTH_TOKEN", "") +ORGANIZATION_ID = os.environ.get("MAT3RA_ORGANIZATION_ID", "") +HOST = os.environ.get("MAT3RA_API_HOST", "platform.mat3ra.com") +PORT = int(os.environ.get("MAT3RA_API_PORT", "443")) +VERSION = os.environ.get("MAT3RA_API_VERSION", "2018-10-01") +SECURE = os.environ.get("MAT3RA_API_SECURE", "true").lower() != "false" + +ENDPOINT_ARGS = [HOST, PORT, ACCOUNT_ID, AUTH_TOKEN, VERSION, SECURE] + +# Owner ID: prefer org, fall back to personal account +OWNER_ID = ORGANIZATION_ID or ACCOUNT_ID + + +def _endpoint_args() -> list: + """Return endpoint args, re-reading env vars at call time so the server + can be started before credentials are set (useful during development).""" + account_id = os.environ.get("MAT3RA_ACCOUNT_ID", ACCOUNT_ID) + auth_token = os.environ.get("MAT3RA_AUTH_TOKEN", AUTH_TOKEN) + return [HOST, PORT, account_id, auth_token, VERSION, SECURE] + + +def _owner() -> str: + org = os.environ.get("MAT3RA_ORGANIZATION_ID", ORGANIZATION_ID) + account = os.environ.get("MAT3RA_ACCOUNT_ID", ACCOUNT_ID) + return org or account + + +# --------------------------------------------------------------------------- +# Helper +# --------------------------------------------------------------------------- +def _json(obj) -> str: + return json.dumps(obj, indent=2, default=str) + + +# --------------------------------------------------------------------------- +# MCP server +# --------------------------------------------------------------------------- +server = Server("mat3ra-mcp") + + +# ── Tool catalogue ────────────────────────────────────────────────────────── + +@server.list_tools() +async def list_tools() -> list[Tool]: + return [ + Tool( + name="list_materials", + description=( + "Search materials in the Mat3ra account. " + "Pass a 'formula' key to filter by chemical formula (e.g. 'Si', 'GaAs'). " + "Any additional key-value pairs are forwarded directly to the Mongo-style query." + ), + inputSchema={ + "type": "object", + "properties": { + "formula": { + "type": "string", + "description": "Chemical formula to search for, e.g. 'Si' or 'GaAs'.", + }, + "query": { + "type": "object", + "description": "Additional Mongo-style query fields (optional).", + }, + }, + }, + ), + Tool( + name="get_material", + description="Fetch a single material by its Mat3ra ID.", + inputSchema={ + "type": "object", + "required": ["material_id"], + "properties": { + "material_id": { + "type": "string", + "description": "The _id of the material.", + } + }, + }, + ), + Tool( + name="create_material", + description=( + "Create a new material on Mat3ra. " + "Provide a JSON config following https://docs.mat3ra.com/materials/data/." + ), + inputSchema={ + "type": "object", + "required": ["config"], + "properties": { + "config": { + "type": "object", + "description": "Material config JSON (name, basis, lattice, tags, …).", + } + }, + }, + ), + Tool( + name="list_workflows", + description="List workflows in the Mat3ra account.", + inputSchema={ + "type": "object", + "properties": { + "query": { + "type": "object", + "description": "Optional Mongo-style query filter.", + } + }, + }, + ), + Tool( + name="list_jobs", + description="List jobs in the Mat3ra account.", + inputSchema={ + "type": "object", + "properties": { + "query": { + "type": "object", + "description": "Optional Mongo-style query filter.", + } + }, + }, + ), + Tool( + name="create_and_submit_job", + description=( + "Create and immediately submit a calculation job on Mat3ra. " + "Requires a material_id and a workflow_id." + ), + inputSchema={ + "type": "object", + "required": ["material_id", "workflow_id"], + "properties": { + "material_id": { + "type": "string", + "description": "The _id of the material to use.", + }, + "workflow_id": { + "type": "string", + "description": "The _id of the workflow to use.", + }, + "job_name": { + "type": "string", + "description": "Human-readable name for the job (optional).", + }, + }, + }, + ), + Tool( + name="get_job", + description="Fetch a single job by its Mat3ra ID, including its current status.", + inputSchema={ + "type": "object", + "required": ["job_id"], + "properties": { + "job_id": { + "type": "string", + "description": "The _id of the job.", + } + }, + }, + ), + ] + + +# ── Tool handlers ──────────────────────────────────────────────────────────── + +@server.call_tool() +async def call_tool(name: str, arguments: dict) -> list[TextContent]: + try: + result = await _dispatch(name, arguments) + except Exception as exc: + return [TextContent(type="text", text=f"Error: {exc}")] + return [TextContent(type="text", text=_json(result))] + + +async def _dispatch(name: str, args: dict): + ep = _endpoint_args() + owner = _owner() + + if name == "list_materials": + q: dict = args.get("query", {}) + if formula := args.get("formula"): + q["formula"] = formula + q.setdefault("owner._id", owner) + endpoint = MaterialEndpoints(*ep) + return endpoint.list(q) + + if name == "get_material": + endpoint = MaterialEndpoints(*ep) + return endpoint.get(args["material_id"]) + + if name == "create_material": + endpoint = MaterialEndpoints(*ep) + return endpoint.create(args["config"], owner) + + if name == "list_workflows": + q = args.get("query", {}) + q.setdefault("owner._id", owner) + endpoint = WorkflowEndpoints(*ep) + return endpoint.list(q) + + if name == "list_jobs": + q = args.get("query", {}) + q.setdefault("owner._id", owner) + endpoint = JobEndpoints(*ep) + return endpoint.list(q) + + if name == "create_and_submit_job": + endpoint = JobEndpoints(*ep) + config = { + "owner": {"_id": owner}, + "_material": {"_id": args["material_id"]}, + "workflow": {"_id": args["workflow_id"]}, + "name": args.get("job_name", "MCP Job"), + } + job = endpoint.create(config) + endpoint.submit(job["_id"]) + # re-fetch to include the updated status + return endpoint.get(job["_id"]) + + if name == "get_job": + endpoint = JobEndpoints(*ep) + return endpoint.get(args["job_id"]) + + raise ValueError(f"Unknown tool: {name}") + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- +async def main(): + import mcp.server.stdio + + async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): + await server.run( + read_stream, + write_stream, + InitializationOptions( + server_name="mat3ra-mcp", + server_version="0.1.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) + + +if __name__ == "__main__": + import asyncio + + asyncio.run(main())