claude-code/infer.py

"""Call the local mlx_lm.server for one hint. Returns hint str or None."""
import json, urllib.request
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent))
from clean_output import clean_hint

URL = "http://127.0.0.1:8080/v1/chat/completions"


def get_hint(prompt, timeout=8.0):
    body = json.dumps({
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 24,
        "temperature": 0.0,
    }).encode()
    req = urllib.request.Request(URL, data=body, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            out = json.load(r)
        text = out["choices"][0]["message"]["content"]
    except Exception:
        return None  # server down / timeout so no hint, never break the tool
    return clean_hint(text) 


if __name__ == "__main__":
    print(get_hint(sys.stdin.read()))