feat: Add HTTP proxy support for OpenAI API and CLI options (#174)

- Add support for HTTP proxy configuration

Co-authored-by: Jovy <[email protected]>

README.md

@@ -130,6 +130,7 @@ python3 xiaogpt.py
 | start_conversation | 开始持续对话关键词                            | `开始持续对话`         |
 | end_conversation   | 结束持续对话关键词                            | `结束持续对话`         |
 | stream             | 使用流式响应，获得更快的响应                  | `false`                |
+| proxy              | 支持 HTTP 代理，传入 http proxy URL         | ""                     |
 | gpt_options        | OpenAI API 的参数字典                         | `{}`                   |
 
 ## 注意

xiao_config.json.example

@@ -15,5 +15,6 @@
     "start_conversation": "开始持续对话",
     "end_conversation": "结束持续对话",
     "stream": false,
+    "proxy": "",
     "gpt_options": {}
 }

xiaogpt/bot/chatgptapi_bot.py

@@ -6,11 +6,13 @@
 
 
 class ChatGPTBot(BaseBot):
-    def __init__(self, openai_key, api_base=None):
+    def __init__(self, openai_key, api_base=None, proxy=None):
         self.history = []
         openai.api_key = openai_key
         if api_base:
             openai.api_base = api_base
+        if proxy:
+            openai.proxy = proxy
 
     async def ask(self, query, **options):
         ms = []

xiaogpt/bot/gpt3_bot.py

@@ -5,10 +5,12 @@
 
 
 class GPT3Bot:
-    def __init__(self, openai_key, api_base=None):
+    def __init__(self, openai_key, api_base=None, proxy=None):
         openai.api_key = openai_key
         if api_base:
             openai.api_base = api_base
+        if proxy:
+            openai.proxy = proxy
         self.history = []
 
     async def ask(self, query, **options):

xiaogpt/cli.py

@@ -27,6 +27,11 @@ def main():
         dest="openai_key",
         help="openai api key",
     )
+    parser.add_argument(
+        "--proxy",
+        dest="proxy",
+        help="http proxy url like http://localhost:8080",
+    )
     parser.add_argument(
         "--cookie",
         dest="cookie",

xiaogpt/config.py

@@ -6,6 +6,8 @@
 from dataclasses import dataclass, field
 from typing import Iterable, Any
 
+from xiaogpt.utils import validate_proxy
+
 LATEST_ASK_API = "https://userprofile.mina.mi.com/device_profile/v2/conversation?source=dialogu&hardware={hardware}&timestamp={timestamp}&limit=2"
 COOKIE_TEMPLATE = "deviceId={device_id}; serviceToken={service_token}; userId={user_id}"
 WAKEUP_KEYWORD = "小爱同学"
@@ -55,6 +57,7 @@ class Config:
     account: str = os.getenv("MI_USER", "")
     password: str = os.getenv("MI_PASS", "")
     openai_key: str = os.getenv("OPENAI_API_KEY", "")
+    proxy: str | None = None
     mi_did: str = os.getenv("MI_DID", "")
     keyword: Iterable[str] = KEY_WORD
     change_prompt_keyword: Iterable[str] = CHANGE_PROMPT_KEY_WORD
@@ -105,4 +108,6 @@ def read_from_config(self, config_path: str) -> None:
                         key, value = "bot", "chatgptapi"
                     elif key == "use_gpt3":
                         key, value = "bot", "gpt3"
+                    elif key == "proxy":
+                        validate_proxy(value)
                     setattr(self, key, value)

xiaogpt/utils.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import re
 from http.cookies import SimpleCookie
+from urllib.parse import urlparse
 
 from requests.utils import cookiejar_from_dict
 
@@ -46,3 +47,15 @@ def find_key_by_partial_string(dictionary, partial_key):
     for key, value in dictionary.items():
         if key in partial_key:
             return value
+
+
+def validate_proxy(proxy_str: str) -> bool:
+    """Do a simple validation of the http proxy string."""
+
+    parsed = urlparse(proxy_str)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError("Proxy scheme must be http or https")
+    if not (parsed.hostname and parsed.port):
+        raise ValueError("Proxy hostname and port must be set")
+
+    return True