From e33994657bae9e6cb33c8e246481fb0095a31904 Mon Sep 17 00:00:00 2001
From: cavities <cavities@sina.com>
Date: Thu, 7 Mar 2024 10:55:28 +0800
Subject: [PATCH] ".. add proxies kwargs in duckduckgo tool run func .."

---
 promptulate/tools/duckduckgo/api_wrapper.py | 6 +++---
 promptulate/tools/duckduckgo/tools.py       | 7 ++++++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/promptulate/tools/duckduckgo/api_wrapper.py b/promptulate/tools/duckduckgo/api_wrapper.py
index 905e7350..4c8a9636 100644
--- a/promptulate/tools/duckduckgo/api_wrapper.py
+++ b/promptulate/tools/duckduckgo/api_wrapper.py
@@ -10,7 +10,7 @@ class DuckDuckGoSearchAPIWrapper(BaseModel):
     safe_search: str = "moderate"
     time: Optional[str] = "y"
     max_num_of_results: int = 5
-
+    proxies: Optional[str] = None
     class Config:
         """Configuration for this pydantic object."""
 
@@ -37,7 +37,7 @@ def query(
         if not num_results:
             num_results = self.max_num_of_results
 
-        with DDGS() as ddgs:
+        with DDGS(self.proxies) as ddgs:
             results = ddgs.text(
                 keyword,
                 region=self.region,
@@ -74,7 +74,7 @@ def query_by_formatted_results(
         if not num_results:
             num_results = self.max_num_of_results
 
-        with DDGS() as ddgs:
+        with DDGS(self.proxies) as ddgs:
             results = ddgs.text(
                 query,
                 region=self.region,
diff --git a/promptulate/tools/duckduckgo/tools.py b/promptulate/tools/duckduckgo/tools.py
index 7c75061f..6ebc24f2 100644
--- a/promptulate/tools/duckduckgo/tools.py
+++ b/promptulate/tools/duckduckgo/tools.py
@@ -34,7 +34,9 @@ def _run(self, keyword: str, **kwargs) -> Union[str, List[str]]:
             result_type="original"
         """
         from duckduckgo_search.exceptions import RateLimitException
-
+        if "proxies" in kwargs:
+            if kwargs["proxies"]:
+                self.api_wrapper.proxies = kwargs["proxies"]
         attempt = 0
         while attempt < self.max_retry:
             try:
@@ -80,6 +82,9 @@ def _run(self, keyword: str, **kwargs) -> Union[str, List[Dict[str, str]]]:
             default is string. Return List[Dict[str, str]] type data if you pass
             result_type="original"
         """
+        if "proxies" in kwargs:
+            if kwargs["proxies"]:
+                self.api_wrapper.proxies = kwargs["proxies"]
         result = self.api_wrapper.query_by_formatted_results(keyword, **kwargs)
         if "return_type" in kwargs and kwargs["return_type"] == "original":
             return result