172 lines
5.4 KiB
Python
172 lines
5.4 KiB
Python
"""
|
||
测试脚本：上传 PDF / DOCX / TXT 文件到阿里云 OSS → 获取 URL → 发送给 GLM-4.6V 识别
|
||
|
||
支持的文件类型:
|
||
- .pdf → 上传 OSS 后以 file_url 类型发送 URL 给 GLM
|
||
- .docx → 上传 OSS 后以 file_url 类型发送 URL 给 GLM
|
||
- .txt → 上传 OSS 后以 file_url 类型发送 URL 给 GLM
|
||
|
||
用法:
|
||
cd server
|
||
source ~/.bashrc && source .venv/bin/activate
|
||
python -m utils.test_oss_doc_glm --file <本地文件路径> [--prompt "请总结这份文件"]
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 确保 server 目录在 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||
|
||
from utils.oss_uploader import upload_file
|
||
from utils.glm_adapter import glm_chat_sync
|
||
|
||
|
||
# 文件类型分组
|
||
DOC_EXTS = {".pdf", ".doc", ".docx", ".xlsx", ".xls", ".pptx", ".ppt"}
|
||
TXT_EXTS = {".txt", ".md", ".csv", ".json", ".log"}
|
||
IMG_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"}
|
||
# 所有可通过 file_url 发送的类型
|
||
FILE_URL_EXTS = DOC_EXTS | TXT_EXTS
|
||
|
||
|
||
def detect_file_type(suffix: str) -> str:
|
||
"""根据后缀判断文件类别: 'file_url' / 'image' / 'unknown'"""
|
||
suffix = suffix.lower()
|
||
if suffix in FILE_URL_EXTS:
|
||
return "file_url"
|
||
elif suffix in IMG_EXTS:
|
||
return "image"
|
||
return "unknown"
|
||
|
||
|
||
def build_messages_for_file_url(file_url: str, prompt: str) -> list:
|
||
"""
|
||
为文档/文本文件构建消息。
|
||
使用 file_url 类型，直接传递 OSS 的 URL 给 GLM。
|
||
"""
|
||
return [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "file_url",
|
||
"file_url": {"url": file_url},
|
||
},
|
||
{
|
||
"type": "text",
|
||
"text": prompt,
|
||
},
|
||
],
|
||
}
|
||
]
|
||
|
||
|
||
def build_messages_for_image(file_url: str, prompt: str) -> list:
|
||
"""为图片文件构建消息，使用 image_url 类型。"""
|
||
return [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image_url", "image_url": {"url": file_url}},
|
||
{"type": "text", "text": prompt},
|
||
],
|
||
}
|
||
]
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="上传 PDF/DOCX/TXT 文件到 OSS 并让 GLM-4.6V 识别"
|
||
)
|
||
parser.add_argument("--file", required=True, help="要上传的本地文件路径")
|
||
parser.add_argument(
|
||
"--prompt", default="请总结这份文件的主要内容", help="发给 GLM 的提示词"
|
||
)
|
||
parser.add_argument(
|
||
"--model", default="glm-4.6v", help="GLM 模型名称（默认: glm-4.6v）"
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
file_path = Path(args.file).resolve()
|
||
if not file_path.exists():
|
||
print(f"❌ 文件不存在: {file_path}")
|
||
sys.exit(1)
|
||
|
||
suffix = file_path.suffix.lower()
|
||
file_type = detect_file_type(suffix)
|
||
print(f"📂 文件信息:")
|
||
print(f" 路径: {file_path}")
|
||
print(f" 后缀: {suffix}")
|
||
print(f" 类型: {file_type}")
|
||
print(f" 大小: {file_path.stat().st_size / 1024:.1f} KB")
|
||
print()
|
||
|
||
# ── 第一步：上传文件到 OSS ────────────────────────────────
|
||
print(f"📤 正在上传文件到阿里云 OSS...")
|
||
oss_result = upload_file(str(file_path))
|
||
file_url = oss_result["url"]
|
||
print(f"✅ OSS 上传成功!")
|
||
print(f" URL: {file_url}")
|
||
print(f" ETag: {oss_result['etag']}")
|
||
print()
|
||
|
||
# ── 第二步：根据文件类型构建消息 ──────────────────────────
|
||
print(f"🔧 正在构建 GLM 消息...")
|
||
|
||
if file_type == "file_url":
|
||
# PDF / DOCX / TXT 等：使用 file_url 类型发送 OSS URL
|
||
print(f" 策略: 使用 file_url 发送 OSS 链接")
|
||
messages = build_messages_for_file_url(file_url, args.prompt)
|
||
elif file_type == "image":
|
||
# 图片：使用 image_url
|
||
print(f" 策略: 使用 image_url 发送 OSS 链接")
|
||
messages = build_messages_for_image(file_url, args.prompt)
|
||
else:
|
||
print(f"❌ 不支持的文件类型: {suffix}")
|
||
print(f" 支持: {', '.join(sorted(FILE_URL_EXTS | IMG_EXTS))}")
|
||
sys.exit(1)
|
||
|
||
print()
|
||
|
||
# ── 第三步：发送给 GLM 识别 ──────────────────────────────
|
||
print(f"🤖 正在请求 GLM ({args.model}) 识别文件...")
|
||
print(f" 提示词: {args.prompt}")
|
||
print()
|
||
|
||
try:
|
||
result = glm_chat_sync(
|
||
messages=messages,
|
||
model=args.model,
|
||
temperature=0.7,
|
||
max_tokens=4096,
|
||
)
|
||
|
||
print("─" * 60)
|
||
print("📝 GLM 回复:")
|
||
print("─" * 60)
|
||
print(result["content"])
|
||
print("─" * 60)
|
||
|
||
if result.get("usage"):
|
||
usage = result["usage"]
|
||
print(
|
||
f"\n📊 Token 用量: 输入 {usage['promptTokens']} | "
|
||
f"输出 {usage['completionTokens']} | "
|
||
f"总计 {usage['totalTokens']}"
|
||
)
|
||
|
||
print(f"\n✅ 测试完成! 使用模型: {result.get('model', args.model)}")
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ GLM 请求失败: {e}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|