From 1dd5a7885732d3f2fa90a869173a980c5c070cbd Mon Sep 17 00:00:00 2001 From: bs258q Date: Wed, 13 May 2026 12:43:32 -0700 Subject: [PATCH 1/4] feat: add tool descriptions, confidence tracking, and threshold fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix constrained decoding for flat parameter format {"key": "string"} — previously only JSON Schema {"properties": {"key": {...}}} worked; now both formats populate the param trie correctly - Add confidence tracking to ConstrainedDecoder: captures softmax max probability over valid name tokens at the first IN_NAME step - Add threshold/no-match fallback to generate(): returns {"match":false,"confidence":0.31} when confidence < threshold - Add return_confidence param to generate() for tuple return - Tool descriptions already work via encoder (no structural change needed) - Add comprehensive unit tests covering Trie, ToolConstraints, JsonStateMachine, apply_constraints, and ConstrainedDecoder confidence Signed-off-by: bs258q --- needle/model/constrained.py | 26 +++- needle/model/run.py | 26 +++- tests/__init__.py | 0 tests/test_constrained.py | 295 ++++++++++++++++++++++++++++++++++++ 4 files changed, 343 insertions(+), 4 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/test_constrained.py 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/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_constrained.py b/tests/test_constrained.py new file mode 100644 index 0000000..0d03a13 --- /dev/null +++ b/tests/test_constrained.py @@ -0,0 +1,295 @@ +"""Unit tests for grammar-constrained decoding. + +Covers: +- ToolConstraints param trie: flat, JSON Schema, and description-bearing formats +- Trie operations +- JsonStateMachine state transitions +- apply_constraints logit masking +- _check_token_valid token validation +- ConstrainedDecoder confidence tracking +""" + +import numpy as np +import pytest + +from needle.model.constrained import ( + JsonState, + JsonStateMachine, + Trie, + ToolConstraints, + ConstrainedDecoder, + TokenIndex, + _check_token_valid, + apply_constraints, + build_token_strings, +) + + +# ── Trie ────────────────────────────────────────────────────────────────────── + +class TestTrie: + def test_insert_and_get_node(self): + t = Trie() + t.insert("location") + assert t.get_node("loc") is not None + + def test_terminal_flag(self): + t = Trie() + t.insert("location") + assert t.get_node("location").is_terminal + + def test_prefix_not_terminal(self): + t = Trie() + t.insert("location") + assert not t.get_node("loc").is_terminal + + def test_missing_prefix_returns_none(self): + t = Trie() + t.insert("location") + assert t.get_node("xyz") is None + + def test_empty_prefix_returns_root(self): + t = Trie() + t.insert("ab") + node = t.get_node("") + assert node is not None + assert "a" in node.children + + def test_words_property(self): + t = Trie() + for w in ["alpha", "beta", "gamma"]: + t.insert(w) + assert sorted(t.words) == ["alpha", "beta", "gamma"] + + +# ── ToolConstraints ──────────────────────────────────────────────────────────── + +class TestToolConstraints: + def test_flat_string_params(self): + """Regression: flat {\"key\": \"string\"} must populate param trie.""" + tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' + tc = ToolConstraints(tools) + trie = tc.get_param_trie("get_weather") + assert trie is not None + assert trie.get_node("location").is_terminal + + def test_flat_multiple_params(self): + tools = '[{"name":"set_alarm","parameters":{"time":"string","label":"string"}}]' + tc = ToolConstraints(tools) + trie = tc.get_param_trie("set_alarm") + assert trie.get_node("time").is_terminal + assert trie.get_node("label").is_terminal + + def test_json_schema_properties_format(self): + tools = '[{"name":"search","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]' + tc = ToolConstraints(tools) + trie = tc.get_param_trie("search") + assert trie is not None + assert trie.get_node("query").is_terminal + + def test_description_field_ignored_in_param_trie(self): + """description at tool level must not leak into param trie.""" + tools = '[{"name":"get_weather","description":"Get current weather","parameters":{"location":"string"}}]' + tc = ToolConstraints(tools) + trie = tc.get_param_trie("get_weather") + assert trie.get_node("location").is_terminal + assert trie.get_node("description") is None + + def test_description_does_not_break_name_trie(self): + tools = '[{"name":"get_weather","description":"Get current weather","parameters":{"location":"string"}}]' + tc = ToolConstraints(tools) + assert tc.name_trie.get_node("get_weather").is_terminal + + def test_name_trie_populated(self): + tools = '[{"name":"get_weather","parameters":{"location":"string"}},{"name":"set_alarm","parameters":{"time":"string"}}]' + tc = ToolConstraints(tools) + assert tc.name_trie.get_node("get_weather").is_terminal + assert tc.name_trie.get_node("set_alarm").is_terminal + + def test_unknown_tool_returns_none(self): + tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' + tc = ToolConstraints(tools) + assert tc.get_param_trie("nonexistent") is None + + def test_invalid_json_does_not_raise(self): + tc = ToolConstraints("not json at all") + assert tc.name_trie.get_node("x") is None + + def test_empty_tools_list(self): + tc = ToolConstraints("[]") + assert tc.name_trie.words == [] + + def test_multi_tool_playground_format(self): + """Matches exact format used in playground and test.py.""" + tools = '''[ + {"name": "get_weather", "description": "Get weather", "parameters": {"location": "string"}}, + {"name": "set_alarm", "description": "Set an alarm", "parameters": {"time": "string", "label": "string"}}, + {"name": "search_web", "description": "Search the web", "parameters": {"query": "string"}}, + {"name": "send_message", "description": "Send a message", "parameters": {"to": "string", "body": "string"}} + ]''' + tc = ToolConstraints(tools) + assert tc.get_param_trie("get_weather").get_node("location").is_terminal + assert tc.get_param_trie("set_alarm").get_node("time").is_terminal + assert tc.get_param_trie("set_alarm").get_node("label").is_terminal + assert tc.get_param_trie("search_web").get_node("query").is_terminal + assert tc.get_param_trie("send_message").get_node("to").is_terminal + assert tc.get_param_trie("send_message").get_node("body").is_terminal + + +# ── _check_token_valid ──────────────────────────────────────────────────────── + +class TestCheckTokenValid: + def setup_method(self): + self.trie = Trie() + self.trie.insert("location") + self.trie.insert("lang") + + def test_valid_prefix_token(self): + node = self.trie.get_node("") + assert _check_token_valid("loc", node) + + def test_valid_full_word_plus_quote(self): + node = self.trie.get_node("") + assert _check_token_valid('location"', node) + + def test_invalid_token_not_in_trie(self): + node = self.trie.get_node("") + assert not _check_token_valid("xyz", node) + + def test_quote_at_nonterminal_invalid(self): + node = self.trie.get_node("loc") + assert not _check_token_valid('"', node) + + def test_quote_at_terminal_valid(self): + node = self.trie.get_node("location") + assert _check_token_valid('"', node) + + +# ── apply_constraints ──────────────────────────────────────────────────────── + +class TestApplyConstraints: + def _make_token_strings(self): + return ["", "[", "{", '"', "l", "o", "c", "a", "t", "i", "n", "x", 'location"'] + + def _make_token_index(self, token_strings): + return TokenIndex(token_strings) + + def test_valid_tokens_unmasked(self): + trie = Trie() + trie.insert("location") + token_strings = self._make_token_strings() + token_index = self._make_token_index(token_strings) + logits = np.zeros(len(token_strings)) + result = apply_constraints(logits, JsonState.IN_ARG_KEY, trie.root, token_strings, token_index) + assert result[4] == 0.0 # "l" valid + assert result[11] == -np.inf # "x" invalid + + def test_fallback_on_empty_trie(self): + trie = Trie() + token_strings = self._make_token_strings() + token_index = self._make_token_index(token_strings) + logits = np.ones(len(token_strings)) + result = apply_constraints(logits, JsonState.IN_ARG_KEY, trie.root, token_strings, token_index) + np.testing.assert_array_equal(result, logits) + + +# ── JsonStateMachine ────────────────────────────────────────────────────────── + +class TestJsonStateMachine: + def test_initial_state_is_free(self): + assert JsonStateMachine().state == JsonState.FREE + + def test_enters_in_name_after_name_prefix(self): + m = JsonStateMachine() + m.feed('[{"name":"') + assert m.state == JsonState.IN_NAME + + def test_captures_tool_name(self): + m = JsonStateMachine() + m.feed('[{"name":"get_weather"') + assert m.current_function == "get_weather" + assert m.state == JsonState.FREE + + def test_enters_in_arg_key_after_arguments(self): + m = JsonStateMachine() + m.feed('[{"name":"get_weather","arguments":{"') + assert m.state == JsonState.IN_ARG_KEY + + def test_captures_arg_key(self): + m = JsonStateMachine() + m.feed('[{"name":"get_weather","arguments":{"location"') + assert m.state == JsonState.FREE + + def test_second_arg_key_triggers_in_arg_key(self): + m = JsonStateMachine() + m.feed('[{"name":"set_alarm","arguments":{"time":"7am","') + assert m.state == JsonState.IN_ARG_KEY + + def test_constrained_buf_resets_on_close_quote(self): + m = JsonStateMachine() + m.feed('[{"name":"get_weather","arguments":{"location"') + assert m.constrained_buf == "" + + +# ── ConstrainedDecoder confidence ──────────────────────────────────────────── + +class TestConstrainedDecoderConfidence: + def _make_decoder(self, tools_json: str, vocab: list[str]) -> ConstrainedDecoder: + token_index = TokenIndex(vocab) + tc = ToolConstraints(tools_json) + return ConstrainedDecoder([tc], vocab, token_index) + + def test_confidence_starts_at_zero(self): + vocab = ["", "g", "s", '"', "et_weather", 'et_weather"'] + decoder = self._make_decoder('[{"name":"get_weather","parameters":{"location":"string"}}]', vocab) + assert decoder.get_confidence(0) == 0.0 + + def test_confidence_set_after_in_name_step(self): + tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' + # vocab includes 'g' (start of get_weather) and 's' (start of search) + vocab = ["", "g", "s", "e", "t", "_", "w", "a", "r", "h", "o", '"', + "et_weather", 'et_weather"', "earch_web", 'earch_web"'] + decoder = self._make_decoder(tools, vocab) + + # Drive machine to IN_NAME state + decoder.machines[0].feed('[{"name":"') + assert decoder.machines[0].state == JsonState.IN_NAME + assert decoder.machines[0].constrained_buf == "" + + # Apply constraints — confidence should be captured + logits = np.zeros(len(vocab)) + logits[1] = 2.0 # 'g' — high probability + logits[2] = 0.5 # 's' — lower + decoder.constrain_logits(logits, 0) + + conf = decoder.get_confidence(0) + assert 0.0 < conf <= 1.0 + + def test_confidence_is_higher_when_model_more_certain(self): + tools = '[{"name":"get_weather","parameters":{"location":"string"}},{"name":"search_web","parameters":{"query":"string"}}]' + vocab = ["", "g", "s", '"'] + decoder_certain = self._make_decoder(tools, vocab) + decoder_uncertain = self._make_decoder(tools, vocab) + + for decoder in (decoder_certain, decoder_uncertain): + decoder.machines[0].feed('[{"name":"') + + # Certain: model heavily favors 'g' + logits_certain = np.array([0.0, 10.0, 0.1, 0.0]) + decoder_certain.constrain_logits(logits_certain, 0) + + # Uncertain: roughly equal between 'g' and 's' + logits_uncertain = np.array([0.0, 1.0, 1.0, 0.0]) + decoder_uncertain.constrain_logits(logits_uncertain, 0) + + assert decoder_certain.get_confidence(0) > decoder_uncertain.get_confidence(0) + + def test_confidence_in_valid_range(self): + tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' + vocab = ["", "g", "s", '"'] + decoder = self._make_decoder(tools, vocab) + decoder.machines[0].feed('[{"name":"') + logits = np.random.randn(len(vocab)) + decoder.constrain_logits(logits, 0) + conf = decoder.get_confidence(0) + assert 0.0 <= conf <= 1.0 From 7f983295d7ec8c59116128ed985fb77b5442728a Mon Sep 17 00:00:00 2001 From: bs258q Date: Wed, 13 May 2026 12:45:06 -0700 Subject: [PATCH 2/4] chore: remove test files Signed-off-by: bs258q --- tests/__init__.py | 0 tests/test_constrained.py | 295 -------------------------------------- 2 files changed, 295 deletions(-) delete mode 100644 tests/__init__.py delete mode 100644 tests/test_constrained.py diff --git a/tests/__init__.py b/tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/test_constrained.py b/tests/test_constrained.py deleted file mode 100644 index 0d03a13..0000000 --- a/tests/test_constrained.py +++ /dev/null @@ -1,295 +0,0 @@ -"""Unit tests for grammar-constrained decoding. - -Covers: -- ToolConstraints param trie: flat, JSON Schema, and description-bearing formats -- Trie operations -- JsonStateMachine state transitions -- apply_constraints logit masking -- _check_token_valid token validation -- ConstrainedDecoder confidence tracking -""" - -import numpy as np -import pytest - -from needle.model.constrained import ( - JsonState, - JsonStateMachine, - Trie, - ToolConstraints, - ConstrainedDecoder, - TokenIndex, - _check_token_valid, - apply_constraints, - build_token_strings, -) - - -# ── Trie ────────────────────────────────────────────────────────────────────── - -class TestTrie: - def test_insert_and_get_node(self): - t = Trie() - t.insert("location") - assert t.get_node("loc") is not None - - def test_terminal_flag(self): - t = Trie() - t.insert("location") - assert t.get_node("location").is_terminal - - def test_prefix_not_terminal(self): - t = Trie() - t.insert("location") - assert not t.get_node("loc").is_terminal - - def test_missing_prefix_returns_none(self): - t = Trie() - t.insert("location") - assert t.get_node("xyz") is None - - def test_empty_prefix_returns_root(self): - t = Trie() - t.insert("ab") - node = t.get_node("") - assert node is not None - assert "a" in node.children - - def test_words_property(self): - t = Trie() - for w in ["alpha", "beta", "gamma"]: - t.insert(w) - assert sorted(t.words) == ["alpha", "beta", "gamma"] - - -# ── ToolConstraints ──────────────────────────────────────────────────────────── - -class TestToolConstraints: - def test_flat_string_params(self): - """Regression: flat {\"key\": \"string\"} must populate param trie.""" - tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' - tc = ToolConstraints(tools) - trie = tc.get_param_trie("get_weather") - assert trie is not None - assert trie.get_node("location").is_terminal - - def test_flat_multiple_params(self): - tools = '[{"name":"set_alarm","parameters":{"time":"string","label":"string"}}]' - tc = ToolConstraints(tools) - trie = tc.get_param_trie("set_alarm") - assert trie.get_node("time").is_terminal - assert trie.get_node("label").is_terminal - - def test_json_schema_properties_format(self): - tools = '[{"name":"search","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]' - tc = ToolConstraints(tools) - trie = tc.get_param_trie("search") - assert trie is not None - assert trie.get_node("query").is_terminal - - def test_description_field_ignored_in_param_trie(self): - """description at tool level must not leak into param trie.""" - tools = '[{"name":"get_weather","description":"Get current weather","parameters":{"location":"string"}}]' - tc = ToolConstraints(tools) - trie = tc.get_param_trie("get_weather") - assert trie.get_node("location").is_terminal - assert trie.get_node("description") is None - - def test_description_does_not_break_name_trie(self): - tools = '[{"name":"get_weather","description":"Get current weather","parameters":{"location":"string"}}]' - tc = ToolConstraints(tools) - assert tc.name_trie.get_node("get_weather").is_terminal - - def test_name_trie_populated(self): - tools = '[{"name":"get_weather","parameters":{"location":"string"}},{"name":"set_alarm","parameters":{"time":"string"}}]' - tc = ToolConstraints(tools) - assert tc.name_trie.get_node("get_weather").is_terminal - assert tc.name_trie.get_node("set_alarm").is_terminal - - def test_unknown_tool_returns_none(self): - tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' - tc = ToolConstraints(tools) - assert tc.get_param_trie("nonexistent") is None - - def test_invalid_json_does_not_raise(self): - tc = ToolConstraints("not json at all") - assert tc.name_trie.get_node("x") is None - - def test_empty_tools_list(self): - tc = ToolConstraints("[]") - assert tc.name_trie.words == [] - - def test_multi_tool_playground_format(self): - """Matches exact format used in playground and test.py.""" - tools = '''[ - {"name": "get_weather", "description": "Get weather", "parameters": {"location": "string"}}, - {"name": "set_alarm", "description": "Set an alarm", "parameters": {"time": "string", "label": "string"}}, - {"name": "search_web", "description": "Search the web", "parameters": {"query": "string"}}, - {"name": "send_message", "description": "Send a message", "parameters": {"to": "string", "body": "string"}} - ]''' - tc = ToolConstraints(tools) - assert tc.get_param_trie("get_weather").get_node("location").is_terminal - assert tc.get_param_trie("set_alarm").get_node("time").is_terminal - assert tc.get_param_trie("set_alarm").get_node("label").is_terminal - assert tc.get_param_trie("search_web").get_node("query").is_terminal - assert tc.get_param_trie("send_message").get_node("to").is_terminal - assert tc.get_param_trie("send_message").get_node("body").is_terminal - - -# ── _check_token_valid ──────────────────────────────────────────────────────── - -class TestCheckTokenValid: - def setup_method(self): - self.trie = Trie() - self.trie.insert("location") - self.trie.insert("lang") - - def test_valid_prefix_token(self): - node = self.trie.get_node("") - assert _check_token_valid("loc", node) - - def test_valid_full_word_plus_quote(self): - node = self.trie.get_node("") - assert _check_token_valid('location"', node) - - def test_invalid_token_not_in_trie(self): - node = self.trie.get_node("") - assert not _check_token_valid("xyz", node) - - def test_quote_at_nonterminal_invalid(self): - node = self.trie.get_node("loc") - assert not _check_token_valid('"', node) - - def test_quote_at_terminal_valid(self): - node = self.trie.get_node("location") - assert _check_token_valid('"', node) - - -# ── apply_constraints ──────────────────────────────────────────────────────── - -class TestApplyConstraints: - def _make_token_strings(self): - return ["", "[", "{", '"', "l", "o", "c", "a", "t", "i", "n", "x", 'location"'] - - def _make_token_index(self, token_strings): - return TokenIndex(token_strings) - - def test_valid_tokens_unmasked(self): - trie = Trie() - trie.insert("location") - token_strings = self._make_token_strings() - token_index = self._make_token_index(token_strings) - logits = np.zeros(len(token_strings)) - result = apply_constraints(logits, JsonState.IN_ARG_KEY, trie.root, token_strings, token_index) - assert result[4] == 0.0 # "l" valid - assert result[11] == -np.inf # "x" invalid - - def test_fallback_on_empty_trie(self): - trie = Trie() - token_strings = self._make_token_strings() - token_index = self._make_token_index(token_strings) - logits = np.ones(len(token_strings)) - result = apply_constraints(logits, JsonState.IN_ARG_KEY, trie.root, token_strings, token_index) - np.testing.assert_array_equal(result, logits) - - -# ── JsonStateMachine ────────────────────────────────────────────────────────── - -class TestJsonStateMachine: - def test_initial_state_is_free(self): - assert JsonStateMachine().state == JsonState.FREE - - def test_enters_in_name_after_name_prefix(self): - m = JsonStateMachine() - m.feed('[{"name":"') - assert m.state == JsonState.IN_NAME - - def test_captures_tool_name(self): - m = JsonStateMachine() - m.feed('[{"name":"get_weather"') - assert m.current_function == "get_weather" - assert m.state == JsonState.FREE - - def test_enters_in_arg_key_after_arguments(self): - m = JsonStateMachine() - m.feed('[{"name":"get_weather","arguments":{"') - assert m.state == JsonState.IN_ARG_KEY - - def test_captures_arg_key(self): - m = JsonStateMachine() - m.feed('[{"name":"get_weather","arguments":{"location"') - assert m.state == JsonState.FREE - - def test_second_arg_key_triggers_in_arg_key(self): - m = JsonStateMachine() - m.feed('[{"name":"set_alarm","arguments":{"time":"7am","') - assert m.state == JsonState.IN_ARG_KEY - - def test_constrained_buf_resets_on_close_quote(self): - m = JsonStateMachine() - m.feed('[{"name":"get_weather","arguments":{"location"') - assert m.constrained_buf == "" - - -# ── ConstrainedDecoder confidence ──────────────────────────────────────────── - -class TestConstrainedDecoderConfidence: - def _make_decoder(self, tools_json: str, vocab: list[str]) -> ConstrainedDecoder: - token_index = TokenIndex(vocab) - tc = ToolConstraints(tools_json) - return ConstrainedDecoder([tc], vocab, token_index) - - def test_confidence_starts_at_zero(self): - vocab = ["", "g", "s", '"', "et_weather", 'et_weather"'] - decoder = self._make_decoder('[{"name":"get_weather","parameters":{"location":"string"}}]', vocab) - assert decoder.get_confidence(0) == 0.0 - - def test_confidence_set_after_in_name_step(self): - tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' - # vocab includes 'g' (start of get_weather) and 's' (start of search) - vocab = ["", "g", "s", "e", "t", "_", "w", "a", "r", "h", "o", '"', - "et_weather", 'et_weather"', "earch_web", 'earch_web"'] - decoder = self._make_decoder(tools, vocab) - - # Drive machine to IN_NAME state - decoder.machines[0].feed('[{"name":"') - assert decoder.machines[0].state == JsonState.IN_NAME - assert decoder.machines[0].constrained_buf == "" - - # Apply constraints — confidence should be captured - logits = np.zeros(len(vocab)) - logits[1] = 2.0 # 'g' — high probability - logits[2] = 0.5 # 's' — lower - decoder.constrain_logits(logits, 0) - - conf = decoder.get_confidence(0) - assert 0.0 < conf <= 1.0 - - def test_confidence_is_higher_when_model_more_certain(self): - tools = '[{"name":"get_weather","parameters":{"location":"string"}},{"name":"search_web","parameters":{"query":"string"}}]' - vocab = ["", "g", "s", '"'] - decoder_certain = self._make_decoder(tools, vocab) - decoder_uncertain = self._make_decoder(tools, vocab) - - for decoder in (decoder_certain, decoder_uncertain): - decoder.machines[0].feed('[{"name":"') - - # Certain: model heavily favors 'g' - logits_certain = np.array([0.0, 10.0, 0.1, 0.0]) - decoder_certain.constrain_logits(logits_certain, 0) - - # Uncertain: roughly equal between 'g' and 's' - logits_uncertain = np.array([0.0, 1.0, 1.0, 0.0]) - decoder_uncertain.constrain_logits(logits_uncertain, 0) - - assert decoder_certain.get_confidence(0) > decoder_uncertain.get_confidence(0) - - def test_confidence_in_valid_range(self): - tools = '[{"name":"get_weather","parameters":{"location":"string"}}]' - vocab = ["", "g", "s", '"'] - decoder = self._make_decoder(tools, vocab) - decoder.machines[0].feed('[{"name":"') - logits = np.random.randn(len(vocab)) - decoder.constrain_logits(logits, 0) - conf = decoder.get_confidence(0) - assert 0.0 <= conf <= 1.0 From 10619b9325b423fd1a2fa299baebb88d2fe8e481 Mon Sep 17 00:00:00 2001 From: bs258q Date: Wed, 13 May 2026 12:58:34 -0700 Subject: [PATCH 3/4] fix: for showing confidence scores Signed-off-by: bs258q --- needle/ui/server.py | 11 ++++++----- needle/ui/static/app.js | 7 +++++++ needle/ui/static/index.html | 1 + 3 files changed, 14 insertions(+), 5 deletions(-) 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 @@ 

+        

       

     
   

From c1816906f3823cea027033265e8733366d7b26fa Mon Sep 17 00:00:00 2001
From: bs258q 
Date: Wed, 13 May 2026 13:36:32 -0700
Subject: [PATCH 4/4] HTTP MCP endpoint and MCP server to be used by
 orchestrators for tool calling

Signed-off-by: bs258q 
---
 README.md            |  56 +++++++++
 needle/cli.py        |  15 +++
 needle/mcp_server.py | 280 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 351 insertions(+)
 create mode 100644 needle/mcp_server.py

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()