From 3ff52f1809d62f056b2c929c0ffbc2b8b1b67ea7 Mon Sep 17 00:00:00 2001
From: Guangdong Liu <liugddx@gmail.com>
Date: Thu, 7 Aug 2025 09:04:51 +0800
Subject: [PATCH] feat: Enhance response validation and parsing in tool.py
 (#23456)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 api/core/tools/custom_tool/tool.py | 65 +++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 14 deletions(-)

diff --git a/api/core/tools/custom_tool/tool.py b/api/core/tools/custom_tool/tool.py
index 333ef2834..e112de957 100644
--- a/api/core/tools/custom_tool/tool.py
+++ b/api/core/tools/custom_tool/tool.py
@@ -1,7 +1,8 @@
 import json
 from collections.abc import Generator
+from dataclasses import dataclass
 from os import getenv
-from typing import Any, Optional
+from typing import Any, Optional, Union
 from urllib.parse import urlencode
 
 import httpx
@@ -20,6 +21,20 @@ API_TOOL_DEFAULT_TIMEOUT = (
 )
 
 
+@dataclass
+class ParsedResponse:
+    """Represents a parsed HTTP response with type information"""
+
+    content: Union[str, dict]
+    is_json: bool
+
+    def to_string(self) -> str:
+        """Convert response to string format for credential validation"""
+        if isinstance(self.content, dict):
+            return json.dumps(self.content, ensure_ascii=False)
+        return str(self.content)
+
+
 class ApiTool(Tool):
     """
     Api tool
@@ -58,7 +73,9 @@ class ApiTool(Tool):
 
         response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters)
         # validate response
-        return self.validate_and_parse_response(response)
+        parsed_response = self.validate_and_parse_response(response)
+        # For credential validation, always return as string
+        return parsed_response.to_string()
 
     def tool_provider_type(self) -> ToolProviderType:
         return ToolProviderType.API
@@ -112,23 +129,36 @@ class ApiTool(Tool):
 
         return headers
 
-    def validate_and_parse_response(self, response: httpx.Response) -> str:
+    def validate_and_parse_response(self, response: httpx.Response) -> ParsedResponse:
         """
-        validate the response
+        validate the response and return parsed content with type information
+
+        :return: ParsedResponse with content and is_json flag
         """
         if isinstance(response, httpx.Response):
             if response.status_code >= 400:
                 raise ToolInvokeError(f"Request failed with status code {response.status_code} and {response.text}")
             if not response.content:
-                return "Empty response from the tool, please check your parameters and try again."
+                return ParsedResponse(
+                    "Empty response from the tool, please check your parameters and try again.", False
+                )
+
+            # Check content type
+            content_type = response.headers.get("content-type", "").lower()
+            is_json_content_type = "application/json" in content_type
+
+            # Try to parse as JSON
             try:
-                response = response.json()
-                try:
-                    return json.dumps(response, ensure_ascii=False)
-                except Exception:
-                    return json.dumps(response)
+                json_response = response.json()
+                # If content-type indicates JSON, return as JSON object
+                if is_json_content_type:
+                    return ParsedResponse(json_response, True)
+                else:
+                    # If content-type doesn't indicate JSON, treat as text regardless of content
+                    return ParsedResponse(response.text, False)
             except Exception:
-                return response.text
+                # Not valid JSON, return as text
+                return ParsedResponse(response.text, False)
         else:
             raise ValueError(f"Invalid response type {type(response)}")
 
@@ -369,7 +399,14 @@ class ApiTool(Tool):
         response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, tool_parameters)
 
         # validate response
-        response = self.validate_and_parse_response(response)
+        parsed_response = self.validate_and_parse_response(response)
 
-        # assemble invoke message
-        yield self.create_text_message(response)
+        # assemble invoke message based on response type
+        if parsed_response.is_json and isinstance(parsed_response.content, dict):
+            yield self.create_json_message(parsed_response.content)
+        else:
+            # Convert to string if needed and create text message
+            text_response = (
+                parsed_response.content if isinstance(parsed_response.content, str) else str(parsed_response.content)
+            )
+            yield self.create_text_message(text_response)