File size: 2,966 Bytes
172caf6 e9078f4 d3d83ec e9078f4 d3d83ec e9078f4 172caf6 fea9976 e9078f4 73a56f9 e9078f4 fea9976 c2d189b 678763e cd7d2b9 172caf6 1f1ca89 fef3045 172caf6 c2d189b 172caf6 c2d189b e9078f4 172caf6 c2d189b 172caf6 e9078f4 fea9976 e9078f4 fea9976 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 70 71 72 73 |
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?"}]
for key,value in res_dict.items():
if key =="chat_id" and value is not None:
self.chat_id = None
self.__session_type = "chat"
if key == "agent_id" and value is not None:
self.agent_id = None
self.__session_type = "agent"
super().__init__(rag, res_dict)
def ask(self, question,stream=True,**kwargs):
if self.__session_type == "agent":
res=self._ask_agent(question,stream)
elif self.__session_type == "chat":
res=self._ask_chat(question,stream,**kwargs)
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"] is not True:
if json_data["data"].get("running_status"):
continue
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,**kwargs):
json_data={"question": question, "stream": True,"session_id":self.id}
json_data.update(kwargs)
res = self.post(f"/chats/{self.chat_id}/completions",
json_data, 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) |