Files
DronePlanning/tools/test_validate/modules/api_client.py
2026-02-26 19:37:55 +08:00

116 lines
4.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import requests
import time
import json
class APIClient:
def __init__(self, base_url="http://127.0.0.1:8000"):
self.base_url = base_url
self.endpoint = "/generate_plan"
def send_debug_stage(
self, prompt: str, drone_state: str = "on_ground", target_stage: int = 1, timeout=120
):
"""
Stage 分阶段调试请求。
target_stage: 1-6
Returns: 同 send_request 结构data 为 debug 端返回的 {target_stage, stage_name, output}
"""
url = f"{self.base_url}/debug_stage"
payload = {
"user_prompt": prompt,
"drone_state": drone_state,
"target_stage": target_stage,
}
headers = {"Content-Type": "application/json"}
start_time = time.time()
try:
response = requests.post(url, json=payload, headers=headers, timeout=timeout)
latency = time.time() - start_time
response.raise_for_status()
try:
data = response.json()
if "error" in data and len(data) == 1:
return {
"success": False,
"data": data,
"latency": latency,
"error": data["error"],
"http_status": response.status_code,
}
return {
"success": True,
"data": data,
"latency": latency,
"error": None,
"http_status": response.status_code,
}
except json.JSONDecodeError:
return {
"success": False,
"data": None,
"latency": latency,
"error": f"Invalid JSON: {response.text[:200]}",
"http_status": response.status_code,
}
except requests.exceptions.RequestException as e:
latency = time.time() - start_time
return {
"success": False,
"data": None,
"latency": latency,
"error": str(e),
"http_status": getattr(e.response, "status_code", None) if hasattr(e, "response") and e.response is not None else None,
}
def send_request(self, prompt, drone_state="on_ground", timeout=60):
"""
Sends a request to the API and returns a structured result.
Returns:
dict: {
"success": bool,
"data": dict or None,
"latency": float (seconds),
"error": str or None,
"http_status": int or None
}
"""
url = f"{self.base_url}{self.endpoint}"
payload = {"user_prompt": prompt, "drone_state": drone_state}
headers = {"Content-Type": "application/json"}
start_time = time.time()
try:
response = requests.post(url, json=payload, headers=headers, timeout=timeout)
latency = time.time() - start_time
response.raise_for_status()
try:
data = response.json()
return {
"success": True,
"data": data,
"latency": latency,
"error": None,
"http_status": response.status_code
}
except json.JSONDecodeError:
return {
"success": False,
"data": None,
"latency": latency,
"error": f"Invalid JSON response: {response.text[:200]}",
"http_status": response.status_code
}
except requests.exceptions.RequestException as e:
latency = time.time() - start_time
return {
"success": False,
"data": None,
"latency": latency,
"error": str(e),
"http_status": getattr(e.response, 'status_code', None)
}