add quota manager

app/embedding.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 import numpy as np
 from loguru import logger
 import httpx
@@ -18,10 +18,7 @@ class EmbeddingClient:
         settings = get_settings()
         self.provider = getattr(settings, 'embedding_provider', 'default')
         self.model = getattr(settings, 'embedding_model', 'models/embedding-001')
-        if self.provider == 'gemini':
-            self.gemini_client = GeminiClient(settings.gemini_api_key, model=self.model)
-        else:
-            self.gemini_client = None
+        self.gemini_client: Optional[GeminiClient] = GeminiClient() if self.provider == 'gemini' else None
 
     @timing_decorator_async
     async def create_embedding(self, text: str) -> List[float]:
@@ -30,12 +27,15 @@ class EmbeddingClient:
         Input: text (str)
         Output: list[float] embedding vector.
         """
-        if self.provider == 'gemini' and self.gemini_client:
+        if self.provider == 'gemini':
+            if not self.gemini_client:
+                raise RuntimeError("GeminiClient is not initialized")
             try:
                 # GeminiClient.create_embedding là hàm sync, chạy trong executor
                 import asyncio
                 loop = asyncio.get_event_loop()
-                embedding = await loop.run_in_executor(None, self.gemini_client.create_embedding, text)
+                gemini_client = self.gemini_client  # type: ignore
+                embedding = await loop.run_in_executor(None, lambda: gemini_client.create_embedding(text, model=self.model))
                 # Kiểm tra kiểu dữ liệu trả về
                 if isinstance(embedding, list):
                     preview = f"{embedding[:10]}...{embedding[-10:]}" if len(embedding) > 20 else str(embedding)

app/gemini_client.py

@@ -3,6 +3,7 @@ from google.generativeai.client import configure
 from google.generativeai.generative_models import GenerativeModel
 from loguru import logger
 from .request_limit_manager import RequestLimitManager
+from typing import List, Optional
 
 class GeminiClient:
     def __init__(self):
@@ -49,17 +50,22 @@ class GeminiClient:
                 continue
         return 0
 
-    def create_embedding(self, text: str) -> list:
+    def create_embedding(self, text: str, model: Optional[str] = None) -> list:
         last_error = None
-        for key, model in self.limit_manager.iterate_key_model():
+        for key, m in self.limit_manager.iterate_key_model():
+            m = m or ""
+            use_model = model if model not in (None, "") else m
+            if not use_model:
+                continue
+            use_model = str(use_model)
             try:
                 configure(api_key=key)
                 response = embed_content(
-                    model=model,
+                    model=use_model,
                     content=text,
                     task_type="retrieval_query"
                 )
-                self.limit_manager.log_request(key, model, success=True)
+                self.limit_manager.log_request(key, use_model, success=True)
                 logger.info(f"[GEMINI][EMBEDDING][RAW_RESPONSE] {response['embedding'][:10]} ..... {response['embedding'][-10:]}")
                 return response['embedding']
             except Exception as e:
@@ -67,10 +73,10 @@ class GeminiClient:
                 msg = str(e)
                 if "429" in msg or "rate limit" in msg.lower():
                     retry_delay = 60
-                    m = re.search(r'retry_delay.*?seconds: (\d+)', msg)
-                    if m:
-                        retry_delay = int(m.group(1))
-                    self.limit_manager.log_request(key, model, success=False, retry_delay=retry_delay)
+                    m_retry = re.search(r'retry_delay.*?seconds: (\d+)', msg)
+                    if m_retry:
+                        retry_delay = int(m_retry.group(1))
+                    self.limit_manager.log_request(key, use_model, success=False, retry_delay=retry_delay)
                 last_error = e
                 continue
         raise last_error or RuntimeError("No available Gemini API key/model")