liuhua
liuhua
commited on
Commit
·
95f8bbb
1
Parent(s):
4cb6b27
Test Cases (#2993)
Browse files### What problem does this PR solve?
Test Cases
### Type of change
- [x] Refactoring
Co-authored-by: liuhua <[email protected]>
- api/apps/sdk/dataset.py +9 -11
- api/apps/sdk/doc.py +1 -0
- api/utils/api_utils.py +1 -1
- sdk/python/ragflow/ragflow.py +3 -1
- sdk/python/test/common.py +0 -4
- sdk/python/test/conftest.py +52 -0
- sdk/python/test/t_chat.py +66 -56
- sdk/python/test/t_dataset.py +54 -53
- sdk/python/test/t_document.py +317 -294
- sdk/python/test/t_session.py +108 -50
- sdk/python/test/test_sdkbase.py +0 -3
api/apps/sdk/dataset.py
CHANGED
@@ -37,9 +37,9 @@ def create(tenant_id):
|
|
37 |
language = req.get("language")
|
38 |
chunk_method = req.get("chunk_method")
|
39 |
parser_config = req.get("parser_config")
|
40 |
-
valid_permission =
|
41 |
-
valid_language =
|
42 |
-
valid_chunk_method =
|
43 |
check_validation=valid(permission,valid_permission,language,valid_language,chunk_method,valid_chunk_method)
|
44 |
if check_validation:
|
45 |
return check_validation
|
@@ -47,10 +47,8 @@ def create(tenant_id):
|
|
47 |
if "tenant_id" in req:
|
48 |
return get_error_data_result(
|
49 |
retmsg="`tenant_id` must not be provided")
|
50 |
-
chunk_count
|
51 |
-
|
52 |
-
if chunk_count or document_count:
|
53 |
-
return get_error_data_result(retmsg="`chunk_count` or `document_count` must be 0 or not be provided")
|
54 |
if "name" not in req:
|
55 |
return get_error_data_result(
|
56 |
retmsg="`name` is not empty!")
|
@@ -123,10 +121,10 @@ def update(tenant_id,dataset_id):
|
|
123 |
language = req.get("language")
|
124 |
chunk_method = req.get("chunk_method")
|
125 |
parser_config = req.get("parser_config")
|
126 |
-
valid_permission =
|
127 |
-
valid_language =
|
128 |
-
valid_chunk_method =
|
129 |
-
"knowledge_graph", "email"
|
130 |
check_validation = valid(permission, valid_permission, language, valid_language, chunk_method, valid_chunk_method)
|
131 |
if check_validation:
|
132 |
return check_validation
|
|
|
37 |
language = req.get("language")
|
38 |
chunk_method = req.get("chunk_method")
|
39 |
parser_config = req.get("parser_config")
|
40 |
+
valid_permission = ["me", "team"]
|
41 |
+
valid_language =["Chinese", "English"]
|
42 |
+
valid_chunk_method = ["naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email"]
|
43 |
check_validation=valid(permission,valid_permission,language,valid_language,chunk_method,valid_chunk_method)
|
44 |
if check_validation:
|
45 |
return check_validation
|
|
|
47 |
if "tenant_id" in req:
|
48 |
return get_error_data_result(
|
49 |
retmsg="`tenant_id` must not be provided")
|
50 |
+
if "chunk_count" in req or "document_count" in req:
|
51 |
+
return get_error_data_result(retmsg="`chunk_count` or `document_count` must not be provided")
|
|
|
|
|
52 |
if "name" not in req:
|
53 |
return get_error_data_result(
|
54 |
retmsg="`name` is not empty!")
|
|
|
121 |
language = req.get("language")
|
122 |
chunk_method = req.get("chunk_method")
|
123 |
parser_config = req.get("parser_config")
|
124 |
+
valid_permission = ["me", "team"]
|
125 |
+
valid_language = ["Chinese", "English"]
|
126 |
+
valid_chunk_method = ["naive", "manual", "qa", "table", "paper", "book", "laws", "presentation", "picture", "one",
|
127 |
+
"knowledge_graph", "email"]
|
128 |
check_validation = valid(permission, valid_permission, language, valid_language, chunk_method, valid_chunk_method)
|
129 |
if check_validation:
|
130 |
return check_validation
|
api/apps/sdk/doc.py
CHANGED
@@ -44,6 +44,7 @@ from rag.nlp import search
|
|
44 |
from rag.utils import rmSpace
|
45 |
from rag.utils.es_conn import ELASTICSEARCH
|
46 |
from rag.utils.storage_factory import STORAGE_IMPL
|
|
|
47 |
|
48 |
MAXIMUM_OF_UPLOADING_FILES = 256
|
49 |
|
|
|
44 |
from rag.utils import rmSpace
|
45 |
from rag.utils.es_conn import ELASTICSEARCH
|
46 |
from rag.utils.storage_factory import STORAGE_IMPL
|
47 |
+
import os
|
48 |
|
49 |
MAXIMUM_OF_UPLOADING_FILES = 256
|
50 |
|
api/utils/api_utils.py
CHANGED
@@ -337,7 +337,7 @@ def valid(permission,valid_permission,language,valid_language,chunk_method,valid
|
|
337 |
|
338 |
def valid_parameter(parameter,valid_values):
|
339 |
if parameter and parameter not in valid_values:
|
340 |
-
return get_error_data_result(f"{parameter}
|
341 |
|
342 |
def get_parser_config(chunk_method,parser_config):
|
343 |
if parser_config:
|
|
|
337 |
|
338 |
def valid_parameter(parameter,valid_values):
|
339 |
if parameter and parameter not in valid_values:
|
340 |
+
return get_error_data_result(f"{parameter} is not in {valid_values}")
|
341 |
|
342 |
def get_parser_config(chunk_method,parser_config):
|
343 |
if parser_config:
|
sdk/python/ragflow/ragflow.py
CHANGED
@@ -51,6 +51,8 @@ class RAGFlow:
|
|
51 |
def create_dataset(self, name: str, avatar: str = "", description: str = "", language: str = "English",
|
52 |
permission: str = "me",chunk_method: str = "naive",
|
53 |
parser_config: DataSet.ParserConfig = None) -> DataSet:
|
|
|
|
|
54 |
res = self.post("/dataset",
|
55 |
{"name": name, "avatar": avatar, "description": description, "language": language,
|
56 |
"permission": permission, "chunk_method": chunk_method,
|
@@ -91,7 +93,7 @@ class RAGFlow:
|
|
91 |
llm: Chat.LLM = None, prompt: Chat.Prompt = None) -> Chat:
|
92 |
dataset_list = []
|
93 |
for dataset in datasets:
|
94 |
-
dataset_list.append(dataset.
|
95 |
|
96 |
if llm is None:
|
97 |
llm = Chat.LLM(self, {"model_name": None,
|
|
|
51 |
def create_dataset(self, name: str, avatar: str = "", description: str = "", language: str = "English",
|
52 |
permission: str = "me",chunk_method: str = "naive",
|
53 |
parser_config: DataSet.ParserConfig = None) -> DataSet:
|
54 |
+
if parser_config:
|
55 |
+
parser_config = parser_config.to_json()
|
56 |
res = self.post("/dataset",
|
57 |
{"name": name, "avatar": avatar, "description": description, "language": language,
|
58 |
"permission": permission, "chunk_method": chunk_method,
|
|
|
93 |
llm: Chat.LLM = None, prompt: Chat.Prompt = None) -> Chat:
|
94 |
dataset_list = []
|
95 |
for dataset in datasets:
|
96 |
+
dataset_list.append(dataset.id)
|
97 |
|
98 |
if llm is None:
|
99 |
llm = Chat.LLM(self, {"model_name": None,
|
sdk/python/test/common.py
DELETED
@@ -1,4 +0,0 @@
|
|
1 |
-
|
2 |
-
|
3 |
-
API_KEY = 'ragflow-NiYmZjNTVjODYwNzExZWZiODEwMDI0Mm'
|
4 |
-
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
|
|
|
|
|
|
|
|
|
sdk/python/test/conftest.py
ADDED
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import pytest
|
2 |
+
import requests
|
3 |
+
import string
|
4 |
+
import random
|
5 |
+
|
6 |
+
|
7 |
+
|
8 |
+
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
9 |
+
|
10 |
+
def generate_random_email():
|
11 |
+
return 'user_' + ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))+'@1.com'
|
12 |
+
|
13 |
+
EMAIL = generate_random_email()
|
14 |
+
# password is "123"
|
15 |
+
PASSWORD='''ctAseGvejiaSWWZ88T/m4FQVOpQyUvP+x7sXtdv3feqZACiQleuewkUi35E16wSd5C5QcnkkcV9cYc8TKPTRZlxappDuirxghxoOvFcJxFU4ixLsD
|
16 |
+
fN33jCHRoDUW81IH9zjij/vaw8IbVyb6vuwg6MX6inOEBRRzVbRYxXOu1wkWY6SsI8X70oF9aeLFp/PzQpjoe/YbSqpTq8qqrmHzn9vO+yvyYyvmDsphXe
|
17 |
+
X8f7fp9c7vUsfOCkM+gHY3PadG+QHa7KI7mzTKgUTZImK6BZtfRBATDTthEUbbaTewY4H0MnWiCeeDhcbeQao6cFy1To8pE3RpmxnGnS8BsBn8w=='''
|
18 |
+
|
19 |
+
def get_email():
|
20 |
+
return EMAIL
|
21 |
+
|
22 |
+
def register():
|
23 |
+
url = HOST_ADDRESS + "/v1/user/register"
|
24 |
+
name = "user"
|
25 |
+
register_data = {"email":EMAIL,"nickname":name,"password":PASSWORD}
|
26 |
+
res = requests.post(url=url,json=register_data)
|
27 |
+
res = res.json()
|
28 |
+
if res.get("retcode") != 0:
|
29 |
+
raise Exception(res.get("retmsg"))
|
30 |
+
|
31 |
+
def login():
|
32 |
+
url = HOST_ADDRESS + "/v1/user/login"
|
33 |
+
login_data = {"email":EMAIL,"password":PASSWORD}
|
34 |
+
response=requests.post(url=url,json=login_data)
|
35 |
+
res = response.json()
|
36 |
+
if res.get("retcode")!=0:
|
37 |
+
raise Exception(res.get("retmsg"))
|
38 |
+
auth = response.headers["Authorization"]
|
39 |
+
return auth
|
40 |
+
|
41 |
+
@pytest.fixture(scope="session")
|
42 |
+
def get_api_key_fixture():
|
43 |
+
register()
|
44 |
+
auth = login()
|
45 |
+
url = HOST_ADDRESS + "/v1/system/new_token"
|
46 |
+
auth = {"Authorization": auth}
|
47 |
+
response = requests.post(url=url,headers=auth)
|
48 |
+
res = response.json()
|
49 |
+
if res.get("retcode") != 0:
|
50 |
+
raise Exception(res.get("retmsg"))
|
51 |
+
return res["data"].get("token")
|
52 |
+
|
sdk/python/test/t_chat.py
CHANGED
@@ -1,57 +1,67 @@
|
|
1 |
from ragflow import RAGFlow, Chat
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
from ragflow import RAGFlow, Chat
|
2 |
+
import time
|
3 |
+
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
4 |
+
|
5 |
+
def test_create_chat_with_name(get_api_key_fixture):
|
6 |
+
API_KEY = get_api_key_fixture
|
7 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
8 |
+
kb = rag.create_dataset(name="test_create_chat")
|
9 |
+
displayed_name = "ragflow.txt"
|
10 |
+
with open("./ragflow.txt","rb") as file:
|
11 |
+
blob = file.read()
|
12 |
+
document = {"displayed_name":displayed_name,"blob":blob}
|
13 |
+
documents = []
|
14 |
+
documents.append(document)
|
15 |
+
doc_ids = []
|
16 |
+
docs= kb.upload_documents(documents)
|
17 |
+
for doc in docs:
|
18 |
+
doc_ids.append(doc.id)
|
19 |
+
kb.async_parse_documents(doc_ids)
|
20 |
+
time.sleep(60)
|
21 |
+
rag.create_chat("test_create", datasets=[kb])
|
22 |
+
|
23 |
+
|
24 |
+
def test_update_chat_with_name(get_api_key_fixture):
|
25 |
+
API_KEY = get_api_key_fixture
|
26 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
27 |
+
kb = rag.create_dataset(name="test_update_chat")
|
28 |
+
displayed_name = "ragflow.txt"
|
29 |
+
with open("./ragflow.txt", "rb") as file:
|
30 |
+
blob = file.read()
|
31 |
+
document = {"displayed_name": displayed_name, "blob": blob}
|
32 |
+
documents = []
|
33 |
+
documents.append(document)
|
34 |
+
doc_ids = []
|
35 |
+
docs = kb.upload_documents(documents)
|
36 |
+
for doc in docs:
|
37 |
+
doc_ids.append(doc.id)
|
38 |
+
kb.async_parse_documents(doc_ids)
|
39 |
+
time.sleep(60)
|
40 |
+
chat = rag.create_chat("test_update", datasets=[kb])
|
41 |
+
chat.update({"name": "new_chat"})
|
42 |
+
|
43 |
+
|
44 |
+
def test_delete_chats_with_success(get_api_key_fixture):
|
45 |
+
API_KEY = get_api_key_fixture
|
46 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
47 |
+
kb = rag.create_dataset(name="test_delete_chat")
|
48 |
+
displayed_name = "ragflow.txt"
|
49 |
+
with open("./ragflow.txt", "rb") as file:
|
50 |
+
blob = file.read()
|
51 |
+
document = {"displayed_name": displayed_name, "blob": blob}
|
52 |
+
documents = []
|
53 |
+
documents.append(document)
|
54 |
+
doc_ids = []
|
55 |
+
docs = kb.upload_documents(documents)
|
56 |
+
for doc in docs:
|
57 |
+
doc_ids.append(doc.id)
|
58 |
+
kb.async_parse_documents(doc_ids)
|
59 |
+
time.sleep(60)
|
60 |
+
chat = rag.create_chat("test_delete", datasets=[kb])
|
61 |
+
rag.delete_chats(ids=[chat.id])
|
62 |
+
|
63 |
+
API_KEY = get_api_key_fixture
|
64 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
65 |
+
rag.list_chats()
|
66 |
+
|
67 |
+
|
sdk/python/test/t_dataset.py
CHANGED
@@ -1,53 +1,54 @@
|
|
1 |
-
from ragflow import RAGFlow
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
""
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
|
|
|
1 |
+
from ragflow import RAGFlow
|
2 |
+
import random
|
3 |
+
import pytest
|
4 |
+
|
5 |
+
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
6 |
+
|
7 |
+
def test_create_dataset_with_name(get_api_key_fixture):
|
8 |
+
API_KEY = get_api_key_fixture
|
9 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
10 |
+
rag.create_dataset("test_create_dataset_with_name")
|
11 |
+
|
12 |
+
def test_create_dataset_with_duplicated_name(get_api_key_fixture):
|
13 |
+
API_KEY = get_api_key_fixture
|
14 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
15 |
+
with pytest.raises(Exception) as exc_info:
|
16 |
+
rag.create_dataset("test_create_dataset_with_name")
|
17 |
+
assert str(exc_info.value) == "Duplicated dataset name in creating dataset."
|
18 |
+
|
19 |
+
def test_create_dataset_with_random_chunk_method(get_api_key_fixture):
|
20 |
+
API_KEY = get_api_key_fixture
|
21 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
22 |
+
valid_chunk_methods = ["naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email"]
|
23 |
+
random_chunk_method = random.choice(valid_chunk_methods)
|
24 |
+
rag.create_dataset("test_create_dataset_with_random_chunk_method",chunk_method=random_chunk_method)
|
25 |
+
|
26 |
+
def test_create_dataset_with_invalid_parameter(get_api_key_fixture):
|
27 |
+
API_KEY = get_api_key_fixture
|
28 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
29 |
+
valid_chunk_methods = ["naive", "manual", "qa", "table", "paper", "book", "laws", "presentation", "picture", "one",
|
30 |
+
"knowledge_graph", "email"]
|
31 |
+
chunk_method = "invalid_chunk_method"
|
32 |
+
with pytest.raises(Exception) as exc_info:
|
33 |
+
rag.create_dataset("test_create_dataset_with_name",chunk_method=chunk_method)
|
34 |
+
assert str(exc_info.value) == f"{chunk_method} is not in {valid_chunk_methods}"
|
35 |
+
|
36 |
+
|
37 |
+
def test_update_dataset_with_name(get_api_key_fixture):
|
38 |
+
API_KEY = get_api_key_fixture
|
39 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
40 |
+
ds = rag.create_dataset("test_update_dataset")
|
41 |
+
ds.update({"name": "updated_dataset"})
|
42 |
+
|
43 |
+
|
44 |
+
def test_delete_datasets_with_success(get_api_key_fixture):
|
45 |
+
API_KEY = get_api_key_fixture
|
46 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
47 |
+
ds = rag.create_dataset("MA")
|
48 |
+
rag.delete_datasets(ids=[ds.id])
|
49 |
+
|
50 |
+
|
51 |
+
def test_list_datasets_with_success(get_api_key_fixture):
|
52 |
+
API_KEY = get_api_key_fixture
|
53 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
54 |
+
rag.list_datasets()
|
sdk/python/test/t_document.py
CHANGED
@@ -1,298 +1,321 @@
|
|
1 |
from ragflow import RAGFlow, DataSet, Document, Chunk
|
2 |
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
#
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
""
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
""
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
|
123 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
for progress, msg in doc.join(interval=5, timeout=10):
|
153 |
-
print(progress, msg)
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
# Test cancelling the parsing operation
|
159 |
-
doc.cancel()
|
160 |
-
# Print message to confirm parsing has been cancelled successfully
|
161 |
-
print("Parsing cancelled successfully")
|
162 |
-
|
163 |
-
def test_bulk_parse_and_cancel_documents(self):
|
164 |
-
# Initialize RAGFlow with API key and host address
|
165 |
-
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
166 |
-
|
167 |
-
# Create a dataset
|
168 |
-
ds = rag.create_dataset(name="God5")
|
169 |
-
assert ds is not None, "Dataset creation failed"
|
170 |
-
assert ds.name == "God5", "Dataset name does not match"
|
171 |
-
|
172 |
-
# Prepare a list of file names and paths
|
173 |
-
documents = [
|
174 |
-
{'name': 'test1.txt', 'path': 'test_data/test1.txt'},
|
175 |
-
{'name': 'test2.txt', 'path': 'test_data/test2.txt'},
|
176 |
-
{'name': 'test3.txt', 'path': 'test_data/test3.txt'}
|
177 |
-
]
|
178 |
-
|
179 |
-
# Create documents in bulk
|
180 |
-
for doc_info in documents:
|
181 |
-
with open(doc_info['path'], "rb") as file:
|
182 |
-
created_doc = rag.create_document(ds, name=doc_info['name'], blob=file.read())
|
183 |
-
assert created_doc is not None, f"Failed to create document {doc_info['name']}"
|
184 |
-
|
185 |
-
# Retrieve document objects in bulk
|
186 |
-
docs = [rag.get_document(name=doc_info['name']) for doc_info in documents]
|
187 |
-
ids = [doc.id for doc in docs]
|
188 |
-
assert len(docs) == len(documents), "Mismatch between created documents and fetched documents"
|
189 |
-
|
190 |
-
# Initiate asynchronous parsing for all documents
|
191 |
-
rag.async_parse_documents(ids)
|
192 |
-
print("Async bulk parsing initiated")
|
193 |
-
|
194 |
-
# Wait for all documents to finish parsing and check progress
|
195 |
-
for doc in docs:
|
196 |
-
for progress, msg in doc.join(interval=5, timeout=10):
|
197 |
-
print(f"{doc.name}: Progress: {progress}, Message: {msg}")
|
198 |
-
|
199 |
-
# Assert that progress is within the valid range
|
200 |
-
assert 0 <= progress <= 100, f"Invalid progress: {progress} for document {doc.name}"
|
201 |
-
|
202 |
-
# Assert that the message is not empty
|
203 |
-
assert msg, f"Message should not be empty for document {doc.name}"
|
204 |
-
|
205 |
-
# If progress reaches 100%, assert that parsing is completed successfully
|
206 |
-
if progress == 100:
|
207 |
-
assert "completed" in msg.lower(), f"Document {doc.name} did not complete successfully"
|
208 |
-
|
209 |
-
# Cancel parsing for all documents in bulk
|
210 |
-
cancel_result = rag.async_cancel_parse_documents(ids)
|
211 |
-
assert cancel_result is None or isinstance(cancel_result, type(None)), "Failed to cancel document parsing"
|
212 |
-
print("Async bulk parsing cancelled")
|
213 |
-
|
214 |
-
def test_parse_document_and_chunk_list(self):
|
215 |
-
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
216 |
-
ds = rag.create_dataset(name="God7")
|
217 |
-
name = 'story.txt'
|
218 |
-
path = 'test_data/story.txt'
|
219 |
-
# name = "Test Document rag.txt"
|
220 |
-
# blob = " Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps.Sample document content for rag test66. rag wonderful apple os documents apps.Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps."
|
221 |
-
rag.create_document(ds, name=name, blob=open(path, "rb").read())
|
222 |
-
doc = rag.get_document(name=name)
|
223 |
-
doc.async_parse()
|
224 |
-
|
225 |
-
# Wait for parsing to complete and get progress updates using join
|
226 |
-
for progress, msg in doc.join(interval=5, timeout=30):
|
227 |
-
print(progress, msg)
|
228 |
-
# Assert that progress is within 0 to 100
|
229 |
-
assert 0 <= progress <= 100, f"Invalid progress: {progress}"
|
230 |
# Assert that the message is not empty
|
231 |
-
assert msg, "Message should not be empty"
|
232 |
-
|
233 |
-
|
234 |
-
|
235 |
-
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
244 |
-
|
245 |
-
|
246 |
-
|
247 |
-
|
248 |
-
|
249 |
-
|
250 |
-
|
251 |
-
|
252 |
-
|
253 |
-
|
254 |
-
|
255 |
-
|
256 |
-
|
257 |
-
|
258 |
-
|
259 |
-
|
260 |
-
|
261 |
-
assert
|
262 |
-
|
263 |
-
|
264 |
-
|
265 |
-
assert
|
266 |
-
|
267 |
-
|
268 |
-
|
269 |
-
|
270 |
-
|
271 |
-
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
281 |
-
|
282 |
-
|
283 |
-
|
284 |
-
|
285 |
-
|
286 |
-
|
287 |
-
|
288 |
-
|
289 |
-
|
290 |
-
|
291 |
-
|
292 |
-
|
293 |
-
|
294 |
-
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
from ragflow import RAGFlow, DataSet, Document, Chunk
|
2 |
|
3 |
+
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
4 |
+
|
5 |
+
|
6 |
+
def test_upload_document_with_success(get_api_key_fixture):
|
7 |
+
"""
|
8 |
+
Test ingesting a document into a dataset with success.
|
9 |
+
"""
|
10 |
+
# Initialize RAGFlow instance
|
11 |
+
API_KEY = get_api_key_fixture
|
12 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
13 |
+
|
14 |
+
# Step 1: Create a new dataset
|
15 |
+
ds = rag.create_dataset(name="God")
|
16 |
+
|
17 |
+
# Ensure dataset creation was successful
|
18 |
+
assert isinstance(ds, DataSet), f"Failed to create dataset, error: {ds}"
|
19 |
+
assert ds.name == "God", "Dataset name does not match."
|
20 |
+
|
21 |
+
# Step 2: Create a new document
|
22 |
+
# The blob is the actual file content or a placeholder in this case
|
23 |
+
blob = b"Sample document content for ingestion test."
|
24 |
+
blob_2 = b"test_2."
|
25 |
+
list_1 = []
|
26 |
+
list_1.append({"name": "Test_1.txt",
|
27 |
+
"blob": blob})
|
28 |
+
list_1.append({"name": "Test_2.txt",
|
29 |
+
"blob": blob_2})
|
30 |
+
res = ds.upload_documents(list_1)
|
31 |
+
# Ensure document ingestion was successful
|
32 |
+
assert res is None, f"Failed to create document, error: {res}"
|
33 |
+
|
34 |
+
|
35 |
+
def test_update_document_with_success(get_api_key_fixture):
|
36 |
+
"""
|
37 |
+
Test updating a document with success.
|
38 |
+
Update name or chunk_method are supported
|
39 |
+
"""
|
40 |
+
API_KEY = get_api_key_fixture
|
41 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
42 |
+
ds = rag.list_datasets(name="God")
|
43 |
+
ds = ds[0]
|
44 |
+
doc = ds.list_documents()
|
45 |
+
doc = doc[0]
|
46 |
+
if isinstance(doc, Document):
|
47 |
+
res = doc.update({"chunk_method": "manual", "name": "manual.txt"})
|
48 |
+
assert res is None, f"Failed to update document, error: {res}"
|
49 |
+
else:
|
50 |
+
assert False, f"Failed to get document, error: {doc}"
|
51 |
+
|
52 |
+
|
53 |
+
def test_download_document_with_success(get_api_key_fixture):
|
54 |
+
"""
|
55 |
+
Test downloading a document with success.
|
56 |
+
"""
|
57 |
+
API_KEY = get_api_key_fixture
|
58 |
+
# Initialize RAGFlow instance
|
59 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
60 |
+
|
61 |
+
# Retrieve a document
|
62 |
+
ds = rag.list_datasets(name="God")
|
63 |
+
ds = ds[0]
|
64 |
+
doc = ds.list_documents(name="manual.txt")
|
65 |
+
doc = doc[0]
|
66 |
+
# Check if the retrieved document is of type Document
|
67 |
+
if isinstance(doc, Document):
|
68 |
+
# Download the document content and save it to a file
|
69 |
+
with open("./ragflow.txt", "wb+") as file:
|
70 |
+
file.write(doc.download())
|
71 |
+
# Print the document object for debugging
|
72 |
+
print(doc)
|
73 |
+
|
74 |
+
# Assert that the download was successful
|
75 |
+
assert True, f"Failed to download document, error: {doc}"
|
76 |
+
else:
|
77 |
+
# If the document retrieval fails, assert failure
|
78 |
+
assert False, f"Failed to get document, error: {doc}"
|
79 |
+
|
80 |
+
|
81 |
+
def test_list_documents_in_dataset_with_success(get_api_key_fixture):
|
82 |
+
"""
|
83 |
+
Test list all documents into a dataset with success.
|
84 |
+
"""
|
85 |
+
API_KEY = get_api_key_fixture
|
86 |
+
# Initialize RAGFlow instance
|
87 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
88 |
+
|
89 |
+
# Step 1: Create a new dataset
|
90 |
+
ds = rag.create_dataset(name="God2")
|
91 |
+
|
92 |
+
# Ensure dataset creation was successful
|
93 |
+
assert isinstance(ds, DataSet), f"Failed to create dataset, error: {ds}"
|
94 |
+
assert ds.name == "God2", "Dataset name does not match."
|
95 |
+
|
96 |
+
# Step 2: Create a new document
|
97 |
+
# The blob is the actual file content or a placeholder in this case
|
98 |
+
name1 = "Test Document111.txt"
|
99 |
+
blob1 = b"Sample document content for ingestion test111."
|
100 |
+
name2 = "Test Document222.txt"
|
101 |
+
blob2 = b"Sample document content for ingestion test222."
|
102 |
+
list_1 = [{"name": name1, "blob": blob1}, {"name": name2, "blob": blob2}]
|
103 |
+
ds.upload_documents(list_1)
|
104 |
+
for d in ds.list_documents(keywords="test", offset=0, limit=12):
|
105 |
+
assert isinstance(d, Document), "Failed to upload documents"
|
106 |
+
|
107 |
+
|
108 |
+
def test_delete_documents_in_dataset_with_success(get_api_key_fixture):
|
109 |
+
"""
|
110 |
+
Test list all documents into a dataset with success.
|
111 |
+
"""
|
112 |
+
API_KEY = get_api_key_fixture
|
113 |
+
# Initialize RAGFlow instance
|
114 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
115 |
+
|
116 |
+
# Step 1: Create a new dataset
|
117 |
+
ds = rag.create_dataset(name="God3")
|
118 |
+
|
119 |
+
# Ensure dataset creation was successful
|
120 |
+
assert isinstance(ds, DataSet), f"Failed to create dataset, error: {ds}"
|
121 |
+
assert ds.name == "God3", "Dataset name does not match."
|
122 |
+
|
123 |
+
# Step 2: Create a new document
|
124 |
+
# The blob is the actual file content or a placeholder in this case
|
125 |
+
name1 = "Test Document333.txt"
|
126 |
+
blob1 = b"Sample document content for ingestion test333."
|
127 |
+
name2 = "Test Document444.txt"
|
128 |
+
blob2 = b"Sample document content for ingestion test444."
|
129 |
+
ds.upload_documents([{"name": name1, "blob": blob1}, {"name": name2, "blob": blob2}])
|
130 |
+
for d in ds.list_documents(keywords="document", offset=0, limit=12):
|
131 |
+
assert isinstance(d, Document)
|
132 |
+
ds.delete_documents([d.id])
|
133 |
+
remaining_docs = ds.list_documents(keywords="rag", offset=0, limit=12)
|
134 |
+
assert len(remaining_docs) == 0, "Documents were not properly deleted."
|
135 |
+
|
136 |
+
|
137 |
+
def test_parse_and_cancel_document(get_api_key_fixture):
|
138 |
+
API_KEY = get_api_key_fixture
|
139 |
+
# Initialize RAGFlow with API key and host address
|
140 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
141 |
+
|
142 |
+
# Create a dataset with a specific name
|
143 |
+
ds = rag.create_dataset(name="God4")
|
144 |
+
|
145 |
+
# Define the document name and path
|
146 |
+
name3 = 'westworld.pdf'
|
147 |
+
path = './test_data/westworld.pdf'
|
148 |
+
|
149 |
+
# Create a document in the dataset using the file path
|
150 |
+
ds.upload_documents({"name": name3, "blob": open(path, "rb").read()})
|
151 |
+
|
152 |
+
# Retrieve the document by name
|
153 |
+
doc = rag.list_documents(name="westworld.pdf")
|
154 |
+
doc = doc[0]
|
155 |
+
ds.async_parse_documents(document_ids=[])
|
156 |
+
|
157 |
+
# Print message to confirm asynchronous parsing has been initiated
|
158 |
+
print("Async parsing initiated")
|
159 |
+
|
160 |
+
# Use join to wait for parsing to complete and get progress updates
|
161 |
+
for progress, msg in doc.join(interval=5, timeout=10):
|
162 |
+
print(progress, msg)
|
163 |
+
# Assert that the progress is within the valid range (0 to 100)
|
164 |
+
assert 0 <= progress <= 100, f"Invalid progress: {progress}"
|
165 |
+
# Assert that the message is not empty
|
166 |
+
assert msg, "Message should not be empty"
|
167 |
+
# Test cancelling the parsing operation
|
168 |
+
doc.cancel()
|
169 |
+
# Print message to confirm parsing has been cancelled successfully
|
170 |
+
print("Parsing cancelled successfully")
|
171 |
+
|
172 |
+
|
173 |
+
def test_bulk_parse_and_cancel_documents(get_api_key_fixture):
|
174 |
+
API_KEY = get_api_key_fixture
|
175 |
+
# Initialize RAGFlow with API key and host address
|
176 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
177 |
+
|
178 |
+
# Create a dataset
|
179 |
+
ds = rag.create_dataset(name="God5")
|
180 |
+
assert ds is not None, "Dataset creation failed"
|
181 |
+
assert ds.name == "God5", "Dataset name does not match"
|
182 |
+
|
183 |
+
# Prepare a list of file names and paths
|
184 |
+
documents = [
|
185 |
+
{'name': 'test1.txt', 'path': 'test_data/test1.txt'},
|
186 |
+
{'name': 'test2.txt', 'path': 'test_data/test2.txt'},
|
187 |
+
{'name': 'test3.txt', 'path': 'test_data/test3.txt'}
|
188 |
+
]
|
189 |
+
|
190 |
+
# Create documents in bulk
|
191 |
+
for doc_info in documents:
|
192 |
+
with open(doc_info['path'], "rb") as file:
|
193 |
+
created_doc = rag.create_document(ds, name=doc_info['name'], blob=file.read())
|
194 |
+
assert created_doc is not None, f"Failed to create document {doc_info['name']}"
|
195 |
+
|
196 |
+
# Retrieve document objects in bulk
|
197 |
+
docs = [rag.get_document(name=doc_info['name']) for doc_info in documents]
|
198 |
+
ids = [doc.id for doc in docs]
|
199 |
+
assert len(docs) == len(documents), "Mismatch between created documents and fetched documents"
|
200 |
+
|
201 |
+
# Initiate asynchronous parsing for all documents
|
202 |
+
rag.async_parse_documents(ids)
|
203 |
+
print("Async bulk parsing initiated")
|
204 |
+
|
205 |
+
# Wait for all documents to finish parsing and check progress
|
206 |
+
for doc in docs:
|
207 |
for progress, msg in doc.join(interval=5, timeout=10):
|
208 |
+
print(f"{doc.name}: Progress: {progress}, Message: {msg}")
|
209 |
+
|
210 |
+
# Assert that progress is within the valid range
|
211 |
+
assert 0 <= progress <= 100, f"Invalid progress: {progress} for document {doc.name}"
|
212 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
213 |
# Assert that the message is not empty
|
214 |
+
assert msg, f"Message should not be empty for document {doc.name}"
|
215 |
+
|
216 |
+
# If progress reaches 100%, assert that parsing is completed successfully
|
217 |
+
if progress == 100:
|
218 |
+
assert "completed" in msg.lower(), f"Document {doc.name} did not complete successfully"
|
219 |
+
|
220 |
+
# Cancel parsing for all documents in bulk
|
221 |
+
cancel_result = rag.async_cancel_parse_documents(ids)
|
222 |
+
assert cancel_result is None or isinstance(cancel_result, type(None)), "Failed to cancel document parsing"
|
223 |
+
print("Async bulk parsing cancelled")
|
224 |
+
|
225 |
+
|
226 |
+
def test_parse_document_and_chunk_list(get_api_key_fixture):
|
227 |
+
API_KEY = get_api_key_fixture
|
228 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
229 |
+
ds = rag.create_dataset(name="God7")
|
230 |
+
name = 'story.txt'
|
231 |
+
path = 'test_data/story.txt'
|
232 |
+
# name = "Test Document rag.txt"
|
233 |
+
# blob = " Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps.Sample document content for rag test66. rag wonderful apple os documents apps.Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps. Sample document content for rag test66. rag wonderful apple os documents apps."
|
234 |
+
rag.create_document(ds, name=name, blob=open(path, "rb").read())
|
235 |
+
doc = rag.get_document(name=name)
|
236 |
+
doc.async_parse()
|
237 |
+
|
238 |
+
# Wait for parsing to complete and get progress updates using join
|
239 |
+
for progress, msg in doc.join(interval=5, timeout=30):
|
240 |
+
print(progress, msg)
|
241 |
+
# Assert that progress is within 0 to 100
|
242 |
+
assert 0 <= progress <= 100, f"Invalid progress: {progress}"
|
243 |
+
# Assert that the message is not empty
|
244 |
+
assert msg, "Message should not be empty"
|
245 |
+
|
246 |
+
for c in doc.list_chunks(keywords="rag", offset=0, limit=12):
|
247 |
+
print(c)
|
248 |
+
assert c is not None, "Chunk is None"
|
249 |
+
assert "rag" in c['content_with_weight'].lower(), f"Keyword 'rag' not found in chunk content: {c.content}"
|
250 |
+
|
251 |
+
|
252 |
+
def test_add_chunk_to_chunk_list(get_api_key_fixture):
|
253 |
+
API_KEY = get_api_key_fixture
|
254 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
255 |
+
doc = rag.get_document(name='story.txt')
|
256 |
+
chunk = doc.add_chunk(content="assssdd")
|
257 |
+
assert chunk is not None, "Chunk is None"
|
258 |
+
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
259 |
+
|
260 |
+
|
261 |
+
def test_delete_chunk_of_chunk_list(get_api_key_fixture):
|
262 |
+
API_KEY = get_api_key_fixture
|
263 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
264 |
+
doc = rag.get_document(name='story.txt')
|
265 |
+
chunk = doc.add_chunk(content="assssdd")
|
266 |
+
assert chunk is not None, "Chunk is None"
|
267 |
+
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
268 |
+
doc = rag.get_document(name='story.txt')
|
269 |
+
chunk_count_before = doc.chunk_count
|
270 |
+
chunk.delete()
|
271 |
+
doc = rag.get_document(name='story.txt')
|
272 |
+
assert doc.chunk_count == chunk_count_before - 1, "Chunk was not deleted"
|
273 |
+
|
274 |
+
|
275 |
+
def test_update_chunk_content(get_api_key_fixture):
|
276 |
+
API_KEY = get_api_key_fixture
|
277 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
278 |
+
doc = rag.get_document(name='story.txt')
|
279 |
+
chunk = doc.add_chunk(content="assssddd")
|
280 |
+
assert chunk is not None, "Chunk is None"
|
281 |
+
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
282 |
+
chunk.content = "ragflow123"
|
283 |
+
res = chunk.save()
|
284 |
+
assert res is True, f"Failed to update chunk content, error: {res}"
|
285 |
+
|
286 |
+
|
287 |
+
def test_update_chunk_available(get_api_key_fixture):
|
288 |
+
API_KEY = get_api_key_fixture
|
289 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
290 |
+
doc = rag.get_document(name='story.txt')
|
291 |
+
chunk = doc.add_chunk(content="ragflow")
|
292 |
+
assert chunk is not None, "Chunk is None"
|
293 |
+
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
294 |
+
chunk.available = 0
|
295 |
+
res = chunk.save()
|
296 |
+
assert res is True, f"Failed to update chunk status, error: {res}"
|
297 |
+
|
298 |
+
|
299 |
+
def test_retrieval_chunks(get_api_key_fixture):
|
300 |
+
API_KEY = get_api_key_fixture
|
301 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
302 |
+
ds = rag.create_dataset(name="God8")
|
303 |
+
name = 'ragflow_test.txt'
|
304 |
+
path = 'test_data/ragflow_test.txt'
|
305 |
+
rag.create_document(ds, name=name, blob=open(path, "rb").read())
|
306 |
+
doc = rag.get_document(name=name)
|
307 |
+
doc.async_parse()
|
308 |
+
# Wait for parsing to complete and get progress updates using join
|
309 |
+
for progress, msg in doc.join(interval=5, timeout=30):
|
310 |
+
print(progress, msg)
|
311 |
+
assert 0 <= progress <= 100, f"Invalid progress: {progress}"
|
312 |
+
assert msg, "Message should not be empty"
|
313 |
+
for c in rag.retrieval(question="What's ragflow?",
|
314 |
+
datasets=[ds.id], documents=[doc],
|
315 |
+
offset=0, limit=6, similarity_threshold=0.1,
|
316 |
+
vector_similarity_weight=0.3,
|
317 |
+
top_k=1024
|
318 |
+
):
|
319 |
+
print(c)
|
320 |
+
assert c is not None, "Chunk is None"
|
321 |
+
assert "ragflow" in c.content.lower(), f"Keyword 'rag' not found in chunk content: {c.content}"
|
sdk/python/test/t_session.py
CHANGED
@@ -1,52 +1,110 @@
|
|
1 |
from ragflow import RAGFlow,Session
|
|
|
|
|
2 |
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
from ragflow import RAGFlow,Session
|
2 |
+
import time
|
3 |
+
HOST_ADDRESS = 'http://127.0.0.1:9380'
|
4 |
|
5 |
+
|
6 |
+
def test_create_session_with_success(get_api_key_fixture):
|
7 |
+
API_KEY = get_api_key_fixture
|
8 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
9 |
+
kb = rag.create_dataset(name="test_create_session")
|
10 |
+
displayed_name = "ragflow.txt"
|
11 |
+
with open("./ragflow.txt", "rb") as file:
|
12 |
+
blob = file.read()
|
13 |
+
document = {"displayed_name": displayed_name, "blob": blob}
|
14 |
+
documents = []
|
15 |
+
documents.append(document)
|
16 |
+
doc_ids = []
|
17 |
+
docs = kb.upload_documents(documents)
|
18 |
+
for doc in docs:
|
19 |
+
doc_ids.append(doc.id)
|
20 |
+
kb.async_parse_documents(doc_ids)
|
21 |
+
time.sleep(60)
|
22 |
+
assistant = rag.create_chat(name="test_create_session", datasets=[kb])
|
23 |
+
assistant.create_session()
|
24 |
+
|
25 |
+
|
26 |
+
def test_create_conversation_with_success(get_api_key_fixture):
|
27 |
+
API_KEY = get_api_key_fixture
|
28 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
29 |
+
kb = rag.create_dataset(name="test_create_conversation")
|
30 |
+
displayed_name = "ragflow.txt"
|
31 |
+
with open("./ragflow.txt","rb") as file:
|
32 |
+
blob = file.read()
|
33 |
+
document = {"displayed_name":displayed_name,"blob":blob}
|
34 |
+
documents = []
|
35 |
+
documents.append(document)
|
36 |
+
doc_ids = []
|
37 |
+
docs= kb.upload_documents(documents)
|
38 |
+
for doc in docs:
|
39 |
+
doc_ids.append(doc.id)
|
40 |
+
kb.async_parse_documents(doc_ids)
|
41 |
+
time.sleep(60)
|
42 |
+
assistant = rag.create_chat(name="test_create_conversation", datasets=[kb])
|
43 |
+
session = assistant.create_session()
|
44 |
+
question = "What is AI"
|
45 |
+
for ans in session.ask(question, stream=True):
|
46 |
+
pass
|
47 |
+
assert not ans.content.startswith("**ERROR**"), "Please check this error."
|
48 |
+
|
49 |
+
|
50 |
+
def test_delete_sessions_with_success(get_api_key_fixture):
|
51 |
+
API_KEY = get_api_key_fixture
|
52 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
53 |
+
kb = rag.create_dataset(name="test_delete_session")
|
54 |
+
displayed_name = "ragflow.txt"
|
55 |
+
with open("./ragflow.txt","rb") as file:
|
56 |
+
blob = file.read()
|
57 |
+
document = {"displayed_name":displayed_name,"blob":blob}
|
58 |
+
documents = []
|
59 |
+
documents.append(document)
|
60 |
+
doc_ids = []
|
61 |
+
docs= kb.upload_documents(documents)
|
62 |
+
for doc in docs:
|
63 |
+
doc_ids.append(doc.id)
|
64 |
+
kb.async_parse_documents(doc_ids)
|
65 |
+
time.sleep(60)
|
66 |
+
assistant = rag.create_chat(name="test_delete_session", datasets=[kb])
|
67 |
+
session = assistant.create_session()
|
68 |
+
assistant.delete_sessions(ids=[session.id])
|
69 |
+
|
70 |
+
def test_update_session_with_name(get_api_key_fixture):
|
71 |
+
API_KEY = get_api_key_fixture
|
72 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
73 |
+
kb = rag.create_dataset(name="test_update_session")
|
74 |
+
displayed_name = "ragflow.txt"
|
75 |
+
with open("./ragflow.txt","rb") as file:
|
76 |
+
blob = file.read()
|
77 |
+
document = {"displayed_name":displayed_name,"blob":blob}
|
78 |
+
documents = []
|
79 |
+
documents.append(document)
|
80 |
+
doc_ids = []
|
81 |
+
docs= kb.upload_documents(documents)
|
82 |
+
for doc in docs:
|
83 |
+
doc_ids.append(doc.id)
|
84 |
+
kb.async_parse_documents(doc_ids)
|
85 |
+
time.sleep(60)
|
86 |
+
assistant = rag.create_chat(name="test_update_session", datasets=[kb])
|
87 |
+
session = assistant.create_session(name="old session")
|
88 |
+
session.update({"name": "new session"})
|
89 |
+
|
90 |
+
|
91 |
+
def test_list_sessions_with_success(get_api_key_fixture):
|
92 |
+
API_KEY = get_api_key_fixture
|
93 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
94 |
+
kb = rag.create_dataset(name="test_list_session")
|
95 |
+
displayed_name = "ragflow.txt"
|
96 |
+
with open("./ragflow.txt","rb") as file:
|
97 |
+
blob = file.read()
|
98 |
+
document = {"displayed_name":displayed_name,"blob":blob}
|
99 |
+
documents = []
|
100 |
+
documents.append(document)
|
101 |
+
doc_ids = []
|
102 |
+
docs= kb.upload_documents(documents)
|
103 |
+
for doc in docs:
|
104 |
+
doc_ids.append(doc.id)
|
105 |
+
kb.async_parse_documents(doc_ids)
|
106 |
+
time.sleep(60)
|
107 |
+
assistant = rag.create_chat(name="test_list_session", datasets=[kb])
|
108 |
+
assistant.create_session("test_1")
|
109 |
+
assistant.create_session("test_2")
|
110 |
+
assistant.list_sessions()
|
sdk/python/test/test_sdkbase.py
DELETED
@@ -1,3 +0,0 @@
|
|
1 |
-
class TestSdk():
|
2 |
-
def test_version(self):
|
3 |
-
print("test_sdk")
|
|
|
|
|
|
|
|