add for settings tokens limit
This commit is contained in:
parent
5aab77a753
commit
04d19844aa
@ -1,3 +0,0 @@
|
||||
|
||||
|
||||
GEMINI_TOKENS_LIMIT: int = 500_000
|
@ -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
|
||||
|
@ -5,3 +5,4 @@ from pydantic import BaseModel
|
||||
class GeminiSettings(BaseModel):
|
||||
API_KEY: str
|
||||
MODEL_NAME: str = "gemini-1.5-flash"
|
||||
TOKENS_LIMIT: int
|
Loading…
x
Reference in New Issue
Block a user