| 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) | |