348 lines
14 KiB
Python
348 lines
14 KiB
Python
"""HTTP plumbing shared by every feature module.
|
||
|
||
Holds the transport concerns only: credentials resolution, headers, retries for
|
||
transient failures, automatic JWT refresh from the account's session cookie,
|
||
progress polling and file download. Feature specific request bodies live in
|
||
``tripo/features/*``.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
from . import auth, glb, s3
|
||
from .accounts import Accounts
|
||
from .config import Config
|
||
from .errors import TripoAPIError, TripoTimeout
|
||
|
||
# Transient transport/server failures worth retrying (the studio occasionally
|
||
# returns a bare 502 mid-generation, which must not abort a running poll).
|
||
TRANSIENT_STATUS = {429, 500, 502, 503, 504}
|
||
RETRY_ATTEMPTS = 4
|
||
RETRY_BACKOFF = 1.5
|
||
|
||
ORIGIN = "https://studio.tripo3d.ai"
|
||
USER_AGENT = "tripocli/2.1.0 (Mozilla/5.0; AI Agent Client)"
|
||
|
||
# Model formats the studio's "上传 3D 模型" importer accepts (UI: OBJ/FBX/STL/GLB ≤150MB).
|
||
MODEL_FORMATS = ("glb", "obj", "fbx", "stl")
|
||
|
||
NO_CREDENTIALS = (
|
||
"未找到可用凭据。请先添加账号（粘贴浏览器 cookie 即可，之后 CLI 自动续期 JWT）：\n"
|
||
" tripocli account add --cookie \"<整条 Cookie 或 ory_kratos_session 的值>\"\n"
|
||
"也可用环境变量 TRIPO_COOKIE 提供。"
|
||
)
|
||
|
||
|
||
def _jwt_expiry(token):
|
||
import base64
|
||
try:
|
||
payload = token.split(".")[1]
|
||
payload += "=" * (-len(payload) % 4)
|
||
return int(json.loads(base64.urlsafe_b64decode(payload).decode()).get("exp") or 0)
|
||
except Exception:
|
||
return 0
|
||
|
||
|
||
class TripoClient:
|
||
def __init__(self, token=None, device_id=None, api_base=None, account=None,
|
||
auto_refresh=True, accounts=None, config=None):
|
||
self.config = config or Config()
|
||
self.accounts = accounts or Accounts()
|
||
|
||
env_token = os.environ.get("TRIPO_TOKEN") or os.environ.get("TRIPO_API_KEY")
|
||
env_cookie = os.environ.get("TRIPO_COOKIE")
|
||
self.account_name, record = self.accounts.get(account or os.environ.get("TRIPO_ACCOUNT"))
|
||
record = record or {}
|
||
|
||
self.token = token or env_token or record.get("token")
|
||
self.cookie = env_cookie or record.get("cookie")
|
||
self.device_id = device_id or self.config.device_id
|
||
self.api_base = (api_base or self.config.api_base).rstrip("/")
|
||
self.auto_refresh = auto_refresh and not token and not env_token
|
||
self.token_source = "env" if env_token else ("account" if record.get("token") else None)
|
||
|
||
# --- credentials ---
|
||
|
||
def token_expiring(self):
|
||
exp = _jwt_expiry(self.token or "")
|
||
return not exp or exp - time.time() < 120
|
||
|
||
def refresh_token(self):
|
||
"""Mint a fresh JWT from the session cookie (pure HTTP, no browser)."""
|
||
if not self.cookie:
|
||
raise TripoAPIError(NO_CREDENTIALS, code="no-credentials")
|
||
token, session_expires = auth.mint_token(self.api_base, self.cookie, self.device_id)
|
||
self.token = token
|
||
if self.account_name:
|
||
self.accounts.update_token(self.account_name, token, session_expires)
|
||
return True
|
||
|
||
def ensure_credentials(self):
|
||
if self.token and not self.token_expiring():
|
||
return
|
||
if self.cookie:
|
||
self.refresh_token()
|
||
return
|
||
if not self.token:
|
||
raise TripoAPIError(NO_CREDENTIALS, code="no-credentials")
|
||
|
||
# --- transport ---
|
||
|
||
def _headers(self, extra=None):
|
||
headers = {
|
||
"Accept": "application/json",
|
||
"Content-Type": "application/json",
|
||
"Origin": ORIGIN,
|
||
"Referer": ORIGIN + "/",
|
||
"User-Agent": USER_AGENT,
|
||
"x-tripo-device-id": self.device_id,
|
||
}
|
||
if self.token:
|
||
headers["Authorization"] = f"Bearer {self.token}"
|
||
if extra:
|
||
headers.update(extra)
|
||
return headers
|
||
|
||
def _raw_request(self, method, endpoint, params=None, body=None, headers=None, timeout=30):
|
||
url = f"{self.api_base}{endpoint}"
|
||
if params:
|
||
query = urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})
|
||
if query:
|
||
url = f"{url}?{query}"
|
||
|
||
payload = json.dumps(body).encode("utf-8") if body is not None else None
|
||
request = urllib.request.Request(
|
||
url, data=payload, headers=self._headers(headers), method=method.upper()
|
||
)
|
||
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
text = response.read().decode("utf-8")
|
||
if not text:
|
||
return None
|
||
return json.loads(text)
|
||
except urllib.error.HTTPError as error:
|
||
err_body = error.read().decode("utf-8", "replace")
|
||
try:
|
||
err_json = json.loads(err_body)
|
||
except json.JSONDecodeError:
|
||
raise TripoAPIError(f"HTTP {error.code}: {error.reason}", code=error.code) from None
|
||
raise TripoAPIError(
|
||
err_json.get("message", f"HTTP {error.code}"),
|
||
code=err_json.get("code", error.code),
|
||
suggestion=err_json.get("suggestion"),
|
||
raw=err_json,
|
||
) from None
|
||
except urllib.error.URLError as error:
|
||
raise TripoAPIError(f"Network error: {error.reason}") from None
|
||
|
||
def request(self, method, endpoint, params=None, body=None, headers=None, timeout=30):
|
||
"""Perform an API call, retrying transient failures and expired tokens."""
|
||
attempt = 0
|
||
refreshed = False
|
||
while True:
|
||
attempt += 1
|
||
try:
|
||
result = self._raw_request(method, endpoint, params, body, headers, timeout)
|
||
if isinstance(result, dict) and result.get("code") not in (None, 0):
|
||
raise TripoAPIError(
|
||
result.get("message", "API request failed"),
|
||
code=result.get("code"),
|
||
suggestion=result.get("suggestion"),
|
||
raw=result,
|
||
)
|
||
return result
|
||
except TripoAPIError as error:
|
||
if error.is_auth_error and self.auto_refresh and not refreshed:
|
||
refreshed = True
|
||
try:
|
||
if self.refresh_token():
|
||
continue
|
||
except TripoAPIError:
|
||
pass # fall through and surface the original error
|
||
transient = error.code in TRANSIENT_STATUS or error.code is None
|
||
if transient and attempt < RETRY_ATTEMPTS:
|
||
time.sleep(RETRY_BACKOFF * attempt)
|
||
continue
|
||
raise
|
||
|
||
# --- shared API calls ---
|
||
|
||
def get_data(self, method, endpoint, **kwargs):
|
||
result = self.request(method, endpoint, **kwargs)
|
||
return result.get("data") if result else None
|
||
|
||
def progress(self, operator_ids):
|
||
if isinstance(operator_ids, str):
|
||
operator_ids = [operator_ids]
|
||
return self.get_data("POST", "/v2/studio/progress", body={"ids": operator_ids}) or []
|
||
|
||
def project_detail(self, project_id, locale="zh-CN"):
|
||
return self.get_data("GET", f"/v2/studio/project/detail/v3/{project_id}", params={"locale": locale})
|
||
|
||
def project_operator(self, project_id):
|
||
"""Current operator block of a project (carries operator_id + model_version)."""
|
||
detail = self.project_detail(project_id) or {}
|
||
return detail.get("operator") or {}
|
||
|
||
def part_names(self, project_id):
|
||
"""Node names inside the project's current mesh (the ``part_name_list``).
|
||
|
||
Read from the model file itself, since the detail API does not expose
|
||
them. Returns an empty list when the model cannot be read.
|
||
"""
|
||
detail = self.project_detail(project_id) or {}
|
||
url = (detail.get("operator") or {}).get("model_url")
|
||
if not url:
|
||
return []
|
||
request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
|
||
with urllib.request.urlopen(request, timeout=120) as response:
|
||
data = response.read()
|
||
return glb.node_names(data)
|
||
|
||
def payment_profile(self):
|
||
return self.get_data("GET", "/v2/studio/user/profile/payment")
|
||
|
||
def whoami(self):
|
||
return self.get_data("GET", "/v2/studio/studio/whoami", params={"tokenizeAs": "default_jwt"})
|
||
|
||
def quota(self):
|
||
return self.get_data("POST", "/v2/studio/operation/quota", body={})
|
||
|
||
def assets(self, asset_type="mine", type="model", page=1, size=20, offset=0, locale="zh-CN"):
|
||
return self.get_data(
|
||
"GET",
|
||
"/v2/studio/assets/v2",
|
||
params={
|
||
"asset_type": asset_type,
|
||
"type": type,
|
||
"page": page,
|
||
"size": size,
|
||
"offset": offset,
|
||
"locale": locale,
|
||
},
|
||
)
|
||
|
||
def download_url(self, operator_id, file_name):
|
||
data = self.get_data(
|
||
"POST", "/v2/studio/operation/download_with_name",
|
||
body={"operator_id": operator_id, "file_name": file_name},
|
||
)
|
||
return (data or {}).get("model_url")
|
||
|
||
def export(self, project_id, operator_id, export_format, name="model", model_version="default"):
|
||
fmt = {"glb": "gltf", "gltf": "gltf"}.get(export_format.lower(), export_format.lower())
|
||
return self.get_data(
|
||
"POST", "/v2/studio/operation/export",
|
||
body={
|
||
"project_id": project_id,
|
||
"operator_id": operator_id,
|
||
"format": fmt,
|
||
"name": name,
|
||
"model_version": model_version,
|
||
},
|
||
)
|
||
|
||
# --- images ---
|
||
|
||
def upload_image(self, path):
|
||
"""Upload a local image and return an image object ready to be submitted.
|
||
|
||
Verified flow: temporary_token -> SigV4 PUT -> ``/audit/image`` -> the
|
||
``image`` object must carry ``image_audit_result`` and ``image_source``,
|
||
otherwise every generation endpoint answers 11002.
|
||
"""
|
||
path = Path(path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"图片文件不存在: {path}")
|
||
|
||
fmt = path.suffix.lower().lstrip(".") or "png"
|
||
if fmt == "jpeg":
|
||
fmt = "jpg"
|
||
if fmt not in ("png", "jpg", "webp"):
|
||
raise ValueError(f"不支持的图片格式: .{fmt} (支持 png/jpg/webp)")
|
||
content_type = "image/jpeg" if fmt == "jpg" else f"image/{fmt}"
|
||
|
||
sts = self.get_data(
|
||
"POST", "/v2/studio/storage/temporary_token", body={"client": "aws", "format": fmt}
|
||
)
|
||
s3.upload_bytes(sts, path.read_bytes(), content_type=content_type)
|
||
|
||
image = {"bucket": sts["resource_bucket"], "key": sts["resource_uri"]}
|
||
audit = self.get_data("POST", "/v2/studio/audit/image", body={"image": image}) or {}
|
||
image["image_audit_result"] = audit.get("result", "pass")
|
||
image["image_source"] = "upload"
|
||
return image
|
||
|
||
# --- models ---
|
||
|
||
def upload_model(self, path):
|
||
"""Upload a local 3D model and return the ``import_user_model`` payload.
|
||
|
||
Verified flow: ``temporary_token`` (format = file extension) -> SigV4 PUT
|
||
-> the import body is ``{format, model:{bucket,key}, name,
|
||
transform_matrix, use_original_uv}``. The studio UI rotates the mesh
|
||
before importing (that is what ``transform_matrix`` carries); a CLI
|
||
import uses the identity matrix.
|
||
"""
|
||
path = Path(path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"模型文件不存在: {path}")
|
||
|
||
fmt = path.suffix.lower().lstrip(".")
|
||
if fmt not in MODEL_FORMATS:
|
||
raise ValueError(f"不支持的模型格式: .{fmt} (支持 glb/obj/fbx/stl)")
|
||
|
||
sts = self.get_data(
|
||
"POST", "/v2/studio/storage/temporary_token", body={"client": "aws", "format": fmt}
|
||
)
|
||
s3.upload_bytes(sts, path.read_bytes(), content_type="application/octet-stream")
|
||
return {
|
||
"format": fmt,
|
||
"model": {"bucket": sts["resource_bucket"], "key": sts["resource_uri"]},
|
||
"name": path.name,
|
||
}
|
||
|
||
# --- waiting & downloading ---
|
||
|
||
def wait_task(self, operator_id, timeout=900, poll_interval=3, on_progress=None):
|
||
"""Poll until the operator succeeds, fails or the timeout expires."""
|
||
if not operator_id:
|
||
# Guard: polling a missing id would otherwise spin until the timeout.
|
||
raise TripoAPIError("没有可等待的 operator_id（接口未返回任务 id）", code="no-operator")
|
||
deadline = time.time() + timeout
|
||
last_progress = -1
|
||
while time.time() < deadline:
|
||
items = self.progress([operator_id])
|
||
if items:
|
||
task = items[0]
|
||
status = task.get("status")
|
||
progress = task.get("progress", 0)
|
||
if progress != last_progress:
|
||
last_progress = progress
|
||
if on_progress:
|
||
on_progress(status, progress, task.get("left_time", -1), task)
|
||
if status == "success":
|
||
return task
|
||
if status in ("failed", "banned", "expired"):
|
||
raise TripoAPIError(f"任务 {operator_id} 失败: {json.dumps(task, ensure_ascii=False)}")
|
||
time.sleep(poll_interval)
|
||
raise TripoTimeout(f"任务 {operator_id} 在 {timeout}s 内未完成")
|
||
|
||
def download(self, url, output_path):
|
||
output_path = Path(output_path)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
|
||
with urllib.request.urlopen(request, timeout=300) as response, open(output_path, "wb") as handle:
|
||
while True:
|
||
chunk = response.read(64 * 1024)
|
||
if not chunk:
|
||
break
|
||
handle.write(chunk)
|
||
return output_path
|