Auto commit: 2026-03-28 14:57:14

2026-03-28 14:57:14 +08:00
parent fefc7239eb
commit 5f7cdac9be
6 changed files with 652 additions and 0 deletions
@@ -0,0 +1,158 @@
+"""Popiart AI OpenAPI 公共模块：HTTP 请求、鉴权、图片保存等工具函数"""
+
+import base64
+import datetime
+import json
+import os
+import re
+import struct
+import sys
+import time
+from typing import Optional, Tuple
+import urllib.request
+import urllib.error
+
+# ========== 配置 ==========
+POPI_OPENAPI_URL = os.environ.get("POPI_OPENAPI_URL", "https://llmapitest.popi.art")
+API_KEY = os.environ.get("POPI_OPENAPI_KEY", "")
+
+if not API_KEY:
+    print("错误：请设置 POPI_OPENAPI_KEY 环境变量", file=sys.stderr)
+    print("  export POPI_OPENAPI_KEY=sk-xxxx", file=sys.stderr)
+    sys.exit(1)
+
+# 轮询配置
+POLL_INTERVAL = 30   # 秒
+MAX_WAIT_TIME = 600  # 秒
+# ==========================
+
+
+def headers(content_type="application/json"):
+    h = {"Authorization": f"Bearer {API_KEY}"}
+    if content_type:
+        h["Content-Type"] = content_type
+    return h
+
+
+def post_json(path: str, body: dict, timeout: int = 120) -> dict:
+    url = f"{POPI_OPENAPI_URL.rstrip('/')}{path}"
+    data = json.dumps(body).encode("utf-8")
+    req = urllib.request.Request(url, data=data, method="POST", headers=headers())
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except urllib.error.HTTPError as e:
+        body_text = e.read().decode("utf-8") if e.fp else ""
+        print(f"❌ HTTP {e.code}: {body_text}", file=sys.stderr)
+        sys.exit(1)
+    except urllib.error.URLError as e:
+        print(f"❌ 网络错误: {e.reason}", file=sys.stderr)
+        sys.exit(1)
+
+
+def get_json(path: str, timeout: int = 30) -> dict:
+    url = f"{POPI_OPENAPI_URL.rstrip('/')}{path}"
+    req = urllib.request.Request(url, method="GET", headers=headers(content_type=None))
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except urllib.error.HTTPError as e:
+        body_text = e.read().decode("utf-8") if e.fp else ""
+        print(f"❌ HTTP {e.code}: {body_text}", file=sys.stderr)
+        sys.exit(1)
+    except urllib.error.URLError as e:
+        print(f"❌ 网络错误: {e.reason}", file=sys.stderr)
+        sys.exit(1)
+
+
+def post_multipart(path: str, fields: dict, file_fields: list = None, timeout: int = 30) -> dict:
+    """
+    multipart/form-data POST。
+    fields: {name: value}
+    file_fields: [(field_name, filename, bytes, mime_type), ...] 或 None，支持多个同名字段
+    """
+    import requests as _requests
+    url = f"{POPI_OPENAPI_URL.rstrip('/')}{path}"
+    # requests 支持 list of tuples 以传递同名多字段
+    files = [(k, (None, v)) for k, v in fields.items()]
+    for fname, filename, data, mime in (file_fields or []):
+        files.append((fname, (filename, data, mime)))
+    resp = _requests.post(url, headers={"Authorization": f"Bearer {API_KEY}"}, files=files, timeout=timeout)
+    if resp.status_code != 200:
+        print(f"❌ HTTP {resp.status_code}: {resp.text}", file=sys.stderr)
+        sys.exit(1)
+    return resp.json()
+
+
+def download_stream(path: str, output_path: str, timeout: int = 120) -> bool:
+    """流式下载文件到 output_path，返回是否成功"""
+    import requests as _requests
+    url = f"{POPI_OPENAPI_URL.rstrip('/')}{path}"
+    resp = _requests.get(url, headers={"Authorization": f"Bearer {API_KEY}"}, stream=True, timeout=timeout)
+    if resp.status_code != 200:
+        print(f"❌ 下载失败 HTTP {resp.status_code}: {resp.text}", file=sys.stderr)
+        return False
+    total = int(resp.headers.get("content-length", 0))
+    downloaded = 0
+    with open(output_path, "wb") as f:
+        for chunk in resp.iter_content(chunk_size=8192):
+            if chunk:
+                f.write(chunk)
+                downloaded += len(chunk)
+                if total:
+                    print(f"\r   {downloaded/total*100:.1f}%  ({downloaded}/{total} bytes)", end="")
+    print()
+    return True
+
+
+def save_base64_image(b64_data: str, prefix: str = "image") -> str:
+    """将 base64 图片数据保存为 PNG 文件，返回文件路径"""
+    ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"{prefix}_{ts}.png"
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    generation_dir = os.path.dirname(script_dir)
+    filepath = os.path.join(generation_dir,"generation", filename)
+    with open(filepath, "wb") as f:
+        f.write(base64.b64decode(b64_data))
+    # 输出超链接格式（使用 file:// 协议）
+    print(f"\n🔗 点击查看图片: file://{filepath}")
+    return filepath
+
+
+
+def extract_image_urls(markdown_text: str) -> list:
+    """从 markdown 文本中提取图片 URL"""
+    return re.findall(r'!\[.*?\]\((https?://[^)]+)\)', markdown_text)
+
+
+def poll_video(video_id: str) -> Optional[str]:
+    """
+    轮询视频生成状态，返回最终 task_id（用于下载），失败返回 None。
+    网关返回结构：{ code, data: { status, progress, result_url, ... } }
+    """
+    url = f"/v1/video/generations/{video_id}"
+    start = time.time()
+    while True:
+        elapsed = int(time.time() - start)
+        if elapsed > MAX_WAIT_TIME:
+            print(f"⏱️  超时：已等待 {MAX_WAIT_TIME}s")
+            return None
+        result = get_json(url)
+        data = result.get("data") or result
+        status = data.get("status", "")
+        progress = data.get("progress", 0)
+        print(f"  [{elapsed:>4}s] status={status}  progress={progress}%")
+        if status == "SUCCESS":
+            url_result = data.get("result_url", "")
+            if url_result:
+                print(f"  result_url: {url_result}")
+            return data.get("task_id") or video_id
+        elif status == "FAILURE":
+            print(f"❌ 生成失败: {data.get('fail_reason', '')}")
+            return None
+        print(f"  等待 {POLL_INTERVAL}s...")
+        time.sleep(POLL_INTERVAL)
+
+
+def ts() -> str:
+    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")