65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
|
import json
|
|||
|
|
|
|||
|
|
from .base import Base
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Session(Base):
|
|||
|
|
def __init__(self, rag, res_dict):
|
|||
|
|
self.id = None
|
|||
|
|
self.name = "New session"
|
|||
|
|
self.messages = [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
|
|||
|
|
|
|||
|
|
self.assistant_id = None
|
|||
|
|
super().__init__(rag, res_dict)
|
|||
|
|
|
|||
|
|
def chat(self, question: str, stream: bool = False):
|
|||
|
|
res = self.post("/session/completion",
|
|||
|
|
{"id": self.id, "question": question, "stream": stream, "messages": self.messages})
|
|||
|
|
res = res.text
|
|||
|
|
response_lines = res.splitlines()
|
|||
|
|
message_list = []
|
|||
|
|
for line in response_lines:
|
|||
|
|
if line.startswith("data:"):
|
|||
|
|
json_data = json.loads(line[5:])
|
|||
|
|
if json_data["data"] != True:
|
|||
|
|
answer = json_data["data"]["answer"]
|
|||
|
|
reference = json_data["data"]["reference"]
|
|||
|
|
temp_dict = {
|
|||
|
|
"content": answer,
|
|||
|
|
"role": "assistant",
|
|||
|
|
"reference": reference
|
|||
|
|
}
|
|||
|
|
message = Message(self.rag, temp_dict)
|
|||
|
|
message_list.append(message)
|
|||
|
|
return message_list
|
|||
|
|
|
|||
|
|
def save(self):
|
|||
|
|
res = self.post("/session/save",
|
|||
|
|
{"id": self.id, "dialog_id": self.assistant_id, "name": self.name, "message": self.messages})
|
|||
|
|
res = res.json()
|
|||
|
|
if res.get("retmsg") == "success": return True
|
|||
|
|
raise Exception(res.get("retmsg"))
|
|||
|
|
|
|||
|
|
class Message(Base):
|
|||
|
|
def __init__(self, rag, res_dict):
|
|||
|
|
self.content = "您好，我是您的助手小樱，长得可爱又善良，can I help you?"
|
|||
|
|
self.reference = []
|
|||
|
|
self.role = "assistant"
|
|||
|
|
self.prompt=None
|
|||
|
|
super().__init__(rag, res_dict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Chunk(Base):
|
|||
|
|
def __init__(self, rag, res_dict):
|
|||
|
|
self.id = None
|
|||
|
|
self.content = None
|
|||
|
|
self.document_id = None
|
|||
|
|
self.document_name = None
|
|||
|
|
self.knowledgebase_id = None
|
|||
|
|
self.image_id = None
|
|||
|
|
self.similarity = None
|
|||
|
|
self.vector_similarity = None
|
|||
|
|
self.term_similarity = None
|
|||
|
|
self.positions = None
|
|||
|
|
super().__init__(rag, res_dict)
|