Switch to agent arch

scripts/diagnose_embeddings.py

@@ -1,10 +1,9 @@
 import asyncio
-import os
 import logging
-import typer
+import os
 import random
-from google import genai
-from google.genai import types
+
+import typer
 from dotenv import load_dotenv
 from embedder.vertex_ai import VertexAIEmbedder
 
@@ -27,7 +26,7 @@ CONTENT_LIST = [
 ]
 TASK_TYPE = "RETRIEVAL_DOCUMENT"
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 app = typer.Typer()
@@ -43,14 +42,14 @@ async def embed_content_task():
 async def run_test(concurrency: int):
     """Continuously calls the embedding API and tracks requests."""
     total_requests = 0
-    
+
     logger.info(f"Starting diagnostic test with {concurrency} concurrent requests on model '{MODEL_NAME}'.")
     logger.info("Press Ctrl+C to stop.")
 
     while True:
         # Create tasks, passing project_id and location
         tasks = [embed_content_task() for _ in range(concurrency)]
-        
+
         try:
             await asyncio.gather(*tasks)
             total_requests += concurrency