File size: 3,441 Bytes
172caf6 cd7d2b9 172caf6 cd7d2b9 c2d189b 9a8dfa4 3d9274d c2d189b 678763e cd7d2b9 172caf6 c2d189b 172caf6 c2d189b 95da4bf c2d189b 172caf6 c2d189b 172caf6 cd7d2b9 9a8dfa4 cd7d2b9 c2d189b cd7d2b9 c2d189b 172caf6 c2d189b 172caf6 c2d189b 678763e 172caf6 678763e 95da4bf 678763e 172caf6 3d9274d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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.chat_id = None
super().__init__(rag, res_dict)
def ask(self, question: str, stream: bool = False):
for message in self.messages:
if "reference" in message:
message.pop("reference")
res = self.post(f"/chats/{self.chat_id}/completions",
{"question": question, "stream": True,"session_id":self.id}, stream=stream)
for line in res.iter_lines():
line = line.decode("utf-8")
if line.startswith("{"):
json_data = json.loads(line)
raise Exception(json_data["message"])
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"
}
if "chunks" in reference:
chunks = reference["chunks"]
chunk_list = []
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"],
"dataset_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)
temp_dict["reference"] = chunk_list
message = Message(self.rag, temp_dict)
yield message
def update(self,update_message):
res = self.put(f"/chats/{self.chat_id}/sessions/{self.id}",
update_message)
res = res.json()
if res.get("code") != 0:
raise Exception(res.get("message"))
class Message(Base):
def __init__(self, rag, res_dict):
self.content = "Hi! I am your assistant，can I help you?"
self.reference = None
self.role = "assistant"
self.prompt = None
self.id = None
super().__init__(rag, res_dict)
class Chunk(Base):
def __init__(self, rag, res_dict):
self.id = None
self.content = None
self.document_id = ""
self.document_name = ""
self.dataset_id = ""
self.image_id = ""
self.similarity = None
self.vector_similarity = None
self.term_similarity = None
self.positions = None
super().__init__(rag, res_dict)
|