소스 검색

add tidb spend limit config (#9999)

Jyong 5 달 전
부모
커밋
5580bcf870
2개의 변경된 파일7개의 추가작업 그리고 1개의 파일을 삭제
  1. 5 0
      api/configs/middleware/vdb/tidb_on_qdrant_config.py
  2. 2 1
      api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

+ 5 - 0
api/configs/middleware/vdb/tidb_on_qdrant_config.py

@@ -63,3 +63,8 @@ class TidbOnQdrantConfig(BaseSettings):
         description="Tidb project id",
         default=None,
     )
+
+    TIDB_SPEND_LIMIT: Optional[int] = Field(
+        description="Tidb spend limit",
+        default=100,
+    )

+ 2 - 1
api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -4,6 +4,7 @@ import uuid
 import requests
 from requests.auth import HTTPDigestAuth
 
+from configs import dify_config
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import TidbAuthBinding
@@ -208,7 +209,7 @@ class TidbService:
             }
 
             spending_limit = {
-                "monthly": 10,
+                "monthly": dify_config.TIDB_SPEND_LIMIT,
             }
             password = str(uuid.uuid4()).replace("-", "")[:16]
             display_name = str(uuid.uuid4()).replace("-", "")