diff --git a/README.md b/README.md index 351ddfc..0fc6bd5 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,7 @@ needle finetune data.jsonl needle playground Test and finetune via web UI needle finetune Finetune on your own data needle run --query "..." --tools Single inference +needle mcp-server Run Needle as a local MCP tool router needle train Full training run needle pretrain Pretrain on PleIAs/SYNTH needle eval --checkpoint Evaluate a checkpoint @@ -113,6 +114,61 @@ needle generate-data Synthesize training data via Gemini needle tpu TPU management (see docs/tpu.md) ``` +## Local MCP Server + +Needle can be run as a local MCP tool router for agent workflows, with either stdio transport or a lightweight HTTP transport. + +### Start the local MCP server + +```bash +needle mcp-server --checkpoint checkpoints/needle.pkl +``` + +### Run over HTTP + +```bash +needle mcp-server --checkpoint checkpoints/needle.pkl --transport http --host 127.0.0.1 --port 8765 +``` + +### Claude Desktop / MCP orchestrator configuration + +Use the local Needle process as a tool router for any MCP-compatible orchestrator. + +```json +{ + "mcpServers": { + "needle": { + "command": "needle", + "args": [ + "mcp-server", + "--checkpoint", + "/path/to/needle/checkpoints/needle.pkl", + "--transport", + "http", + "--host", + "127.0.0.1", + "--port", + "8765" + ] + } + } +} +``` + +With `--transport http`, the server listens on `http://127.0.0.1:8765` and accepts JSON-RPC 2.0 requests at the root path. + +### Health check + +```bash +curl http://127.0.0.1:8765/health +``` + +### How the router works + +1. The orchestrator sends a natural language query and the available tool metadata. +2. Needle returns a `tool_call`, `confidence`, and `match` flag. +3. The orchestrator can execute the chosen tool locally if confidence is high, or escalate to a cloud LLM if not. + ``` @misc{ndubuaku2026needle, title={Needle}, diff --git a/needle/cli.py b/needle/cli.py index 705ba92..15c1079 100644 --- a/needle/cli.py +++ b/needle/cli.py @@ -243,6 +243,18 @@ def main(): p.add_argument("--port", type=int, default=7860) p.add_argument("--host", type=str, default="127.0.0.1") + p = sub.add_parser("mcp-server", add_help=False) + p.add_argument("--checkpoint", type=str, default="checkpoints/needle.pkl") + p.add_argument( + "--transport", + type=str, + choices=["stdio", "http"], + default="stdio", + help="MCP transport to run: stdio for JSON-RPC over stdin/stdout, http for local HTTP server", + ) + p.add_argument("--host", type=str, default="127.0.0.1", help="Host to bind for HTTP transport") + p.add_argument("--port", type=int, default=8765, help="Port to bind for HTTP transport") + p = sub.add_parser("tpu", add_help=False) tpu_sub = p.add_subparsers(dest="tpu_action") @@ -336,6 +348,9 @@ def main(): elif args.command == "playground": from .ui.server import main as ui_main ui_main(args) + elif args.command == "mcp-server": + from .mcp_server import main as mcp_main + mcp_main() elif args.command == "tpu": from .utils.tpu import tpu_dispatch tpu_dispatch(args) diff --git a/needle/mcp_server.py b/needle/mcp_server.py new file mode 100644 index 0000000..e451981 --- /dev/null +++ b/needle/mcp_server.py @@ -0,0 +1,280 @@ +"""Needle MCP server — exposes Needle as a local tool-dispatch MCP tool. + +Any MCP-compatible orchestrator (Claude Desktop, custom agents) can use +Needle as a fast local router: Needle picks the tool call, the orchestrator +executes it. If Needle's confidence is below the caller's threshold the +response signals a no-match so the orchestrator can escalate to a cloud LLM. + +Usage: + needle mcp-server --checkpoint checkpoints/needle.pkl + +Claude Desktop config (~/.claude/claude_desktop_config.json): + { + "mcpServers": { + "needle": { + "command": "needle", + "args": ["mcp-server", "--checkpoint", "/path/to/needle.pkl"] + } + } + } +""" + +import argparse +import json +import logging +import sys +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any + +logger = logging.getLogger(__name__) + +_DISPATCH_TOOL = { + "name": "dispatch", + "description": ( + "Route a natural language query to the correct tool call using the " + "Needle on-device model. Returns a tool call JSON and a confidence " + "score. Use threshold to decide whether to trust the result or " + "escalate to a cloud LLM." + ), + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Natural language user query to route.", + }, + "tools": { + "type": "array", + "description": "Available tool definitions (name, description, parameters).", + "items": {"type": "object"}, + }, + "threshold": { + "type": "number", + "description": ( + "Confidence threshold in [0, 1]. When confidence is below " + "this value the response sets match=false. Default 0 " + "(always return best guess)." + ), + "default": 0.0, + }, + }, + "required": ["query", "tools"], + }, +} + + +def _load_model(checkpoint_path: str): + from .dataset.dataset import get_tokenizer + from .model.architecture import SimpleAttentionNetwork + from .model.run import load_checkpoint + + params, config = load_checkpoint(checkpoint_path) + model = SimpleAttentionNetwork(config) + tokenizer = get_tokenizer() + return model, params, tokenizer + + +def _dispatch(model, params, tokenizer, query: str, tools: list, threshold: float): + from .model.run import generate + + tools_json = json.dumps(tools, separators=(",", ":")) + result, confidence = generate( + model, + params, + tokenizer, + query=query, + tools=tools_json, + stream=False, + threshold=threshold, + return_confidence=True, + ) + try: + tool_call = json.loads(result) + except (json.JSONDecodeError, TypeError): + tool_call = result + + return { + "tool_call": tool_call, + "confidence": round(confidence, 4), + "match": not (isinstance(tool_call, dict) and tool_call.get("match") is False), + } + + +def _handle_mcp_request(req: dict, model: Any, params: Any, tokenizer: Any): + method = req.get("method", "") + req_id = req.get("id") + + if method == "initialize": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "serverInfo": {"name": "needle", "version": "1.0.0"}, + "capabilities": {"tools": {}}, + }, + } + + if method == "tools/list": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": [_DISPATCH_TOOL]}, + } + + if method == "tools/call": + params = req.get("params", {}) + name = params.get("name") + args = params.get("arguments", {}) + if name != "dispatch": + return _error(req_id, -32601, f"Unknown tool: {name}") + try: + result = _dispatch( + model, + params, + tokenizer, + query=args["query"], + tools=args.get("tools", []), + threshold=float(args.get("threshold", 0.0)), + ) + except Exception as exc: + logger.exception("dispatch failed") + return _error(req_id, -32603, str(exc)) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": json.dumps(result)}], + "isError": False, + }, + } + + if method == "notifications/initialized": + return None + + return _error(req_id, -32601, f"Method not found: {method}") + + +class _StdioMCPServer: + """Minimal stdio MCP server (JSON-RPC 2.0, MCP 2024-11 spec).""" + + def __init__(self, model, params, tokenizer): + self._model = model + self._params = params + self._tokenizer = tokenizer + + def run(self): + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except json.JSONDecodeError: + continue + resp = _handle_mcp_request(req, self._model, self._params, self._tokenizer) + if resp is not None: + sys.stdout.write(json.dumps(resp) + "\n") + sys.stdout.flush() + + +class _HTTPMCPHandler(BaseHTTPRequestHandler): + server_version = "NeedleMCP/1.0" + protocol_version = "HTTP/1.1" + + def do_POST(self): + content_length = int(self.headers.get("Content-Length", 0) or 0) + body = self.rfile.read(content_length) + try: + req = json.loads(body.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + self.send_response(400) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps({"error": "Invalid JSON"}).encode("utf-8")) + return + + resp = self.server.handle_mcp(req) + if resp is None: + resp = {} + encoded = json.dumps(resp).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def do_GET(self): + if self.path == "/health": + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps({"status": "ok"}).encode("utf-8")) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, format: str, *args: Any) -> None: + return + + +class _HTTPMCPServer(HTTPServer): + def __init__(self, server_address, RequestHandlerClass, model, params, tokenizer): + super().__init__(server_address, RequestHandlerClass) + self.model = model + self.params = params + self.tokenizer = tokenizer + + def handle_mcp(self, req: dict): + return _handle_mcp_request(req, self.model, self.params, self.tokenizer) + + +def _run_http(model, params, tokenizer, host: str, port: int): + server = _HTTPMCPServer((host, port), _HTTPMCPHandler, model, params, tokenizer) + logger.warning("Needle MCP HTTP server listening on %s:%s", host, port) + try: + server.serve_forever() + except KeyboardInterrupt: + pass + + +def main(): + parser = argparse.ArgumentParser(description="Needle MCP server") + parser.add_argument( + "--checkpoint", + default="checkpoints/needle.pkl", + help="Path to Needle checkpoint (.pkl)", + ) + parser.add_argument( + "--transport", + choices=["stdio", "http"], + default="stdio", + help="Server transport protocol for MCP (stdio or http)", + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="Host to bind when using HTTP transport", + ) + parser.add_argument( + "--port", + type=int, + default=8765, + help="Port to bind when using HTTP transport", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.WARNING, stream=sys.stderr) + print(f"[needle-mcp] Loading checkpoint: {args.checkpoint}", file=sys.stderr) + model, params, tokenizer = _load_model(args.checkpoint) + print("[needle-mcp] Ready", file=sys.stderr) + + if args.transport == "http": + print(f"[needle-mcp] Listening on http://{args.host}:{args.port}", file=sys.stderr) + _run_http(model, params, tokenizer, args.host, args.port) + else: + _StdioMCPServer(model, params, tokenizer).run() + + +if __name__ == "__main__": + main() diff --git a/needle/model/constrained.py b/needle/model/constrained.py index b6afbd7..2d10480 100644 --- a/needle/model/constrained.py +++ b/needle/model/constrained.py @@ -96,8 +96,9 @@ def __init__(self, tools_json: str): params = tool.get("parameters", {}) if isinstance(params, dict): param_trie = Trie() - for key, val in params.items(): - if isinstance(val, dict): + properties = params.get("properties", params) + for key in properties: + if isinstance(key, str) and key: param_trie.insert(key) self.param_tries[name] = param_trie @@ -347,11 +348,21 @@ def __init__(self, tool_constraints_list: list[ToolConstraints], self.machines = [JsonStateMachine() for _ in range(self.batch_size)] self.token_strings = token_strings self.token_index = token_index + self._confidence: list[float] = [0.0] * self.batch_size def is_active(self, batch_idx: int) -> bool: """Return True if this batch element is currently in a constrained state.""" return self.machines[batch_idx].state != JsonState.FREE + def get_confidence(self, batch_idx: int = 0) -> float: + """Return confidence [0, 1] from the last tool name selection. + + Computed as the max softmax probability over valid name-starting tokens + at the first constrained step. Returns 0.0 if no constrained name + generation has occurred yet. + """ + return self._confidence[batch_idx] + def constrain_logits(self, logits: np.ndarray, batch_idx: int) -> np.ndarray: """Apply grammar constraints to logits for a single batch element.""" machine = self.machines[batch_idx] @@ -374,7 +385,16 @@ def constrain_logits(self, logits: np.ndarray, batch_idx: int) -> np.ndarray: logger.warning("Constrained decoding: off-trie at %r, falling back", machine.constrained_buf) return logits - return apply_constraints(logits, machine.state, node, self.token_strings, self.token_index) + masked = apply_constraints(logits, machine.state, node, self.token_strings, self.token_index) + + if machine.state == JsonState.IN_NAME and machine.constrained_buf == "": + valid = masked > -np.inf + if valid.any(): + subset = masked[valid] + exp_l = np.exp(subset - subset.max()) + self._confidence[batch_idx] = float(exp_l.max() / exp_l.sum()) + + return masked def update(self, batch_idx: int, token_id: int): """Advance the state machine for *batch_idx* with the selected token.""" diff --git a/needle/model/run.py b/needle/model/run.py index 9216045..06b3b05 100644 --- a/needle/model/run.py +++ b/needle/model/run.py @@ -103,11 +103,27 @@ def _build_encoder_input(tokenizer, query, tools, max_enc_len=DEFAULT_MAX_ENC_LE return q_toks + [tools_sep_id] + t_toks -def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, seed=0, stream=True, task_token_id=None, normalize=True, constrained=True): +def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, seed=0, stream=True, task_token_id=None, normalize=True, constrained=True, threshold: float = 0.0, return_confidence: bool = False): """Generate tool-call output. Encoder: [query_tokens..., , tools_tokens...] truncated to max_enc_len. Decoder: prefilled with [EOS], model predicts first, then answer tokens. + + Tool descriptions are supported — include a ``"description"`` field in each + tool object and the encoder will incorporate the text for better semantic + matching:: + + [{"name": "get_weather", + "description": "Get current weather for a location", + "parameters": {"location": "string"}}] + + Args: + threshold: confidence threshold in [0, 1]. When > 0 and the model's + confidence is below this value, returns a no-match sentinel instead + of a tool call: ``{"match":false,"confidence":0.31}``. + Set to 0.0 (default) to disable and always return the best guess. + return_confidence: if True, returns a ``(result, confidence)`` tuple + instead of just the result string. """ name_map = {} if normalize: @@ -176,6 +192,14 @@ def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MA result = result[len(""):] if normalize and name_map: result = restore_tool_names(result, name_map) + + confidence = constrained_decoder.get_confidence(0) if constrained_decoder else 1.0 + + if threshold > 0.0 and confidence < threshold: + result = _json.dumps({"match": False, "confidence": round(confidence, 4)}, separators=(",", ":")) + + if return_confidence: + return result, confidence return result diff --git a/needle/ui/server.py b/needle/ui/server.py index d75bbc9..d531a58 100644 --- a/needle/ui/server.py +++ b/needle/ui/server.py @@ -377,11 +377,11 @@ def _handle_generate(self): self._json_response(400, {"error": str(exc)}) return - result, error = _run_generate(query, tools, seed, max_gen_len, constrained) + result, confidence, error = _run_generate(query, tools, seed, max_gen_len, constrained) if error: self._json_response(500, {"error": error}) else: - self._json_response(200, {"result": result}) + self._json_response(200, {"result": result, "confidence": confidence}) def _handle_finetune(self): if not _is_local_request(self.client_address[0]): @@ -437,7 +437,7 @@ def _run_generate(query, tools, seed, max_gen_len, constrained): if _model is None or _params is None or _tokenizer is None: return None, "model is not loaded" try: - result = generate( + result, confidence = generate( _model, _params, _tokenizer, @@ -447,11 +447,12 @@ def _run_generate(query, tools, seed, max_gen_len, constrained): seed=seed, stream=False, constrained=constrained, + return_confidence=True, ) - return result, None + return result, confidence, None except Exception as exc: print(f"[generate] {exc}", file=sys.stderr) - return None, "Generation failed" + return None, None, "Generation failed" def _load_checkpoint(path, display_name=None): diff --git a/needle/ui/static/app.js b/needle/ui/static/app.js index cb06b58..e7c8efb 100644 --- a/needle/ui/static/app.js +++ b/needle/ui/static/app.js @@ -77,6 +77,8 @@ async function send() { btn.disabled = true; result.className = "result-box"; result.textContent = "Running..."; + var confEl = document.getElementById("confidence"); + if (confEl) confEl.textContent = ""; try { var r = await fetch("/generate", { @@ -101,6 +103,11 @@ async function send() { } else { result.className = "result-box has-result"; result.textContent = data.result; + if (data.confidence !== null && data.confidence !== undefined) { + var pct = Math.round(data.confidence * 100); + var confEl = document.getElementById("confidence"); + if (confEl) confEl.textContent = "Confidence: " + pct + "%"; + } } } catch (e) { result.textContent = ""; diff --git a/needle/ui/static/index.html b/needle/ui/static/index.html index a06db16..f65b4d0 100644 --- a/needle/ui/static/index.html +++ b/needle/ui/static/index.html @@ -65,6 +65,7 @@ 

+