VietCat commited on
Commit
f0e68b1
·
1 Parent(s): c0907f8

add quota manager

Browse files
Files changed (3) hide show
  1. app/config.py +8 -14
  2. app/main.py +1 -1
  3. app/request_limit_manager.py +2 -2
app/config.py CHANGED
@@ -34,8 +34,8 @@ class Settings(BaseSettings):
34
  # Gemini Configuration
35
  # Hỗ trợ nhiều API key và model cho Gemini
36
  # Định nghĩa biến môi trường: GEMINI_API_KEYS="key1,key2,..."; GEMINI_MODELS="model1,model2,..."
37
- gemini_api_keys: list[str] = []
38
- gemini_models: list[str] = []
39
  gemini_base_url: str = os.getenv("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1") or ""
40
 
41
  # LLM (chat/completion) provider/model
@@ -48,19 +48,13 @@ class Settings(BaseSettings):
48
  rerank_provider: str = os.getenv("RERANK_PROVIDER", "") or llm_provider
49
  rerank_model: str = os.getenv("RERANK_MODEL", "") or llm_model
50
 
51
- @field_validator("gemini_api_keys", mode="before")
52
- @classmethod
53
- def split_gemini_api_keys(cls, v):
54
- if isinstance(v, str):
55
- return [k.strip() for k in v.split(",") if k.strip()]
56
- return v
57
 
58
- @field_validator("gemini_models", mode="before")
59
- @classmethod
60
- def split_gemini_models(cls, v):
61
- if isinstance(v, str):
62
- return [m.strip() for m in v.split(",") if m.strip()]
63
- return v
64
 
65
  class Config:
66
  env_file = ".env"
 
34
  # Gemini Configuration
35
  # Hỗ trợ nhiều API key và model cho Gemini
36
  # Định nghĩa biến môi trường: GEMINI_API_KEYS="key1,key2,..."; GEMINI_MODELS="model1,model2,..."
37
+ gemini_api_keys: str = ""
38
+ gemini_models: str = ""
39
  gemini_base_url: str = os.getenv("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1") or ""
40
 
41
  # LLM (chat/completion) provider/model
 
48
  rerank_provider: str = os.getenv("RERANK_PROVIDER", "") or llm_provider
49
  rerank_model: str = os.getenv("RERANK_MODEL", "") or llm_model
50
 
51
+ @property
52
+ def gemini_api_keys_list(self):
53
+ return [k.strip() for k in self.gemini_api_keys.split(",") if k.strip()]
 
 
 
54
 
55
+ @property
56
+ def gemini_models_list(self):
57
+ return [m.strip() for m in self.gemini_models.split(",") if m.strip()]
 
 
 
58
 
59
  class Config:
60
  env_file = ".env"
app/main.py CHANGED
@@ -64,7 +64,7 @@ VEHICLE_KEYWORDS = ["xe máy", "ô tô", "xe đạp", "xe hơi"]
64
  llm_client = create_llm_client(
65
  provider="gemini",
66
  base_url=settings.gemini_base_url,
67
- model=settings.gemini_models[0] if settings.gemini_models else "gemini-2.5-flash"
68
  )
69
 
70
  reranker = Reranker()
 
64
  llm_client = create_llm_client(
65
  provider="gemini",
66
  base_url=settings.gemini_base_url,
67
+ model=settings.gemini_models_list[0] if settings.gemini_models_list else "gemini-2.5-flash"
68
  )
69
 
70
  reranker = Reranker()
app/request_limit_manager.py CHANGED
@@ -12,8 +12,8 @@ class RequestLimitManager:
12
  def _init_keys_models(self):
13
  settings = get_settings()
14
  if self.provider == "gemini":
15
- self.api_keys: List[str] = getattr(settings, 'gemini_api_keys', [])
16
- self.models: List[str] = getattr(settings, 'gemini_models', [])
17
  # Có thể mở rộng cho provider khác ở đây
18
  self.status: Dict[str, Dict[str, Dict[str, Union[str, float]]]] = {}
19
  now = time.time()
 
12
  def _init_keys_models(self):
13
  settings = get_settings()
14
  if self.provider == "gemini":
15
+ self.api_keys: List[str] = getattr(settings, 'gemini_api_keys_list', [])
16
+ self.models: List[str] = getattr(settings, 'gemini_models_list', [])
17
  # Có thể mở rộng cho provider khác ở đây
18
  self.status: Dict[str, Dict[str, Dict[str, Union[str, float]]]] = {}
19
  now = time.time()