From 9faf7b0464b5bfd7fd2d3c255b6fc770e0c2da2b Mon Sep 17 00:00:00 2001 From: Harsh23Kashyap <55448981+Harsh23Kashyap@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:55:19 +0530 Subject: [PATCH 1/2] fix(tianditu): add request timeouts and credential guards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the same three-bug reliability pattern from PR #3456 (tools/google), and the matching PRs #3466 (tools/searchapi) and #3468 (tools/serper): - requests.get had no timeout= argument in any of the three tool files (geocoder.py, staticmap.py, poisearch.py). Stalled hosts blocked until the framework-level 120s ceiling. Add timeout=10 to every call site. - self.runtime.credentials["tianditu_api_key"] raised KeyError when the key was absent or empty. Replace with an early guard that yields a friendly error message and returns; use .get() for the value lookup. - Wrap the HTTP calls in try/except requests.RequestException so a timeout, connection reset, DNS failure, or 5xx surfaces as a friendly message instead of a stack trace. tianditu does not take user-supplied hl/gl (the geocoder API uses a keyWord and the user's region comes from upstream), so the invoke fall-through bug from PR #3456 does not apply. No public API change; successful 200-path responses are byte-identical. Adds tools/tianditu/tests/test_tianditu.py with 11 in-process test cases covering all three tool files. Bumps manifest 0.0.7 -> 0.0.8. Refactors nothing else; the staticmap.PoiSearchTool class name is kept as-is to avoid unrelated churn. (It is mis-named — the file is the static-map image tool, not POI search — but renaming is out of scope.) Refs #3465 Refs #3456 --- tools/tianditu/manifest.yaml | 2 +- tools/tianditu/tests/__init__.py | 0 tools/tianditu/tests/test_tianditu.py | 282 ++++++++++++++++++++++++++ tools/tianditu/tools/geocoder.py | 20 +- tools/tianditu/tools/poisearch.py | 63 +++--- tools/tianditu/tools/staticmap.py | 57 ++++-- 6 files changed, 372 insertions(+), 52 deletions(-) create mode 100644 tools/tianditu/tests/__init__.py create mode 100644 tools/tianditu/tests/test_tianditu.py diff --git a/tools/tianditu/manifest.yaml b/tools/tianditu/manifest.yaml index 903db0308..78f781699 100644 --- a/tools/tianditu/manifest.yaml +++ b/tools/tianditu/manifest.yaml @@ -36,4 +36,4 @@ tags: - utilities - travel type: plugin -version: 0.0.7 +version: 0.0.8 diff --git a/tools/tianditu/tests/__init__.py b/tools/tianditu/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tools/tianditu/tests/test_tianditu.py b/tools/tianditu/tests/test_tianditu.py new file mode 100644 index 000000000..a24b195a1 --- /dev/null +++ b/tools/tianditu/tests/test_tianditu.py @@ -0,0 +1,282 @@ +"""In-process tests for tools/tianditu (no pytest required). + +These tests cover two reliability fixes shipped in this PR: + +- Bug 1: missing `timeout=` on `requests.get` calls. +- Bug 2: `KeyError` when `tianditu_api_key` is missing from credentials. + +`tianditu` does not take user-supplied `hl` / `gl` (the geocoder API takes a +`keyWord` and an optional `region` already from upstream, not through us), +so the third bug from PR #3456 (invoke fall-through on invalid hl/gl) +does not apply. + +They use only Python stdlib (`unittest.mock`) and tiny in-test stubs for +`requests`, `dify_plugin`, and `pydantic` so they can run in environments +where the plugin dependencies are not pre-installed. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +_HERE = Path(__file__).resolve().parent # tests/ +_PLUGIN_ROOT = _HERE.parent # tools/tianditu/ +sys.path.insert(0, str(_PLUGIN_ROOT)) # so `tools.geocoder` etc. resolve + + +# ---- Stub modules ---------------------------------------------------------- + +def _ensure_stub_modules() -> None: + if "requests" not in sys.modules: + requests_stub = types.ModuleType("requests") + requests_stub.exceptions = types.SimpleNamespace( + RequestException=type("RequestException", (Exception,), {}) + ) + + def _default_get(*args, **kwargs): + raise AssertionError( + "requests.get was called in a test that did not patch it" + ) + + requests_stub.get = _default_get + sys.modules["requests"] = requests_stub + + if "dify_plugin" not in sys.modules: + dify_plugin_stub = types.ModuleType("dify_plugin") + + class _BaseTool: + def __init__(self): + self.runtime = SimpleNamespace(credentials={}) + + def create_text_message(self, text): + return ("text", text) + + def create_json_message(self, json): + return ("json", json) + + def create_blob_message(self, blob, meta=None): + return ("blob", blob, meta) + + dify_plugin_stub.Tool = _BaseTool + + dify_plugin_entities_stub = types.ModuleType("dify_plugin.entities") + tool_module_stub = types.ModuleType("dify_plugin.entities.tool") + + class _ToolInvokeMessage: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + tool_module_stub.ToolInvokeMessage = _ToolInvokeMessage + + dify_plugin_entities_stub.tool = tool_module_stub + sys.modules["dify_plugin"] = dify_plugin_stub + sys.modules["dify_plugin.entities"] = dify_plugin_entities_stub + sys.modules["dify_plugin.entities.tool"] = tool_module_stub + + +_ensure_stub_modules() + +# Force-fresh imports so we get the patched stubs above. +for _name in list(sys.modules): + if _name.startswith("tools.tianditu") or _name.startswith("tools.geocoder") \ + or _name.startswith("tools.staticmap") or _name.startswith("tools.poisearch"): + sys.modules.pop(_name, None) + +from importlib import util as _importlib_util # noqa: E402 + + +def _load_module(file_name: str): + path = _PLUGIN_ROOT / "tools" / file_name + spec = _importlib_util.spec_from_file_location(f"tianditu.{file_name[:-3]}", path) + mod = _importlib_util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +geocoder = _load_module("geocoder.py") +staticmap = _load_module("staticmap.py") +poisearch = _load_module("poisearch.py") + + +# ---- Helpers --------------------------------------------------------------- + +class _FakeResponse: + def __init__(self, payload=None, *, raw_bytes=b"\x89PNG\r\n"): + if payload is not None: + self._payload = payload + else: + self._payload = { + "status": "0", + "message": "ok", + "location": {"lon": 116.404, "lat": 39.915}, + } + self._raw_bytes = raw_bytes + self.calls = 0 # some files make multiple requests.get calls + + def json(self): + self.calls += 1 + return self._payload + + @property + def content(self): + return self._raw_bytes + + +def _make_tool(tool_cls, credentials): + tool = object.__new__(tool_cls) + + def _text(text): + return ("text", text) + + def _json(json=None, **_kw): + return ("json", json) + + def _blob(blob, meta=None): + return ("blob", blob, meta) + + tool.create_text_message = _text + tool.create_json_message = _json + tool.create_blob_message = _blob + tool.runtime = SimpleNamespace(credentials=credentials) + return tool + + +def _flatten(messages): + return list(messages) + + +# ============================================================================= +# Bug 2: missing-credential guard +# ============================================================================= + +def test_geocoder_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(geocoder.GeocoderTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert len(messages) == 1 + assert messages[0][0] == "text" + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_geocoder_empty_api_key_returns_message_and_no_http_call(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": ""}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_staticmap_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(staticmap.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +def test_poisearch_missing_api_key_returns_message_and_no_http_call(): + tool = _make_tool(poisearch.PoiSearchTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(poisearch.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing", "baseAddress": "Beijing"})) + assert messages[0][1] == "Tianditu API key is required." + assert fake_get.call_count == 0 + + +# ============================================================================= +# Bug 2 corollary: missing key did NOT raise KeyError pre-fix. +# ============================================================================= + +def test_geocoder_missing_api_key_does_not_raise_keyerror(): + tool = _make_tool(geocoder.GeocoderTool, credentials={}) + fake_get = mock.Mock(side_effect=AssertionError("requests.get must not be called")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert messages[0] == ("text", "Tianditu API key is required.") + + +# ============================================================================= +# Bug 1: requests.get forwards a 10s timeout on every call. +# ============================================================================= + +def test_geocoder_requests_get_uses_10s_timeout(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse(payload={"status": "0", "message": "ok"}) + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing"})) + assert fake_get.call_count == 1 + _, kwargs = fake_get.call_args + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +def test_staticmap_requests_get_uses_10s_timeout_on_every_call(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(staticmap.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing"})) + # staticmap makes 2 calls: geocoder look-up + static image fetch. + assert fake_get.call_count == 2 + for call in fake_get.call_args_list: + _, kwargs = call + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +def test_poisearch_requests_get_uses_10s_timeout_on_every_call(): + tool = _make_tool(poisearch.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(poisearch.requests, "get", fake_get): + _flatten(tool._invoke({"keyword": "Beijing", "baseAddress": "Beijing"})) + # poisearch makes 2 calls: geocoder look-up + v2/search + assert fake_get.call_count == 2 + for call in fake_get.call_args_list: + _, kwargs = call + assert kwargs.get("timeout") == 10, f"expected timeout=10 kwarg, got {kwargs!r}" + + +# ============================================================================= +# Happy-path smoke +# ============================================================================= + +def test_geocoder_happy_path_yields_json_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "json" for m in messages) + + +def test_staticmap_happy_path_yields_blob_message(): + tool = _make_tool(staticmap.PoiSearchTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse() + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(staticmap.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "blob" for m in messages) + + +# ============================================================================= +# Network failure surfaces as a friendly message, not a stack trace. +# ============================================================================= + +def test_geocoder_request_exception_yields_friendly_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_get = mock.Mock(side_effect=geocoder.requests.exceptions.RequestException("boom")) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any( + m[0] == "text" and "boom" in m[1] for m in messages + ), f"expected friendly error message containing 'boom', got {messages!r}" diff --git a/tools/tianditu/tools/geocoder.py b/tools/tianditu/tools/geocoder.py index fc07e0e52..27223c624 100644 --- a/tools/tianditu/tools/geocoder.py +++ b/tools/tianditu/tools/geocoder.py @@ -17,9 +17,21 @@ def _invoke( if not keyword: yield self.create_text_message("Invalid parameter keyword") return - tk = self.runtime.credentials["tianditu_api_key"] + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + params = {"keyWord": keyword} - result = requests.get( - base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk - ).json() + try: + result = requests.get( + base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk, + timeout=10, + ).json() + except requests.exceptions.RequestException as exc: + yield self.create_text_message( + f"An error occurred while invoking the tool: {exc}." + ) + return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/poisearch.py b/tools/tianditu/tools/poisearch.py index efd08aa7f..0d0d795ee 100644 --- a/tools/tianditu/tools/poisearch.py +++ b/tools/tianditu/tools/poisearch.py @@ -22,29 +22,42 @@ def _invoke( if not baseAddress: yield self.create_text_message("Invalid parameter baseAddress") return - tk = self.runtime.credentials["tianditu_api_key"] - base_coords = requests.get( - geocoder_base_url - + "?ds=" - + json.dumps({"keyWord": baseAddress}, ensure_ascii=False) - + "&tk=" - + tk - ).json() - params = { - "keyWord": keyword, - "queryRadius": 5000, - "queryType": 3, - "pointLonlat": str(base_coords["location"]["lon"]) - + "," - + str(base_coords["location"]["lat"]), - "start": 0, - "count": 100, - } - result = requests.get( - base_url - + "?postStr=" - + json.dumps(params, ensure_ascii=False) - + "&type=query&tk=" - + tk - ).json() + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + + try: + base_coords = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps({"keyWord": baseAddress}, ensure_ascii=False) + + "&tk=" + + tk, + timeout=10, + ).json() + params = { + "keyWord": keyword, + "queryRadius": 5000, + "queryType": 3, + "pointLonlat": str(base_coords["location"]["lon"]) + + "," + + str(base_coords["location"]["lat"]), + "start": 0, + "count": 100, + } + result = requests.get( + base_url + + "?postStr=" + + json.dumps(params, ensure_ascii=False) + + "&type=query&tk=" + + tk, + timeout=10, + ).json() + except requests.exceptions.RequestException as exc: + yield self.create_text_message( + f"An error occurred while invoking the tool: {exc}." + ) + return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/staticmap.py b/tools/tianditu/tools/staticmap.py index 00a1e83f6..7715c926c 100644 --- a/tools/tianditu/tools/staticmap.py +++ b/tools/tianditu/tools/staticmap.py @@ -18,28 +18,41 @@ def _invoke( if not keyword: yield self.create_text_message("Invalid parameter keyword") return - tk = self.runtime.credentials["tianditu_api_key"] - keyword_coords = requests.get( - geocoder_base_url - + "?ds=" - + json.dumps({"keyWord": keyword}, ensure_ascii=False) - + "&tk=" - + tk - ).json() - coords = ( - str(keyword_coords["location"]["lon"]) - + "," - + str(keyword_coords["location"]["lat"]) - ) - result = requests.get( - base_url - + "?center=" - + coords - + "&markers=" - + coords - + "&width=400&height=300&zoom=14&tk=" - + tk - ).content + + tk = self.runtime.credentials.get("tianditu_api_key") + if not tk: + yield self.create_text_message("Tianditu API key is required.") + return + + try: + keyword_coords = requests.get( + geocoder_base_url + + "?ds=" + + json.dumps({"keyWord": keyword}, ensure_ascii=False) + + "&tk=" + + tk, + timeout=10, + ).json() + coords = ( + str(keyword_coords["location"]["lon"]) + + "," + + str(keyword_coords["location"]["lat"]) + ) + result = requests.get( + base_url + + "?center=" + + coords + + "&markers=" + + coords + + "&width=400&height=300&zoom=14&tk=" + + tk, + timeout=10, + ).content + except requests.exceptions.RequestException as exc: + yield self.create_text_message( + f"An error occurred while invoking the tool: {exc}." + ) + return yield self.create_blob_message( blob=result, meta={"mime_type": "image/png"}, From 85d109fe3e9ad71d553be56c7615424c34d3ae6d Mon Sep 17 00:00:00 2001 From: Harsh23Kashyap <55448981+Harsh23Kashyap@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:38:11 +0530 Subject: [PATCH 2/2] fix(tianditu): call raise_for_status and stop leaking tk in error messages Copilot review flagged two issues on PR #3469: - requests.get skipped raise_for_status(), so 4xx/5xx were treated as success. - str(exc) in the yielded message can include the request URL with tk. Add raise_for_status() before parsing each response. Use a generic network/upstream failure message that never forwards URLs or credentials. Refs #3465 --- tools/tianditu/tests/test_tianditu.py | 40 +++++++++++++++++++++++++-- tools/tianditu/tools/geocoder.py | 15 ++++++---- tools/tianditu/tools/poisearch.py | 21 +++++++++----- tools/tianditu/tools/staticmap.py | 21 +++++++++----- 4 files changed, 75 insertions(+), 22 deletions(-) diff --git a/tools/tianditu/tests/test_tianditu.py b/tools/tianditu/tests/test_tianditu.py index a24b195a1..57ae77feb 100644 --- a/tools/tianditu/tests/test_tianditu.py +++ b/tools/tianditu/tests/test_tianditu.py @@ -107,7 +107,7 @@ def _load_module(file_name: str): # ---- Helpers --------------------------------------------------------------- class _FakeResponse: - def __init__(self, payload=None, *, raw_bytes=b"\x89PNG\r\n"): + def __init__(self, payload=None, *, raw_bytes=b"\x89PNG\r\n", status_code=200): if payload is not None: self._payload = payload else: @@ -117,7 +117,14 @@ def __init__(self, payload=None, *, raw_bytes=b"\x89PNG\r\n"): "location": {"lon": 116.404, "lat": 39.915}, } self._raw_bytes = raw_bytes - self.calls = 0 # some files make multiple requests.get calls + self.status_code = status_code + self.calls = 0 + + def raise_for_status(self): + if self.status_code >= 400: + raise geocoder.requests.exceptions.RequestException( + f"HTTP {self.status_code}" + ) def json(self): self.calls += 1 @@ -278,5 +285,32 @@ def test_geocoder_request_exception_yields_friendly_message(): with mock.patch.object(geocoder.requests, "get", fake_get): messages = _flatten(tool._invoke({"keyword": "Beijing"})) assert any( + m[0] == "text" and "network or upstream failure" in m[1] for m in messages + ), f"expected generic friendly error, got {messages!r}" + assert not any( m[0] == "text" and "boom" in m[1] for m in messages - ), f"expected friendly error message containing 'boom', got {messages!r}" + ), "exception text must not be forwarded to the user" + + +def test_geocoder_request_does_not_leak_tk_in_message(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret-key"}) + fake_get = mock.Mock( + side_effect=geocoder.requests.exceptions.RequestException( + "GET http://api.tianditu.gov.cn/geocoder?ds=%7B%22keyWord%22%3A%22Beijing%22%7D&tk=secret-key" + ) + ) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + for m in messages: + if m[0] == "text": + assert "secret-key" not in m[1], f"API key leaked in message: {m[1]!r}" + + +def test_geocoder_http_4xx_returns_message_and_no_json_dump(): + tool = _make_tool(geocoder.GeocoderTool, credentials={"tianditu_api_key": "secret"}) + fake_response = _FakeResponse(status_code=500) + fake_get = mock.Mock(return_value=fake_response) + with mock.patch.object(geocoder.requests, "get", fake_get): + messages = _flatten(tool._invoke({"keyword": "Beijing"})) + assert any(m[0] == "text" and "network or upstream failure" in m[1] for m in messages) + assert not any(m[0] == "json" for m in messages) diff --git a/tools/tianditu/tools/geocoder.py b/tools/tianditu/tools/geocoder.py index 27223c624..a11b1d4b7 100644 --- a/tools/tianditu/tools/geocoder.py +++ b/tools/tianditu/tools/geocoder.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class GeocoderTool(Tool): def _invoke( @@ -25,13 +30,13 @@ def _invoke( params = {"keyWord": keyword} try: - result = requests.get( + response = requests.get( base_url + "?ds=" + json.dumps(params, ensure_ascii=False) + "&tk=" + tk, timeout=10, - ).json() - except requests.exceptions.RequestException as exc: - yield self.create_text_message( - f"An error occurred while invoking the tool: {exc}." ) + response.raise_for_status() + result = response.json() + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/poisearch.py b/tools/tianditu/tools/poisearch.py index 0d0d795ee..325231af0 100644 --- a/tools/tianditu/tools/poisearch.py +++ b/tools/tianditu/tools/poisearch.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class PoiSearchTool(Tool): def _invoke( @@ -29,14 +34,16 @@ def _invoke( return try: - base_coords = requests.get( + geo_response = requests.get( geocoder_base_url + "?ds=" + json.dumps({"keyWord": baseAddress}, ensure_ascii=False) + "&tk=" + tk, timeout=10, - ).json() + ) + geo_response.raise_for_status() + base_coords = geo_response.json() params = { "keyWord": keyword, "queryRadius": 5000, @@ -47,17 +54,17 @@ def _invoke( "start": 0, "count": 100, } - result = requests.get( + search_response = requests.get( base_url + "?postStr=" + json.dumps(params, ensure_ascii=False) + "&type=query&tk=" + tk, timeout=10, - ).json() - except requests.exceptions.RequestException as exc: - yield self.create_text_message( - f"An error occurred while invoking the tool: {exc}." ) + search_response.raise_for_status() + result = search_response.json() + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) return yield self.create_json_message(result) diff --git a/tools/tianditu/tools/staticmap.py b/tools/tianditu/tools/staticmap.py index 7715c926c..af1844866 100644 --- a/tools/tianditu/tools/staticmap.py +++ b/tools/tianditu/tools/staticmap.py @@ -4,6 +4,11 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin import Tool +_NETWORK_ERROR = ( + "An error occurred while invoking the tool (network or upstream failure). " + "Please retry shortly." +) + class PoiSearchTool(Tool): def _invoke( @@ -25,20 +30,22 @@ def _invoke( return try: - keyword_coords = requests.get( + geo_response = requests.get( geocoder_base_url + "?ds=" + json.dumps({"keyWord": keyword}, ensure_ascii=False) + "&tk=" + tk, timeout=10, - ).json() + ) + geo_response.raise_for_status() + keyword_coords = geo_response.json() coords = ( str(keyword_coords["location"]["lon"]) + "," + str(keyword_coords["location"]["lat"]) ) - result = requests.get( + img_response = requests.get( base_url + "?center=" + coords @@ -47,11 +54,11 @@ def _invoke( + "&width=400&height=300&zoom=14&tk=" + tk, timeout=10, - ).content - except requests.exceptions.RequestException as exc: - yield self.create_text_message( - f"An error occurred while invoking the tool: {exc}." ) + img_response.raise_for_status() + result = img_response.content + except requests.exceptions.RequestException: + yield self.create_text_message(_NETWORK_ERROR) return yield self.create_blob_message( blob=result,