规范化vertex ai的调用方式
* GoogleAuthUtils.GoogleGenaiClient.generate_content() 调用增加了timeout参数 默认为300秒 * 规范化VertexAI调用方式 --------- Merge request URL: https://g-ldyi2063.coding.net/p/dev/d/modalDeploy/git/merge/4827?initial=true Co-authored-by: shuohigh@gmail.com
This commit is contained in:
4
.gitignore
vendored
4
.gitignore
vendored
@@ -3,4 +3,6 @@
|
|||||||
.runtime.env
|
.runtime.env
|
||||||
.idea
|
.idea
|
||||||
.venv
|
.venv
|
||||||
test*
|
tests/*
|
||||||
|
!tests/**/*.py
|
||||||
|
tests/generate_appendix.py
|
||||||
@@ -12,7 +12,7 @@ dependencies = [
|
|||||||
"backoff>=2.2.1",
|
"backoff>=2.2.1",
|
||||||
"crcmod>=1.7",
|
"crcmod>=1.7",
|
||||||
"loguru>=0.7.3",
|
"loguru>=0.7.3",
|
||||||
"httpx==0.27.0",
|
"httpx>=0.28.1",
|
||||||
"pydantic>=2.11.3",
|
"pydantic>=2.11.3",
|
||||||
"pydantic-settings>=2.9.1",
|
"pydantic-settings>=2.9.1",
|
||||||
"webvtt-py>=0.5.1",
|
"webvtt-py>=0.5.1",
|
||||||
@@ -34,7 +34,7 @@ dependencies = [
|
|||||||
"m3u8>=6.0.0",
|
"m3u8>=6.0.0",
|
||||||
"aiofiles==23.2.1",
|
"aiofiles==23.2.1",
|
||||||
"annotated-types==0.6.0",
|
"annotated-types==0.6.0",
|
||||||
"anyio==4.3.0",
|
"anyio>=4.8.0",
|
||||||
"browser-cookie3==0.19.1",
|
"browser-cookie3==0.19.1",
|
||||||
"certifi==2024.2.2",
|
"certifi==2024.2.2",
|
||||||
"click==8.1.7",
|
"click==8.1.7",
|
||||||
@@ -59,14 +59,15 @@ dependencies = [
|
|||||||
"tornado==6.4",
|
"tornado==6.4",
|
||||||
"ua-parser==0.18.0",
|
"ua-parser==0.18.0",
|
||||||
"user-agents==2.2.0",
|
"user-agents==2.2.0",
|
||||||
"websockets==12.0",
|
"websockets>=12.0",
|
||||||
"gmssl==3.2.2",
|
"gmssl==3.2.2",
|
||||||
"tenacity~=9.0.0",
|
"tenacity",
|
||||||
"retry>=0.9.2",
|
"retry>=0.9.2",
|
||||||
"ffmpy>=0.5.0",
|
"ffmpy>=0.5.0",
|
||||||
"watchdog>=6.0.0",
|
"watchdog>=6.0.0",
|
||||||
"pyfiglet>=1.0.3",
|
"pyfiglet>=1.0.3",
|
||||||
"google-api-python-client>=2.172.0",
|
"google-api-python-client>=2.172.0",
|
||||||
|
"google-genai==1.21.0",
|
||||||
]
|
]
|
||||||
classifiers = [
|
classifiers = [
|
||||||
"Programming Language :: Python :: 3",
|
"Programming Language :: Python :: 3",
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
from typing import Union, Dict, Any, BinaryIO
|
import random
|
||||||
from urllib.parse import urlencode, quote
|
|
||||||
|
from typing import Union, Dict, BinaryIO, List, Optional
|
||||||
|
from urllib.parse import quote
|
||||||
|
|
||||||
import backoff
|
import backoff
|
||||||
import httpx
|
import httpx
|
||||||
@@ -7,9 +9,7 @@ import asyncio
|
|||||||
from loguru import logger
|
from loguru import logger
|
||||||
import aiofiles
|
import aiofiles
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from pydantic import BaseModel, computed_field, Field, field_validator
|
||||||
from pydantic import BaseModel
|
|
||||||
|
|
||||||
from BowongModalFunctions.models.web_model import GoogleUploadResponse
|
from BowongModalFunctions.models.web_model import GoogleUploadResponse
|
||||||
|
|
||||||
|
|
||||||
@@ -106,6 +106,7 @@ class HTTPDownloadUtils:
|
|||||||
|
|
||||||
|
|
||||||
from google.oauth2 import service_account
|
from google.oauth2 import service_account
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
|
||||||
class GoogleAuthUtils:
|
class GoogleAuthUtils:
|
||||||
@@ -114,6 +115,58 @@ class GoogleAuthUtils:
|
|||||||
expires_in: int
|
expires_in: int
|
||||||
token_type: str
|
token_type: str
|
||||||
|
|
||||||
|
class VertexAIRequestModel(types._common.BaseModel):
|
||||||
|
contents: Union[types.ContentListUnion, types.ContentListUnionDict] = Field()
|
||||||
|
generation_config: Optional[types.GenerateContentConfig] = Field(default=None, alias="generationConfig",
|
||||||
|
serialization_alias="generationConfig")
|
||||||
|
|
||||||
|
@field_validator("generation_config", mode="before")
|
||||||
|
def validate_generation_config(cls, v):
|
||||||
|
if not v:
|
||||||
|
return None
|
||||||
|
elif isinstance(v, types.GenerateContentConfig):
|
||||||
|
if issubclass(v.response_schema, BaseModel):
|
||||||
|
v.response_schema = v.response_schema.model_json_schema()
|
||||||
|
return v
|
||||||
|
elif isinstance(v.response_schema, Dict):
|
||||||
|
return v
|
||||||
|
raise TypeError("generation_config must be of type 'GenerateContentConfig'")
|
||||||
|
|
||||||
|
@computed_field(alias="safetySettings")
|
||||||
|
@property
|
||||||
|
def safety_settings(self) -> Optional[List[types.SafetySetting]]:
|
||||||
|
return self.generation_config.safety_settings
|
||||||
|
|
||||||
|
class GoogleGenaiClient(BaseModel):
|
||||||
|
cloudflare_project_id: str
|
||||||
|
cloudflare_gateway_id: str
|
||||||
|
google_project_id: str
|
||||||
|
regions: List[str]
|
||||||
|
access_token: str
|
||||||
|
|
||||||
|
@computed_field(description="通过Cloudflare gateway调用的url")
|
||||||
|
@property
|
||||||
|
def gateway_url(self) -> str:
|
||||||
|
return f"https://gateway.ai.cloudflare.com/v1/{self.cloudflare_project_id}/{self.cloudflare_gateway_id}/google-vertex-ai/v1/projects/{self.google_project_id}/locations/{random.choice(self.regions)}/publishers/google/models"
|
||||||
|
|
||||||
|
def generate_content(self, model_id: str, contents: Union[types.ContentListUnion, types.ContentListUnionDict],
|
||||||
|
config: Optional[types.GenerateContentConfig] = None,
|
||||||
|
timeout: int = 30) -> types.GenerateContentResponse:
|
||||||
|
parameter_model = GoogleAuthUtils.VertexAIRequestModel(contents=contents, generationConfig=config)
|
||||||
|
json_body = parameter_model.model_dump_json(indent=2, exclude_none=True, by_alias=True,
|
||||||
|
exclude={"generation_config": {"safety_settings"}})
|
||||||
|
logger.info(json_body)
|
||||||
|
url = f"{self.gateway_url}/{model_id}:generateContent"
|
||||||
|
logger.info(f"Authorization : Bearer {self.access_token}")
|
||||||
|
logger.info(f"POST {url}")
|
||||||
|
response = httpx.post(url=url, timeout=timeout,
|
||||||
|
headers={"Content-Type": "application/json",
|
||||||
|
"Authorization": f"Bearer {self.access_token}"},
|
||||||
|
content=json_body)
|
||||||
|
response.raise_for_status()
|
||||||
|
logger.info(response.text)
|
||||||
|
return types.GenerateContentResponse.model_validate_json(response.text)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
async def get_google_auth_jwt(service_account_info: Dict[str, str], scopes: list[str]) -> GoogleAuthResponse:
|
async def get_google_auth_jwt(service_account_info: Dict[str, str], scopes: list[str]) -> GoogleAuthResponse:
|
||||||
credentials: service_account.Credentials = service_account.Credentials.from_service_account_info(
|
credentials: service_account.Credentials = service_account.Credentials.from_service_account_info(
|
||||||
|
|||||||
@@ -39,9 +39,13 @@ class TokenManager:
|
|||||||
ttwid_conf = tiktok_manager.get("ttwid", None)
|
ttwid_conf = tiktok_manager.get("ttwid", None)
|
||||||
odin_tt_conf = tiktok_manager.get("odin_tt", None)
|
odin_tt_conf = tiktok_manager.get("odin_tt", None)
|
||||||
proxies_conf = tiktok_manager.get("proxies", None)
|
proxies_conf = tiktok_manager.get("proxies", None)
|
||||||
|
async_proxies = {
|
||||||
|
"http://": httpx.AsyncHTTPTransport(proxy=proxies_conf.get("http", None)),
|
||||||
|
"https://": httpx.AsyncHTTPTransport(proxy=proxies_conf.get("https", None)),
|
||||||
|
}
|
||||||
proxies = {
|
proxies = {
|
||||||
"http://": proxies_conf.get("http", None),
|
"http://": httpx.HTTPTransport(proxy=proxies_conf.get("http", None)),
|
||||||
"https://": proxies_conf.get("https", None),
|
"https://": httpx.HTTPTransport(proxy=proxies_conf.get("https", None)),
|
||||||
}
|
}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -67,7 +71,8 @@ class TokenManager:
|
|||||||
}
|
}
|
||||||
|
|
||||||
transport = httpx.HTTPTransport(retries=5)
|
transport = httpx.HTTPTransport(retries=5)
|
||||||
with httpx.Client(transport=transport, proxies=cls.proxies) as client:
|
|
||||||
|
with httpx.Client(transport=transport, mounts=cls.proxies) as client:
|
||||||
try:
|
try:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
cls.token_conf["url"], headers=headers, content=payload
|
cls.token_conf["url"], headers=headers, content=payload
|
||||||
@@ -118,7 +123,7 @@ class TokenManager:
|
|||||||
生成请求必带的ttwid (Generate the essential ttwid for requests)
|
生成请求必带的ttwid (Generate the essential ttwid for requests)
|
||||||
"""
|
"""
|
||||||
transport = httpx.HTTPTransport(retries=5)
|
transport = httpx.HTTPTransport(retries=5)
|
||||||
with httpx.Client(transport=transport, proxies=cls.proxies) as client:
|
with httpx.Client(transport=transport, mounts=cls.proxies) as client:
|
||||||
try:
|
try:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
cls.ttwid_conf["url"],
|
cls.ttwid_conf["url"],
|
||||||
@@ -166,7 +171,7 @@ class TokenManager:
|
|||||||
生成请求必带的odin_tt (Generate the essential odin_tt for requests)
|
生成请求必带的odin_tt (Generate the essential odin_tt for requests)
|
||||||
"""
|
"""
|
||||||
transport = httpx.HTTPTransport(retries=5)
|
transport = httpx.HTTPTransport(retries=5)
|
||||||
with httpx.Client(transport=transport, proxies=cls.proxies) as client:
|
with httpx.Client(transport=transport, mounts=cls.proxies) as client:
|
||||||
try:
|
try:
|
||||||
response = client.get(cls.odin_tt_conf["url"])
|
response = client.get(cls.odin_tt_conf["url"])
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
@@ -272,7 +277,7 @@ class SecUserIdFetcher:
|
|||||||
|
|
||||||
transport = httpx.AsyncHTTPTransport(retries=5)
|
transport = httpx.AsyncHTTPTransport(retries=5)
|
||||||
async with httpx.AsyncClient(
|
async with httpx.AsyncClient(
|
||||||
transport=transport, proxies=TokenManager.proxies, timeout=10
|
transport=transport, mounts=TokenManager.async_proxies, timeout=10
|
||||||
) as client:
|
) as client:
|
||||||
try:
|
try:
|
||||||
response = await client.get(url, follow_redirects=True)
|
response = await client.get(url, follow_redirects=True)
|
||||||
@@ -363,7 +368,7 @@ class SecUserIdFetcher:
|
|||||||
|
|
||||||
transport = httpx.AsyncHTTPTransport(retries=5)
|
transport = httpx.AsyncHTTPTransport(retries=5)
|
||||||
async with httpx.AsyncClient(
|
async with httpx.AsyncClient(
|
||||||
transport=transport, proxies=TokenManager.proxies, timeout=10
|
transport=transport, mounts=TokenManager.async_proxies, timeout=10
|
||||||
) as client:
|
) as client:
|
||||||
try:
|
try:
|
||||||
response = await client.get(url, follow_redirects=True)
|
response = await client.get(url, follow_redirects=True)
|
||||||
@@ -477,7 +482,7 @@ class AwemeIdFetcher:
|
|||||||
print(f"输入的URL需要重定向: {url}")
|
print(f"输入的URL需要重定向: {url}")
|
||||||
transport = httpx.AsyncHTTPTransport(retries=10)
|
transport = httpx.AsyncHTTPTransport(retries=10)
|
||||||
async with httpx.AsyncClient(
|
async with httpx.AsyncClient(
|
||||||
transport=transport, proxies=TokenManager.proxies, timeout=10
|
transport=transport, mounts=TokenManager.async_proxies, timeout=10
|
||||||
) as client:
|
) as client:
|
||||||
try:
|
try:
|
||||||
response = await client.get(url, follow_redirects=True)
|
response = await client.get(url, follow_redirects=True)
|
||||||
|
|||||||
176
tests/CDN_speed_test_case.py
Normal file
176
tests/CDN_speed_test_case.py
Normal file
@@ -0,0 +1,176 @@
|
|||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
import statistics
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
import m3u8
|
||||||
|
import time
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import List, Dict
|
||||||
|
from loguru import logger
|
||||||
|
from BowongModalFunctions.utils.HTTPUtils import HTTPDownloadUtils
|
||||||
|
import modal
|
||||||
|
|
||||||
|
class CloudFrontCDNTestCase(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_modal_restore(self):
|
||||||
|
fn_id = ""
|
||||||
|
fn = modal.Function.from_name("bowong-ai-video","ffmpeg_stream_record_restore", environment_name="test")
|
||||||
|
restore_id = fn.spawn(fn_id=fn_id).object_id
|
||||||
|
logger.info(f"restore id: {restore_id}")
|
||||||
|
|
||||||
|
async def test_hls_speed(self):
|
||||||
|
# 生成唯一ID和临时目录
|
||||||
|
id = uuid.uuid4()
|
||||||
|
temp_dir = f"./videos/{id}"
|
||||||
|
os.makedirs(temp_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# HLS URL
|
||||||
|
# hls_url = "https://d2nj71io21vkj2.cloudfront.net/test/records/hls/fc-01JXBVTMMYSFNZ4SN1D225NMYV/playlist.m3u8"
|
||||||
|
hls_url = "https://cdn.roasmax.cn/test/records/hls/fc-01JWZS8954RZP4B13H3TA6TKMM/playlist.m3u8"
|
||||||
|
|
||||||
|
# 加载播放列表
|
||||||
|
playlist = m3u8.load(hls_url)
|
||||||
|
total_sample_size = len(playlist.segments)
|
||||||
|
used_sample_size = min(100, total_sample_size)
|
||||||
|
logger.info(f"total samples = {total_sample_size}, used sample = {used_sample_size}")
|
||||||
|
# 准备下载任务
|
||||||
|
urls = [segment.absolute_uri for segment in playlist.segments]
|
||||||
|
output_paths = [f"{temp_dir}/{segment.uri}" for segment in playlist.segments]
|
||||||
|
|
||||||
|
# 记录下载开始时间
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
# 执行下载并收集结果
|
||||||
|
results = await HTTPDownloadUtils.batch_download(
|
||||||
|
urls=urls[:used_sample_size],
|
||||||
|
output_paths=output_paths[:used_sample_size]
|
||||||
|
)
|
||||||
|
|
||||||
|
# 计算总下载时间
|
||||||
|
total_time = time.time() - start_time
|
||||||
|
|
||||||
|
# 生成报告
|
||||||
|
report = self._generate_speed_report(
|
||||||
|
results=results,
|
||||||
|
urls=urls,
|
||||||
|
total_time=total_time,
|
||||||
|
playlist=playlist
|
||||||
|
)
|
||||||
|
|
||||||
|
# 打印报告
|
||||||
|
self._print_report(report)
|
||||||
|
|
||||||
|
# 保存报告到文件
|
||||||
|
self._save_report(report, temp_dir)
|
||||||
|
|
||||||
|
# 清理临时文件
|
||||||
|
self._cleanup(temp_dir)
|
||||||
|
|
||||||
|
return report
|
||||||
|
|
||||||
|
def _generate_speed_report(
|
||||||
|
self,
|
||||||
|
results: List[str],
|
||||||
|
urls: List[str],
|
||||||
|
total_time: float,
|
||||||
|
playlist: m3u8.M3U8
|
||||||
|
) -> Dict:
|
||||||
|
"""生成速度测试报告"""
|
||||||
|
|
||||||
|
# 计算每个分片的大小和下载时间
|
||||||
|
segment_stats = []
|
||||||
|
total_size = 0
|
||||||
|
|
||||||
|
for i, result in enumerate(results):
|
||||||
|
if isinstance(result, Exception):
|
||||||
|
logger.error(f"Failed to download segment {i}: {str(result)}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
segment_size = os.path.getsize(result)
|
||||||
|
total_size += segment_size
|
||||||
|
|
||||||
|
segment_stats.append({
|
||||||
|
'index': i,
|
||||||
|
'url': urls[i],
|
||||||
|
'size': segment_size,
|
||||||
|
'duration': playlist.segments[i].duration,
|
||||||
|
'bandwidth': segment_size / playlist.segments[i].duration if playlist.segments[i].duration else 0
|
||||||
|
})
|
||||||
|
|
||||||
|
# 计算统计指标
|
||||||
|
sizes = [stat['size'] for stat in segment_stats]
|
||||||
|
bandwidths = [stat['bandwidth'] for stat in segment_stats]
|
||||||
|
|
||||||
|
report = {
|
||||||
|
'timestamp': datetime.now().isoformat(),
|
||||||
|
'total_segments': len(playlist.segments),
|
||||||
|
'successful_downloads': len(segment_stats),
|
||||||
|
'total_size': total_size,
|
||||||
|
'total_time': total_time,
|
||||||
|
'average_speed': total_size / total_time if total_time > 0 else 0,
|
||||||
|
'statistics': {
|
||||||
|
'size': {
|
||||||
|
'mean': statistics.mean(sizes),
|
||||||
|
'median': statistics.median(sizes),
|
||||||
|
'stdev': statistics.stdev(sizes) if len(sizes) > 1 else 0,
|
||||||
|
'min': min(sizes),
|
||||||
|
'max': max(sizes)
|
||||||
|
},
|
||||||
|
'bandwidth': {
|
||||||
|
'mean': statistics.mean(bandwidths),
|
||||||
|
'median': statistics.median(bandwidths),
|
||||||
|
'stdev': statistics.stdev(bandwidths) if len(bandwidths) > 1 else 0,
|
||||||
|
'min': min(bandwidths),
|
||||||
|
'max': max(bandwidths)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'segments': segment_stats
|
||||||
|
}
|
||||||
|
|
||||||
|
return report
|
||||||
|
|
||||||
|
def _print_report(self, report: Dict):
|
||||||
|
"""打印报告到控制台"""
|
||||||
|
logger.info("=== CloudFront CDN Speed Test Report ===")
|
||||||
|
logger.info(f"Timestamp: {report['timestamp']}")
|
||||||
|
logger.info(f"Total Segments: {report['total_segments']}")
|
||||||
|
logger.info(f"Successful Downloads: {report['successful_downloads']}")
|
||||||
|
logger.info(f"Total Size: {report['total_size'] / 1024 / 1024:.2f} MB")
|
||||||
|
logger.info(f"Total Time: {report['total_time']:.2f} seconds")
|
||||||
|
logger.info(f"Average Speed: {report['average_speed'] / 1024 / 1024:.2f} MB/s")
|
||||||
|
|
||||||
|
logger.info("Size Statistics:")
|
||||||
|
logger.info(f" Mean: {report['statistics']['size']['mean'] / 1024:.2f} KB")
|
||||||
|
logger.info(f" Median: {report['statistics']['size']['median'] / 1024:.2f} KB")
|
||||||
|
logger.info(f" StdDev: {report['statistics']['size']['stdev'] / 1024:.2f} KB")
|
||||||
|
logger.info(f" Min: {report['statistics']['size']['min'] / 1024:.2f} KB")
|
||||||
|
logger.info(f" Max: {report['statistics']['size']['max'] / 1024:.2f} KB")
|
||||||
|
|
||||||
|
logger.info("Bandwidth Statistics:")
|
||||||
|
logger.info(f" Mean: {report['statistics']['bandwidth']['mean'] / 1024:.2f} KB/s")
|
||||||
|
logger.info(f" Median: {report['statistics']['bandwidth']['median'] / 1024:.2f} KB/s")
|
||||||
|
logger.info(f" StdDev: {report['statistics']['bandwidth']['stdev'] / 1024:.2f} KB/s")
|
||||||
|
logger.info(f" Min: {report['statistics']['bandwidth']['min'] / 1024:.2f} KB/s")
|
||||||
|
logger.info(f" Max: {report['statistics']['bandwidth']['max'] / 1024:.2f} KB/s")
|
||||||
|
|
||||||
|
def _save_report(self, report: Dict, temp_dir: str):
|
||||||
|
"""保存报告到文件"""
|
||||||
|
import json
|
||||||
|
report_path = os.path.join(temp_dir, 'speed_test_report.json')
|
||||||
|
with open(report_path, 'w') as f:
|
||||||
|
json.dump(report, f, indent=2)
|
||||||
|
logger.info(f"\nReport saved to: {report_path}")
|
||||||
|
|
||||||
|
def _cleanup(self, temp_dir: str):
|
||||||
|
"""清理临时文件"""
|
||||||
|
import shutil
|
||||||
|
try:
|
||||||
|
shutil.rmtree(temp_dir)
|
||||||
|
logger.info(f"Cleaned up temporary directory: {temp_dir}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to cleanup temporary directory: {str(e)}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
261
tests/ffmpeg_test_case.py
Normal file
261
tests/ffmpeg_test_case.py
Normal file
@@ -0,0 +1,261 @@
|
|||||||
|
import json
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
from typing import List, Optional
|
||||||
|
import httpx
|
||||||
|
from loguru import logger
|
||||||
|
from tqdm import tqdm
|
||||||
|
from BowongModalFunctions.models.ffmpeg_worker_model import FFMpegSliceSegment
|
||||||
|
from BowongModalFunctions.utils.VideoUtils import VideoUtils, FFMPEGSliceOptions
|
||||||
|
|
||||||
|
|
||||||
|
class FFMPEGTestCase(unittest.IsolatedAsyncioTestCase):
|
||||||
|
temp_dir = f"./videos"
|
||||||
|
douyin_cookies = "ttwid=1%7CB1qls3GdnZhUov9o2NxOMxxYS2ff6OSvEWbv0ytbES4%7C1680522049%7C280d802d6d478e3e78d0c807f7c487e7ffec0ae4e5fdd6a0fe74c3c6af149511; my_rd=1; passport_csrf_token=3ab34460fa656183fccfb904b16ff742; passport_csrf_token_default=3ab34460fa656183fccfb904b16ff742; d_ticket=9f562383ac0547d0b561904513229d76c9c21; n_mh=hvnJEQ4Q5eiH74-84kTFUyv4VK8xtSrpRZG1AhCeFNI; store-region=cn-fj; store-region-src=uid; LOGIN_STATUS=1; __security_server_data_status=1; FORCE_LOGIN=%7B%22videoConsumedRemainSeconds%22%3A180%7D; pwa2=%223%7C0%7C3%7C0%22; download_guide=%223%2F20230729%2F0%22; volume_info=%7B%22isUserMute%22%3Afalse%2C%22isMute%22%3Afalse%2C%22volume%22%3A0.6%7D; strategyABtestKey=%221690824679.923%22; stream_recommend_feed_params=%22%7B%5C%22cookie_enabled%5C%22%3Atrue%2C%5C%22screen_width%5C%22%3A1536%2C%5C%22screen_height%5C%22%3A864%2C%5C%22browser_online%5C%22%3Atrue%2C%5C%22cpu_core_num%5C%22%3A8%2C%5C%22device_memory%5C%22%3A8%2C%5C%22downlink%5C%22%3A10%2C%5C%22effective_type%5C%22%3A%5C%224g%5C%22%2C%5C%22round_trip_time%5C%22%3A150%7D%22; VIDEO_FILTER_MEMO_SELECT=%7B%22expireTime%22%3A1691443863751%2C%22type%22%3Anull%7D; home_can_add_dy_2_desktop=%221%22; __live_version__=%221.1.1.2169%22; device_web_cpu_core=8; device_web_memory_size=8; xgplayer_user_id=346045893336; csrf_session_id=2e00356b5cd8544d17a0e66484946f28; odin_tt=724eb4dd23bc6ffaed9a1571ac4c757ef597768a70c75fef695b95845b7ffcd8b1524278c2ac31c2587996d058e03414595f0a4e856c53bd0d5e5f56dc6d82e24004dc77773e6b83ced6f80f1bb70627; __ac_nonce=064caded4009deafd8b89; __ac_signature=_02B4Z6wo00f01HLUuwwAAIDBh6tRkVLvBQBy9L-AAHiHf7; ttcid=2e9619ebbb8449eaa3d5a42d8ce88ec835; webcast_leading_last_show_time=1691016922379; webcast_leading_total_show_times=1; webcast_local_quality=sd; live_can_add_dy_2_desktop=%221%22; msToken=1JDHnVPw_9yTvzIrwb7cQj8dCMNOoesXbA_IooV8cezcOdpe4pzusZE7NB7tZn9TBXPr0ylxmv-KMs5rqbNUBHP4P7VBFUu0ZAht_BEylqrLpzgt3y5ne_38hXDOX8o=; msToken=jV_yeN1IQKUd9PlNtpL7k5vthGKcHo0dEh_QPUQhr8G3cuYv-Jbb4NnIxGDmhVOkZOCSihNpA2kvYtHiTW25XNNX_yrsv5FN8O6zm3qmCIXcEe0LywLn7oBO2gITEeg=; tt_scid=mYfqpfbDjqXrIGJuQ7q-DlQJfUSG51qG.KUdzztuGP83OjuVLXnQHjsz-BRHRJu4e986"
|
||||||
|
|
||||||
|
def setUp(self) -> None:
|
||||||
|
os.makedirs(self.temp_dir, exist_ok=True)
|
||||||
|
|
||||||
|
async def test_ffmpeg_concat(self):
|
||||||
|
videos = {
|
||||||
|
"urls": [
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/vod/ap-seoul/1500034234/1397757909715120723.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/vod/ap-seoul/1500034234/1397757909695745707.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP2H5B547XSNM2VC9ZGF4/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP2H5B547XSNM2VC9ZGF4/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP2H5B547XSNM2VC9ZGF4/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP0KJPKHBSEFFPT9KY4DY/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP1XYBQ1A34XJM566WPZ8/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP1C13J3CDQESJF99QT5K/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP1PXBQ0MQGM9YBNPZQ1E/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP4A7J7RPRFHA1KGWW3S6/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP2RD3MTRJ5D54KX6T9D3/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP3Y6JZMFK2WDP2A4KF3N/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP3PFTM1A2Y7PFQ96HQJJ/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP4JQVHFWBBN0JEJ5HBHX/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP54TK35VP8E659AP2201/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP1PXBQ0MQGM9YBNPZQ1E/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP6GM06MZEPB5Y12CA2ZX/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP3A7MSSSAKF4EJAG8H70/output_0.mp4",
|
||||||
|
"https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JWDYP4TPNGHW9W12JCTTA80V/output_0.mp4"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
outputs: List[str] = []
|
||||||
|
with httpx.Client() as client:
|
||||||
|
for i, video in enumerate(videos.get("urls")):
|
||||||
|
output_path = f"{self.temp_dir}/input_{i}.mp4"
|
||||||
|
with client.stream("GET", video) as response:
|
||||||
|
response.raise_for_status()
|
||||||
|
file_size = int(response.headers.get("content-length", 0))
|
||||||
|
progress_bar = tqdm(
|
||||||
|
total=file_size,
|
||||||
|
unit='iB',
|
||||||
|
unit_scale=True,
|
||||||
|
desc='Downloading'
|
||||||
|
)
|
||||||
|
# 以二进制写模式打开文件
|
||||||
|
with open(output_path, 'wb') as file:
|
||||||
|
# 分块下载,每次读取1MB
|
||||||
|
chunk_size = 1024 * 1024 # 1MB
|
||||||
|
downloaded_size = 0
|
||||||
|
for chunk in response.iter_bytes(chunk_size=chunk_size):
|
||||||
|
if chunk:
|
||||||
|
file.write(chunk)
|
||||||
|
downloaded_size += len(chunk)
|
||||||
|
|
||||||
|
progress_bar.update(len(chunk))
|
||||||
|
# 每下载100MB记录一次日志
|
||||||
|
if downloaded_size % (100 * 1024 * 1024) == 0:
|
||||||
|
logger.info(
|
||||||
|
f"Downloaded: {downloaded_size / (1024 * 1024 * 1024):.2f} GB")
|
||||||
|
progress_bar.close()
|
||||||
|
self.assertTrue(os.path.exists(output_path))
|
||||||
|
outputs.append(output_path)
|
||||||
|
final_output_path = f"{self.temp_dir}/output.mp4"
|
||||||
|
await VideoUtils.ffmpeg_concat_medias(media_paths=outputs,
|
||||||
|
output_path=final_output_path)
|
||||||
|
self.assertTrue(os.path.exists(final_output_path))
|
||||||
|
|
||||||
|
async def test_ffmpeg_extract_frame(self):
|
||||||
|
video_url = "https://d2nj71io21vkj2.cloudfront.net/test/slice/outputs/fc-01JW84BFYAH0F9BS9T9DHDRKKH/output_0.mp4"
|
||||||
|
with httpx.Client() as client:
|
||||||
|
output_path = f"{self.temp_dir}/input.mp4"
|
||||||
|
with client.stream("GET", video_url) as response:
|
||||||
|
response.raise_for_status()
|
||||||
|
file_size = int(response.headers.get("content-length", 0))
|
||||||
|
progress_bar = tqdm(
|
||||||
|
total=file_size,
|
||||||
|
unit='iB',
|
||||||
|
unit_scale=True,
|
||||||
|
desc='Downloading'
|
||||||
|
)
|
||||||
|
# 以二进制写模式打开文件
|
||||||
|
with open(output_path, 'wb') as file:
|
||||||
|
# 分块下载,每次读取1MB
|
||||||
|
chunk_size = 1024 * 1024 # 1MB
|
||||||
|
downloaded_size = 0
|
||||||
|
for chunk in response.iter_bytes(chunk_size=chunk_size):
|
||||||
|
if chunk:
|
||||||
|
file.write(chunk)
|
||||||
|
downloaded_size += len(chunk)
|
||||||
|
|
||||||
|
progress_bar.update(len(chunk))
|
||||||
|
# 每下载100MB记录一次日志
|
||||||
|
if downloaded_size % (100 * 1024 * 1024) == 0:
|
||||||
|
logger.info(
|
||||||
|
f"Downloaded: {downloaded_size / (1024 * 1024 * 1024):.2f} GB")
|
||||||
|
progress_bar.close()
|
||||||
|
self.assertTrue(os.path.exists(output_path))
|
||||||
|
image_path, metadata = await VideoUtils.ffmpeg_extract_frame_image(video_path=output_path, frame_index=1)
|
||||||
|
logger.info(f"Extracted {image_path} and metadata: {metadata}")
|
||||||
|
|
||||||
|
async def test_ffmpeg_slice_stream(self):
|
||||||
|
from pydantic import TypeAdapter
|
||||||
|
|
||||||
|
stream_url = "https://cdn.roasmax.cn/test/records/hls/fc-01JXY1AS1HDGS300EQ54ATAKHF/playlist.m3u8"
|
||||||
|
adapter = TypeAdapter(List[FFMpegSliceSegment])
|
||||||
|
segments = adapter.validate_json("""[
|
||||||
|
{
|
||||||
|
"start": 920.425,
|
||||||
|
"end": 921.688
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"start": 922.425,
|
||||||
|
"end": 923.688
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"start": 940.425,
|
||||||
|
"end": 941.688
|
||||||
|
}
|
||||||
|
]""")
|
||||||
|
for segment in segments:
|
||||||
|
logger.info(f"{segment.start.toFormatStr()} --> {segment.end.toFormatStr()}")
|
||||||
|
options = FFMPEGSliceOptions()
|
||||||
|
outputs = await VideoUtils.ffmpeg_slice_media(media_path=stream_url,
|
||||||
|
media_markers=segments,
|
||||||
|
options=options,
|
||||||
|
is_streams=True,
|
||||||
|
output_path=f"{self.temp_dir}/output_stream.mp4")
|
||||||
|
for (output, metadata) in outputs:
|
||||||
|
self.assertTrue(os.path.exists(output))
|
||||||
|
|
||||||
|
async def test_ffmpeg_streamlink(self):
|
||||||
|
webcast_id = "333555252930"
|
||||||
|
segment_duration = 5
|
||||||
|
title = None
|
||||||
|
live_url: Optional[str] = None
|
||||||
|
quality: Optional[str] = None
|
||||||
|
|
||||||
|
with httpx.Client() as client:
|
||||||
|
cookie_getter_response = client.get(url='https://live.douyin.com/')
|
||||||
|
cookie_getter_response.raise_for_status()
|
||||||
|
stream_url_response = client.get(url='https://live.douyin.com/webcast/room/web/enter/',
|
||||||
|
params={
|
||||||
|
'aid': 6383,
|
||||||
|
'device_platform': 'web',
|
||||||
|
'browser_language': 'zh-CN',
|
||||||
|
'browser_platform': 'Win32',
|
||||||
|
'browser_name': 'Chrome',
|
||||||
|
'browser_version': '100.0.0.0',
|
||||||
|
'web_rid': webcast_id
|
||||||
|
})
|
||||||
|
stream_url_response.raise_for_status()
|
||||||
|
stream_info_json = stream_url_response.json()
|
||||||
|
if data := stream_info_json['data']['data']:
|
||||||
|
data = data[0]
|
||||||
|
if data['status'] == 2:
|
||||||
|
title = data['title']
|
||||||
|
live_url = ''
|
||||||
|
stream_data = json.loads(data['stream_url']['live_core_sdk_data']['pull_data']['stream_data'])
|
||||||
|
for quality_code in ('origin', 'uhd', 'hd', 'sd', 'md', 'ld'):
|
||||||
|
if quality_data := stream_data['data'].get(quality_code):
|
||||||
|
logger.info(f"质量={quality_code}")
|
||||||
|
live_url = quality_data['main']['flv']
|
||||||
|
quality = quality_code
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
raise ValueError("直播间没有开播")
|
||||||
|
logger.info(f"{title} url = {live_url}")
|
||||||
|
if live_url:
|
||||||
|
output_dir = f"{self.temp_dir}/{title}"
|
||||||
|
os.makedirs(output_dir, exist_ok=True)
|
||||||
|
# output_file_pattern = "%10d.ts"
|
||||||
|
# output_pattern = f"{output_dir}/{output_file_pattern}"
|
||||||
|
logger.info(f"开始录制 [{title}|{quality}] {live_url} 到 {output_dir}")
|
||||||
|
# await VideoUtils.ffmpeg_stream_record_as_hls(stream_url=live_url,
|
||||||
|
# segment_duration=segment_duration,
|
||||||
|
# stream_content_timeout=stream_content_timeout,
|
||||||
|
# segments_output_dir=output_dir,
|
||||||
|
# playlist_output_dir=manifest_output_dir,
|
||||||
|
# playlist_output_method='PUT',
|
||||||
|
# playlist_output_headers={
|
||||||
|
# "Content-Type": "application/vnd.apple.mpegurl",
|
||||||
|
# "Authorization": f"Bearer bowong7777"
|
||||||
|
# },
|
||||||
|
# manifest_segment_prefix=manifest_segment_prefix)
|
||||||
|
logger.info(f'停止录制:{output_dir}')
|
||||||
|
|
||||||
|
async def test_ffprobe_read_video(self):
|
||||||
|
# video = "./videos/input_0.mp4"
|
||||||
|
video = "./videos/fc-01JXKK5VWK6B924TDB9WZETACW/playlist.m3u8"
|
||||||
|
video = "https://cdn.roasmax.cn/prod/records/hls/fc-01JXKK5VWK6B924TDB9WZETACW/playlist.m3u8"
|
||||||
|
video_metadata = VideoUtils.ffprobe_media_metadata(video)
|
||||||
|
logger.info(video_metadata)
|
||||||
|
self.assertIsNotNone(video_metadata)
|
||||||
|
|
||||||
|
async def test_ffprobe_read_audio(self):
|
||||||
|
audio = "./videos/output.wav"
|
||||||
|
audio_metadata = VideoUtils.ffprobe_media_metadata(audio)
|
||||||
|
self.assertIsNotNone(audio_metadata)
|
||||||
|
|
||||||
|
async def test_ffprobe_read_hls(self):
|
||||||
|
# hls = "./fc-01JWX1CD01JTV0DECWVWX1X9FM/playlist.m3u8"
|
||||||
|
hls = "./fc-01JWZS8954RZP4B13H3TA6TKMM/playlist.m3u8"
|
||||||
|
hls_metadata = VideoUtils.ffprobe_media_metadata(hls)
|
||||||
|
self.assertIsNotNone(hls_metadata)
|
||||||
|
|
||||||
|
async def test_ffprobe_read_image(self):
|
||||||
|
image = "./1928302183527354368.png"
|
||||||
|
image_metadata = VideoUtils.ffprobe_media_metadata(image)
|
||||||
|
self.assertIsNotNone(image_metadata)
|
||||||
|
|
||||||
|
async def test_ffprobe_read_gif(self):
|
||||||
|
gif = "./flow_frame.gif"
|
||||||
|
gif_metadata = VideoUtils.ffprobe_media_metadata(gif)
|
||||||
|
self.assertIsNotNone(gif_metadata)
|
||||||
|
|
||||||
|
async def test_ffmpeg_convert_stream(self):
|
||||||
|
stream_url = "https://cdn.roasmax.cn/prod/records/hls/fc-01JXKJ2JCTDKWJ6J97FN6ABQPM/playlist.m3u8"
|
||||||
|
output_path = f"{self.temp_dir}/fc-01JXKJ2JCTDKWJ6J97FN6ABQPM/output.mp4"
|
||||||
|
options = FFMPEGSliceOptions()
|
||||||
|
local_output, metadata = await VideoUtils.ffmpeg_convert_stream_media(media_stream_url=stream_url,
|
||||||
|
options=options,
|
||||||
|
output_path=output_path)
|
||||||
|
self.assertIsNotNone(local_output)
|
||||||
|
|
||||||
|
async def test_ffmpeg_best_convert(self):
|
||||||
|
media_stream_url = "https://cdn.roasmax.cn/test/records/hls/fc-01JY0ZRKCG6K7WSD9H4Q0X922K/playlist.m3u8"
|
||||||
|
local_m3u8_path, temp_dir, diff = await VideoUtils.convert_m3u8_to_local_source(
|
||||||
|
media_stream_url=media_stream_url, temp_dir=f"./videos/large")
|
||||||
|
|
||||||
|
ffmpeg_cmd = VideoUtils.async_ffmpeg_init()
|
||||||
|
ffmpeg_cmd.input(local_m3u8_path,
|
||||||
|
protocol_whitelist="file,http,https,tcp,tls")
|
||||||
|
h264_options = {
|
||||||
|
"reset_timestamps": "1",
|
||||||
|
"vcodec": "libx264",
|
||||||
|
"movflags": "+faststart", # 优化MP4文件结构
|
||||||
|
"pix_fmt": "yuv420p",
|
||||||
|
"acodec": "aac",
|
||||||
|
"aac_coder": "fast",
|
||||||
|
"ar": 44100,
|
||||||
|
"ac": 2,
|
||||||
|
"ab": "192k",
|
||||||
|
"crf": 18,
|
||||||
|
"r": 30
|
||||||
|
}
|
||||||
|
ffmpeg_cmd.output("./videos/output_large_h264.mp4", options=h264_options)
|
||||||
|
await ffmpeg_cmd.execute()
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
110
tests/google_test_case.py
Normal file
110
tests/google_test_case.py
Normal file
@@ -0,0 +1,110 @@
|
|||||||
|
import json
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
from loguru import logger
|
||||||
|
from google.genai import types
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from BowongModalFunctions.utils.HTTPUtils import GoogleAuthUtils
|
||||||
|
|
||||||
|
|
||||||
|
class BaseResponse(BaseModel):
|
||||||
|
message: str
|
||||||
|
|
||||||
|
|
||||||
|
class GoogleTestCase(unittest.IsolatedAsyncioTestCase):
|
||||||
|
service_account_info: dict = {
|
||||||
|
"type": "service_account",
|
||||||
|
"project_id": "gen-lang-client-0413414134",
|
||||||
|
"private_key_id": "48c91fc4cae8158edaad1f52577e4c98143a8cd9",
|
||||||
|
"private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCzNfkzjGOSAv+e\nHSWOEq87sE8cNdt0AXdAyRL66rMuerGjGpOoP5Ok/LfZrx7DdGg7f9w1DZmw8P81\nvj7s2ZchEGfRrDVQNigaogJzDWQBnCUZBMmaFBcnMndPDb9gqM9fP4gWJoAcRoxw\nFzBi7sPdl5C5Y24UdoHky6z+YKHtLqo3kdB+qXCsJR8U4eqJG16EW/OlS26L/hSP\n8tLNFI3SgcJiRWeCO5pRRpX6nfGf5wju0KMaJKzBRbDJwF3NEj3nmoXSfyoD+itV\nuv5DDCwojB/4nLT2EuxAr5vyY+JY6LmCZhWgPcXy60nsDcUxJzcvRaCsb+C2exZR\n9q4jgXUdAgMBAAECggEAOfle6Zcj6us/aBYDvSc8OvH5VaXynV+QBYxGsJdWadXV\nO29wjwAqMjhy/V/ScuZohb8CLMN+kagU13z4/EQTyOV2wHSWNqGebac1ZaTSUlcC\nBUrwMQEI0GxZ/l/zJkDV/PkffBLuZLdJ3UUTKR4WjMvoTKDmzoXb1XkyOIRoPcLN\nQXqGRl2A/BLgL0mxZsnvBXavcp0o4TfIxC73+ZEnJmrbuoHlDGbXWQvSOGJbi3gE\nLLSJ/+Sn2o9nhHPJI3M9xfMHnU7Fwo5Dt+vSl/Vn4+dvNd2djEjefQTSU19yE6n8\nW5/QzriG4IClBEjqTxYxnL+VQNmUm5dwXqB0C2ph/QKBgQDZ8gKJgU4xH0Woxh+d\nh5AjjnKVk0YlCS9MNA9VlGu2O8ohVj9LZ4azNfYyZMp5DP4gXW21elZWwRyfvSv0\nRlo1CrQIZwY/ETw8NOp9L8+OXQinjL8pLuqYo8rWU4jwdyHrBlwic0KSSOq785Xk\nmdSdU3NOPqTnnUHnDrJLFyEFowKBgQDSgI9v798XAgNhwRLVAwyXTf1N39R8jAl6\nm+m3xzPEblnOKp04cMcjjhV8AqNadg5bZ1Io5Qwy30PofcQmNLCCXs31gCn+0c0i\nEehSXz+QgkNmSxWLXl3SODWY4XN7ThmLJcL36iebKF8t50xvzKbER0xEMwzrrmVq\nW0YwpjfGPwKBgF7Plyb2Z2ubLRSUy+AdvyiYqWREYzltW3QNGbajEJCARhhmirZk\n3QZNLUMS8bnjWxH9UuKly7WF4Mvk4aAsksWMWHFnUCJTfx657mBzUhmeg0tQQUDL\nNicc6fp+8I2bZdf2NlKOTaGRsvv8pXKDMSkXyot5WQehM7AuhoWAFE99AoGAWJ8F\nREQBcQdI8zO8wO8asuyDkvCD3bd7GiJfwB5eXflzV4e7TxKz0/UyeFYH/cKsArE5\n9ruPai9ywIOKO+d81DYjkZLWm1Aqg4h0fZFaCnW8+Gjt9hHRf/poHif0XVohCOLp\n9UOgTwMtJv80v/Cx2PqHUkMH0oVGbwNkRoEEBDMCgYEAvf7a3Xxjl3Ymmy+15oOW\n+iX0/3Ntmr6TUjxFzRRnamO3CO7Vm3qCLOceE2C5/TCi07NxGXkY/NIEUywBGrLe\nSmm2ny5/u6vrDygZMGSB59RVnrAiX7zkqaIy6pY6cgPRQhclHZ9s34Nnd1J2GM4v\nxfdWj16ZNTMAaaWd9u+nZhg=\n-----END PRIVATE KEY-----\n",
|
||||||
|
"client_email": "gemini-api@gen-lang-client-0413414134.iam.gserviceaccount.com",
|
||||||
|
"client_id": "116149182781835050625",
|
||||||
|
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
|
||||||
|
"token_uri": "https://oauth2.googleapis.com/token",
|
||||||
|
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
|
||||||
|
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/gemini-api%40gen-lang-client-0413414134.iam.gserviceaccount.com",
|
||||||
|
"universe_domain": "googleapis.com"
|
||||||
|
}
|
||||||
|
bucket_name = "dy-media-storage"
|
||||||
|
cloudflare_project_id = "67720b647ff2b55cf37ba3ef9e677083"
|
||||||
|
cloudflare_gateway_id = "bowong-dev"
|
||||||
|
|
||||||
|
async def test_google_new_token(self):
|
||||||
|
cred = await GoogleAuthUtils.get_google_auth_jwt(service_account_info=self.service_account_info,
|
||||||
|
scopes=[
|
||||||
|
'https://www.googleapis.com/auth/cloud-platform'])
|
||||||
|
self.assertIsNotNone(cred)
|
||||||
|
|
||||||
|
async def test_google_cloud_storage_upload(self):
|
||||||
|
filepath = "./videos/input_1.mp4"
|
||||||
|
cred = await GoogleAuthUtils.get_google_auth_jwt(service_account_info=self.service_account_info,
|
||||||
|
scopes=[
|
||||||
|
'https://www.googleapis.com/auth/cloud-platform'])
|
||||||
|
prefix = "test/123"
|
||||||
|
filename = os.path.basename(filepath)
|
||||||
|
with open(filepath, 'rb') as file:
|
||||||
|
response = await GoogleAuthUtils.google_upload_file(file_stream=file,
|
||||||
|
content_type="video/mp4",
|
||||||
|
google_api_key=cred.access_token,
|
||||||
|
bucket_name=self.bucket_name,
|
||||||
|
filename=f"{prefix}/{filename}")
|
||||||
|
logger.info(response.model_dump_json(indent=2))
|
||||||
|
self.assertIsNotNone(response)
|
||||||
|
|
||||||
|
async def test_google_inference_with_sdk(self):
|
||||||
|
cred = await GoogleAuthUtils.get_google_auth_jwt(service_account_info=self.service_account_info,
|
||||||
|
scopes=[
|
||||||
|
'https://www.googleapis.com/auth/cloud-platform'])
|
||||||
|
|
||||||
|
logger.info(cred.model_dump_json(indent=2))
|
||||||
|
client = GoogleAuthUtils.GoogleGenaiClient(
|
||||||
|
cloudflare_project_id=self.cloudflare_project_id,
|
||||||
|
cloudflare_gateway_id=self.cloudflare_gateway_id,
|
||||||
|
google_project_id=self.service_account_info.get('project_id'),
|
||||||
|
regions=['us-central1'], access_token=cred.access_token,
|
||||||
|
)
|
||||||
|
result = client.generate_content(model_id="gemini-2.5-flash",
|
||||||
|
contents=[types.Content(role='user',
|
||||||
|
parts=[
|
||||||
|
types.Part(file_data=types.FileData(
|
||||||
|
mime_type="video/mp4",
|
||||||
|
file_uri="gs://dy-media-storage/videos/035b3053-73f8-45b7-9bf8-428df9025608.mp4"
|
||||||
|
)),
|
||||||
|
types.Part.from_text(
|
||||||
|
text="帮我总结一下这个视频里有什么"),
|
||||||
|
])],
|
||||||
|
config=types.GenerateContentConfig(
|
||||||
|
temperature=0.1,
|
||||||
|
top_p=0.7,
|
||||||
|
safety_settings=[
|
||||||
|
types.SafetySetting(
|
||||||
|
category=types.HarmCategory.HARM_CATEGORY_HARASSMENT,
|
||||||
|
threshold=types.HarmBlockThreshold.BLOCK_NONE,
|
||||||
|
),
|
||||||
|
types.SafetySetting(
|
||||||
|
category=types.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY,
|
||||||
|
threshold=types.HarmBlockThreshold.BLOCK_NONE,
|
||||||
|
),
|
||||||
|
types.SafetySetting(
|
||||||
|
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
||||||
|
threshold=types.HarmBlockThreshold.BLOCK_NONE,
|
||||||
|
),
|
||||||
|
types.SafetySetting(
|
||||||
|
category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
|
||||||
|
threshold=types.HarmBlockThreshold.BLOCK_NONE,
|
||||||
|
),
|
||||||
|
types.SafetySetting(
|
||||||
|
category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
|
||||||
|
threshold=types.HarmBlockThreshold.BLOCK_NONE,
|
||||||
|
)
|
||||||
|
],
|
||||||
|
response_mime_type="application/json",
|
||||||
|
response_schema=BaseResponse
|
||||||
|
))
|
||||||
|
logger.info(result.model_dump_json(indent=2, exclude_none=True))
|
||||||
|
self.assertIsNotNone(result)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
30
tests/pydantic_model_validation_test.py
Normal file
30
tests/pydantic_model_validation_test.py
Normal file
@@ -0,0 +1,30 @@
|
|||||||
|
import base64
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
|
||||||
|
class PydanticModelTestCase(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_nakama_login(self):
|
||||||
|
nakama_endpoint = 'http://43.143.58.201:7350'
|
||||||
|
email = "ybj.yege@gmail.com"
|
||||||
|
password = "Bowong7777"
|
||||||
|
basic_auth_str = base64.b64encode(f"{email}:{password}".encode('utf8')).decode('utf8')
|
||||||
|
logger.info(basic_auth_str)
|
||||||
|
with httpx.Client() as client:
|
||||||
|
login_response = client.post(url=f"{nakama_endpoint}/v2/account/authenticate/email?create=false",
|
||||||
|
headers={
|
||||||
|
"Authorization": f"Basic {basic_auth_str}",
|
||||||
|
},
|
||||||
|
json={"email": "ybj.yege@gmail.com", "password": "Bowong7777"}
|
||||||
|
)
|
||||||
|
login_response.raise_for_status()
|
||||||
|
session_json = login_response.json()
|
||||||
|
logger.info(f"Login response: {session_json}")
|
||||||
|
self.assertEqual(True, True)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
138
tests/s3_test_case.py
Normal file
138
tests/s3_test_case.py
Normal file
@@ -0,0 +1,138 @@
|
|||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import unittest
|
||||||
|
import xml.etree.ElementTree as ET
|
||||||
|
|
||||||
|
import boto3
|
||||||
|
import httpx
|
||||||
|
from botocore.client import BaseClient
|
||||||
|
from botocore.config import Config
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from BowongModalFunctions.utils.PathUtils import FileUtils
|
||||||
|
|
||||||
|
|
||||||
|
class S3TestCase(unittest.IsolatedAsyncioTestCase):
|
||||||
|
AWS_ACCESS_KEY_ID = "AKIAYRH5NGRSWHN2L4M6"
|
||||||
|
AWS_SECRET_ACCESS_KEY = "kfAqoOmIiyiywi25xaAkJUQbZ/EKDnzvI6NRCW1l"
|
||||||
|
S3_BUCKET = "modal-media-cache"
|
||||||
|
S3_REGION = "ap-northeast-2"
|
||||||
|
client: BaseClient
|
||||||
|
semaphore = asyncio.Semaphore(20)
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.client = boto3.client("s3",
|
||||||
|
aws_access_key_id=self.AWS_ACCESS_KEY_ID,
|
||||||
|
aws_secret_access_key=self.AWS_SECRET_ACCESS_KEY,
|
||||||
|
region_name=self.S3_REGION,
|
||||||
|
endpoint_url="https://s3-accelerate.amazonaws.com",
|
||||||
|
config=Config(s3={'addressing_style': 'virtual'}, signature_version='s3v4'))
|
||||||
|
|
||||||
|
async def test_cloudfront_presign_upload(self):
|
||||||
|
key = "upload/test.mp4"
|
||||||
|
signed_url = self.client.generate_presigned_url("put_object",
|
||||||
|
Params={
|
||||||
|
'Bucket': self.S3_BUCKET,
|
||||||
|
'Key': key,
|
||||||
|
"ContentType": "video/mp4",
|
||||||
|
}, ExpiresIn=3600)
|
||||||
|
logger.info(signed_url)
|
||||||
|
|
||||||
|
async def test_cloudfront_presign_multipart_upload(self):
|
||||||
|
local_filename = "./videos/fc-01JWDYP1PXBQ0MQGM9YBNPZQ1E.mp4"
|
||||||
|
chunk_size = 1024 * 1024 * 5
|
||||||
|
total_file_size = FileUtils.get_file_size(local_filename)
|
||||||
|
chunk_count = math.ceil(total_file_size / chunk_size)
|
||||||
|
logger.info(f"{chunk_count} * {chunk_size} of {total_file_size} bytes")
|
||||||
|
key = "upload/test2.mp4"
|
||||||
|
content_type = "video/mp4"
|
||||||
|
multipart_upload_response = self.client.create_multipart_upload(Bucket=self.S3_BUCKET, Key=key,
|
||||||
|
ContentType=content_type, )
|
||||||
|
logger.info(multipart_upload_response)
|
||||||
|
upload_id = multipart_upload_response.get("UploadId")
|
||||||
|
signed_urls = []
|
||||||
|
for i in range(chunk_count):
|
||||||
|
signed_url = self.client.generate_presigned_url("upload_part",
|
||||||
|
Params={
|
||||||
|
'Bucket': self.S3_BUCKET,
|
||||||
|
'Key': key,
|
||||||
|
'PartNumber': i + 1,
|
||||||
|
'UploadId': upload_id,
|
||||||
|
}, ExpiresIn=3600)
|
||||||
|
signed_urls.append(signed_url)
|
||||||
|
# logger.info(signed_url)
|
||||||
|
logger.info(signed_urls)
|
||||||
|
signed_completed_url = self.client.generate_presigned_url("complete_multipart_upload",
|
||||||
|
Params={
|
||||||
|
'Bucket': self.S3_BUCKET,
|
||||||
|
'Key': key,
|
||||||
|
'UploadId': upload_id,
|
||||||
|
}, ExpiresIn=3600)
|
||||||
|
signed_list_url = self.client.generate_presigned_url("list_parts",
|
||||||
|
Params={
|
||||||
|
'Bucket': self.S3_BUCKET,
|
||||||
|
'Key': key,
|
||||||
|
'UploadId': upload_id,
|
||||||
|
}, ExpiresIn=3600)
|
||||||
|
json_data = {
|
||||||
|
"parts": signed_urls,
|
||||||
|
"list": signed_list_url,
|
||||||
|
"completed": signed_completed_url
|
||||||
|
}
|
||||||
|
logger.info(json.dumps(json_data, ensure_ascii=False, indent=2))
|
||||||
|
|
||||||
|
async def upload_task(part_number: int, url: str, local_file: str, start: int, end: int,
|
||||||
|
client: httpx.AsyncClient):
|
||||||
|
async with self.semaphore:
|
||||||
|
try:
|
||||||
|
logger.info(f"[{part_number}] Uploading {start} -> {end}")
|
||||||
|
with open(local_file, "rb") as f:
|
||||||
|
f.seek(start)
|
||||||
|
data = f.read(end - start)
|
||||||
|
response = await client.put(url=url, content=data)
|
||||||
|
response.raise_for_status()
|
||||||
|
Etag = response.headers.get("ETag").strip('"')
|
||||||
|
# logger.info(f"headers: {response.headers}")
|
||||||
|
logger.info(Etag)
|
||||||
|
return {"ETag": Etag, "PartNumber": part_number}
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception(e)
|
||||||
|
raise e
|
||||||
|
|
||||||
|
upload_tasks = []
|
||||||
|
async with httpx.AsyncClient(timeout=30) as client:
|
||||||
|
async with self.semaphore:
|
||||||
|
for i, signed_url in enumerate(signed_urls):
|
||||||
|
start = i * chunk_size
|
||||||
|
end = min(start + chunk_size, total_file_size)
|
||||||
|
curr_chunk_size = end - start
|
||||||
|
logger.info(f"No.{i + 1} chunks of {curr_chunk_size} bytes")
|
||||||
|
upload_tasks.append(
|
||||||
|
upload_task(part_number=i + 1, url=signed_url,
|
||||||
|
local_file=local_filename, start=start, end=end,
|
||||||
|
client=client))
|
||||||
|
# 等待所有上传任务完成
|
||||||
|
results = await asyncio.gather(*upload_tasks)
|
||||||
|
logger.info(results)
|
||||||
|
results.sort(key=lambda x: x["PartNumber"])
|
||||||
|
logger.info(json.dumps(results, ensure_ascii=False, indent=2))
|
||||||
|
# 生成XML格式的请求体
|
||||||
|
root = ET.Element('CompleteMultipartUpload')
|
||||||
|
for part in results:
|
||||||
|
part_elem = ET.SubElement(root, 'Part')
|
||||||
|
part_number = ET.SubElement(part_elem, 'PartNumber')
|
||||||
|
part_number.text = str(part['PartNumber'])
|
||||||
|
etag = ET.SubElement(part_elem, 'ETag')
|
||||||
|
etag_str = part['ETag']
|
||||||
|
etag.text = f'"{etag_str}"'
|
||||||
|
# 将XML转换为字符串
|
||||||
|
xml_data = ET.tostring(root, encoding='utf-8')
|
||||||
|
logger.info(f"xml : {xml_data}")
|
||||||
|
completed_response = await client.post(url=signed_completed_url, content=xml_data,
|
||||||
|
headers={"Content-Type": "application/xml"})
|
||||||
|
logger.info(f"[HTTP {completed_response.status_code}] {completed_response.text}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user