AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-06-04 11:14:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add program_name attribute to TiDB connection (#5499)

Browse Source 
Signed-off-by: Xiaoguang Sun <sunxiaoguang@gmail.com>
This commit is contained in:
Xiaoguang Sun 2024-06-24 14:41:07 +08:00 committed by GitHub
parent 756d9a4bc2
commit 13fcd7a901
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/config.py
View File

@ -33,6 +33,7 @@ class Config:
dotenv.load_dotenv() dotenv.load_dotenv()
self.TESTING = False self.TESTING = False
self.APPLICATION_NAME = "langgenius/dify"
# cors settings # cors settings
self.CONSOLE_CORS_ALLOW_ORIGINS = get_cors_allow_origins( self.CONSOLE_CORS_ALLOW_ORIGINS = get_cors_allow_origins(

6
api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py
View File

@ -26,6 +26,7 @@ class TiDBVectorConfig(BaseModel):
user: str user: str
password: str password: str
database: str database: str
program_name: str
@model_validator(mode='before') @model_validator(mode='before')
def validate_config(cls, values: dict) -> dict: def validate_config(cls, values: dict) -> dict:
@ -39,6 +40,8 @@ class TiDBVectorConfig(BaseModel):
raise ValueError("config TIDB_VECTOR_PASSWORD is required") raise ValueError("config TIDB_VECTOR_PASSWORD is required")
if not values['database']: if not values['database']:
raise ValueError("config TIDB_VECTOR_DATABASE is required") raise ValueError("config TIDB_VECTOR_DATABASE is required")
if not values['program_name']:
raise ValueError("config APPLICATION_NAME is required")
return values return values
@ -65,7 +68,7 @@ class TiDBVector(BaseVector):
super().__init__(collection_name) super().__init__(collection_name)
self._client_config = config self._client_config = config
self._url = (f"mysql+pymysql://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}?" self._url = (f"mysql+pymysql://{config.user}:{config.password}@{config.host}:{config.port}/{config.database}?"
f"ssl_verify_cert=true&ssl_verify_identity=true") f"ssl_verify_cert=true&ssl_verify_identity=true&program_name={config.program_name}")
self._distance_func = distance_func.lower() self._distance_func = distance_func.lower()
self._engine = create_engine(self._url) self._engine = create_engine(self._url)
self._orm_base = declarative_base() self._orm_base = declarative_base()
@ -245,5 +248,6 @@ class TiDBVectorFactory(AbstractVectorFactory):
user=config.get('TIDB_VECTOR_USER'), user=config.get('TIDB_VECTOR_USER'),
password=config.get('TIDB_VECTOR_PASSWORD'), password=config.get('TIDB_VECTOR_PASSWORD'),
database=config.get('TIDB_VECTOR_DATABASE'), database=config.get('TIDB_VECTOR_DATABASE'),
program_name=config.get('APPLICATION_NAME'),
), ),
) )

3
api/tests/integration_tests/vdb/tidb_vector/test_tidb_vector.py
View File

@ -16,7 +16,8 @@ def tidb_vector():
port="4000", port="4000",
user="xxx.root", user="xxx.root",
password="xxxxxx", password="xxxxxx",
database="dify" database="dify",
program_name="langgenius/dify"
) )
) )