A proxy for copilot labs. pip install rich

copilot_proxy.py

#!/usr/bin/env python3
"""
Very simple HTTP server in python for logging requests
Usage::
    ./server.py [<port>]
"""
import json
from http.server import BaseHTTPRequestHandler, HTTPServer

import requests
from rich import print as print
from rich import print_json

URL = 'https://copilot-proxy.githubusercontent.com'


def decode_response(response: str):
    lines = [line.strip()
             for line in response.splitlines() if len(line.strip()) > 0]

    # remove "data: " from the start of the line
    lines = [line[5:] for line in lines]

    # remove the last line
    lines = lines[:-1]

    return [json.loads(line) for line in lines]


class Proxy(BaseHTTPRequestHandler):
    def _set_response(self):
        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()

    def do_GET(self):
        print("GET request,\nPath: %s\nHeaders:\n%s\n",
              str(self.path), str(self.headers))

    def do_POST(self):
        # <--- Gets the size of data
        content_length = int(self.headers['Content-Length'])
        # <--- Gets the data itself
        post_data = self.rfile.read(content_length).decode()
        post_data = json.loads(post_data)

        post_data['temperature'] = 0.75
        # post_data['top_p'] = 0

        print("POST Body:")
        print(post_data)

        self._set_response()

        # relay the request to the server
        r = requests.post(f'{URL}{self.path}',
                          headers=self.headers, data=json.dumps(post_data))

        # relay the response back to the client
        data = decode_response(r.content.decode())
        print(data)

        joined = ''.join(x['choices'][0]['text'] for x in data)
        print(joined)

        self.wfile.write(r.content)


def run(server_class=HTTPServer, handler_class=Proxy, port=8080):
    server_address = ('', port)

    httpd = server_class(server_address, handler_class)

    print('Starting httpd...\n')

    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass

    httpd.server_close()

    print('Stopping httpd...\n')


if __name__ == '__main__':
    from sys import argv

    if len(argv) == 2:
        run(port=int(argv[1]))
    else:
        run()