117 lines
4.0 KiB
Python
117 lines
4.0 KiB
Python
|
import asyncio
|
|||
|
|
from google.cloud import storage
|
|||
|
|
from google.oauth2 import service_account
|
|||
|
|
from vertexai.generative_models import Part
|
|||
|
|
import os
|
|||
|
|
import mimetypes
|
|||
|
|
from typing import List
|
|||
|
|
from google.cloud import aiplatform
|
|||
|
|
|
|||
|
|
# 替换为您的 GCS Bucket 名称和目标文件夹
|
|||
|
|
GCS_BUCKET_NAME = "aida-test-vertex-ai-bucket"
|
|||
|
|
GCS_BLOB_FOLDER = "user_456/"
|
|||
|
|
KEY_FILE_PATH = "/workspace/lc_stylist_agent/request.json"
|
|||
|
|
# 1. 从 JSON 文件加载凭证
|
|||
|
|
aiplatform.init(
|
|||
|
|
project='aida-461108',
|
|||
|
|
location='us-central1'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MyVertexAIClient:
|
|||
|
|
def __init__(self):
|
|||
|
|
try:
|
|||
|
|
self.credentials = service_account.Credentials.from_service_account_file(KEY_FILE_PATH)
|
|||
|
|
except Exception as e:
|
|||
|
|
# 这里的异常处理应根据实际情况调整
|
|||
|
|
raise RuntimeError(f"Failed to load credentials from file {KEY_FILE_PATH}: {e}")
|
|||
|
|
# 初始化 GCS 客户端 (可以在 setup 或 __init__ 中完成)
|
|||
|
|
self.gcs_client = storage.Client(
|
|||
|
|
project=self.credentials.project_id,
|
|||
|
|
credentials=self.credentials
|
|||
|
|
)
|
|||
|
|
self.uploaded_gcs_uris: List[str] = [] # 用于记录和后续清理
|
|||
|
|
|
|||
|
|
# --- 辅助方法：安全地执行同步 GCS 上传 ---
|
|||
|
|
def _upload_to_gcs_sync(self, bucket_name: str, blob_name: str, local_file_path: str) -> str:
|
|||
|
|
"""同步方法：将文件上传到 GCS 并返回 GCS URI。"""
|
|||
|
|
bucket = self.gcs_client.bucket(bucket_name)
|
|||
|
|
blob = bucket.blob(blob_name)
|
|||
|
|
blob.upload_from_filename(local_file_path)
|
|||
|
|
|
|||
|
|
gcs_uri = f"gs://{bucket_name}/{blob_name}"
|
|||
|
|
return gcs_uri
|
|||
|
|
|
|||
|
|
# --- 目标方法：改写 gemini_client.aio.files.upload ---
|
|||
|
|
async def upload_file_to_vertex_part(self, merged_image_path: str) -> Part:
|
|||
|
|
"""
|
|||
|
|
[改写目标] 异步地将本地文件上传到 GCS，并返回 Vertex AI Part 对象。
|
|||
|
|
|
|||
|
|
:param merged_image_path: 本地文件的路径。
|
|||
|
|
:return: 包含 GCS URI 的 Vertex AI Part 对象。
|
|||
|
|
"""
|
|||
|
|
# 1. 确定 GCS 路径
|
|||
|
|
file_name = os.path.basename(merged_image_path)
|
|||
|
|
# 使用唯一标识符以防冲突，这里简化为文件名
|
|||
|
|
blob_name = GCS_BLOB_FOLDER + file_name
|
|||
|
|
|
|||
|
|
# 2. 异步执行 GCS 上传 (将同步 I/O 放入线程池)
|
|||
|
|
loop = asyncio.get_event_loop()
|
|||
|
|
gcs_uri = await loop.run_in_executor(
|
|||
|
|
None, # 使用默认线程池
|
|||
|
|
self._upload_to_gcs_sync,
|
|||
|
|
GCS_BUCKET_NAME,
|
|||
|
|
blob_name,
|
|||
|
|
merged_image_path
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 3. 记录 GCS URI 以供后续清理
|
|||
|
|
self.uploaded_gcs_uris.append(gcs_uri)
|
|||
|
|
|
|||
|
|
# 4. 确定 MIME 类型
|
|||
|
|
mime_type, _ = mimetypes.guess_type(merged_image_path)
|
|||
|
|
if not mime_type:
|
|||
|
|
mime_type = 'application/octet-stream' # 提供默认值
|
|||
|
|
|
|||
|
|
# 5. 创建 Vertex AI Part 对象
|
|||
|
|
# 这是 Vertex AI SDK 传递文件的标准方式
|
|||
|
|
vertex_part = Part.from_uri(
|
|||
|
|
uri=gcs_uri,
|
|||
|
|
mime_type=mime_type
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return vertex_part
|
|||
|
|
|
|||
|
|
def _list_clear_object(self, bucket_name: str, prefix: str, delimiter=None):
|
|||
|
|
blobs = self.gcs_client.list_blobs(
|
|||
|
|
bucket_name, prefix=prefix, delimiter=delimiter
|
|||
|
|
)
|
|||
|
|
for blob in blobs:
|
|||
|
|
blob.delete()
|
|||
|
|
print(f" ✅ 已删除: {blob.name}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# --- 原始代码的调用改写 ---
|
|||
|
|
|
|||
|
|
# 假设 self.vertex_client 是 MyVertexAIClient 的实例
|
|||
|
|
# original: myfile = await self.gemini_client.aio.files.upload(file=merged_image_path)
|
|||
|
|
# original: content_parts.append(myfile)
|
|||
|
|
|
|||
|
|
# 改写后:
|
|||
|
|
merged_image_path = "/workspace/lc_stylist_agent/app/core/data/outfit_output/35e8626c-943f-4a3f-a0cc-1280d8bcf84d.jpg"
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def run():
|
|||
|
|
vertex_ai_server = MyVertexAIClient()
|
|||
|
|
# file_name = os.path.basename(merged_image_path)
|
|||
|
|
# blob_name = GCS_BLOB_FOLDER + file_name
|
|||
|
|
|
|||
|
|
# vertex_part = vertex_ai_server._upload_to_gcs_sync(GCS_BUCKET_NAME, blob_name, merged_image_path)
|
|||
|
|
# print(vertex_part)
|
|||
|
|
|
|||
|
|
vertex_ai_server._list_clear_object(GCS_BUCKET_NAME, "user_456")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# content_parts.append(vertex_part)
|
|||
|
|
asyncio.run(run())
|