""" AI 助手 — 调用 DeepSeek 分析 OCR 文本结果 只负责:接收文本 → 调 AI → 返回结构化数据 不负责 OCR 识别 """ import json import os import re import time from pathlib import Path from typing import Optional import requests # ── 配置 ──────────────────────────────────────────────── DEEPSEEK_API_URL = "https://ai-api.dfwy.tech/v1/messages" DEEPSEEK_MODEL = "deepseek-v4-flash-nothinking" # 项目根目录(用于读取 config.json) PROJECT_ROOT = Path(__file__).parent class AIParser: """调用 DeepSeek 分析 OCR 文本,返回结构化店铺列表""" # 默认 API Key(写死,也可通过环境变量或 config.json 覆盖) _DEFAULT_KEY = "" def __init__(self, api_key: Optional[str] = None): """ Args: api_key: DeepSeek API Key,不传则从环境变量或 config.json 读取 """ self.api_key = api_key or self._load_api_key() if not self.api_key: raise RuntimeError( "未设置 DeepSeek API Key。请:\n" "1. 设环境变量 DEEPSEEK_API_KEY\n" "2. 或在 config.json 中添加 deepseek_api_key 字段\n" "3. 或在初始化时传入 api_key" ) @staticmethod def _load_api_key() -> str: """从环境变量、config.json 或默认值读取 API Key""" key = os.environ.get("DEEPSEEK_API_KEY", "") if key: return key cfg_path = PROJECT_ROOT / "config.json" if cfg_path.exists(): try: with open(cfg_path) as f: cfg = json.load(f) key = cfg.get("deepseek_api_key", "") if key: return key except Exception: pass return AIParser._DEFAULT_KEY # ── 核心:分析 OCR 结果 → 店铺列表 ────────────────── def parse_shops( self, ocr_results: list, screen_size: tuple = (1220, 2712), ) -> list: """ 将 OCR 识别结果传给 DeepSeek,返回结构化店铺列表 Args: ocr_results: OCR.recognize(detail="all") 返回值 [{"text": str, "confidence": float, "bbox": [[x1,y1],...]}, ...] screen_size: (宽, 高),默认 1220x2712 Returns: [[店铺名, 商品标题, 价格, [点击x, 点击y], 配送距离y], ...] 点击坐标是店铺卡片中心点 配送距离y = 配送时间/距离文字的y坐标,用于滑动到其上方 """ # 1. 格式化 OCR 数据 — 按 y 坐标排序(从上到下) sorted_results = sorted(ocr_results, key=lambda r: r["bbox"][0][1]) ocr_lines = [] for r in sorted_results: box = r["bbox"] cy = (box[0][1] + box[2][1]) // 2 cx = (box[0][0] + box[2][0]) // 2 ocr_lines.append( f"[x={cx:4d}, y={cy:4d}, conf={r['confidence']:.2f}] {r['text']}" ) ocr_text = "\n".join(ocr_lines) screen_w, screen_h = screen_size # 2. system prompt system_prompt = """你是OCR文本结构化助手。从OCR文字列表中找出每个"店铺卡片"的4个字段,必须严格使用OCR原文。 【OCR数据格式】 每行: [x=坐标x, y=坐标y, conf=置信度] OCR识别出的文字 y越大越靠下,文字按从上到下排列。 【判断店铺卡片方法】 - 同卡片的文字y坐标连续相近,卡片间有明显y间隔(>80px) - 卡片从上到下:商品标题(多行) → 价格(¥/¥开头) → 店铺名(以店/超市/旗舰店等结尾) 【字段规则 — 严格使用OCR原文,不得修改】 1. 店铺名 = 卡片最下方像店名的文字。 如果文字含"(",只取"("之前的部分(括号内是分店名,不重要)。 店铺名通常以"店/超市/旗舰店"结尾。 2. 商品标题 = 卡片上方描述商品的文字。 多行时要完整拼接在一起,不要截断。 3. 价格 = 卡片中间带¥或¥符号的文字。 只取**数值+货币符号**部分(如"¥12.8"、"¥67"),不要含"预估价"等额外文字。 如果找不到价格,留空字符串""。 4. 点击坐标[x,y] = OCR数据中该卡片区域中间位置的实际坐标,用整数。 【约束】 - 所有字段值必须严格等于或取自OCR原文,不得修改、拼凑、推断 - 店铺名如果找不到,留空字符串"" - 找不到价格的卡片也要列出,价格留空 - 有多少个卡片就输出多少个,不要遗漏 - 输出纯JSON数组,不要任何多余文字 返回格式: [[店铺名(去括号), 商品标题(完整), 价格(纯数字), [点击x, 点击y], 店铺名y], ...] 店铺名y = 店铺名文字在OCR数据中的y坐标,必须从OCR数据中找""" # 3. user prompt user_prompt = ( f"屏幕OCR识别结果(屏幕尺寸 {screen_w}x{screen_h}):\n\n" f"{ocr_text}\n\n请按格式返回店铺列表。" ) # 4. 调 AI print(f"[ai] 分析 {len(ocr_results)} 个文本块...") response = self._call(system_prompt, user_prompt) # 5. 解析返回 return self._parse_response(response, screen_w, screen_h) # ── 页面类型检测 ────────────────────────────────── def check_page(self, ocr_results: list) -> dict: """ 分析进入商品页后的页面类型 Returns: {"type": "normal"|"risk"|"qrcode"|"unknown", "shop": "店铺名(如有)", "detail": "额外描述"} """ # 1. 格式化 sorted_r = sorted(ocr_results, key=lambda r: r["bbox"][0][1]) lines = [] for r in sorted_r: box = r["bbox"] cy = (box[0][1] + box[2][1]) // 2 cx = (box[0][0] + box[2][0]) // 2 lines.append(f"[x={cx:4d}, y={cy:4d}] {r['text']}") ocr_text = "\n".join(lines) system_prompt = """你是手机屏幕分析助手。收到OCR文字列表,判断当前屏幕是什么类型的页面。 【页面类型】 1. normal = 正常商品/店铺详情页(可以看到商品信息、价格、评价、店铺名等) 2. risk = 风控/验证码/拦截页(包含"验证"、"拖动滑块"、"安全验证"、"禁止搜索"、"操作频繁"、"风险"、"请输入验证码"等) 3. qrcode = 二维码/扫码页(包含"二维码"、"扫码"、"请用微信"等) 4. unknown = 其他无法识别的页面 **优先检测 risk 类型**:只要有类似验证码、风控拦截的文字,就返回risk 【输出格式】 只返回一个JSON对象,不要任何多余文字: {"type": "normal|risk|qrcode|unknown", "shop": "店铺名或空字符串", "detail": "简要说明"} """ user_prompt = f"屏幕OCR文字:\n{ocr_text}\n\n请判断页面类型。" print(f"[ai] 检测页面类型...") response = self._call(system_prompt, user_prompt) import json cleaned = response.strip() if cleaned.startswith("```"): lines = cleaned.split("\n") if lines[0].startswith("```"): lines = lines[1:] if lines and lines[-1].strip() == "```": lines = lines[:-1] cleaned = "\n".join(lines).strip() try: data = json.loads(cleaned) return { "type": data.get("type", "unknown"), "shop": data.get("shop", ""), "detail": data.get("detail", ""), } except json.JSONDecodeError: if not response.strip(): return {"type": "normal", "shop": "", "detail": "API调用失败,默认继续"} return {"type": "unknown", "shop": "", "detail": response[:100]} # ── API 调用 ──────────────────────────────────────── def _call(self, system_prompt: str, user_content: str, timeout: int = 120, retries: int = 3) -> str: """调用 API(Anthropic 格式),失败自动重试""" headers = { "x-api-key": self.api_key, "anthropic-version": "2023-06-01", "Content-Type": "application/json", } payload = { "model": DEEPSEEK_MODEL, "max_tokens": 8192, "system": system_prompt, "messages": [ {"role": "user", "content": user_content}, ], } last_err = None for attempt in range(retries): try: resp = requests.post(DEEPSEEK_API_URL, json=payload, headers=headers, timeout=timeout) resp.raise_for_status() data = resp.json() # 从 content 中提取所有 text 类型块 texts = [c["text"] for c in data.get("content", []) if c.get("type") == "text"] return "".join(texts) except Exception as e: last_err = e if attempt < retries - 1: print(f" [ai] API调用失败(第{attempt+1}次),重试... ({e})") time.sleep(2) print(f" [ai] API调用{retries}次均失败,跳过: {last_err}") return "" # ── 解析 ──────────────────────────────────────────── @staticmethod def _parse_response(response: str, screen_w: int, screen_h: int) -> list: """解析 AI 返回的 JSON""" cleaned = response.strip() # 去掉 markdown 代码块包裹 if cleaned.startswith("```"): lines = cleaned.split("\n") if lines[0].startswith("```"): lines = lines[1:] if lines and lines[-1].strip() == "```": lines = lines[:-1] cleaned = "\n".join(lines).strip() # 解析 JSON try: data = json.loads(cleaned) except json.JSONDecodeError: print(f"[ai] WARNING: JSON 解析失败,尝试提取数组片段") match = re.search(r'\[[\s\S]*\]', cleaned) if match: try: data = json.loads(match.group(0)) except json.JSONDecodeError: print(f"[ai] ERROR: JSON 解析失败") return [] else: print(f"[ai] ERROR: 未找到 JSON 数组") return [] if not isinstance(data, list): print(f"[ai] ERROR: 返回不是数组") return [] validated = [] for item in data: # 兼容 dict 格式(模型偶发不按 list 模板输出时自动转换) if isinstance(item, dict): shop = str(item.get("name") or item.get("shop") or item.get("shop_name") or "") title = str(item.get("title") or item.get("description") or item.get("product_name") or "") price = str(item.get("price") or item.get("min_price") or "") coords = item.get("coords") or item.get("click_pos") or item.get("click") if not isinstance(coords, list) and item.get("x") is not None and item.get("y") is not None: coords = [item.get("x"), item.get("y")] shop_name_y = item.get("shop_name_y") or item.get("y") item = [shop, title, price, coords, shop_name_y] if not isinstance(item, list) or len(item) < 4: continue shop = str(item[0]) if item[0] else "" title = str(item[1]) if item[1] else "" price = str(item[2]) if item[2] else "" coords = item[3] if isinstance(coords, list) and len(coords) == 2: cx, cy = int(coords[0]), int(coords[1]) if 0 <= cx <= screen_w and 0 <= cy <= screen_h: # 第5个字段=店铺名文字的y坐标 shop_name_y = int(item[4]) if len(item) >= 5 and isinstance(item[4], (int, float)) else min(cy + 200, screen_h) validated.append([shop, title, price, [cx, cy], shop_name_y]) print(f"[ai] 解析到 {len(validated)} 个店铺") return validated # ── 独立测试 ──────────────────────────────────────────── if __name__ == "__main__": import sys from steps.ocr import OCR # 读取已有截图 test_img = str(PROJECT_ROOT / "screenshots" / "step2_result.png") if not os.path.exists(test_img): print(f"测试图片不存在: {test_img}") print("请先跑一次 main.py 生成截图") sys.exit(1) # OCR 识别 ocr = OCR() raw = ocr.recognize(test_img, detail="all") print(f"OCR 识别到 {len(raw)} 个文本块\n") # AI 分析 parser = AIParser() shops = parser.parse_shops(raw) print("\n" + "=" * 60) print(f" 店铺列表 ({len(shops)} 个)") print("=" * 60) for s in shops: print(f" 店铺: {s[0]}") print(f" 商品: {s[1]}") print(f" 价格: {s[2]}") print(f" 点击: ({s[3][0]}, {s[3][1]})") print(f" 配送y: {s[4]}") print()