fm_proxy.py

gistfile1.txt

from http.server import BaseHTTPRequestHandler, HTTPServer
from http.client import HTTPConnection
import time
import json
import subprocess
import sys
import threading

UPSTREAM_HOST = "127.0.0.1"
UPSTREAM_PORT = 1976
PROXY_PORT = 1977


def wait_for_fm(max_attempts=60, interval=0.5):
    """Wait for fm serve to be reachable on UPSTREAM_PORT."""
    import httpx
    for i in range(max_attempts):
        try:
            r = httpx.get(f"http://127.0.0.1:{UPSTREAM_PORT}/health", timeout=0.5)
            if r.status_code == 200:
                try:
                    r.json()
                    print(f"[proxy] fm serve is ready (attempt {i + 1})")
                    return True
                except Exception:
                    pass
        except Exception:
            pass
        time.sleep(interval)
    return False


def spawn_fm():
    """Start fm serve in the background, capturing stdout/stderr."""
    proc = subprocess.Popen(
                            ["fm", "serve", "--port", str(UPSTREAM_PORT)],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT,
                            text=True,
                            bufsize=1,  # line-buffered
                            )
    return proc


def log_console(stdout_pipe):
    """Read fm serve stdout and print text chunks to console."""
    for line in stdout_pipe:
        line = line.rstrip("\n")
        if not line:
            continue
        print(f"[fm] {line}", flush=True)


class Proxy(BaseHTTPRequestHandler):
    def do_GET(self):
        path = self.path
        if path == "/v1/models?":
            path = "/v1/models"
        
        conn = HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT, timeout=30)
        try:
            conn.request("GET", path, headers=self.headers)
            resp = conn.getresponse()
            
            self.send_response(resp.status)
            for k, v in resp.getheaders():
                if k.lower() != "transfer-encoding":
                    self.send_header(k, v)
            self.end_headers()
            
            # Only wrap errors if not SSE and body is text-like
            if resp.status >= 400:
                try:
                    body = resp.read().decode("utf-8", errors="ignore")
                    content_type = resp.getheader("Content-Type", "")
                    
                    wrap_error = "text/eventStream" not in content_type
                    
                    if wrap_error:
                        try:
                            parsed = json.loads(body)
                            if isinstance(parsed, dict) and "error" in parsed:
                                self.wfile.write(body.encode())
                                self.wfile.flush()
                                return
                        except json.JSONDecodeError:
                            pass
                        
                        error_obj = {
                            "error": {
                                "message": f"Upstream error: {resp.status}",
                                "code": resp.status,
                                "detail": body[:500]  if body else "No upstream response"
                            }
                        }
                        print(f"[proxy] Upstream error {resp.status} during SSE stream (not wrapped).")
                        self.wfile.write(body.encode())
                        self.wfile.flush()
                    else:
                        pass
                except Exception as e:
                    fallback = json.dumps({
                        "error": {
                            "message": f"Proxy error processing {resp.status}: {str(e)}",
                            "code": 500
                        }
                    })
                    
                    self.wfile.write(json.dumps(fallback).encode())
                    self.wfile.flush()
                return
            
            while True:
                try:
                    chunk = resp.read(4096)
                except socket.timeout:
                    break
                    
                if not chunk:
                    break
                self.wfile.write(chunk)
                self.wfile.flush()
                
        finally:
            conn.close()
    
    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length)
        
        path = self.path
        
        conn = HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT, timeout=30)
        try:
            conn.request("POST", path, body=body, headers=self.headers)
            resp = conn.getresponse()
            
            self.send_response(resp.status)
            for k, v in resp.getheaders():
                if k.lower() != "transfer-encoding":
                    self.send_header(k, v)
            self.end_headers()
            
            # Only wrap errors if not SSE and body is text-like
            if resp.status >= 400:
                try:
                    resp_body = resp.read().decode("utf-8", errors="ignore")
                    content_type = resp.getheader("Content-Type", "")
                    
                    wrap_error = "text/eventStream" not in content_type
                    
                    if wrap_error:
                        try:
                            parsed = json.loads(resp_body)
                            if isinstance(parsed, dict) and "error" in parsed:
                                self.wfile.write(resp_body.encode())
                                self.wfile.flush()
                                return
                        except json.JSONDecodeError:
                            pass
                        
                        error_obj = {
                            "error": {
                                "message": f"Upstream error: {resp.status}",
                                "code": resp.status,
                                "detail": resp_body[:500]  if resp_body else "No upstream response"
                            }
                        }
                        
                        self.wfile.write(json.dumps(error_obj).encode())
                        self.wfile.flush()
                    return
                except Exception as e:
                    fallback = json.dumps({
                        "error": {
                            "message": f"Proxy error processing {resp.status}: {str(e)}",
                            "code": 500
                        }
                    })
                    self.wfile.write(json.dumps(fallback).encode())
                    self.wfile.flush()
                return
            
            #SSE processing
            while True:
                try:
                    chunk = resp.read(4096)
                except socket.timeout:
                    break
                    
                if not chunk:
                    break
                
                text = chunk.decode("utf-8", errors="ignore")
                
                for line in text.splitlines():
                    if not line.startswith("data:"):
                        continue
                    
                    payload = line[6:].strip()
                    
                    if payload == "[DONE]":
                        continue
                    
                    try:
                        obj = json.loads(payload)
                    except json.JSONDecodeError:
                        continue
                    
                    # 🔥 PUT IT HERE
                    obj.setdefault("object", "chat.completion.chunk")
                    obj.setdefault("created", int(time.time()))
                    
                    choice = obj["choices"][0]
                    choice.setdefault("index", 0)
                    choice.setdefault("finish_reason", None)
                    
                    delta = choice.setdefault("delta", {})
                    delta.setdefault("content", None)
                    
                    # re-emit SSE
                    self.wfile.write(f"data: {json.dumps(obj)}\n\n".encode())
                    self.wfile.flush()
        finally:
            conn.close()

def main():
    print("[proxy] Starting fm serve...")
    
    proc = None
    while True:
        proc = spawn_fm()
        print(f"[proxy] fm serve PID: {proc.pid}")
        
        # Start logger thread for this process
        logger_thread = threading.Thread(target=log_console, args=(proc.stdout,), daemon=True)
        logger_thread.start()
        
        # Wait for readiness
        if not wait_for_fm():
            if proc.poll() is not None:
                stdout = proc.stdout.read()
                print(f"[proxy] fm serve exited immediately. Output:\n{stdout}")
            else:
                print("[proxy] fm serve did not become ready within timeout.")
            sys.exit(1)
        
        print(f"[proxy] Starting HTTP proxy on port {PROXY_PORT}...")
        
        server = HTTPServer(("127.0.0.1", PROXY_PORT), Proxy)
        
        # Monitor thread: watch for fm serve crash
        def monitor():
            proc.wait()
            code = proc.returncode
            print(f"[proxy] fm serve (PID {proc.pid}) exited with code {code}. Shutting down.")
            server.shutdown()  # stop accepting new requests
        
        monitor_thread = threading.Thread(target=monitor, daemon=True)
        monitor_thread.start()
        
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            print("\n[proxy] Shutting down...")
            break
        
        # After shutdown, wait a bit before restarting
        print("[proxy] fm serve stopped. Restarting in 5s...")
        time.sleep(5)
    
    if proc:
        proc.terminate()
        proc.wait()

if __name__ == "__main__":
    main()