Skip to content

Commit

Permalink
SDK for session (#2354)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?

Includes SDK for creating, updating sessions, getting sessions, listing
sessions, and dialogues
#1102 
### Type of change

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

---------

Co-authored-by: liuhua <[email protected]>
  • Loading branch information
Feiue and liuhua authored Sep 11, 2024
1 parent 7fad48f commit 1fc14ff
Show file tree
Hide file tree
Showing 7 changed files with 298 additions and 131 deletions.
153 changes: 124 additions & 29 deletions api/apps/sdk/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,13 @@
# limitations under the License.
#
import json
from copy import deepcopy
from uuid import uuid4

from flask import request, Response

from api.db import StatusEnum
from api.db.services.dialog_service import DialogService, ConversationService, chat
from api.settings import RetCode
from api.utils import get_uuid
from api.utils.api_utils import get_data_error_result
from api.utils.api_utils import get_json_result, token_required
Expand All @@ -31,12 +31,6 @@
def set_conversation(tenant_id):
req = request.json
conv_id = req.get("id")
if "messages" in req:
req["message"] = req.pop("messages")
if req["message"]:
for message in req["message"]:
if "reference" in message:
req["reference"] = message.pop("reference")
if "assistant_id" in req:
req["dialog_id"] = req.pop("assistant_id")
if "id" in req:
Expand All @@ -52,10 +46,12 @@ def set_conversation(tenant_id):
return get_data_error_result(retmsg="You do not own the assistant")
if "dialog_id" in req and not req.get("dialog_id"):
return get_data_error_result(retmsg="assistant_id can not be empty.")
if "message" in req:
return get_data_error_result(retmsg="message can not be change")
if "reference" in req:
return get_data_error_result(retmsg="reference can not be change")
if "name" in req and not req.get("name"):
return get_data_error_result(retmsg="name can not be empty.")
if "message" in req and not req.get("message"):
return get_data_error_result(retmsg="messages can not be empty")
if not ConversationService.update_by_id(conv_id, req):
return get_data_error_result(retmsg="Session updates error")
return get_json_result(data=True)
Expand All @@ -69,22 +65,17 @@ def set_conversation(tenant_id):
"id": get_uuid(),
"dialog_id": req["dialog_id"],
"name": req.get("name", "New session"),
"message": req.get("message", [{"role": "assistant", "content": dia[0].prompt_config["prologue"]}]),
"reference": req.get("reference", [])
"message": [{"role": "assistant", "content": "Hi! I am your assistant,can I help you?"}]
}
if not conv.get("name"):
return get_data_error_result(retmsg="name can not be empty.")
if not conv.get("message"):
return get_data_error_result(retmsg="messages can not be empty")
ConversationService.save(**conv)
e, conv = ConversationService.get_by_id(conv["id"])
if not e:
return get_data_error_result(retmsg="Fail to new session!")
conv = conv.to_dict()
conv["messages"] = conv.pop("message")
conv['messages'] = conv.pop("message")
conv["assistant_id"] = conv.pop("dialog_id")
for message in conv["messages"]:
message["reference"] = conv.get("reference")
del conv["reference"]
return get_json_result(data=conv)

Expand All @@ -96,31 +87,28 @@ def completion(tenant_id):
# req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
# {"role": "user", "content": "上海有吗?"}
# ]}
if "id" not in req:
return get_data_error_result(retmsg="id is required")
conv = ConversationService.query(id=req["id"])
if not conv:
return get_data_error_result(retmsg="Session does not exist")
conv = conv[0]
if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_data_error_result(retmsg="You do not own the session")
msg = []
question = {
"content": req.get("question"),
"role": "user",
"id": str(uuid4())
}
req["messages"].append(question)
for m in req["messages"]:
conv.message.append(question)
for m in conv.message:
if m["role"] == "system": continue
if m["role"] == "assistant" and not msg: continue
m["id"] = m.get("id", str(uuid4()))
msg.append(m)
message_id = msg[-1].get("id")
conv = ConversationService.query(id=req["id"])
conv = conv[0]
if not conv:
return get_data_error_result(retmsg="Session does not exist")
if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_data_error_result(retmsg="You do not own the session")
conv.message = deepcopy(req["messages"])
e, dia = DialogService.get_by_id(conv.dialog_id)
if not e:
return get_data_error_result(retmsg="Dialog not found!")
del req["id"]
del req["messages"]

if not conv.reference:
conv.reference = []
Expand Down Expand Up @@ -166,3 +154,110 @@ def stream():
ConversationService.update_by_id(conv.id, conv.to_dict())
break
return get_json_result(data=answer)


