mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
Feat: change default models (#7777)
### What problem does this PR solve? change default models to buildin models https://github.com/infiniflow/ragflow/issues/7774 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
@ -20,7 +20,9 @@ import pytest
|
||||
import requests
|
||||
|
||||
HOST_ADDRESS = os.getenv("HOST_ADDRESS", "http://127.0.0.1:9380")
|
||||
|
||||
ZHIPU_AI_API_KEY = os.getenv("ZHIPU_AI_API_KEY", "ca148e43209c40109e2bc2f56281dd11.BltyA2N1B043B7Ra")
|
||||
if ZHIPU_AI_API_KEY is None:
|
||||
pytest.exit("Error: Environment variable ZHIPU_AI_API_KEY must be set")
|
||||
|
||||
# def generate_random_email():
|
||||
# return 'user_' + ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))+'@1.com'
|
||||
@ -87,3 +89,64 @@ def get_auth():
|
||||
@pytest.fixture(scope="session")
|
||||
def get_email():
|
||||
return EMAIL
|
||||
|
||||
|
||||
def get_my_llms(auth, name):
|
||||
url = HOST_ADDRESS + "/v1/llm/my_llms"
|
||||
authorization = {"Authorization": auth}
|
||||
response = requests.get(url=url, headers=authorization)
|
||||
res = response.json()
|
||||
if res.get("code") != 0:
|
||||
raise Exception(res.get("message"))
|
||||
if name in res.get("data"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def add_models(auth):
|
||||
url = HOST_ADDRESS + "/v1/llm/set_api_key"
|
||||
authorization = {"Authorization": auth}
|
||||
models_info = {
|
||||
"ZHIPU-AI": {"llm_factory": "ZHIPU-AI", "api_key": ZHIPU_AI_API_KEY},
|
||||
}
|
||||
|
||||
for name, model_info in models_info.items():
|
||||
if not get_my_llms(auth, name):
|
||||
response = requests.post(url=url, headers=authorization, json=model_info)
|
||||
res = response.json()
|
||||
if res.get("code") != 0:
|
||||
pytest.exit(f"Critical error in add_models: {res.get('message')}")
|
||||
|
||||
|
||||
def get_tenant_info(auth):
|
||||
url = HOST_ADDRESS + "/v1/user/tenant_info"
|
||||
authorization = {"Authorization": auth}
|
||||
response = requests.get(url=url, headers=authorization)
|
||||
res = response.json()
|
||||
if res.get("code") != 0:
|
||||
raise Exception(res.get("message"))
|
||||
return res["data"].get("tenant_id")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def set_tenant_info(get_auth):
|
||||
auth = get_auth
|
||||
try:
|
||||
add_models(auth)
|
||||
tenant_id = get_tenant_info(auth)
|
||||
except Exception as e:
|
||||
pytest.exit(f"Error in set_tenant_info: {str(e)}")
|
||||
url = HOST_ADDRESS + "/v1/user/set_tenant_info"
|
||||
authorization = {"Authorization": get_auth}
|
||||
tenant_info = {
|
||||
"tenant_id": tenant_id,
|
||||
"llm_id": "glm-4-flash@ZHIPU-AI",
|
||||
"embd_id": "BAAI/bge-large-zh-v1.5@BAAI",
|
||||
"img2txt_id": "glm-4v@ZHIPU-AI",
|
||||
"asr_id": "",
|
||||
"tts_id": None,
|
||||
}
|
||||
response = requests.post(url=url, headers=authorization, json=tenant_info)
|
||||
res = response.json()
|
||||
if res.get("code") != 0:
|
||||
raise Exception(res.get("message"))
|
||||
|
||||
Reference in New Issue
Block a user