浏览代码

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

Signed-off-by: Xiaoguang Sun <sunxiaoguang@gmail.com>
tags/0.6.12
Xiaoguang Sun 1年前
父节点
当前提交
13fcd7a901
没有帐户链接到提交者的电子邮件

+ 1
- 0
api/config.py 查看文件

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(

+ 5
- 1
api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py 查看文件

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:
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




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()
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'),
), ),
) )

+ 2
- 1
api/tests/integration_tests/vdb/tidb_vector/test_tidb_vector.py 查看文件

port="4000", port="4000",
user="xxx.root", user="xxx.root",
password="xxxxxx", password="xxxxxx",
database="dify"
database="dify",
program_name="langgenius/dify"
) )
) )



正在加载...
取消
保存