feat: Add proxy support

Author: qingyanbaby

.env.example

@@ -4,4 +4,9 @@ MODEL_MAPPING='{
     "gpt-3.5-turbo": "GPT-3.5-Turbo",
     "gpt-4o": "GPT-4o",
     "gpt-4-turbo": "GPT-4-Turbo"
-}'
+}'
+PROXY_TYPE=socks
+PROXY_HOST=127.0.0.1
+PROXY_PORT=6668
+PROXY_USERNAME=
+PROXY_PASSWORD=

README.md

@@ -3,6 +3,11 @@
 
 This is a project that converts the official `poe.com` API to the OpenAI API. It only supports the `/v1/chat/completions` endpoint.
 
+## Changelog
+1.0.0  `/v1/chat/completions` endpoint support
+1.0.1  Environment variable custom model mapping
+1.1.0  Add proxy support
+
 ## Usage
 ### Running Locally
 
@@ -67,4 +72,14 @@ MODEL_MAPPING='{
     "gpt-4o": "GPT-4o",
     "gpt-4-turbo": "GPT-4-Turbo"
 }'
+```
+
+## Proxy Settings
+```shell
+# Please customize and edit in the .env file
+PROXY_TYPE=socks # socks/http, socks only supports socks5 proxy
+PROXY_HOST=127.0.0.1 # Proxy address
+PROXY_PORT=6668 # Proxy port
+PROXY_USERNAME= # Proxy username, optional
+PROXY_PASSWORD= # Proxy password, optional
 ```

README_ZH.md

@@ -2,6 +2,11 @@
 这是一个`poe.com`官方API转openai的API的项目
 仅支持接口`/v1/chat/completions`
 
+## 更新日志
+1.0.0  `/v1/chat/completions`接口支持
+1.0.1  环境变量自定义模型映射
+1.1.0  添加proxy支持
+
 ## 使用方式
 ### 本地运行
 
@@ -66,4 +71,14 @@ MODEL_MAPPING='{
     "gpt-4o": "GPT-4o",
     "gpt-4-turbo": "GPT-4-Turbo"
 }'
+```
+
+## 代理设置
+```shell
+# 请在 .env 文件自定义编辑
+PROXY_TYPE=socks # socks/http，socks仅支持socks5代理
+PROXY_HOST=127.0.0.1 # 代理地址
+PROXY_PORT=6668 # 代理端口
+PROXY_USERNAME= # 代理用户名，可选
+PROXY_PASSWORD= # 代理密码，可选
 ```

docker-compose.yml

@@ -12,4 +12,9 @@ services:
         "gpt-3.5-turbo": "GPT-3.5-Turbo",
         "gpt-4o": "GPT-4o",
         "gpt-4-turbo": "GPT-4-Turbo"
-      }'
+      }'
+      PROXY_TYPE: socks
+      PROXY_HOST:
+      PROXY_PORT:
+      PROXY_USERNAME:
+      PROXY_PASSWORD:

poe_api.py

@@ -2,7 +2,8 @@
 import logging
 import os
 
-from fastapi import WebSocket, Form
+import httpx
+from fastapi import Form
 from fastapi.responses import JSONResponse
 from fastapi_poe.client import get_bot_response, get_final_response, QueryRequest
 from fastapi_poe.types import ProtocolMessage
@@ -13,6 +14,7 @@
 
 client_dict = {}
 
+
 async def get_responses(api_key, prompt=[], bot="gpt-4"):
     bot_name = get_bot(bot)
     # "system", "user", "bot"
@@ -30,14 +32,17 @@ async def get_responses(api_key, prompt=[], bot="gpt-4"):
         **additional_params
     )
 
-    return await get_final_response(query, bot_name=bot_name, api_key=api_key)
+    session = create_client()
+    return await get_final_response(query, bot_name=bot_name, api_key=api_key, session=session)
 
 
 async def stream_get_responses(api_key, prompt, bot):
     bot_name = get_bot(bot)
     messages = openai_message_to_poe_message(prompt)
+
+    session = create_client()
     async for partial in get_bot_response(messages=messages, bot_name=bot_name, api_key=api_key,
-                                          skip_system_prompt=True):
+                                          skip_system_prompt=True, session=session):
         yield partial.text
 
 
@@ -69,24 +74,6 @@ async def ask(token: str = Form(...), bot: str = Form(...), content: str = Form(
         return JSONResponse(status_code=400, content={"message": errmsg})
 
 
-# @app.websocket("/stream")
-async def websocket_endpoint(websocket: WebSocket):
-    try:
-        await websocket.accept()
-        token = await websocket.receive_text()
-        bot = await websocket.receive_text()
-        content = await websocket.receive_text()
-        add_token(token)
-        async for ret in stream_get_responses(token, content, bot):
-            await websocket.send_text(ret)
-    except Exception as e:
-        errmsg = f"An exception of type {type(e).__name__} occurred. Arguments: {e.args}"
-        logging.info(errmsg)
-        await websocket.send_text(errmsg)
-    finally:
-        await websocket.close()
-
-
 def get_bot(model):
     model_mapping = json.loads(os.environ.get("MODEL_MAPPING", "{}"))
     return model_mapping.get(model, "GPT-4o")
@@ -100,3 +87,48 @@ def openai_message_to_poe_message(messages=[]):
         new_messages.append(ProtocolMessage(role=role, content=message["content"]))
 
     return new_messages
+
+
+def create_client():
+    proxy_config = {
+        "proxy_type": os.environ.get("PROXY_TYPE"),
+        "proxy_host": os.environ.get("PROXY_HOST"),
+        "proxy_port": os.environ.get("PROXY_PORT"),
+        "proxy_username": os.environ.get("PROXY_USERNAME"),
+        "proxy_password": os.environ.get("PROXY_PASSWORD"),
+    }
+
+    proxy = create_proxy(proxy_config)
+    client = httpx.AsyncClient(timeout=600, proxies=proxy)
+    return client
+
+
+def create_proxy(proxy_config):
+    proxy_type = proxy_config["proxy_type"]
+    proxy_url = create_proxy_url(proxy_config)
+
+    if proxy_type in ["http", "socks"] and proxy_url:
+        return {
+            "http://": proxy_url,
+            "https://": proxy_url,
+        }
+    else:
+        return None
+
+
+def create_proxy_url(proxy_config):
+    proxy_type = proxy_config["proxy_type"]
+    proxy_host = proxy_config["proxy_host"]
+    proxy_port = proxy_config["proxy_port"]
+    proxy_username = proxy_config["proxy_username"]
+    proxy_password = proxy_config["proxy_password"]
+
+    if not proxy_host or not proxy_port:
+        return None
+
+    if proxy_type == "http":
+        return f"http://{proxy_username}:{proxy_password}@{proxy_host}:{proxy_port}"
+    elif proxy_type == "socks":
+        return f"socks5://{proxy_username}:{proxy_password}@{proxy_host}:{proxy_port}"
+    else:
+        return None

requirements.txt

@@ -12,7 +12,8 @@ dataclasses-json==0.6.4
 exceptiongroup==1.2.0
 fastapi==0.109.2
 fastapi-cache2==0.2.1
-fastapi_poe==0.0.42
+fastapi_poe==0.0.46
+fp==0.2
 frozenlist==1.4.1
 greenlet==3.0.3
 gunicorn==21.2.0
@@ -43,6 +44,7 @@ regex==2023.12.25
 requests==2.31.0
 six==1.16.0
 sniffio==1.3.0
+socksio==1.0.0
 SQLAlchemy==2.0.27
 sse-starlette==2.1.0
 starlette==0.36.3