diff --git a/README.md b/README.md index 8a49774..d228894 100644 --- a/README.md +++ b/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() diff --git a/README_zh_CN.md b/README_zh_CN.md index 943d352..4395a51 100644 --- a/README_zh_CN.md +++ b/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() diff --git a/openai_api_call/proxy.py b/openai_api_call/proxy.py index e99fcde..658fcac 100644 --- a/openai_api_call/proxy.py +++ b/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!") \ No newline at end of file diff --git a/tests/test_proxy.py b/tests/test_proxy.py index 438b9c0..7318c66 100644 --- a/tests/test_proxy.py +++ b/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