add ipv6 support to the API

extensions/openai/script.py

@@ -353,22 +353,44 @@ async def handle_unload_loras():
 
 
 def run_server():
-    server_addr = '0.0.0.0' if shared.args.listen else '127.0.0.1'
+    server_addrV6 = '[::]' if shared.args.listen else '[::1]'
+    server_addrV4 = '0.0.0.0' if shared.args.listen else '127.0.0.1'
+
     port = int(os.environ.get('OPENEDAI_PORT', shared.args.api_port))
+    server_addrs = []
+
+    disable_ipv6 = os.environ.get('OPENEDAI_DISABLE_IPV6', shared.args.api_disable_ipv6)
+    if not disable_ipv6:
+        server_addrs.append(server_addrV6)
+
+    disable_ipv4 = os.environ.get('OPENEDAI_DISABLE_IPV4', shared.args.api_disable_ipv4)
+    if not disable_ipv4:
+        server_addrs.append(server_addrV4)
+
+    if len(server_addrs) < 1:
+        logger.error('you MUST enable IPv6 or IPv4 for the API to work')
+        raise Exception('you MUST enable IPv6 or IPv4 for the API to work')
 
     ssl_certfile = os.environ.get('OPENEDAI_CERT_PATH', shared.args.ssl_certfile)
     ssl_keyfile = os.environ.get('OPENEDAI_KEY_PATH', shared.args.ssl_keyfile)
 
+    url_proto = 'http://'
+    if ssl_keyfile and ssl_certfile:
+        url_proto = 'https://'
+
     if shared.args.public_api:
         def on_start(public_url: str):
             logger.info(f'OpenAI-compatible API URL:\n\n{public_url}\n')
 
         _start_cloudflared(port, shared.args.public_api_id, max_attempts=3, on_start=on_start)
     else:
-        if ssl_keyfile and ssl_certfile:
-            logger.info(f'OpenAI-compatible API URL:\n\nhttps://{server_addr}:{port}\n')
-        else:
-            logger.info(f'OpenAI-compatible API URL:\n\nhttp://{server_addr}:{port}\n')
+        logger.info('\n\nRunning OpenAI-compatible API on:')
+        for i, addr in enumerate(server_addrs):
+            nl = ''
+            if i == len(server_addrs) - 1:
+                nl = '\n\n'
+
+            logger.info(f'{url_proto}{addr}:{port}{nl}')
 
     if shared.args.api_key:
         if not shared.args.admin_key:
@@ -380,7 +402,7 @@ def on_start(public_url: str):
         logger.info(f'OpenAI API admin key (for loading/unloading models):\n\n{shared.args.admin_key}\n')
 
     logging.getLogger("uvicorn.error").propagate = False
-    uvicorn.run(app, host=server_addr, port=port, ssl_certfile=ssl_certfile, ssl_keyfile=ssl_keyfile)
+    uvicorn.run(app, host=server_addrs, port=port, ssl_certfile=ssl_certfile, ssl_keyfile=ssl_keyfile)
 
 
 def setup():

modules/shared.py

@@ -195,6 +195,8 @@
 # API
 group = parser.add_argument_group('API')
 group.add_argument('--api', action='store_true', help='Enable the API extension.')
+group.add_argument('--api-disable-ipv6', action='store_true', help='Disable IPv6 for the API')
+group.add_argument('--api-disable-ipv4', action='store_true', help='Disable IPv4 for the API')
 group.add_argument('--public-api', action='store_true', help='Create a public URL for the API using Cloudfare.')
 group.add_argument('--public-api-id', type=str, help='Tunnel ID for named Cloudflare Tunnel. Use together with public-api option.', default=None)
 group.add_argument('--api-port', type=int, default=5000, help='The listening port for the API.')