127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
||
# /// script
|
||
# requires-python = ">=3.10"
|
||
# dependencies = [
|
||
# "requests>=2.31.0",
|
||
# ]
|
||
# ///
|
||
"""使用 Agnes AI 的多模态模型理解图片。
|
||
传入图片（公网 URL 或本地文件）+ 文本问题，模型基于图片进行描述、分析、问答或信息提取。
|
||
"""
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import requests
|
||
|
||
API_URL = "https://apihub.agnes-ai.com/v1/chat/completions"
|
||
DEFAULT_MODEL = "agnes-2.0-flash"
|
||
|
||
|
||
def file_to_data_uri(path):
|
||
"""把本地图片文件转成 data:image Base64。"""
|
||
mime = mimetypes.guess_type(path)[0] or "image/png"
|
||
with open(path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
return f"data:{mime};base64,{b64}"
|
||
|
||
|
||
def call_api(messages, model, api_key, temperature=None, max_tokens=None, retries=3):
|
||
"""调用 Agnes chat/completions，返回解析后的 JSON。瞬时网络抖动自动重试。"""
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {api_key}",
|
||
}
|
||
payload = {"model": model, "messages": messages}
|
||
if temperature is not None:
|
||
payload["temperature"] = temperature
|
||
if max_tokens is not None:
|
||
payload["max_tokens"] = max_tokens
|
||
|
||
last_err = None
|
||
for attempt in range(1, retries + 1):
|
||
try:
|
||
resp = requests.post(API_URL, headers=headers, json=payload, timeout=180)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
except (requests.exceptions.SSLError,
|
||
requests.exceptions.ConnectionError,
|
||
requests.exceptions.Timeout) as e:
|
||
last_err = e
|
||
if attempt < retries:
|
||
time.sleep(attempt * 2)
|
||
continue
|
||
raise
|
||
raise last_err
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="使用 Agnes AI 理解/分析图片。")
|
||
parser.add_argument("--prompt", required=True, help="对图片的问题或指令（必选）")
|
||
parser.add_argument("--image", action="append", default=None,
|
||
help="图片公网 URL 或 data:image Base64（可多次指定）")
|
||
parser.add_argument("--image-file", action="append", default=None,
|
||
help="本地图片文件路径，自动转 Base64（可多次指定）")
|
||
parser.add_argument("--system", help="可选的系统提示（system 消息）")
|
||
parser.add_argument("--model", default=DEFAULT_MODEL, help=f"模型 ID，默认 {DEFAULT_MODEL}")
|
||
parser.add_argument("--temperature", type=float, help="采样温度，0~1，越低越确定")
|
||
parser.add_argument("--max-tokens", type=int, help="最多生成的 token 数")
|
||
parser.add_argument("--api-key", help="Agnes API Key（也可用 AGNES_API_KEY 环境变量）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
api_key = args.api_key or os.environ.get("AGNES_API_KEY")
|
||
if not api_key:
|
||
print("ERROR: 缺少 API Key，请用 --api-key 或设置 AGNES_API_KEY 环境变量。")
|
||
sys.exit(1)
|
||
|
||
# 收集所有图片输入
|
||
image_urls = list(args.image or [])
|
||
for p in (args.image_file or []):
|
||
if not os.path.isfile(p):
|
||
print(f"ERROR: 本地图片不存在: {p}")
|
||
sys.exit(1)
|
||
image_urls.append(file_to_data_uri(p))
|
||
|
||
# 组装多模态 content：文本 + 若干图片
|
||
content = [{"type": "text", "text": args.prompt}]
|
||
for url in image_urls:
|
||
content.append({"type": "image_url", "image_url": {"url": url}})
|
||
|
||
messages = []
|
||
if args.system:
|
||
messages.append({"role": "system", "content": args.system})
|
||
# 没有图片时退化为纯文本，content 用字符串更稳妥
|
||
messages.append({"role": "user", "content": content if image_urls else args.prompt})
|
||
|
||
try:
|
||
result = call_api(
|
||
messages=messages,
|
||
model=args.model,
|
||
api_key=api_key,
|
||
temperature=args.temperature,
|
||
max_tokens=args.max_tokens,
|
||
)
|
||
except requests.exceptions.RequestException as e:
|
||
print(f"ERROR: API 请求失败: {e}")
|
||
if getattr(e, "response", None) is not None:
|
||
print(f"Response body: {e.response.text[:500]}")
|
||
sys.exit(1)
|
||
|
||
try:
|
||
answer = result["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError):
|
||
print(f"ERROR: 无法解析响应。完整响应: {json.dumps(result, ensure_ascii=False)[:500]}")
|
||
sys.exit(1)
|
||
|
||
print(answer)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|