From b1d36c9aee722b5fae348859892ba229fdc3f573 Mon Sep 17 00:00:00 2001
From: duhan <duhan@189.cn>
Date: Sat, 20 Apr 2024 16:07:29 +0800
Subject: [PATCH] add proxy

---
 replicate/client.py |  6 ++++--
 replicate/proxy.py  | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)
 create mode 100644 replicate/proxy.py

diff --git a/replicate/client.py b/replicate/client.py
index 2441a4c7..8a9d033c 100644
--- a/replicate/client.py
+++ b/replicate/client.py
@@ -30,6 +30,7 @@
 from replicate.run import async_run, run
 from replicate.stream import async_stream, stream
 from replicate.training import Trainings
+from replicate.proxy import get_proxy
 
 if TYPE_CHECKING:
     from replicate.stream import ServerSentEvent
@@ -347,10 +348,11 @@ def _build_httpx_client(
         5.0, read=30.0, write=30.0, connect=5.0, pool=10.0
     )
 
+    proxy = get_proxy()
     transport = kwargs.pop("transport", None) or (
-        httpx.HTTPTransport()
+        httpx.HTTPTransport(proxy=proxy)
         if client_type is httpx.Client
-        else httpx.AsyncHTTPTransport()
+        else httpx.AsyncHTTPTransport(proxy=proxy)
     )
 
     return client_type(
diff --git a/replicate/proxy.py b/replicate/proxy.py
new file mode 100644
index 00000000..bdd2d871
--- /dev/null
+++ b/replicate/proxy.py
@@ -0,0 +1,14 @@
+import os
+import httpx
+
+def get_proxy():
+
+    http_proxy = os.environ.get('HTTP_PROXY') or os.environ.get('http_proxy')
+    if http_proxy:
+        return httpx.Proxy(url=http_proxy)
+
+    https_proxy = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
+    if https_proxy:
+        return httpx.Proxy(url=https_proxy)
+
+    return None