diff --git a/scratch/verify_real_data.py b/scratch/verify_real_data.py index 6432240..8969122 100644 --- a/scratch/verify_real_data.py +++ b/scratch/verify_real_data.py @@ -1,4 +1,3 @@ -import os import sys import time import base64 @@ -11,8 +10,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) import torch -from app.models import get_model, VisualizedBGEEmbeddingModel -from app.config import EMBEDDING_MODELS, RERANK_MODELS +from app.models import get_model def cosine_similarity(v1: list[float], v2: list[float]) -> float: @@ -24,12 +22,16 @@ def cosine_similarity(v1: list[float], v2: list[float]) -> float: return dot / (norm1 * norm2) -def create_sample_image(text: str = "システム構成図", color: tuple = (70, 130, 180)) -> Image.Image: +def create_sample_image( + text: str = "システム構成図", color: tuple = (70, 130, 180) +) -> Image.Image: """実データのテスト用画像を動的に生成""" img = Image.new("RGB", (256, 256), color=(240, 244, 248)) draw = ImageDraw.Draw(img) draw.rectangle([20, 20, 236, 100], fill=color, outline=(30, 60, 90), width=2) - draw.rectangle([40, 140, 216, 220], fill=(220, 230, 242), outline=(30, 60, 90), width=2) + draw.rectangle( + [40, 140, 216, 220], fill=(220, 230, 242), outline=(30, 60, 90), width=2 + ) draw.line([(128, 100), (128, 140)], fill=(30, 60, 90), width=3) return img @@ -38,7 +40,7 @@ def run_text_embedding_verification(): print("\n========================================================") print("1. テキスト埋め込み実データ検証 (RURI, BGE-M3)") print("========================================================") - + test_queries = [ "検索クエリ: 日本の首都はどこですか?", "検索ドキュメント: 日本の首都は東京都であり、政治・経済・文化の中心地です。", @@ -55,26 +57,32 @@ def run_text_embedding_verification(): t0 = time.perf_counter() embeddings = model.encode(test_queries, normalize_embeddings=True) infer_time = time.perf_counter() - t0 - + dim = len(embeddings[0]) sim_relevant = cosine_similarity(embeddings[0].tolist(), embeddings[1].tolist()) - sim_irrelevant = cosine_similarity(embeddings[0].tolist(), embeddings[2].tolist()) - + sim_irrelevant = cosine_similarity( + embeddings[0].tolist(), embeddings[2].tolist() + ) + print(f" ✓ 埋め込み次元数: {dim}") - print(f" ✓ 推論時間 (3文): {infer_time*1000:.1f}ms") + print(f" ✓ 推論時間 (3文): {infer_time * 1000:.1f}ms") print(f" ✓ 関連ドキュメントとの類似度: {sim_relevant:.4f}") print(f" ✓ 無関係ドキュメントとの類似度: {sim_irrelevant:.4f}") - + assert dim > 0, "次元数が不正です" - assert sim_relevant > sim_irrelevant, f"関連文書の類似度({sim_relevant})が無関係文書({sim_irrelevant})を下回っています" - print(f" 🎯 判定: 合格 (関連度判定正常: {sim_relevant:.4f} > {sim_irrelevant:.4f})") + assert sim_relevant > sim_irrelevant, ( + f"関連文書の類似度({sim_relevant})が無関係文書({sim_irrelevant})を下回っています" + ) + print( + f" 🎯 判定: 合格 (関連度判定正常: {sim_relevant:.4f} > {sim_irrelevant:.4f})" + ) def run_multimodal_verification(): print("\n========================================================") print("2. マルチモーダル実データ検証 (bge-visualized-m3)") print("========================================================") - + t0 = time.perf_counter() model = get_model("bge-visualized-m3", device="cpu") print(f" ✓ ロード完了 ({time.perf_counter() - t0:.2f}秒)") @@ -97,30 +105,38 @@ def run_multimodal_verification(): t0 = time.perf_counter() embeddings = model.encode_multimodal(items) infer_time = time.perf_counter() - t0 - + dim = len(embeddings[0]) print(f" ✓ マルチモーダル埋め込み次元数: {dim}") - print(f" ✓ 推論時間 ({len(items)}アイテム): {infer_time*1000:.1f}ms") - + print(f" ✓ 推論時間 ({len(items)}アイテム): {infer_time * 1000:.1f}ms") + # 類似度評価 # システム構成図(画像+テキスト) と クラウドインフラ構成図(テキスト) sim_diagram = cosine_similarity(embeddings[0], embeddings[2]) # システム構成図(画像+テキスト) と 青空と緑の草原(テキスト) sim_mismatch = cosine_similarity(embeddings[0], embeddings[3]) - print(f" ✓ アーキテクチャ図(画像+文) vs クラウドインフラ(文) 類似度: {sim_diagram:.4f}") + print( + f" ✓ アーキテクチャ図(画像+文) vs クラウドインフラ(文) 類似度: {sim_diagram:.4f}" + ) print(f" ✓ アーキテクチャ図(画像+文) vs 草原風景(文) 類似度: {sim_mismatch:.4f}") - - assert dim == 1024, f"bge-visualized-m3 の次元数は 1024 である必要があります (実際: {dim})" - assert sim_diagram > sim_mismatch, f"画像-テキスト間のセマンティック類似度が期待を満たしていません ({sim_diagram} vs {sim_mismatch})" - print(f" 🎯 判定: 合格 (マルチモーダル類似度正常: {sim_diagram:.4f} > {sim_mismatch:.4f})") + + assert dim == 1024, ( + f"bge-visualized-m3 の次元数は 1024 である必要があります (実際: {dim})" + ) + assert sim_diagram > sim_mismatch, ( + f"画像-テキスト間のセマンティック類似度が期待を満たしていません ({sim_diagram} vs {sim_mismatch})" + ) + print( + f" 🎯 判定: 合格 (マルチモーダル類似度正常: {sim_diagram:.4f} > {sim_mismatch:.4f})" + ) def run_reranker_verification(): print("\n========================================================") print("3. リランカー実データ検証 (ruri-v3-reranker-310m)") print("========================================================") - + t0 = time.perf_counter() model = get_model("cl-nagoya/ruri-v3-reranker-310m", device="cpu") print(f" ✓ ロード完了 ({time.perf_counter() - t0:.2f}秒)") @@ -131,18 +147,22 @@ def run_reranker_verification(): "日本の温泉地ランキングでは、草津温泉や別府温泉、有馬温泉などが上位に選ばれています。", "ニューラルネットワークの汎化性能向上のため、学習データのバリデーション分割やクロスバリデーションが推奨されます。", ] - + pairs = [[query, p] for p in passages] t0 = time.perf_counter() scores = model.predict(pairs) infer_time = time.perf_counter() - t0 - - print(f" ✓ 推論時間 ({len(pairs)}ペア): {infer_time*1000:.1f}ms") + + print(f" ✓ 推論時間 ({len(pairs)}ペア): {infer_time * 1000:.1f}ms") for i, (p, score) in enumerate(zip(passages, scores)): - print(f" [{i+1}] スコア: {score:+.4f} | 内容: {p[:35]}...") + print(f" [{i + 1}] スコア: {score:+.4f} | 内容: {p[:35]}...") - assert scores[0] > scores[1], "過学習対策ドキュメントのスコアが温泉ドキュメントを下回っています" - assert scores[2] > scores[1], "汎化性能ドキュメントのスコアが温泉ドキュメントを下回っています" + assert scores[0] > scores[1], ( + "過学習対策ドキュメントのスコアが温泉ドキュメントを下回っています" + ) + assert scores[2] > scores[1], ( + "汎化性能ドキュメントのスコアが温泉ドキュメントを下回っています" + ) print(" 🎯 判定: 合格 (リランキング順位スコア正常)") @@ -150,7 +170,7 @@ def run_device_switching_verification(): print("\n========================================================") print("4. デバイス切り替え・フォールバック検証 (CPU / CUDA)") print("========================================================") - + cuda_available = torch.cuda.is_available() print(f" 現在のCUDA利用可能性: {cuda_available}") @@ -200,7 +220,9 @@ def run_fastapi_endpoints_real_verification(): data = res.json() assert len(data["data"]) == 2 assert len(data["data"][0]["embedding"]) > 0 - print(f" ✓ ステータス 200, 次元数: {len(data['data'][0]['embedding'])}, Usage: {data['usage']}") + print( + f" ✓ ステータス 200, 次元数: {len(data['data'][0]['embedding'])}, Usage: {data['usage']}" + ) # 2. /v1/embeddings (マルチモーダル Base64) print("\n [Endpoint 2] POST /v1/embeddings (マルチモーダル Base64画像)") @@ -223,7 +245,9 @@ def run_fastapi_endpoints_real_verification(): data = res.json() assert len(data["data"]) == 1 assert len(data["data"][0]["embedding"]) == 1024 - print(f" ✓ ステータス 200, マルチモーダル埋め込み次元数: {len(data['data'][0]['embedding'])}") + print( + f" ✓ ステータス 200, マルチモーダル埋め込み次元数: {len(data['data'][0]['embedding'])}" + ) # 3. /v1/rerank (リランキング) print("\n [Endpoint 3] POST /v1/rerank (テキストリランキング)") @@ -247,7 +271,9 @@ def run_fastapi_endpoints_real_verification(): assert len(results) == 2 print(f" ✓ ステータス 200, Top-{len(results)} 返却:") for r in results: - print(f" - Doc {r['document']}: score={r['score']:+.4f} | text={r.get('text', '')[:35]}...") + print( + f" - Doc {r['document']}: score={r['score']:+.4f} | text={r.get('text', '')[:35]}..." + ) assert results[0]["document"] in (0, 2) print(" 🎯 判定: 合格 (全APIエンドポイント実データ推論正常)") @@ -255,14 +281,16 @@ def run_fastapi_endpoints_real_verification(): if __name__ == "__main__": t_start = time.perf_counter() print("🚀 実データ・デバイス動作検証テストを開始します") - + run_text_embedding_verification() run_multimodal_verification() run_reranker_verification() run_device_switching_verification() run_fastapi_endpoints_real_verification() - + total_sec = time.perf_counter() - t_start - print(f"\n========================================================") - print(f"🎉 全ての実データ・デバイス検証テストに合格しました! (総所要時間: {total_sec:.2f}秒)") - print(f"========================================================") + print("\n========================================================") + print( + f"🎉 全ての実データ・デバイス検証テストに合格しました! (総所要時間: {total_sec:.2f}秒)" + ) + print("========================================================") diff --git a/scripts/benchmark_comprehensive.py b/scripts/benchmark_comprehensive.py index b5ee556..6e0509f 100644 --- a/scripts/benchmark_comprehensive.py +++ b/scripts/benchmark_comprehensive.py @@ -27,8 +27,17 @@ def create_mock_image(width: int = 400, height: int = 260) -> str: """Creates a sample mock diagram Base64 Data URL.""" img = Image.new("RGB", (width, height), color=(240, 245, 250)) draw = ImageDraw.Draw(img) - draw.rectangle([20, 20, width - 20, height - 20], fill=(33, 150, 243), outline=(25, 118, 210), width=2) - draw.text((width // 4, height // 2), f"Benchmark Diagram {width}x{height}", fill=(255, 255, 255)) + draw.rectangle( + [20, 20, width - 20, height - 20], + fill=(33, 150, 243), + outline=(25, 118, 210), + width=2, + ) + draw.text( + (width // 4, height // 2), + f"Benchmark Diagram {width}x{height}", + fill=(255, 255, 255), + ) buf = io.BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode("utf-8") @@ -47,13 +56,17 @@ async def run_comprehensive_benchmarks(): print("📊 COMPREHENSIVE BENCHMARK & HARDWARE PROFILING SUITE") print("=" * 80) - device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "Host CPU" - print(f"\n[Environment Information]") + device_name = ( + torch.cuda.get_device_name(0) if torch.cuda.is_available() else "Host CPU" + ) + print("\n[Environment Information]") print(f" • Compute Device : {device_name}") print(f" • PyTorch Version: {torch.__version__}") print(f" • CUDA Available : {torch.cuda.is_available()}") if torch.cuda.is_available(): - total_vram = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024 * 1024) + total_vram = torch.cuda.get_device_properties(0).total_memory / ( + 1024 * 1024 * 1024 + ) print(f" • Total GPU VRAM : {total_vram:.2f} GB") async with httpx.AsyncClient( @@ -62,7 +75,6 @@ async def run_comprehensive_benchmarks(): headers=HEADERS, timeout=180.0, ) as client: - # ====================================================================== # 1. All Models Head-to-Head Comparison (Single Query Latency & VRAM) # ====================================================================== @@ -77,23 +89,28 @@ async def run_comprehensive_benchmarks(): ("bge-visualized-m3", "Multimodal (800M / 1024d)"), ] - print(f"{'Model Name':<28} | {'Type':<22} | {'P50 (ms)':>8} | {'P95 (ms)':>8} | {'P99 (ms)':>8} | {'VRAM (MB)':>9}") + print( + f"{'Model Name':<28} | {'Type':<22} | {'P50 (ms)':>8} | {'P95 (ms)':>8} | {'P99 (ms)':>8} | {'VRAM (MB)':>9}" + ) print("-" * 88) for model_id, model_desc in embedding_models: # Warm up / Load model if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() - + # 50 iterations latencies = [] for i in range(50): t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": model_id, - "input": f"東京都千代田区における自然言語処理技術のベンチマーク測定サンプル {i}", - "input_type": "query" - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": model_id, + "input": f"東京都千代田区における自然言語処理技術のベンチマーク測定サンプル {i}", + "input_type": "query", + }, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200, f"Error {resp.status_code}: {resp.text}" latencies.append(dt) @@ -103,7 +120,9 @@ async def run_comprehensive_benchmarks(): p99 = np.percentile(latencies, 99) vram = get_vram_mb() - print(f"{model_id:<28} | {model_desc:<22} | {p50:8.2f} | {p95:8.2f} | {p99:8.2f} | {vram:9.1f}") + print( + f"{model_id:<28} | {model_desc:<22} | {p50:8.2f} | {p95:8.2f} | {p99:8.2f} | {vram:9.1f}" + ) # Reranker Model print("-" * 88) @@ -111,18 +130,21 @@ async def run_comprehensive_benchmarks(): rerank_lats = [] for i in range(30): t0 = time.perf_counter() - resp = await client.post("/v1/rerank", json={ - "model": rerank_model, - "query": "日本の首都はどこですか?", - "documents": [ - "東京は日本の首都であり、最大の都市です。", - "京都はかつての日本の古都です。", - "大阪は西日本の主要な経済都市です。", - "名古屋は中部地方の中心都市です。", - "福岡は九州地方の主要都市です。" - ], - "top_n": 3 - }) + resp = await client.post( + "/v1/rerank", + json={ + "model": rerank_model, + "query": "日本の首都はどこですか?", + "documents": [ + "東京は日本の首都であり、最大の都市です。", + "京都はかつての日本の古都です。", + "大阪は西日本の主要な経済都市です。", + "名古屋は中部地方の中心都市です。", + "福岡は九州地方の主要都市です。", + ], + "top_n": 3, + }, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200 rerank_lats.append(dt) @@ -131,7 +153,9 @@ async def run_comprehensive_benchmarks(): r_p95 = np.percentile(rerank_lats, 95) r_p99 = np.percentile(rerank_lats, 99) r_vram = get_vram_mb() - print(f"{rerank_model:<28} | {'Reranker (5 docs)':<22} | {r_p50:8.2f} | {r_p95:8.2f} | {r_p99:8.2f} | {r_vram:9.1f}") + print( + f"{rerank_model:<28} | {'Reranker (5 docs)':<22} | {r_p50:8.2f} | {r_p95:8.2f} | {r_p99:8.2f} | {r_vram:9.1f}" + ) # ====================================================================== # 2. Batch Scaling & Peak Throughput (ruri-v3-30m vs ruri-v3-310m) @@ -143,17 +167,20 @@ async def run_comprehensive_benchmarks(): batch_sizes = [1, 8, 32, 64] for m_id in ["cl-nagoya/ruri-v3-30m", "cl-nagoya/ruri-v3-310m"]: print(f"\n--- Model: {m_id} ---") - print(f"{'Batch Size':>10} | {'Total Time (ms)':>15} | {'Throughput (items/s)':>20} | {'Per-Item (ms)':>13}") + print( + f"{'Batch Size':>10} | {'Total Time (ms)':>15} | {'Throughput (items/s)':>20} | {'Per-Item (ms)':>13}" + ) print("-" * 65) - sample_txt = "高度な自然言語処理技術を用いたベクトル検索エンジンの性能テスト。" + sample_txt = ( + "高度な自然言語処理技術を用いたベクトル検索エンジンの性能テスト。" + ) for bs in batch_sizes: b_inputs = [f"{sample_txt} (seq_{j})" for j in range(bs)] t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": m_id, - "input": b_inputs, - "input_type": "document" - }) + resp = await client.post( + "/v1/embeddings", + json={"model": m_id, "input": b_inputs, "input_type": "document"}, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200 qps = bs / (dt / 1000) @@ -168,9 +195,13 @@ async def run_comprehensive_benchmarks(): print("=" * 80) seq_lengths = [32, 128, 512, 1024, 2048] - base_phrase = "自然言語処理における埋め込みベクトルの計算速度とメモリ使用量を検証する。" - - print(f"{'Target Tokens (approx)':<25} | {'Char Length':>12} | {'ruri-30m (ms)':>13} | {'ruri-310m (ms)':>14}") + base_phrase = ( + "自然言語処理における埋め込みベクトルの計算速度とメモリ使用量を検証する。" + ) + + print( + f"{'Target Tokens (approx)':<25} | {'Char Length':>12} | {'ruri-30m (ms)':>13} | {'ruri-310m (ms)':>14}" + ) print("-" * 72) for target_tokens in seq_lengths: @@ -181,17 +212,25 @@ async def run_comprehensive_benchmarks(): # Measure ruri-30m t0 = time.perf_counter() - r1 = await client.post("/v1/embeddings", json={"model": "cl-nagoya/ruri-v3-30m", "input": text_payload}) + r1 = await client.post( + "/v1/embeddings", + json={"model": "cl-nagoya/ruri-v3-30m", "input": text_payload}, + ) t_30m = (time.perf_counter() - t0) * 1000 assert r1.status_code == 200 # Measure ruri-310m t0 = time.perf_counter() - r2 = await client.post("/v1/embeddings", json={"model": "cl-nagoya/ruri-v3-310m", "input": text_payload}) + r2 = await client.post( + "/v1/embeddings", + json={"model": "cl-nagoya/ruri-v3-310m", "input": text_payload}, + ) t_310m = (time.perf_counter() - t0) * 1000 assert r2.status_code == 200 - print(f"{target_tokens:<25d} | {char_len:12d} | {t_30m:13.1f} | {t_310m:14.1f}") + print( + f"{target_tokens:<25d} | {char_len:12d} | {t_30m:13.1f} | {t_310m:14.1f}" + ) # ====================================================================== # 4. Multimodal Modality & Image Resolution Breakdown @@ -209,21 +248,29 @@ async def run_comprehensive_benchmarks(): ("Image Only (Thumbnail 64x64)", {"image_url": img_64}), ("Image Only (Standard 224x224)", {"image_url": img_224}), ("Image Only (Full HD 1080p)", {"image_url": img_1080}), - ("Multimodal (Standard 224 + Text)", {"text": "マイクロサービス構成図", "image_url": img_224}), - ("Multimodal (Full HD 1080p + Text)", {"text": "高解像度インフラ構成図", "image_url": img_1080}), + ( + "Multimodal (Standard 224 + Text)", + {"text": "マイクロサービス構成図", "image_url": img_224}, + ), + ( + "Multimodal (Full HD 1080p + Text)", + {"text": "高解像度インフラ構成図", "image_url": img_1080}, + ), ] - print(f"{'Input Modality / Resolution':<36} | {'Avg Latency (ms)':>16} | {'Output Dim':>10}") + print( + f"{'Input Modality / Resolution':<36} | {'Avg Latency (ms)':>16} | {'Output Dim':>10}" + ) print("-" * 68) for case_name, payload_input in mm_cases: lats = [] for _ in range(5): t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": payload_input - }) + resp = await client.post( + "/v1/embeddings", + json={"model": "bge-visualized-m3", "input": payload_input}, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200 lats.append(dt) diff --git a/scripts/benchmark_suite.py b/scripts/benchmark_suite.py index 253c5ab..46c8975 100644 --- a/scripts/benchmark_suite.py +++ b/scripts/benchmark_suite.py @@ -7,27 +7,38 @@ API_KEY = "test_api_key_secret" HEADERS = {"Authorization": f"Bearer {API_KEY}"} + async def benchmark_endpoint(): print("=" * 80) print("📊 API LATENCY & THROUGHPUT BENCHMARK SUITE") print("=" * 80) - async with httpx.AsyncClient(base_url=BASE_URL, headers=HEADERS, timeout=60.0) as client: + async with httpx.AsyncClient( + base_url=BASE_URL, headers=HEADERS, timeout=60.0 + ) as client: # 1. Warm-up print("\n[Step 1] Warming up model endpoints...") - await client.post("/v1/embeddings", json={"model": "cl-nagoya/ruri-v3-30m", "input": "ウォームアップ"}) + await client.post( + "/v1/embeddings", + json={"model": "cl-nagoya/ruri-v3-30m", "input": "ウォームアップ"}, + ) print(" ✓ Warm-up complete") # 2. Single Query Latency Benchmark (100 sequential requests) - print("\n[Step 2] Measuring Single Text Embedding Latency (100 sequential requests)...") + print( + "\n[Step 2] Measuring Single Text Embedding Latency (100 sequential requests)..." + ) latencies = [] for i in range(100): t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "cl-nagoya/ruri-v3-30m", - "input": f"日本語クエリのベンチマーク測定テスト {i}", - "input_type": "query" - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "cl-nagoya/ruri-v3-30m", + "input": f"日本語クエリのベンチマーク測定テスト {i}", + "input_type": "query", + }, + ) assert resp.status_code == 200 lat = (time.perf_counter() - t0) * 1000 latencies.append(lat) @@ -53,41 +64,51 @@ async def benchmark_endpoint(): for batch_size in [1, 8, 32, 64]: batch_input = [f"{sample_sentence} (idx={j})" for j in range(batch_size)] t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "cl-nagoya/ruri-v3-30m", - "input": batch_input - }) + resp = await client.post( + "/v1/embeddings", + json={"model": "cl-nagoya/ruri-v3-30m", "input": batch_input}, + ) elapsed = time.perf_counter() - t0 assert resp.status_code == 200 data = resp.json() total_tokens = data["usage"]["total_tokens"] - print(f" • Batch Size {batch_size:2d}: {elapsed*1000:6.1f} ms | Throughput: {batch_size/elapsed:6.1f} items/sec | Token Rate: {total_tokens/elapsed:7.1f} tok/sec") + print( + f" • Batch Size {batch_size:2d}: {elapsed * 1000:6.1f} ms | Throughput: {batch_size / elapsed:6.1f} items/sec | Token Rate: {total_tokens / elapsed:7.1f} tok/sec" + ) # 4. Reranking Benchmark - print("\n[Step 4] Measuring Reranker Latency (cl-nagoya/ruri-v3-reranker-310m)...") + print( + "\n[Step 4] Measuring Reranker Latency (cl-nagoya/ruri-v3-reranker-310m)..." + ) docs = [ "東京は日本の首都であり、最大の都市です。", "京都は日本の古都として親しまれています。", "富士山は日本で最も高い山です。", "北海道は日本最北端の島で、広大な自然を有します。", - "沖縄は日本南西部の島々で、美しいサンゴ礁が広がります。" + "沖縄は日本南西部の島々で、美しいサンゴ礁が広がります。", ] rerank_lats = [] for i in range(20): t0 = time.perf_counter() - resp = await client.post("/v1/rerank", json={ - "model": "cl-nagoya/ruri-v3-reranker-310m", - "query": "日本の首都はどこですか?", - "documents": docs - }) + resp = await client.post( + "/v1/rerank", + json={ + "model": "cl-nagoya/ruri-v3-reranker-310m", + "query": "日本の首都はどこですか?", + "documents": docs, + }, + ) assert resp.status_code == 200 rerank_lats.append((time.perf_counter() - t0) * 1000) - print(f" • Rerank (5 docs) Avg Latency: {statistics.mean(rerank_lats):.2f} ms (P50={statistics.median(rerank_lats):.2f} ms, P95={sorted(rerank_lats)[int(len(rerank_lats)*0.95)]:.2f} ms)") + print( + f" • Rerank (5 docs) Avg Latency: {statistics.mean(rerank_lats):.2f} ms (P50={statistics.median(rerank_lats):.2f} ms, P95={sorted(rerank_lats)[int(len(rerank_lats) * 0.95)]:.2f} ms)" + ) print("\n" + "=" * 80) print("✨ BENCHMARK SUITE COMPLETED SUCCESSFULLY") print("=" * 80) + if __name__ == "__main__": asyncio.run(benchmark_endpoint()) diff --git a/scripts/test_multimodal_suite.py b/scripts/test_multimodal_suite.py index 4f6dbb0..ef25c91 100644 --- a/scripts/test_multimodal_suite.py +++ b/scripts/test_multimodal_suite.py @@ -9,7 +9,7 @@ import time import httpx import numpy as np -from PIL import Image, ImageDraw, ImageFont +from PIL import Image, ImageDraw from app.main import app BASE_URL = "http://testserver" @@ -21,25 +21,50 @@ # 1. Real Diagram & Visual Asset Generators # ============================================================================== + def create_architecture_diagram() -> Image.Image: """Microservices / Cloud Architecture Diagram.""" img = Image.new("RGB", (600, 360), color=(245, 247, 250)) draw = ImageDraw.Draw(img) # API Gateway - draw.rounded_rectangle([30, 130, 150, 210], radius=8, fill=(30, 136, 229), outline=(21, 101, 192), width=2) + draw.rounded_rectangle( + [30, 130, 150, 210], + radius=8, + fill=(30, 136, 229), + outline=(21, 101, 192), + width=2, + ) draw.text((45, 160), "API Gateway\n(FastAPI)", fill=(255, 255, 255)) # Service A (Embedding Service) - draw.rounded_rectangle([230, 40, 370, 120], radius=8, fill=(67, 160, 71), outline=(46, 125, 50), width=2) + draw.rounded_rectangle( + [230, 40, 370, 120], + radius=8, + fill=(67, 160, 71), + outline=(46, 125, 50), + width=2, + ) draw.text((245, 70), "Embedding\nService (GPU)", fill=(255, 255, 255)) # Service B (Rerank Service) - draw.rounded_rectangle([230, 220, 370, 300], radius=8, fill=(142, 36, 170), outline=(106, 27, 154), width=2) + draw.rounded_rectangle( + [230, 220, 370, 300], + radius=8, + fill=(142, 36, 170), + outline=(106, 27, 154), + width=2, + ) draw.text((245, 250), "Rerank\nService", fill=(255, 255, 255)) # Vector DB (Milvus / Qdrant) - draw.rounded_rectangle([450, 130, 570, 210], radius=8, fill=(251, 140, 0), outline=(239, 108, 0), width=2) + draw.rounded_rectangle( + [450, 130, 570, 210], + radius=8, + fill=(251, 140, 0), + outline=(239, 108, 0), + width=2, + ) draw.text((465, 160), "Vector Store\n(Qdrant DB)", fill=(255, 255, 255)) # Arrows @@ -86,19 +111,30 @@ def create_flowchart_diagram() -> Image.Image: draw = ImageDraw.Draw(img) # Step 1: Client Request - draw.ellipse([40, 140, 120, 200], fill=(225, 245, 254), outline=(2, 136, 209), width=2) + draw.ellipse( + [40, 140, 120, 200], fill=(225, 245, 254), outline=(2, 136, 209), width=2 + ) draw.text((55, 160), "Client\nLogin", fill=(1, 87, 155)) # Step 2: Auth Check - draw.polygon([(220, 130), (280, 170), (220, 210), (160, 170)], fill=(255, 243, 224), outline=(245, 124, 0), width=2) + draw.polygon( + [(220, 130), (280, 170), (220, 210), (160, 170)], + fill=(255, 243, 224), + outline=(245, 124, 0), + width=2, + ) draw.text((190, 162), "Verify\nToken", fill=(230, 81, 0)) # Step 3: Success Token Granted - draw.rectangle([340, 90, 480, 150], fill=(232, 245, 233), outline=(56, 142, 60), width=2) + draw.rectangle( + [340, 90, 480, 150], fill=(232, 245, 233), outline=(56, 142, 60), width=2 + ) draw.text((360, 110), "200 OK JWT Token\nAccess Granted", fill=(27, 94, 32)) # Step 4: 401 Unauthorized - draw.rectangle([340, 210, 480, 270], fill=(255, 235, 238), outline=(211, 47, 47), width=2) + draw.rectangle( + [340, 210, 480, 270], fill=(255, 235, 238), outline=(211, 47, 47), width=2 + ) draw.text((360, 230), "401 Unauthorized\nInvalid API Key", fill=(183, 28, 28)) # Connecting Lines @@ -170,7 +206,9 @@ def create_extreme_aspect_ratio(mode: str) -> Image.Image: if mode == "ultra_wide": img = Image.new("RGB", (1200, 180), color=(240, 244, 248)) draw = ImageDraw.Draw(img) - draw.text((450, 80), "Ultra Wide Architecture Banner (1200x180)", fill=(33, 33, 33)) + draw.text( + (450, 80), "Ultra Wide Architecture Banner (1200x180)", fill=(33, 33, 33) + ) return img elif mode == "ultra_tall": img = Image.new("RGB", (180, 1200), color=(248, 244, 240)) @@ -180,7 +218,11 @@ def create_extreme_aspect_ratio(mode: str) -> Image.Image: elif mode == "high_res": img = Image.new("RGB", (1920, 1080), color=(230, 238, 245)) draw = ImageDraw.Draw(img) - draw.text((800, 500), "Full HD 1080p High-Resolution Diagram (1920x1080)", fill=(33, 33, 33)) + draw.text( + (800, 500), + "Full HD 1080p High-Resolution Diagram (1920x1080)", + fill=(33, 33, 33), + ) return img elif mode == "thumbnail": img = Image.new("RGB", (64, 64), color=(100, 150, 200)) @@ -210,6 +252,7 @@ def cosine_similarity(a: list[float], b: list[float]) -> float: # 2. Main Test Execution Engine # ============================================================================== + async def run_extended_multimodal_tests(): print("=" * 80) print("🚀 EXTENDED REAL DATA MULTIMODAL TEST & STRESS SUITE (bge-visualized-m3)") @@ -242,7 +285,6 @@ async def run_extended_multimodal_tests(): headers=HEADERS, timeout=180.0, ) as client: - # ====================================================================== # SECTION 1: Image Format Variations & Transparency # ====================================================================== @@ -254,10 +296,16 @@ async def run_extended_multimodal_tests(): for fmt in formats: b64_url = image_to_base64_data_url(assets["architecture"], format=fmt) t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": f"{fmt}形式でエンコードされたシステム構成図", "image_url": b64_url} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": f"{fmt}形式でエンコードされたシステム構成図", + "image_url": b64_url, + }, + }, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200, f"Failed for format {fmt}: {resp.text}" vec = resp.json()["data"][0]["embedding"] @@ -269,10 +317,16 @@ async def run_extended_multimodal_tests(): rgba_draw = ImageDraw.Draw(rgba_img) rgba_draw.rectangle([50, 50, 250, 150], fill=(0, 128, 255, 180)) rgba_b64 = image_to_base64_data_url(rgba_img, format="PNG") - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": "半透明アルファチャンネルを含むRGBA透過PNG画像", "image_url": rgba_b64} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": "半透明アルファチャンネルを含むRGBA透過PNG画像", + "image_url": rgba_b64, + }, + }, + ) assert resp.status_code == 200 print(" ✓ RGBA PNG (Transparent Alpha Channel): 200 OK (dim=1024)") @@ -286,15 +340,23 @@ async def run_extended_multimodal_tests(): for name, img in extreme_assets.items(): b64_url = image_to_base64_data_url(img, format="PNG") t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": f"解像度テスト: {name} (size: {img.size})", "image_url": b64_url} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": f"解像度テスト: {name} (size: {img.size})", + "image_url": b64_url, + }, + }, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200, f"Failed for {name}: {resp.text}" vec = resp.json()["data"][0]["embedding"] assert len(vec) == 1024 - print(f" ✓ {name:12s} ({img.size[0]:4d}x{img.size[1]:4d}): 200 OK ({dt:.1f} ms)") + print( + f" ✓ {name:12s} ({img.size[0]:4d}x{img.size[1]:4d}): 200 OK ({dt:.1f} ms)" + ) # ====================================================================== # SECTION 3: Edge Cases, Schema Variations & Error Handling @@ -304,54 +366,96 @@ async def run_extended_multimodal_tests(): print("=" * 80) # 1. Image only (Empty Text) - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": "", "image_url": image_to_base64_data_url(assets["sketch"])} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": "", + "image_url": image_to_base64_data_url(assets["sketch"]), + }, + }, + ) assert resp.status_code == 200 print(" ✓ Image Only (text=''): 200 OK") # 2. Text only with bge-visualized-m3 - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": "テキストのみの単体クエリエンコード" - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": "テキストのみの単体クエリエンコード", + }, + ) assert resp.status_code == 200 print(" ✓ Text Only with bge-visualized-m3: 200 OK") # 3. Long Japanese Text (>1000 chars) + Diagram - long_jp_text = "このシステムアーキテクチャは、高可用性とスケーラビリティを担保するために設計された最新のマイクロサービス構成です。" * 30 - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": long_jp_text, "image_url": image_to_base64_data_url(assets["architecture"])} - }) + long_jp_text = ( + "このシステムアーキテクチャは、高可用性とスケーラビリティを担保するために設計された最新のマイクロサービス構成です。" + * 30 + ) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": long_jp_text, + "image_url": image_to_base64_data_url(assets["architecture"]), + }, + }, + ) assert resp.status_code == 200 print(f" ✓ Long Japanese Text ({len(long_jp_text)} chars) + Diagram: 200 OK") # 4. Japanese Unicode, Emojis & Symbols - special_text = "🔥【超重要】API 構成図 🚀 (Ver 2.5.0) -> DB 連携 & 高速キャッシュ ⚡️" - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": special_text, "image_url": image_to_base64_data_url(assets["architecture"])} - }) + special_text = ( + "🔥【超重要】API 構成図 🚀 (Ver 2.5.0) -> DB 連携 & 高速キャッシュ ⚡️" + ) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": special_text, + "image_url": image_to_base64_data_url(assets["architecture"]), + }, + }, + ) assert resp.status_code == 200 print(" ✓ Japanese Emojis & Unicode Symbols: 200 OK") # 5. Invalid Base64 Image -> 400 Bad Request - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": "破損した画像データ", "image_url": "data:image/png;base64,invalid_corrupted_base64_!@#$"} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": "破損した画像データ", + "image_url": "data:image/png;base64,invalid_corrupted_base64_!@#$", + }, + }, + ) assert resp.status_code == 400 - print(f" ✓ Invalid Base64 Validation: Correctly returned 400 ({resp.json()['detail'][:40]}...)") + print( + f" ✓ Invalid Base64 Validation: Correctly returned 400 ({resp.json()['detail'][:40]}...)" + ) # 6. Image sent to Text-Only model -> 400 Bad Request - resp = await client.post("/v1/embeddings", json={ - "model": "cl-nagoya/ruri-v3-310m", - "input": {"text": "テキスト専用モデルに画像送信", "image_url": image_to_base64_data_url(assets["sketch"])} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "cl-nagoya/ruri-v3-310m", + "input": { + "text": "テキスト専用モデルに画像送信", + "image_url": image_to_base64_data_url(assets["sketch"]), + }, + }, + ) assert resp.status_code == 400 - print(f" ✓ Text-Only Model Guard: Correctly returned 400 ({resp.json()['detail'][:40]}...)") + print( + f" ✓ Text-Only Model Guard: Correctly returned 400 ({resp.json()['detail'][:40]}...)" + ) # ====================================================================== # SECTION 4: 5x5 Cross-Modal Semantic Retrieval Matrix & Accuracy @@ -364,10 +468,16 @@ async def run_extended_multimodal_tests(): diagram_keys = ["architecture", "performance", "flowchart", "table", "sketch"] diagram_vecs = {} for key in diagram_keys: - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": {"text": f"{key} diagram", "image_url": image_to_base64_data_url(assets[key])} - }) + resp = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": f"{key} diagram", + "image_url": image_to_base64_data_url(assets[key]), + }, + }, + ) diagram_vecs[key] = resp.json()["data"][0]["embedding"] # Domain Text Queries @@ -382,30 +492,41 @@ async def run_extended_multimodal_tests(): # Encode queries and compute similarity matrix query_vecs = {} for q_key, q_text in queries.items(): - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": q_text - }) + resp = await client.post( + "/v1/embeddings", json={"model": "bge-visualized-m3", "input": q_text} + ) query_vecs[q_key] = resp.json()["data"][0]["embedding"] # Print Matrix - print(f"\n{'Query Category':<18} | " + " | ".join([f"{k[:7]:>7}" for k in diagram_keys])) + print( + f"\n{'Query Category':<18} | " + + " | ".join([f"{k[:7]:>7}" for k in diagram_keys]) + ) print("-" * 65) correct_top1_count = 0 for q_key, q_vec in query_vecs.items(): - sims = {d_key: cosine_similarity(q_vec, diagram_vecs[d_key]) for d_key in diagram_keys} - row_str = f"{q_key:<18} | " + " | ".join([f"{sims[k]:7.4f}" for k in diagram_keys]) + sims = { + d_key: cosine_similarity(q_vec, diagram_vecs[d_key]) + for d_key in diagram_keys + } + row_str = f"{q_key:<18} | " + " | ".join( + [f"{sims[k]:7.4f}" for k in diagram_keys] + ) top_match = max(sims, key=sims.get) - is_correct = (top_match == q_key) + is_correct = top_match == q_key if is_correct: correct_top1_count += 1 status = "🎯 Match" if is_correct else "❌ Mismatch" print(f"{row_str} [{status} -> {top_match}]") accuracy = (correct_top1_count / len(queries)) * 100.0 - print(f"\n📊 Top-1 Retrieval Accuracy: {accuracy:.1f}% ({correct_top1_count}/{len(queries)})") - assert correct_top1_count == len(queries), "Cross-modal retrieval failed accuracy check!" + print( + f"\n📊 Top-1 Retrieval Accuracy: {accuracy:.1f}% ({correct_top1_count}/{len(queries)})" + ) + assert correct_top1_count == len(queries), ( + "Cross-modal retrieval failed accuracy check!" + ) # ====================================================================== # SECTION 5: Batch Processing & Throughput Scaling @@ -417,22 +538,24 @@ async def run_extended_multimodal_tests(): batch_sizes = [1, 2, 4, 8] base_item = { "text": "バッチテスト用図面アイテム", - "image_url": image_to_base64_data_url(assets["performance"]) + "image_url": image_to_base64_data_url(assets["performance"]), } for bs in batch_sizes: batch_input = [base_item] * bs t0 = time.perf_counter() - resp = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": batch_input - }) + resp = await client.post( + "/v1/embeddings", + json={"model": "bge-visualized-m3", "input": batch_input}, + ) dt = (time.perf_counter() - t0) * 1000 assert resp.status_code == 200 res_data = resp.json()["data"] assert len(res_data) == bs per_item_ms = dt / bs - print(f" ✓ Batch Size {bs:2d}: Total {dt:7.1f} ms ({per_item_ms:6.1f} ms/item, QPS={bs / (dt/1000):5.1f})") + print( + f" ✓ Batch Size {bs:2d}: Total {dt:7.1f} ms ({per_item_ms:6.1f} ms/item, QPS={bs / (dt / 1000):5.1f})" + ) # ====================================================================== # SECTION 6: High Concurrency & Thread-Safety Stress Test @@ -444,13 +567,16 @@ async def run_extended_multimodal_tests(): async def worker(worker_id: int): diag_name = diagram_keys[worker_id % len(diagram_keys)] t0 = time.perf_counter() - r = await client.post("/v1/embeddings", json={ - "model": "bge-visualized-m3", - "input": { - "text": f"並行ワーカー {worker_id} リクエスト ({diag_name})", - "image_url": image_to_base64_data_url(assets[diag_name]) - } - }) + r = await client.post( + "/v1/embeddings", + json={ + "model": "bge-visualized-m3", + "input": { + "text": f"並行ワーカー {worker_id} リクエスト ({diag_name})", + "image_url": image_to_base64_data_url(assets[diag_name]), + }, + }, + ) dt = (time.perf_counter() - t0) * 1000 assert r.status_code == 200, f"Worker {worker_id} failed: {r.text}" return worker_id, dt @@ -462,10 +588,18 @@ async def worker(worker_id: int): total_time_ms = (time.perf_counter() - t_start) * 1000 latencies = [res[1] for res in results] - print(f" ✓ Processed {num_concurrent} concurrent multimodal requests in {total_time_ms:.1f} ms") - print(f" ✓ Avg Latency: {np.mean(latencies):.1f} ms | Min: {np.min(latencies):.1f} ms | Max: {np.max(latencies):.1f} ms | P95: {np.percentile(latencies, 95):.1f} ms") - print(f" ✓ Concurrency Throughput: {num_concurrent / (total_time_ms / 1000):.2f} req/s") - print(" ✓ Thread-Safety Verified: All 20 workers returned 200 OK without race conditions.") + print( + f" ✓ Processed {num_concurrent} concurrent multimodal requests in {total_time_ms:.1f} ms" + ) + print( + f" ✓ Avg Latency: {np.mean(latencies):.1f} ms | Min: {np.min(latencies):.1f} ms | Max: {np.max(latencies):.1f} ms | P95: {np.percentile(latencies, 95):.1f} ms" + ) + print( + f" ✓ Concurrency Throughput: {num_concurrent / (total_time_ms / 1000):.2f} req/s" + ) + print( + " ✓ Thread-Safety Verified: All 20 workers returned 200 OK without race conditions." + ) print("\n" + "=" * 80) print("🎉 ALL EXTENDED MULTIMODAL TESTS AND LOAD STRESS TESTS PASSED SUCCESSFULLY!") diff --git a/src/app/main.py b/src/app/main.py index a3d2bb1..03ff462 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -22,8 +22,11 @@ EmbeddingResponse, RerankRequest, RerankResponse, + ModelList, + ModelCard, + ErrorResponse, ) -from .models import get_model as get_model +from .models import get_model as get_model, _model_cache from .config import ( EMBEDDING_MODELS, RERANK_MODELS, @@ -188,9 +191,65 @@ def get_rerank_service() -> BaseRerankService: return RerankService(proxy_to_tei_func=_proxy_to_tei, model_loader=get_model) +@app.get( + "/v1/models", + response_model=ModelList, + tags=["Models"], + summary="List available models", + description="Retrieves a list of all currently supported models, including both embedding and reranking models.", + responses={ + 401: {"model": ErrorResponse, "description": "Unauthorized"}, + 429: {"model": ErrorResponse, "description": "Too Many Requests"}, + 503: {"model": ErrorResponse, "description": "Service Unavailable"}, + }, +) +async def list_models(): + """ + Lists available models by dynamically aggregating and deduplicating + from EMBEDDING_MODELS and RERANK_MODELS. + """ + model_ids = sorted(list(set(EMBEDDING_MODELS + RERANK_MODELS))) + models = [ModelCard(id=model_id) for model_id in model_ids] + return ModelList(data=models) + + +@app.post( + "/v1/models/unload", + tags=["Models"], + summary="Unload models", + description="Clears the model cache, releasing any loaded models from memory.", + responses={ + 401: {"model": ErrorResponse, "description": "Unauthorized"}, + 429: {"model": ErrorResponse, "description": "Too Many Requests"}, + 503: {"model": ErrorResponse, "description": "Service Unavailable"}, + }, + dependencies=[Depends(verify_api_key)], +) +async def unload_models(): + """ + Unloads all currently cached models to free memory/VRAM. + """ + _model_cache.clear() + return {"status": "ok", "message": "All models unloaded"} + + @app.post( "/v1/embeddings", response_model=EmbeddingResponse, + tags=["Embeddings"], + summary="Create embeddings", + description="Creates embeddings for the given input text or multimodal item, following OpenAI's API format.", + responses={ + 400: { + "model": ErrorResponse, + "description": "Bad Request (e.g., unsupported model, invalid input)", + }, + 401: {"model": ErrorResponse, "description": "Unauthorized"}, + 413: {"model": ErrorResponse, "description": "Payload Too Large"}, + 429: {"model": ErrorResponse, "description": "Too Many Requests"}, + 500: {"model": ErrorResponse, "description": "Internal Server Error"}, + 503: {"model": ErrorResponse, "description": "Service Unavailable"}, + }, dependencies=[Depends(verify_api_key)], ) async def create_embeddings( @@ -208,6 +267,20 @@ async def create_embeddings( "/v1/rerank", response_model=RerankResponse, response_model_exclude_none=True, + tags=["Rerank"], + summary="Rerank documents", + description="Reranks a list of documents for a given query to determine their relevance.", + responses={ + 400: { + "model": ErrorResponse, + "description": "Bad Request (e.g., unsupported model, invalid input)", + }, + 401: {"model": ErrorResponse, "description": "Unauthorized"}, + 413: {"model": ErrorResponse, "description": "Payload Too Large"}, + 429: {"model": ErrorResponse, "description": "Too Many Requests"}, + 500: {"model": ErrorResponse, "description": "Internal Server Error"}, + 503: {"model": ErrorResponse, "description": "Service Unavailable"}, + }, dependencies=[Depends(verify_api_key)], ) async def create_rerank( diff --git a/src/app/schemas.py b/src/app/schemas.py index a774025..7aeaaef 100644 --- a/src/app/schemas.py +++ b/src/app/schemas.py @@ -56,44 +56,77 @@ class EmbeddingRequest(BaseModel): Annotated[ list[SingleInputItem], Field(min_length=1, max_length=MAX_INPUT_ITEMS) ], - ] - model: str - user: Optional[str] = None + ] = Field( + description="Input text to embed, encoded as a string or array of strings/multimodal items.", + examples=[["Hello world", "How are you?"]], + ) + model: str = Field(description="ID of the model to use.", examples=["BAAI/bge-m3"]) + user: Optional[str] = Field( + default=None, + description="A unique identifier representing your end-user.", + examples=["user-1234"], + ) input_type: Optional[str] = Field( - None, + default=None, description="Type of the input. Maps to Ruri-v3 prefixes: query, document, classification, clustering, sts.", + examples=["document"], ) instruction: Optional[str] = Field( - None, + default=None, description="Specific instruction for the model. For future use with instruction-based models.", + examples=["Represent the document for retrieval: "], ) apply_ruri_prefix: bool = Field( - False, + default=False, description="Automatically apply prefixes based on input shape if true (fallback/compatibility).", ) + model_config = ConfigDict( + json_schema_extra={ + "examples": [ + { + "input": ["Hello world", "How are you?"], + "model": "BAAI/bge-m3", + "input_type": "document", + } + ] + } + ) + class EmbeddingData(BaseModel): - object: str = "embedding" - embedding: list[float] - index: int + object: str = Field( + default="embedding", description="The object type, which is always 'embedding'." + ) + embedding: list[float] = Field( + description="The embedding vector, which is a list of floats." + ) + index: int = Field( + description="The index of the embedding in the list of embeddings." + ) class Usage(BaseModel): - prompt_tokens: int - total_tokens: int + prompt_tokens: int = Field(description="Number of tokens in the prompt.") + total_tokens: int = Field( + description="Total number of tokens used in the request (prompt + completion)." + ) class EmbeddingResponse(BaseModel): - object: str = "list" - data: list[EmbeddingData] - model: str - usage: Usage + object: str = Field( + default="list", description="The object type, which is always 'list'." + ) + data: list[EmbeddingData] = Field(description="A list of embedding objects.") + model: str = Field(description="The ID of the model used.") + usage: Usage = Field(description="Usage statistics for the request.") # --- For /v1/rerank --- class RerankRequest(BaseModel): - query: LimitedString + query: LimitedString = Field( + description="The search query.", examples=["What is the capital of France?"] + ) # Limit list size to prevent memory exhaustion (DoS) documents: Annotated[ list[LimitedString], @@ -101,25 +134,90 @@ class RerankRequest(BaseModel): min_length=1, max_length=MAX_INPUT_ITEMS, description="List of documents to rerank. Limited to MAX_INPUT_ITEMS to prevent DoS.", + examples=[ + ["Paris is the capital of France.", "Berlin is the capital of Germany."] + ], ), ] - model: str + model: str = Field( + description="ID of the model to use.", examples=["BAAI/bge-reranker-v2-m3"] + ) top_n: Optional[int] = Field( - None, validation_alias="top_k", ge=0, le=MAX_INPUT_ITEMS + default=None, + validation_alias="top_k", + ge=0, + le=MAX_INPUT_ITEMS, + description="The number of most relevant documents to return.", + examples=[1], + ) + return_documents: Optional[bool] = Field( + default=None, + description="If true, returns the document text along with the score.", + examples=[True], ) - return_documents: Optional[bool] = None - model_config = ConfigDict(populate_by_name=True) + model_config = ConfigDict( + populate_by_name=True, + json_schema_extra={ + "examples": [ + { + "query": "What is the capital of France?", + "documents": [ + "Paris is the capital of France.", + "Berlin is the capital of Germany.", + ], + "model": "BAAI/bge-reranker-v2-m3", + "top_n": 1, + "return_documents": True, + } + ] + }, + ) class RerankData(BaseModel): - document: int # As per the doc, this is the index - score: float - text: Optional[LimitedString] = None + document: int = Field(description="The index of the document in the original list.") + score: float = Field(description="The relevance score of the document.") + text: Optional[LimitedString] = Field( + default=None, + description="The text of the document, if `return_documents` is true.", + ) class RerankResponse(BaseModel): - query: LimitedString - data: list[RerankData] - model: str - usage: Optional[Usage] = None + query: LimitedString = Field(description="The original search query.") + data: list[RerankData] = Field(description="A list of ranked documents.") + model: str = Field(description="The ID of the model used.") + usage: Optional[Usage] = Field( + default=None, description="Usage statistics for the request." + ) + + +# --- Error Responses --- +class ErrorResponse(BaseModel): + detail: str = Field( + description="A detailed human-readable error message.", + examples=["Invalid API Key"], + ) + + +# --- Models API --- +class ModelCard(BaseModel): + id: str = Field(description="The model identifier.", examples=["BAAI/bge-m3"]) + object: str = Field( + default="model", description="The object type, which is always 'model'." + ) + created: int = Field( + default=0, + description="The Unix timestamp (in seconds) when the model was created.", + ) + owned_by: str = Field( + default="organization", description="The organization that owns the model." + ) + + +class ModelList(BaseModel): + object: str = Field( + default="list", description="The object type, which is always 'list'." + ) + data: list[ModelCard] = Field(description="A list of model objects.")