Muse Code Third Part Provider Test

muse_proxy.py

import json
import http.server
import socketserver
import urllib.request
import urllib.error
import sys
import argparse
import signal
import os

"""
# ~/.config/muse/settings.json

```json
{
  "schema_version": 1,
  "run": {
    "reminder_roster": {
      "interval_steps": 999999, # default 1, use large num to skip
      "skip_if_running": true, # you can remove this
      "max_in_flight_per_agent": 1 #  you can remove this
    },
    "context_compaction": {
      "soft_threshold": 0.8,
      "hard_threshold": 0.9,
      "provider_context_limit_tokens": 32768, # higher priority than the interface
      "tool_result_clearing_enabled": true
    }
  }
}

```bash
# how to start
META_API_KEY="dummy-key" muse --provider meta --base-url http://127.0.0.1:9911 --model glm-5.2 --yolo
```
"""


BACKEND_URL= "http://your.endpoint.com"
PROXY_PORT = 9911
MODEL_ID = "glm-5.2"
MODEL_LABEL = "Third-Party Model"


def make_catalog(model_id, label):
    return {
        "schema_version": 1,
        "data": [
            {
                "provider_id": "meta",
                "profile_id": "default",
                "id": model_id,
                "model_id": model_id,
                "display_label": label,
                "visibility": "visible",
                "release_date": "2025-01-01",
                "display_order": 1,
                "is_current": True,
                "is_default": True,
                "context_limit": 32768,
                "output_limit": 8192,
                "description": f"{label} via muse proxy",
            }
        ],
    }


class ProxyHandler(http.server.BaseHTTPRequestHandler):
    backend_url = BACKEND_URL
    catalog = None

    def log_message(self, fmt, *args):
        if os.environ.get("MUSE_PROXY_DEBUG"):
            super().log_message(fmt, *args)

    def _log_request(self, method, path, body=None):
        if os.environ.get("MUSE_PROXY_DEBUG"):
            import time
            ts = time.strftime("%H:%M:%S")
            extra = ""
            if body and len(body) < 500:
                extra = f" body={body[:200]}"
            elif body:
                extra = f" body_len={len(body)}"
            print(f"[{ts}] {method} {path}{extra}", flush=True)

    def _send_json(self, data, status=200):
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _proxy_to(self, target_url, body=b""):
        is_stream = b'"stream":true' in body or b'"stream": true' in body

        req = urllib.request.Request(target_url, data=body, method=self.command)
        for key in ["Content-Type", "Authorization", "Accept", "User-Agent"]:
            val = self.headers.get(key)
            if val:
                req.add_header(key, val)
        if not req.has_header("Content-Type"):
            req.add_header("Content-Type", "application/json")

        try:
            resp = urllib.request.urlopen(req, timeout=600)

            if is_stream:
                self.send_response(resp.status)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-cache")
                self.send_header("Connection", "keep-alive")
                self.end_headers()
                try:
                    while True:
                        chunk = resp.read(4096)
                        if not chunk:
                            break
                        self.wfile.write(chunk)
                        self.wfile.flush()
                except Exception:
                    pass
                finally:
                    resp.close()
            else:
                resp_body = resp.read()
                self.send_response(resp.status)
                for key, val in resp.getheaders():
                    if key.lower() not in ("transfer-encoding", "connection",
                                            "content-encoding"):
                        self.send_header(key, val)
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                self.wfile.write(resp_body)
                resp.close()
        except urllib.error.HTTPError as e:
            resp_body = e.read()
            self.send_response(e.code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp_body)))
            self.end_headers()
            self.wfile.write(resp_body)
        except Exception as e:
            err = json.dumps({"error": str(e)}).encode()
            self.send_response(502)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(err)))
            self.end_headers()
            self.wfile.write(err)

    def do_GET(self):
        self._log_request("GET", self.path)
        if self.path.startswith("/muse-code/models"):
            self._send_json(self.catalog)
            return
        if self.path.startswith("/muse-code/telemetry"):
            self._send_json({})
            return
        if self.path.startswith("/muse-code/search"):
            self._send_json({"results": []})
            return
        self._send_json({}, 404)

    def do_POST(self):
        path = self.path
        content_len = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(content_len) if content_len > 0 else b""
        self._log_request("POST", path, raw_body)

        if path in ("/responses", "/v1/responses"):
            target = f"{self.backend_url}/v1/responses"
        elif path in ("/chat/completions", "/v1/chat/completions"):
            target = f"{self.backend_url}/v1/chat/completions"
        elif path.startswith("/muse-code/telemetry"):
            self._send_json({})
            return
        elif path.startswith("/muse-code/search"):
            self._send_json({"results": []})
            return
        else:
            target = self.backend_url + path
        self._proxy_to(target, raw_body)

    def do_DELETE(self):
        self._send_json({})

    def do_PUT(self):
        content_len = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(content_len) if content_len > 0 else b""
        self._log_request("PUT", self.path, body)
        self._proxy_to(self.backend_url + self.path, body)


def main():
    parser = argparse.ArgumentParser(description="muse proxy server")
    parser.add_argument("--port", type=int, default=9911)
    parser.add_argument("--backend", default=BACKEND_URL)
    parser.add_argument("--model", default=MODEL_ID)
    parser.add_argument("--label", default=MODEL_LABEL)
    args = parser.parse_args()

    ProxyHandler.backend_url = args.backend
    ProxyHandler.catalog = make_catalog(args.model, args.label)

    server = http.server.ThreadingHTTPServer(("127.0.0.1", args.port), ProxyHandler)
    print(
        f"muse-proxy | port={args.port} backend={args.backend} model={args.model}",
        flush=True,
    )
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()

Thanks