Refactor Dataset API (#2783)

### What problem does this PR solve?

Refactor Dataset API

### Type of change

- [x] Refactoring

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
This commit is contained in:
liuhua
2024-10-11 09:55:27 +08:00
committed by GitHub
parent a2f9c03a95
commit cbd7cd7c4d
11 changed files with 449 additions and 393 deletions

View File

@@ -30,5 +30,9 @@ class Base(object):
res = self.rag.delete(path, params)
return res
def put(self,path, json):
res = self.rag.put(path,json)
return res
def __str__(self):
return str(self.to_json())

View File

@@ -32,24 +32,13 @@ class DataSet(Base):
res_dict.pop(k)
super().__init__(rag, res_dict)
def save(self) -> bool:
res = self.post('/dataset/save',
{"id": self.id, "name": self.name, "avatar": self.avatar, "tenant_id": self.tenant_id,
"description": self.description, "language": self.language, "embedding_model": self.embedding_model,
"permission": self.permission,
"document_count": self.document_count, "chunk_count": self.chunk_count, "parse_method": self.parse_method,
"parser_config": self.parser_config.to_json()
})
def update(self, update_message: dict):
res = self.put(f'/dataset/{self.id}',
update_message)
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])
if res.get("code") != 0:
raise Exception(res["message"])
def delete(self) -> bool:
res = self.rm('/dataset/delete',
{"id": self.id})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])
def list_docs(self, keywords: Optional[str] = None, offset: int = 0, limit: int = -1) -> List[Document]:
"""