fix bug of proxy

README.md

@@ -42,8 +42,8 @@ from openai_api_call import proxy_on, proxy_off, proxy_status
 # Check the current proxy
 proxy_status()
 
-# Set local proxy, port number is 7890 by default
-proxy_on("127.0.0.1", port=7890)
+# Set proxy(example)
+proxy_on(http="127.0.0.1:7890", https="socks://127.0.0.1:7891")
 
 # Check the updated proxy
 proxy_status()

README_zh_CN.md

@@ -40,8 +40,8 @@ from openai_api_call import proxy_on, proxy_off, proxy_status
 # 查看当前代理
 proxy_status()
 
-# 设置本地代理，端口号默认为 7890
-proxy_on("127.0.0.1", port=7890)
+# 设置代理，这里 IP 127.0.0.1 代表本机
+proxy_on(http="127.0.0.1:7890", https="socks://127.0.0.1:7891")
 
 # 查看更新后的代理
 proxy_status()

openai_api_call/proxy.py

@@ -1,15 +1,17 @@
 import os
+from typing import Union
 
-def proxy_on(host:str, port:int=7890):
+def proxy_on(http:Union[str, None]=None, https:Union[str, None]=None):
     """Set proxy for the API call
 
     Args:
-        host (str): proxy host
-        port (int, optional): proxy port. Defaults to 7890.
+        http (str, optional): http proxy. Defaults to None.
+        https (str, optional): https proxy. Defaults to None.
     """
-    host = host.replace("http://", "").replace("https://", "")
-    os.environ['http_proxy'] = f"http://{host}:{port}"
-    os.environ['https_proxy'] = f"https://{host}:{port}"
+    if http is not None:
+        os.environ['http_proxy'] = http
+    if https is not None:
+        os.environ['https_proxy'] = https
 
 def proxy_off():
     """Turn off proxy for the API call"""
@@ -29,3 +31,10 @@ def proxy_status():
         print("`https_proxy` is not set!")
     else:
         print(f"https_proxy:\t{https}")
+
+def proxy_test(url:str="www.facebook.com"):
+    url = url.replace("http://", "").replace("https://", "")
+    if os.system("curl -I https://"+url) != 0:
+        print("Https: Curl to "+url+" failed!")
+    if os.system("curl -I http://"+url) != 0:
+        print("Http: Curl to "+url+" failed!")

tests/test_proxy.py

@@ -5,9 +5,9 @@ def test_proxy():
     proxy_off()
     assert os.environ.get('http_proxy') is None
     assert os.environ.get('https_proxy') is None
-    proxy_on("192.168.0.1")
-    assert os.environ.get('http_proxy') == "http://192.168.0.1:7890"
-    assert os.environ.get('https_proxy') == "https://192.168.0.1:7890"
+    proxy_on(http="127.0.0.1:7890", https="socks://127.0.0.1:7891")
+    assert os.environ.get('http_proxy') == "127.0.0.1:7890"
+    assert os.environ.get('https_proxy') == "socks://127.0.0.1:7891"
     proxy_off()
     assert os.environ.get('http_proxy') is None
     assert os.environ.get('https_proxy') is None