add for settings tokens limit

This commit is contained in:
harold 2025-03-25 10:19:21 +05:00
parent 5aab77a753
commit 04d19844aa
3 changed files with 4 additions and 6 deletions

View File

@ -1,3 +0,0 @@
GEMINI_TOKENS_LIMIT: int = 500_000

View File

@ -2,8 +2,8 @@ import json
from google.generativeai import GenerativeModel
from src.gemini_sdk.constants import GEMINI_TOKENS_LIMIT
from src.service.storage import MESSAGES_STORAGE_SCHEMA
from src.settings.base import settings
def check_current_token_limit(
@ -13,7 +13,7 @@ def check_current_token_limit(
tokens_response_count = model.count_tokens(raw_response).total_tokens
print("tokens", model.count_tokens(raw_response).total_tokens)
if tokens_response_count >= GEMINI_TOKENS_LIMIT:
if tokens_response_count >= settings.GEMINI.TOKENS_LIMIT:
return False
return True

View File

@ -5,3 +5,4 @@ from pydantic import BaseModel
class GeminiSettings(BaseModel):
API_KEY: str
MODEL_NAME: str = "gemini-1.5-flash"
TOKENS_LIMIT: int