103 lines
4.5 KiB
Python
103 lines
4.5 KiB
Python
"""独立视觉客户端：复用既有 OpenAI 兼容视觉配置，但不依赖旧 CLI 或旧 ADB。"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import urllib.error
|
||
import os
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
from .paths import CONFIG
|
||
|
||
|
||
class VisionError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def _config() -> dict:
|
||
path = CONFIG
|
||
if not path.exists():
|
||
raise VisionError("缺少 nextgen/config.local.json 的 vision 配置")
|
||
vision = json.loads(path.read_text(encoding="utf-8")).get("vision", {})
|
||
# 本地 config.local.json 是被 gitignore 的用户私有配置；环境变量仍可覆盖它，便于 CI 或部署注入。
|
||
api_key = os.environ.get(str(vision.get("api_key_env", "AUTOTRAIN_NEXT_VISION_API_KEY")), "") or str(vision.get("api_key", ""))
|
||
if not all(vision.get(key) for key in ("base_url", "model")) or not api_key:
|
||
raise VisionError("vision 配置缺少 base_url/model 或 API key 环境变量")
|
||
return {**vision, "api_key": api_key}
|
||
|
||
|
||
def _json(content: str) -> dict:
|
||
text = content.strip()
|
||
if text.startswith("```"):
|
||
text = text.split("\n", 1)[-1].rsplit("```", 1)[0]
|
||
start, end = text.find("{"), text.rfind("}")
|
||
if start < 0 or end < start:
|
||
raise VisionError(f"视觉模型未返回 JSON：{content[:240]}")
|
||
try:
|
||
return json.loads(text[start:end + 1])
|
||
except json.JSONDecodeError as exc:
|
||
raise VisionError(f"视觉模型 JSON 无效：{exc}") from exc
|
||
|
||
|
||
def ask(image: Path, prompt: str) -> dict:
|
||
cfg = _config()
|
||
encoded = base64.b64encode(image.read_bytes()).decode()
|
||
payload = {
|
||
"model": cfg["model"], "temperature": 0,
|
||
"messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded}"}},
|
||
]}],
|
||
}
|
||
request = urllib.request.Request(
|
||
str(cfg["base_url"]).rstrip("/") + "/chat/completions", data=json.dumps(payload).encode(),
|
||
headers={"Authorization": f"Bearer {cfg['api_key']}", "Content-Type": "application/json"}, method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=45) as response:
|
||
body = json.loads(response.read())
|
||
except urllib.error.HTTPError as exc:
|
||
raise VisionError(f"视觉请求失败 HTTP {exc.code}") from exc
|
||
content = body.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||
return _json(content)
|
||
|
||
|
||
def detect_page(image: Path) -> dict:
|
||
return ask(image, """这是 12306 App 截图。只返回 JSON：
|
||
{"page":"list|verify|invoice_info|invoice_confirm|success|email_confirm|mail_sent|no_ticket|loading|other","summary":"不超过30字","scroll_position":"top|middle|bottom|unknown"}。
|
||
list 是含脱敏姓名和开具按钮的扫码开票单；verify 是身份证后8位核验；invoice_info 有提交；invoice_confirm 有确认；success 是开票成功。""")
|
||
|
||
|
||
def list_rows(image: Path) -> dict:
|
||
parsed = ask(image, """识别这张 12306 扫码开票单列表截图。只返回 JSON：
|
||
{"page_ok":true,"scroll_position":"top|middle|bottom","rows":[{"masked_name":"张*三","id_first":"","id_last":"","button_center":[500,500],"button_text":"开具|已开票|查看"}]}。
|
||
button_center 必须是按钮中心的 0-1000 归一化坐标；只列完整可见行；不是列表页则 page_ok=false、rows=[]。""")
|
||
rows = parsed.get("rows") or []
|
||
valid = []
|
||
for row in rows:
|
||
point = row.get("button_center") or []
|
||
if len(point) != 2:
|
||
continue
|
||
try:
|
||
x, y = float(point[0]) / 1000, float(point[1]) / 1000
|
||
except (TypeError, ValueError):
|
||
continue
|
||
if 0 <= x <= 1 and 0 <= y <= 1:
|
||
valid.append({**row, "point": [x, y]})
|
||
return {"page_ok": bool(parsed.get("page_ok")), "scroll_position": parsed.get("scroll_position", "unknown"), "rows": valid}
|
||
|
||
|
||
def find_button(image: Path, text: str) -> dict:
|
||
parsed = ask(image, f"""在这张 12306 App 截图中找文字为「{text}」的可点击按钮。只返回 JSON：
|
||
{{"found":true,"center":[500,500]}}。
|
||
center 必须是 0-1000 归一化坐标；找不到返回 {{"found":false,"center":[]}}。不要猜测。""")
|
||
center = parsed.get("center") or []
|
||
if not parsed.get("found") or len(center) != 2:
|
||
return {"found": False}
|
||
try:
|
||
point = [float(center[0]) / 1000, float(center[1]) / 1000]
|
||
except (TypeError, ValueError):
|
||
return {"found": False}
|
||
return {"found": all(0 <= value <= 1 for value in point), "point": point}
|