@manager.route('/get', methods=['GET'])
@token_required
def get(tenant_id):
req = request.args
if "id" not in req:
return get_data_error_result(retmsg="id is required")
conv_id = req["id"]
conv = ConversationService.query(id=conv_id)
if not conv:
return get_data_error_result(retmsg="Session does not exist")
if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_data_error_result(retmsg="You do not own the session")
conv = conv[0].to_dict()
conv['messages'] = conv.pop("message")
conv["assistant_id"] = conv.pop("dialog_id")
if conv["reference"]:
messages = conv["messages"]
message_num = 0
chunk_num = 0
while message_num < len(messages):
if message_num != 0 and messages[message_num]["role"] != "user":
chunk_list = []
if "chunks" in conv["reference"][chunk_num]:
chunks = conv["reference"][chunk_num]["chunks"]
for chunk in chunks:
new_chunk = {
"id": chunk["chunk_id"],
"content": chunk["content_with_weight"],
"document_id": chunk["doc_id"],
"document_name": chunk["docnm_kwd"],
"knowledgebase_id": chunk["kb_id"],
"image_id": chunk["img_id"],
"similarity": chunk["similarity"],
"vector_similarity": chunk["vector_similarity"],
"term_similarity": chunk["term_similarity"],
"positions": chunk["positions"],
}
chunk_list.append(new_chunk)
chunk_num += 1
messages[message_num]["reference"] = chunk_list
message_num += 1
del conv["reference"]
return get_json_result(data=conv)


@manager.route('/list', methods=["GET"])
@token_required
def list(tenant_id):
assistant_id = request.args["assistant_id"]
if not DialogService.query(tenant_id=tenant_id, id=assistant_id, status=StatusEnum.VALID.value):
return get_json_result(
data=False, retmsg=f'Only owner of the assistant is authorized for this operation.',
retcode=RetCode.OPERATING_ERROR)
convs = ConversationService.query(
dialog_id=assistant_id,
order_by=ConversationService.model.create_time,
reverse=True)
convs = [d.to_dict() for d in convs]
for conv in convs:
conv['messages'] = conv.pop("message")
conv["assistant_id"] = conv.pop("dialog_id")
if conv["reference"]:
messages = conv["messages"]
message_num = 0
chunk_num = 0
while message_num < len(messages):
if message_num != 0 and messages[message_num]["role"] != "user":
chunk_list = []
if "chunks" in conv["reference"][chunk_num]:
chunks = conv["reference"][chunk_num]["chunks"]
for chunk in chunks:
new_chunk = {
"id": chunk["chunk_id"],
"content": chunk["content_with_weight"],
"document_id": chunk["doc_id"],
"document_name": chunk["docnm_kwd"],
"knowledgebase_id": chunk["kb_id"],
"image_id": chunk["img_id"],
"similarity": chunk["similarity"],
"vector_similarity": chunk["vector_similarity"],
"term_similarity": chunk["term_similarity"],
"positions": chunk["positions"],
}
chunk_list.append(new_chunk)
chunk_num += 1
messages[message_num]["reference"] = chunk_list
message_num += 1
del conv["reference"]
return get_json_result(data=convs)


@manager.route('/delete', methods=["DELETE"])
@token_required
def delete(tenant_id):
id = request.args.get("id")
if not id:
return get_data_error_result(retmsg="`id` is required in deleting operation")
conv = ConversationService.query(id=id)
if not conv:
return get_data_error_result(retmsg="Session doesn't exist")
conv = conv[0]
if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_data_error_result(retmsg="You don't own the session")
ConversationService.delete_by_id(id)
return get_json_result(data=True)
3 changes: 2 additions & 1 deletion sdk/python/ragflow/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@

from .ragflow import RAGFlow
from .modules.dataset import DataSet
from .modules.chat_assistant import Assistant
from .modules.assistant import Assistant
from .modules.session import Session
Original file line number Diff line number Diff line change
@@ -1,71 +1,86 @@
from typing import List

from .base import Base
from .session import Session, Message


class Assistant(Base):
def __init__(self, rag, res_dict):
self.id = ""
self.name = "assistant"
self.avatar = "path/to/avatar"
self.knowledgebases = ["kb1"]
self.llm = Assistant.LLM(rag, {})
self.prompt = Assistant.Prompt(rag, {})
super().__init__(rag, res_dict)

class LLM(Base):
def __init__(self, rag, res_dict):
self.model_name = "deepseek-chat"
self.temperature = 0.1
self.top_p = 0.3
self.presence_penalty = 0.4
self.frequency_penalty = 0.7
self.max_tokens = 512
super().__init__(rag, res_dict)

