100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
|
import json
|
||
|
|
from typing import Any, Dict
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from langchain_core.tools import BaseTool
|
||
|
|
|
||
|
|
from config import Config
|
||
|
|
|
||
|
|
|
||
|
|
class SrApiQueryTool(BaseTool):
|
||
|
|
"""调用 SR API 执行 SQL 查询"""
|
||
|
|
|
||
|
|
name: str = "sr_api_query"
|
||
|
|
description: str = (
|
||
|
|
"调用 SR API 执行 SQL 查询。输入为 JSON 字符串，示例："
|
||
|
|
'{"sql":"SELECT * FROM table","page":1,"rows":10,"orderBySelect":true,"timeout":30}'
|
||
|
|
)
|
||
|
|
|
||
|
|
def _run(self, payload: str) -> str:
|
||
|
|
"""执行 SQL 查询"""
|
||
|
|
try:
|
||
|
|
data = json.loads(payload)
|
||
|
|
except Exception as e:
|
||
|
|
return f"请求参数解析失败: {e}"
|
||
|
|
|
||
|
|
sql = data.get("sql")
|
||
|
|
page = int(data.get("page", 1))
|
||
|
|
rows = int(data.get("rows", 10))
|
||
|
|
order_by_select = bool(data.get("orderBySelect", True))
|
||
|
|
timeout = float(data.get("timeout", 30))
|
||
|
|
|
||
|
|
if not sql:
|
||
|
|
return "缺少 sql"
|
||
|
|
|
||
|
|
cfg = Config.get_section("sr_api")
|
||
|
|
url = cfg.get("url")
|
||
|
|
app_key = cfg.get("llzappkey")
|
||
|
|
secret_key = cfg.get("llzsercret")
|
||
|
|
|
||
|
|
if not url or not app_key or not secret_key:
|
||
|
|
return "sr_api 配置缺失 url/llzAppkey/llzSercret"
|
||
|
|
|
||
|
|
headers = {"llzAppkey": app_key, "llzSercret": secret_key}
|
||
|
|
body = {"sql": sql, "page": page, "rows": rows, "orderBySelect": order_by_select}
|
||
|
|
|
||
|
|
try:
|
||
|
|
with httpx.Client(timeout=timeout) as client:
|
||
|
|
response = client.post(url, json=body, headers=headers)
|
||
|
|
return json.dumps(
|
||
|
|
{
|
||
|
|
"status_code": response.status_code,
|
||
|
|
"headers": dict(response.headers),
|
||
|
|
"text": response.text,
|
||
|
|
},
|
||
|
|
ensure_ascii=False,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
return f"请求失败: {e}"
|
||
|
|
|
||
|
|
async def _arun(self, payload: str) -> str:
|
||
|
|
"""工具的异步版本"""
|
||
|
|
try:
|
||
|
|
data = json.loads(payload)
|
||
|
|
except Exception as e:
|
||
|
|
return f"请求参数解析失败: {e}"
|
||
|
|
|
||
|
|
sql = data.get("sql")
|
||
|
|
page = int(data.get("page", 1))
|
||
|
|
rows = int(data.get("rows", 10))
|
||
|
|
order_by_select = bool(data.get("orderBySelect", True))
|
||
|
|
timeout = float(data.get("timeout", 30))
|
||
|
|
|
||
|
|
if not sql:
|
||
|
|
return "缺少 sql"
|
||
|
|
|
||
|
|
cfg = Config.get_section("sr_api")
|
||
|
|
url = cfg.get("url")
|
||
|
|
app_key = cfg.get("llzappkey")
|
||
|
|
secret_key = cfg.get("llzsercret")
|
||
|
|
|
||
|
|
if not url or not app_key or not secret_key:
|
||
|
|
return "sr_api 配置缺失 url/llzAppkey/llzSercret"
|
||
|
|
|
||
|
|
headers = {"llzAppkey": app_key, "llzSercret": secret_key}
|
||
|
|
body = {"sql": sql, "page": page, "rows": rows, "orderBySelect": order_by_select}
|
||
|
|
|
||
|
|
try:
|
||
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
|
|
response = await client.post(url, json=body, headers=headers)
|
||
|
|
return json.dumps(
|
||
|
|
{
|
||
|
|
"status_code": response.status_code,
|
||
|
|
"headers": dict(response.headers),
|
||
|
|
"text": response.text,
|
||
|
|
},
|
||
|
|
ensure_ascii=False,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
return f"请求失败: {e}"
|