pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: change default models (#7777)

Browse Source 
### 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:
liu an
2025-05-22 11:59:12 +08:00
committed by Yingfeng Zhang
parent 42f4d4dbc8
commit e166f132b3
7 changed files with 221 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
sdk/python/test/conftest.py
View File

@ -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"))

65
sdk/python/test/test_http_api/conftest.py
View File

@ -16,7 +16,6 @@
import os
import pytest
import requests
from common import (
add_chunk,
batch_create_datasets,
@ -49,9 +48,6 @@ MARKER_EXPRESSIONS = {
"p3": "p1 or p2 or p3",
}
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 pytest_addoption(parser: pytest.Parser) -> None:
@ -85,67 +81,6 @@ def get_http_api_auth(get_api_key_fixture):
return RAGFlowHttpApiAuth(get_api_key_fixture)
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"))
@pytest.fixture(scope="function")
def clear_datasets(request, get_http_api_auth):
def cleanup():

57
sdk/python/test/test_sdk_api/t_chat.py
View File

@ -14,8 +14,9 @@
# limitations under the License.
#
from ragflow_sdk import RAGFlow
from common import HOST_ADDRESS
from ragflow_sdk import RAGFlow
from ragflow_sdk.modules.chat import Chat
def test_create_chat_with_name(get_api_key_fixture):
@ -31,7 +32,18 @@ def test_create_chat_with_name(get_api_key_fixture):
docs = kb.upload_documents(documents)
for doc in docs:
doc.add_chunk("This is a test to add chunk")
rag.create_chat("test_create_chat", dataset_ids=[kb.id])
llm = Chat.LLM(
rag,
{
"model_name": "glm-4-flash@ZHIPU-AI",
"temperature": 0.1,
"top_p": 0.3,
"presence_penalty": 0.4,
"frequency_penalty": 0.7,
"max_tokens": 512,
},
)
rag.create_chat("test_create_chat", dataset_ids=[kb.id], llm=llm)
def test_update_chat_with_name(get_api_key_fixture):
@ -47,7 +59,18 @@ def test_update_chat_with_name(get_api_key_fixture):
docs = kb.upload_documents(documents)
for doc in docs:
doc.add_chunk("This is a test to add chunk")
chat = rag.create_chat("test_update_chat", dataset_ids=[kb.id])
llm = Chat.LLM(
rag,
{
"model_name": "glm-4-flash@ZHIPU-AI",
"temperature": 0.1,
"top_p": 0.3,
"presence_penalty": 0.4,
"frequency_penalty": 0.7,
"max_tokens": 512,
},
)
chat = rag.create_chat("test_update_chat", dataset_ids=[kb.id], llm=llm)
chat.update({"name": "new_chat"})
@ -64,7 +87,18 @@ def test_delete_chats_with_success(get_api_key_fixture):
docs = kb.upload_documents(documents)
for doc in docs:
doc.add_chunk("This is a test to add chunk")
chat = rag.create_chat("test_delete_chat", dataset_ids=[kb.id])
llm = Chat.LLM(
rag,
{
"model_name": "glm-4-flash@ZHIPU-AI",
"temperature": 0.1,
"top_p": 0.3,
"presence_penalty": 0.4,
"frequency_penalty": 0.7,
"max_tokens": 512,
},
)
chat = rag.create_chat("test_delete_chat", dataset_ids=[kb.id], llm=llm)
rag.delete_chats(ids=[chat.id])
@ -81,6 +115,17 @@ def test_list_chats_with_success(get_api_key_fixture):
docs = kb.upload_documents(documents)
for doc in docs:
doc.add_chunk("This is a test to add chunk")
rag.create_chat("test_list_1", dataset_ids=[kb.id])
rag.create_chat("test_list_2", dataset_ids=[kb.id])
llm = Chat.LLM(
rag,
{
"model_name": "glm-4-flash@ZHIPU-AI",
"temperature": 0.1,
"top_p": 0.3,
"presence_penalty": 0.4,
"frequency_penalty": 0.7,
"max_tokens": 512,
},
)
rag.create_chat("test_list_1", dataset_ids=[kb.id], llm=llm)
rag.create_chat("test_list_2", dataset_ids=[kb.id], llm=llm)
rag.list_chats()