ragflow-sdk-myfbee-test 1.0.1.dev0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ragflow_sdk_myfbee_test-1.0.1.dev0/PKG-INFO +10 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/pyproject.toml +34 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/__init__.py +39 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/__init__.py +15 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/agent.py +69 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/base.py +62 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/chat.py +67 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/chunk.py +62 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/dataset.py +182 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/document.py +109 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/memory.py +87 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/modules/session.py +203 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk/ragflow.py +398 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk_myfbee_test.egg-info/PKG-INFO +10 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk_myfbee_test.egg-info/SOURCES.txt +17 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk_myfbee_test.egg-info/dependency_links.txt +1 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk_myfbee_test.egg-info/requires.txt +2 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/ragflow_sdk_myfbee_test.egg-info/top_level.txt +1 -0
- ragflow_sdk_myfbee_test-1.0.1.dev0/setup.cfg +4 -0
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ragflow-sdk-myfbee-test
|
|
3
|
+
Version: 1.0.1.dev0
|
|
4
|
+
Summary: Python client sdk of [RAGFlow](https://github.com/infiniflow/ragflow). RAGFlow is an open-source RAG (Retrieval-Augmented Generation) engine based on deep document understanding.
|
|
5
|
+
Author-email: Zhichang Yu <yuzhichang@gmail.com>
|
|
6
|
+
License: Apache License, Version 2.0
|
|
7
|
+
Requires-Python: <3.14,>=3.13
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: requests<3.0.0,>=2.30.0
|
|
10
|
+
Requires-Dist: beartype<1.0.0,>=0.20.0
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "ragflow-sdk-myfbee-test"
|
|
3
|
+
version = "1.0.1-dev"
|
|
4
|
+
description = "Python client sdk of [RAGFlow](https://github.com/infiniflow/ragflow). RAGFlow is an open-source RAG (Retrieval-Augmented Generation) engine based on deep document understanding."
|
|
5
|
+
authors = [{ name = "Zhichang Yu", email = "yuzhichang@gmail.com" }]
|
|
6
|
+
license = { text = "Apache License, Version 2.0" }
|
|
7
|
+
readme = "README.md"
|
|
8
|
+
requires-python = ">=3.13,<3.14"
|
|
9
|
+
dependencies = ["requests>=2.30.0,<3.0.0", "beartype>=0.20.0,<1.0.0"]
|
|
10
|
+
|
|
11
|
+
[dependency-groups]
|
|
12
|
+
test = [
|
|
13
|
+
"hypothesis>=6.131.9",
|
|
14
|
+
"openpyxl>=3.1.5",
|
|
15
|
+
"pillow>=11.1.0",
|
|
16
|
+
"pytest>=8.3.5",
|
|
17
|
+
"python-docx>=1.1.2",
|
|
18
|
+
"python-pptx>=1.0.2",
|
|
19
|
+
"reportlab>=4.3.1",
|
|
20
|
+
"requests>=2.32.3",
|
|
21
|
+
"requests-toolbelt>=1.0.0",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[tool.pytest.ini_options]
|
|
25
|
+
markers = [
|
|
26
|
+
"p1: high priority test cases",
|
|
27
|
+
"p2: medium priority test cases",
|
|
28
|
+
"p3: low priority test cases",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[tool.uv]
|
|
32
|
+
constraint-dependencies = [
|
|
33
|
+
"urllib3>=2.7.0",
|
|
34
|
+
]
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
# Keep runtime validation when beartype is installed, while allowing slim
|
|
18
|
+
# SDK installs to import without the optional checker.
|
|
19
|
+
try:
|
|
20
|
+
from beartype.claw import beartype_this_package
|
|
21
|
+
|
|
22
|
+
beartype_this_package()
|
|
23
|
+
except ImportError:
|
|
24
|
+
pass
|
|
25
|
+
|
|
26
|
+
import importlib.metadata
|
|
27
|
+
|
|
28
|
+
from .ragflow import RAGFlow
|
|
29
|
+
from .modules.dataset import DataSet
|
|
30
|
+
from .modules.chat import Chat
|
|
31
|
+
from .modules.session import Session
|
|
32
|
+
from .modules.document import Document
|
|
33
|
+
from .modules.chunk import Chunk
|
|
34
|
+
from .modules.agent import Agent
|
|
35
|
+
from .modules.memory import Memory
|
|
36
|
+
|
|
37
|
+
__version__ = importlib.metadata.version("ragflow_sdk")
|
|
38
|
+
|
|
39
|
+
__all__ = ["RAGFlow", "DataSet", "Chat", "Session", "Document", "Chunk", "Agent", "Memory"]
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
from .base import Base
|
|
18
|
+
from .session import Session
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class Agent(Base):
|
|
22
|
+
def __init__(self, rag, res_dict):
|
|
23
|
+
self.id = None
|
|
24
|
+
self.avatar = None
|
|
25
|
+
self.canvas_type = None
|
|
26
|
+
self.description = None
|
|
27
|
+
self.dsl = None
|
|
28
|
+
super().__init__(rag, res_dict)
|
|
29
|
+
|
|
30
|
+
class Dsl(Base):
|
|
31
|
+
def __init__(self, rag, res_dict):
|
|
32
|
+
self.answer = []
|
|
33
|
+
self.components = {"begin": {"downstream": ["Answer:China"], "obj": {"component_name": "Begin", "params": {}}, "upstream": []}}
|
|
34
|
+
self.graph = {
|
|
35
|
+
"edges": [],
|
|
36
|
+
"nodes": [{"data": {"label": "Begin", "name": "begin"}, "id": "begin", "position": {"x": 50, "y": 200}, "sourcePosition": "left", "targetPosition": "right", "type": "beginNode"}],
|
|
37
|
+
}
|
|
38
|
+
self.history = []
|
|
39
|
+
self.messages = []
|
|
40
|
+
self.path = []
|
|
41
|
+
self.reference = []
|
|
42
|
+
super().__init__(rag, res_dict)
|
|
43
|
+
|
|
44
|
+
def create_session(self, **kwargs) -> Session:
|
|
45
|
+
res = self.post(f"/agents/{self.id}/sessions", json=kwargs)
|
|
46
|
+
res = res.json()
|
|
47
|
+
if res.get("code") == 0:
|
|
48
|
+
return Session(self.rag, res.get("data"))
|
|
49
|
+
raise Exception(res.get("message"))
|
|
50
|
+
|
|
51
|
+
def list_sessions(self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str = None) -> list[Session]:
|
|
52
|
+
res = self.get(f"/agents/{self.id}/sessions", {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id})
|
|
53
|
+
res = res.json()
|
|
54
|
+
if res.get("code") == 0:
|
|
55
|
+
result_list = []
|
|
56
|
+
for data in res.get("data"):
|
|
57
|
+
temp_agent = Session(self.rag, data)
|
|
58
|
+
result_list.append(temp_agent)
|
|
59
|
+
return result_list
|
|
60
|
+
raise Exception(res.get("message"))
|
|
61
|
+
|
|
62
|
+
def delete_sessions(self, ids: list[str] | None = None, delete_all: bool = False):
|
|
63
|
+
payload = {"ids": ids}
|
|
64
|
+
if delete_all:
|
|
65
|
+
payload["delete_all"] = True
|
|
66
|
+
res = self.rm(f"/agents/{self.id}/sessions", payload)
|
|
67
|
+
res = res.json()
|
|
68
|
+
if res.get("code") != 0:
|
|
69
|
+
raise Exception(res.get("message"))
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class Base:
|
|
19
|
+
def __init__(self, rag, res_dict):
|
|
20
|
+
self.rag = rag
|
|
21
|
+
self._update_from_dict(rag, res_dict)
|
|
22
|
+
|
|
23
|
+
def _update_from_dict(self, rag, res_dict):
|
|
24
|
+
for k, v in res_dict.items():
|
|
25
|
+
if isinstance(v, dict):
|
|
26
|
+
self.__dict__[k] = Base(rag, v)
|
|
27
|
+
else:
|
|
28
|
+
self.__dict__[k] = v
|
|
29
|
+
|
|
30
|
+
def to_json(self):
|
|
31
|
+
pr = {}
|
|
32
|
+
for name in dir(self):
|
|
33
|
+
value = getattr(self, name)
|
|
34
|
+
if not name.startswith("__") and not callable(value) and name != "rag":
|
|
35
|
+
if isinstance(value, Base):
|
|
36
|
+
pr[name] = value.to_json()
|
|
37
|
+
else:
|
|
38
|
+
pr[name] = value
|
|
39
|
+
return pr
|
|
40
|
+
|
|
41
|
+
def post(self, path, json=None, stream=False, files=None):
|
|
42
|
+
res = self.rag.post(path, json, stream=stream, files=files)
|
|
43
|
+
return res
|
|
44
|
+
|
|
45
|
+
def get(self, path, params=None):
|
|
46
|
+
res = self.rag.get(path, params)
|
|
47
|
+
return res
|
|
48
|
+
|
|
49
|
+
def rm(self, path, json):
|
|
50
|
+
res = self.rag.delete(path, json)
|
|
51
|
+
return res
|
|
52
|
+
|
|
53
|
+
def put(self, path, json):
|
|
54
|
+
res = self.rag.put(path, json)
|
|
55
|
+
return res
|
|
56
|
+
|
|
57
|
+
def patch(self, path, json):
|
|
58
|
+
res = self.rag.patch(path, json)
|
|
59
|
+
return res
|
|
60
|
+
|
|
61
|
+
def __str__(self):
|
|
62
|
+
return str(self.to_json())
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
from .base import Base
|
|
19
|
+
from .session import Session
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class Chat(Base):
|
|
23
|
+
def __init__(self, rag, res_dict):
|
|
24
|
+
self.id = ""
|
|
25
|
+
self.name = "assistant"
|
|
26
|
+
self.icon = ""
|
|
27
|
+
self.dataset_ids = []
|
|
28
|
+
self.llm_id = None
|
|
29
|
+
self.llm_setting = {}
|
|
30
|
+
self.prompt_config = {}
|
|
31
|
+
self.similarity_threshold = 0.2
|
|
32
|
+
self.vector_similarity_weight = 0.3
|
|
33
|
+
self.top_n = 6
|
|
34
|
+
self.top_k = 1024
|
|
35
|
+
self.rerank_id = ""
|
|
36
|
+
super().__init__(rag, res_dict)
|
|
37
|
+
|
|
38
|
+
def update(self, update_message: dict):
|
|
39
|
+
if not isinstance(update_message, dict):
|
|
40
|
+
raise Exception("ValueError('`update_message` must be a dict')")
|
|
41
|
+
res = self.patch(f"/chats/{self.id}", update_message)
|
|
42
|
+
res = res.json()
|
|
43
|
+
if res.get("code") != 0:
|
|
44
|
+
raise Exception(res["message"])
|
|
45
|
+
|
|
46
|
+
def create_session(self, name: str = "New session") -> Session:
|
|
47
|
+
res = self.post(f"/chats/{self.id}/sessions", {"name": name})
|
|
48
|
+
res = res.json()
|
|
49
|
+
if res.get("code") == 0:
|
|
50
|
+
return Session(self.rag, res["data"])
|
|
51
|
+
raise Exception(res.get("message"))
|
|
52
|
+
|
|
53
|
+
def list_sessions(self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str = None, name: str = None, user_id: str = None) -> list[Session]:
|
|
54
|
+
res = self.get(f"/chats/{self.id}/sessions", {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id, "name": name, "user_id": user_id})
|
|
55
|
+
res = res.json()
|
|
56
|
+
if res.get("code") == 0:
|
|
57
|
+
result_list = []
|
|
58
|
+
for data in res["data"]:
|
|
59
|
+
result_list.append(Session(self.rag, data))
|
|
60
|
+
return result_list
|
|
61
|
+
raise Exception(res["message"])
|
|
62
|
+
|
|
63
|
+
def delete_sessions(self, ids: list[str] | None = None, delete_all: bool = False):
|
|
64
|
+
res = self.rm(f"/chats/{self.id}/sessions", {"ids": ids, "delete_all": delete_all})
|
|
65
|
+
res = res.json()
|
|
66
|
+
if res.get("code") != 0:
|
|
67
|
+
raise Exception(res.get("message"))
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
from .base import Base
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class ChunkUpdateError(Exception):
|
|
21
|
+
def __init__(self, code=None, message=None, details=None):
|
|
22
|
+
self.code = code
|
|
23
|
+
self.message = message
|
|
24
|
+
self.details = details
|
|
25
|
+
super().__init__(message)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class Chunk(Base):
|
|
29
|
+
def __init__(self, rag, res_dict):
|
|
30
|
+
self.id = ""
|
|
31
|
+
self.content = ""
|
|
32
|
+
self.image_id = ""
|
|
33
|
+
self.important_keywords = []
|
|
34
|
+
self.tag_kwd = []
|
|
35
|
+
self.questions = []
|
|
36
|
+
self.create_time = ""
|
|
37
|
+
self.create_timestamp = 0.0
|
|
38
|
+
self.dataset_id = None
|
|
39
|
+
self.document_name = ""
|
|
40
|
+
self.document_keyword = ""
|
|
41
|
+
self.document_id = ""
|
|
42
|
+
self.available = True
|
|
43
|
+
# Additional fields for retrieval results
|
|
44
|
+
self.similarity = 0.0
|
|
45
|
+
self.vector_similarity = 0.0
|
|
46
|
+
self.term_similarity = 0.0
|
|
47
|
+
self.positions = []
|
|
48
|
+
self.doc_type = ""
|
|
49
|
+
for k in list(res_dict.keys()):
|
|
50
|
+
if k not in self.__dict__:
|
|
51
|
+
res_dict.pop(k)
|
|
52
|
+
super().__init__(rag, res_dict)
|
|
53
|
+
|
|
54
|
+
# for backward compatibility
|
|
55
|
+
if not self.document_name:
|
|
56
|
+
self.document_name = self.document_keyword
|
|
57
|
+
|
|
58
|
+
def update(self, update_message: dict):
|
|
59
|
+
res = self.patch(f"/datasets/{self.dataset_id}/documents/{self.document_id}/chunks/{self.id}", update_message)
|
|
60
|
+
res = res.json()
|
|
61
|
+
if res.get("code") != 0:
|
|
62
|
+
raise ChunkUpdateError(code=res.get("code"), message=res.get("message"), details=res.get("details"))
|
|
@@ -0,0 +1,182 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
from typing import Any
|
|
17
|
+
|
|
18
|
+
from .base import Base
|
|
19
|
+
from .document import Document
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class DataSet(Base):
|
|
23
|
+
class ParserConfig(Base):
|
|
24
|
+
def __init__(self, rag, res_dict):
|
|
25
|
+
super().__init__(rag, res_dict)
|
|
26
|
+
|
|
27
|
+
def __init__(self, rag, res_dict):
|
|
28
|
+
self.id = ""
|
|
29
|
+
self.name = ""
|
|
30
|
+
self.avatar = ""
|
|
31
|
+
self.tenant_id = None
|
|
32
|
+
self.description = ""
|
|
33
|
+
self.embedding_model = ""
|
|
34
|
+
self.permission = "me"
|
|
35
|
+
self.document_count = 0
|
|
36
|
+
self.chunk_count = 0
|
|
37
|
+
self.chunk_method = "naive"
|
|
38
|
+
self.parser_config = None
|
|
39
|
+
self.pagerank = 0
|
|
40
|
+
for k in list(res_dict.keys()):
|
|
41
|
+
if k not in self.__dict__:
|
|
42
|
+
res_dict.pop(k)
|
|
43
|
+
super().__init__(rag, res_dict)
|
|
44
|
+
|
|
45
|
+
def update(self, update_message: dict):
|
|
46
|
+
res = self.put(f"/datasets/{self.id}", update_message)
|
|
47
|
+
res = res.json()
|
|
48
|
+
if res.get("code") != 0:
|
|
49
|
+
raise Exception(res["message"])
|
|
50
|
+
|
|
51
|
+
self._update_from_dict(self.rag, res.get("data", {}))
|
|
52
|
+
return self
|
|
53
|
+
|
|
54
|
+
def upload_documents(self, document_list: list[dict]):
|
|
55
|
+
url = f"/datasets/{self.id}/documents"
|
|
56
|
+
files = [("file", (ele["display_name"], ele["blob"])) for ele in document_list]
|
|
57
|
+
res = self.post(path=url, json=None, files=files)
|
|
58
|
+
res = res.json()
|
|
59
|
+
if res.get("code") == 0:
|
|
60
|
+
doc_list = []
|
|
61
|
+
for doc in res["data"]:
|
|
62
|
+
document = Document(self.rag, doc)
|
|
63
|
+
doc_list.append(document)
|
|
64
|
+
return doc_list
|
|
65
|
+
raise Exception(res.get("message"))
|
|
66
|
+
|
|
67
|
+
def list_documents(
|
|
68
|
+
self,
|
|
69
|
+
id: str | None = None,
|
|
70
|
+
ids: list[str] | None = None,
|
|
71
|
+
name: str | None = None,
|
|
72
|
+
keywords: str | None = None,
|
|
73
|
+
page: int = 1,
|
|
74
|
+
page_size: int = 30,
|
|
75
|
+
orderby: str = "create_time",
|
|
76
|
+
desc: bool = True,
|
|
77
|
+
create_time_from: int = 0,
|
|
78
|
+
create_time_to: int = 0,
|
|
79
|
+
):
|
|
80
|
+
# Validate that id and ids are not used together
|
|
81
|
+
if id and ids:
|
|
82
|
+
raise ValueError("Cannot use both 'id' and 'ids' parameters at the same time.")
|
|
83
|
+
|
|
84
|
+
params = {
|
|
85
|
+
"id": id,
|
|
86
|
+
"name": name,
|
|
87
|
+
"keywords": keywords,
|
|
88
|
+
"page": page,
|
|
89
|
+
"page_size": page_size,
|
|
90
|
+
"orderby": orderby,
|
|
91
|
+
"desc": desc,
|
|
92
|
+
"create_time_from": create_time_from,
|
|
93
|
+
"create_time_to": create_time_to,
|
|
94
|
+
}
|
|
95
|
+
# Handle ids parameter - requests expands list values into multiple query params
|
|
96
|
+
if ids:
|
|
97
|
+
params["ids"] = ids
|
|
98
|
+
res = self.get(f"/datasets/{self.id}/documents", params=params)
|
|
99
|
+
res = res.json()
|
|
100
|
+
documents = []
|
|
101
|
+
if res.get("code") == 0:
|
|
102
|
+
for document in res["data"].get("docs"):
|
|
103
|
+
documents.append(Document(self.rag, document))
|
|
104
|
+
return documents
|
|
105
|
+
raise Exception(res["message"])
|
|
106
|
+
|
|
107
|
+
def delete_documents(self, ids: list[str] | None = None, delete_all: bool = False):
|
|
108
|
+
res = self.rm(f"/datasets/{self.id}/documents", {"ids": ids, "delete_all": delete_all})
|
|
109
|
+
res = res.json()
|
|
110
|
+
if res.get("code") != 0:
|
|
111
|
+
raise Exception(res["message"])
|
|
112
|
+
|
|
113
|
+
def _get_documents_status(self, document_ids):
|
|
114
|
+
import time
|
|
115
|
+
|
|
116
|
+
terminal_states = {"DONE", "FAIL", "CANCEL"}
|
|
117
|
+
interval_sec = 1
|
|
118
|
+
pending = set(document_ids)
|
|
119
|
+
finished = []
|
|
120
|
+
while pending:
|
|
121
|
+
for doc_id in list(pending):
|
|
122
|
+
|
|
123
|
+
def fetch_doc(doc_id: str) -> Document | None:
|
|
124
|
+
try:
|
|
125
|
+
docs = self.list_documents(id=doc_id)
|
|
126
|
+
return docs[0] if docs else None
|
|
127
|
+
except Exception:
|
|
128
|
+
return None
|
|
129
|
+
|
|
130
|
+
doc = fetch_doc(doc_id)
|
|
131
|
+
if doc is None:
|
|
132
|
+
continue
|
|
133
|
+
if isinstance(doc.run, str) and doc.run.upper() in terminal_states:
|
|
134
|
+
finished.append((doc_id, doc.run, doc.chunk_count, doc.token_count))
|
|
135
|
+
pending.discard(doc_id)
|
|
136
|
+
elif float(doc.progress or 0.0) >= 1.0:
|
|
137
|
+
finished.append((doc_id, "DONE", doc.chunk_count, doc.token_count))
|
|
138
|
+
pending.discard(doc_id)
|
|
139
|
+
if pending:
|
|
140
|
+
time.sleep(interval_sec)
|
|
141
|
+
return finished
|
|
142
|
+
|
|
143
|
+
def async_parse_documents(self, document_ids):
|
|
144
|
+
res = self.post(f"/datasets/{self.id}/chunks", {"document_ids": document_ids})
|
|
145
|
+
res = res.json()
|
|
146
|
+
if res.get("code") != 0:
|
|
147
|
+
raise Exception(res.get("message"))
|
|
148
|
+
|
|
149
|
+
def parse_documents(self, document_ids):
|
|
150
|
+
try:
|
|
151
|
+
self.async_parse_documents(document_ids)
|
|
152
|
+
self._get_documents_status(document_ids)
|
|
153
|
+
except KeyboardInterrupt:
|
|
154
|
+
self.async_cancel_parse_documents(document_ids)
|
|
155
|
+
|
|
156
|
+
return self._get_documents_status(document_ids)
|
|
157
|
+
|
|
158
|
+
def async_cancel_parse_documents(self, document_ids):
|
|
159
|
+
res = self.rm(f"/datasets/{self.id}/chunks", {"document_ids": document_ids})
|
|
160
|
+
res = res.json()
|
|
161
|
+
if res.get("code") != 0:
|
|
162
|
+
raise Exception(res.get("message"))
|
|
163
|
+
|
|
164
|
+
def get_auto_metadata(self) -> dict[str, Any]:
|
|
165
|
+
"""
|
|
166
|
+
Retrieve auto-metadata configuration for a dataset via SDK.
|
|
167
|
+
"""
|
|
168
|
+
res = self.get(f"/datasets/{self.id}/metadata/config")
|
|
169
|
+
res = res.json()
|
|
170
|
+
if res.get("code") == 0:
|
|
171
|
+
return res["data"]
|
|
172
|
+
raise Exception(res["message"])
|
|
173
|
+
|
|
174
|
+
def update_auto_metadata(self, **config: Any) -> dict[str, Any]:
|
|
175
|
+
"""
|
|
176
|
+
Update auto-metadata configuration for a dataset via SDK.
|
|
177
|
+
"""
|
|
178
|
+
res = self.put(f"/datasets/{self.id}/metadata/config", config)
|
|
179
|
+
res = res.json()
|
|
180
|
+
if res.get("code") == 0:
|
|
181
|
+
return res["data"]
|
|
182
|
+
raise Exception(res["message"])
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
#
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
|
|
19
|
+
from .base import Base
|
|
20
|
+
from .chunk import Chunk
|
|
21
|
+
|
|
22
|
+
logger = logging.getLogger(__name__)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class Document(Base):
|
|
26
|
+
class ParserConfig(Base):
|
|
27
|
+
def __init__(self, rag, res_dict):
|
|
28
|
+
super().__init__(rag, res_dict)
|
|
29
|
+
|
|
30
|
+
def __init__(self, rag, res_dict):
|
|
31
|
+
self.id = ""
|
|
32
|
+
self.name = ""
|
|
33
|
+
self.thumbnail = None
|
|
34
|
+
self.dataset_id = None
|
|
35
|
+
self.chunk_method = "naive"
|
|
36
|
+
self.parser_config = {"pages": [[1, 1000000]]}
|
|
37
|
+
self.source_type = "local"
|
|
38
|
+
self.type = ""
|
|
39
|
+
self.created_by = ""
|
|
40
|
+
self.size = 0
|
|
41
|
+
self.token_count = 0
|
|
42
|
+
self.chunk_count = 0
|
|
43
|
+
self.progress = 0.0
|
|
44
|
+
self.progress_msg = ""
|
|
45
|
+
self.process_begin_at = None
|
|
46
|
+
self.process_duration = 0.0
|
|
47
|
+
self.run = "0"
|
|
48
|
+
self.status = "1"
|
|
49
|
+
self.meta_fields = {}
|
|
50
|
+
for k in list(res_dict.keys()):
|
|
51
|
+
if k not in self.__dict__:
|
|
52
|
+
res_dict.pop(k)
|
|
53
|
+
super().__init__(rag, res_dict)
|
|
54
|
+
|
|
55
|
+
def update(self, update_message: dict):
|
|
56
|
+
if "meta_fields" in update_message:
|
|
57
|
+
if not isinstance(update_message["meta_fields"], dict):
|
|
58
|
+
raise Exception("meta_fields must be a dictionary")
|
|
59
|
+
res = self.patch(f"/datasets/{self.dataset_id}/documents/{self.id}", update_message)
|
|
60
|
+
res = res.json()
|
|
61
|
+
if res.get("code") != 0:
|
|
62
|
+
raise Exception(res["message"])
|
|
63
|
+
|
|
64
|
+
self._update_from_dict(self.rag, res.get("data", {}))
|
|
65
|
+
return self
|
|
66
|
+
|
|
67
|
+
def download(self):
|
|
68
|
+
res = self.get(f"/datasets/{self.dataset_id}/documents/{self.id}")
|
|
69
|
+
content_disposition = res.headers.get("Content-Disposition", "")
|
|
70
|
+
if content_disposition.lstrip().lower().startswith("attachment"):
|
|
71
|
+
logger.debug("Document.download returning attachment content dataset_id=%s document_id=%s", self.dataset_id, self.id)
|
|
72
|
+
return res.content
|
|
73
|
+
|
|
74
|
+
try:
|
|
75
|
+
response = res.json()
|
|
76
|
+
if isinstance(response, dict) and "code" in response and response["code"] != 0:
|
|
77
|
+
raise Exception(response.get("message"))
|
|
78
|
+
except ValueError:
|
|
79
|
+
pass
|
|
80
|
+
logger.debug("Document.download returning non-attachment content dataset_id=%s document_id=%s", self.dataset_id, self.id)
|
|
81
|
+
return res.content
|
|
82
|
+
|
|
83
|
+
def list_chunks(self, page=1, page_size=30, keywords="", id=""):
|
|
84
|
+
data = {"keywords": keywords, "page": page, "page_size": page_size, "id": id}
|
|
85
|
+
res = self.get(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", data)
|
|
86
|
+
res = res.json()
|
|
87
|
+
if res.get("code") == 0:
|
|
88
|
+
chunks = []
|
|
89
|
+
for data in res["data"].get("chunks"):
|
|
90
|
+
chunk = Chunk(self.rag, data)
|
|
91
|
+
chunks.append(chunk)
|
|
92
|
+
return chunks
|
|
93
|
+
raise Exception(res.get("message"))
|
|
94
|
+
|
|
95
|
+
def add_chunk(self, content: str, important_keywords: list[str] = [], questions: list[str] = [], image_base64: str | None = None, *, tag_kwd: list[str] = []):
|
|
96
|
+
body = {"content": content, "important_keywords": important_keywords, "tag_kwd": tag_kwd, "questions": questions}
|
|
97
|
+
if image_base64 is not None:
|
|
98
|
+
body["image_base64"] = image_base64
|
|
99
|
+
res = self.post(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", body)
|
|
100
|
+
res = res.json()
|
|
101
|
+
if res.get("code") == 0:
|
|
102
|
+
return Chunk(self.rag, res["data"].get("chunk"))
|
|
103
|
+
raise Exception(res.get("message"))
|
|
104
|
+
|
|
105
|
+
def delete_chunks(self, ids: list[str] | None = None, delete_all: bool = False):
|
|
106
|
+
res = self.rm(f"/datasets/{self.dataset_id}/documents/{self.id}/chunks", {"chunk_ids": ids, "delete_all": delete_all})
|
|
107
|
+
res = res.json()
|
|
108
|
+
if res.get("code") != 0:
|
|
109
|
+
raise Exception(res.get("message"))
|