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