class Prompt(Base):
def __init__(self, rag, res_dict):
self.similarity_threshold = 0.2
self.keywords_similarity_weight = 0.7
self.top_n = 8
self.variables = [{"key": "knowledge", "optional": True}]
self.rerank_model = None
self.empty_response = None
self.opener = "Hi! I'm your assistant, what can I do for you?"
self.show_quote = True
self.prompt = (
"You are an intelligent assistant. Please summarize the content of the knowledge base to answer the question. "
"Please list the data in the knowledge base and answer in detail. When all knowledge base content is irrelevant to the question, "
"your answer must include the sentence 'The answer you are looking for is not found in the knowledge base!' "
"Answers need to consider chat history.\nHere is the knowledge base:\n{knowledge}\nThe above is the knowledge base."
)
super().__init__(rag, res_dict)

def save(self) -> bool:
res = self.post('/assistant/save',
{"id": self.id, "name": self.name, "avatar": self.avatar, "knowledgebases": self.knowledgebases,
"llm": self.llm.to_json(), "prompt": self.prompt.to_json()
})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])

def delete(self) -> bool:
res = self.rm('/assistant/delete',
{"id": self.id})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])

def create_session(self, name: str = "New session", messages: List[Message] = [
{"role": "assistant", "reference": [],
"content": "您好,我是您的助手小樱,长得可爱又善良,can I help you?"}]) -> Session:
res = self.post("/session/save", {"name": name, "messages": messages, "assistant_id": self.id, })
res = res.json()
if res.get("retmsg") == "success":
return Session(self.rag, res['data'])
raise Exception(res["retmsg"])

def get_prologue(self):
return self.prompt.opener
from typing import List

from .base import Base
from .session import Session


class Assistant(Base):
def __init__(self, rag, res_dict):
self.id = ""
self.name = "assistant"
self.avatar = "path/to/avatar"
self.knowledgebases = ["kb1"]
self.llm = Assistant.LLM(rag, {})
self.prompt = Assistant.Prompt(rag, {})
super().__init__(rag, res_dict)

class LLM(Base):
def __init__(self, rag, res_dict):
self.model_name = "deepseek-chat"
self.temperature = 0.1
self.top_p = 0.3
self.presence_penalty = 0.4
self.frequency_penalty = 0.7
self.max_tokens = 512
super().__init__(rag, res_dict)

class Prompt(Base):
def __init__(self, rag, res_dict):
self.similarity_threshold = 0.2
self.keywords_similarity_weight = 0.7
self.top_n = 8
self.variables = [{"key": "knowledge", "optional": True}]
self.rerank_model = None
self.empty_response = None
self.opener = "Hi! I'm your assistant, what can I do for you?"
self.show_quote = True
self.prompt = (
"You are an intelligent assistant. Please summarize the content of the knowledge base to answer the question. "
"Please list the data in the knowledge base and answer in detail. When all knowledge base content is irrelevant to the question, "
"your answer must include the sentence 'The answer you are looking for is not found in the knowledge base!' "
"Answers need to consider chat history.\nHere is the knowledge base:\n{knowledge}\nThe above is the knowledge base."
)
super().__init__(rag, res_dict)

def save(self) -> bool:
res = self.post('/assistant/save',
{"id": self.id, "name": self.name, "avatar": self.avatar, "knowledgebases": self.knowledgebases,
"llm": self.llm.to_json(), "prompt": self.prompt.to_json()
})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])

def delete(self) -> bool:
res = self.rm('/assistant/delete',
{"id": self.id})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])

def create_session(self, name: str = "New session") -> Session:
res = self.post("/session/save", {"name": name, "assistant_id": self.id})
res = res.json()
if res.get("retmsg") == "success":
return Session(self.rag, res['data'])
raise Exception(res["retmsg"])

def list_session(self) -> List[Session]:
res = self.get('/session/list', {"assistant_id": self.id})
res = res.json()
if res.get("retmsg") == "success":
result_list = []
for data in res["data"]:
result_list.append(Session(self.rag, data))
return result_list
raise Exception(res["retmsg"])

def get_session(self, id) -> Session:
res = self.get("/session/get", {"id": id})
res = res.json()
if res.get("retmsg") == "success":
return Session(self.rag, res["data"])
raise Exception(res["retmsg"])

def get_prologue(self):
return self.prompt.opener
4 changes: 2 additions & 2 deletions sdk/python/ragflow/modules/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@ def to_json(self):
pr[name] = value
return pr

def post(self, path, param):
res = self.rag.post(path, param)
def post(self, path, param, stream=False):
res = self.rag.post(path, param, stream=stream)
return res

def get(self, path, params):
Expand Down
Loading

0 comments on commit 1fc14ff

Please sign in to comment.