update sdk document (#2374)

### What problem does this PR solve?

_Briefly describe what this PR aims to solve. Include background context
that will help reviewers understand the purpose of the PR._

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
This commit is contained in:
JobSmithManipulation
2024-09-12 14:19:45 +08:00
committed by GitHub
parent 6000c3e304
commit 6722b3d558
7 changed files with 472 additions and 2 deletions

View File

@@ -19,7 +19,7 @@ import requests
from .modules.assistant import Assistant
from .modules.dataset import DataSet
from .modules.document import Document
class RAGFlow:
def __init__(self, user_key, base_url, version='v1'):
@@ -142,3 +142,32 @@ class RAGFlow:
result_list.append(Assistant(self, data))
return result_list
raise Exception(res["retmsg"])
def create_document(self, ds:DataSet, name: str, blob: bytes) -> bool:
url = f"/doc/dataset/{ds.id}/documents/upload"
files = {
'file': (name, blob)
}
data = {
'kb_id': ds.id
}
headers = {
'Authorization': f"Bearer {ds.rag.user_key}"
}
response = requests.post(self.api_url + url, data=data, files=files,
headers=headers)
if response.status_code == 200 and response.json().get('retmsg') == 'success':
return True
else:
raise Exception(f"Upload failed: {response.json().get('retmsg')}")
return False
def get_document(self, id: str = None, name: str = None) -> Document:
res = self.get("/doc/infos", {"id": id, "name": name})
res = res.json()
if res.get("retmsg") == "success":
return Document(self, res['data'])
raise Exception(res["retmsg"])