File size: 2,757 Bytes
172caf6 cd7d2b9 e9078f4 172caf6 73a56f9 e9078f4 73a56f9 e9078f4 73a56f9 c2d189b 678763e cd7d2b9 172caf6 c2d189b 172caf6 c2d189b e9078f4 172caf6 c2d189b 172caf6 e9078f4 73a56f9 e9078f4 73a56f9 e9078f4 cd7d2b9 9a8dfa4 cd7d2b9 c2d189b cd7d2b9 c2d189b 172caf6 c2d189b 172caf6 c2d189b 678763e e9078f4 |
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 |
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
self.agent_id = None
for key,value in res_dict.items():
if key =="chat_id" and value is not None:
self.__session_type = "chat"
if key == "agent_id" and value is not None:
self.__session_type = "agent"
super().__init__(rag, res_dict)
def ask(self, question,stream=True):
if self.__session_type == "agent":
res=self._ask_agent(question,stream)
elif self.__session_type == "chat":
res=self._ask_chat(question,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"]
temp_dict["reference"] = chunks
message = Message(self.rag, temp_dict)
yield message
def _ask_chat(self, question: str, stream: bool):
res = self.post(f"/chats/{self.chat_id}/completions",
{"question": question, "stream": True,"session_id":self.id}, stream=stream)
return res
def _ask_agent(self,question:str,stream:bool):
res = self.post(f"/agents/{self.agent_id}/completions",
{"question": question, "stream": True,"session_id":self.id}, stream=stream)
return res
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) |