Daniel Marques commited on
Commit
d2ccd1f
1 Parent(s): b9ea22a

feat: add redis

Browse files
Files changed (1) hide show
  1. main.py +26 -31
main.py CHANGED
@@ -55,8 +55,30 @@ QA = RetrievalQA.from_chain_type(
55
  },
56
  )
57
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
58
  socket_manager = WebSocketManager()
59
 
 
 
60
  app = FastAPI(title="homepage-app")
61
  api_app = FastAPI(title="api app")
62
 
@@ -143,20 +165,7 @@ def predict(data: Predict):
143
  try:
144
  user_prompt = data.prompt
145
  if user_prompt:
146
- res = QA(user_prompt)
147
-
148
- answer, docs = res["result"], res["source_documents"]
149
-
150
- prompt_response_dict = {
151
- "Prompt": user_prompt,
152
- "Answer": answer,
153
- }
154
-
155
- prompt_response_dict["Sources"] = []
156
- for document in docs:
157
- prompt_response_dict["Sources"].append(
158
- (os.path.basename(str(document.metadata["source"])), str(document.page_content))
159
- )
160
 
161
  return {"response": prompt_response_dict}
162
  else:
@@ -213,7 +222,7 @@ async def websocket_endpoint(websocket: WebSocket, room_id: str, user_id: int):
213
  message = {
214
  "user_id": user_id,
215
  "room_id": room_id,
216
- "message": f"User {user_id} connected to room - {room_id}"
217
  }
218
 
219
  await socket_manager.broadcast_to_room(room_id, json.dumps(message))
@@ -228,23 +237,9 @@ async def websocket_endpoint(websocket: WebSocket, room_id: str, user_id: int):
228
  "message": data
229
  }
230
 
231
- await socket_manager.broadcast_to_room(room_id, json.dumps(message))
232
-
233
- # user_prompt = await websocket.receive_text()
234
- # response = QA(inputs=user_prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True)
235
- # answer, docs = response["result"], response["source_documents"]
236
-
237
- # prompt_response_dict = {
238
- # "Prompt": user_prompt,
239
- # "Answer": answer,
240
- # }
241
 
242
- # prompt_response_dict["Sources"] = []
243
- # for document in docs:
244
- # prompt_response_dict["Sources"].append(
245
- # (os.path.basename(str(document.metadata["source"])), str(document.page_content))
246
- # )
247
- # await websocket.send_json(prompt_response_dict)
248
 
249
  except WebSocketDisconnect:
250
  await socket_manager.remove_user_from_room(room_id, websocket)
 
55
  },
56
  )
57
 
58
+ def sendPromptChain(QA, user_prompt):
59
+
60
+ res = QA(user_prompt)
61
+
62
+ answer, docs = res["result"], res["source_documents"]
63
+
64
+ prompt_response_dict = {
65
+ "Prompt": user_prompt,
66
+ "Answer": answer,
67
+ }
68
+
69
+ prompt_response_dict["Sources"] = []
70
+ for document in docs:
71
+ prompt_response_dict["Sources"].append(
72
+ (os.path.basename(str(document.metadata["source"])), str(document.page_content))
73
+ )
74
+
75
+ return prompt_response_dict;
76
+
77
+
78
  socket_manager = WebSocketManager()
79
 
80
+
81
+
82
  app = FastAPI(title="homepage-app")
83
  api_app = FastAPI(title="api app")
84
 
 
165
  try:
166
  user_prompt = data.prompt
167
  if user_prompt:
168
+ prompt_response_dict = sendPromptChain(QA, user_prompt)
 
 
 
 
 
 
 
 
 
 
 
 
 
169
 
170
  return {"response": prompt_response_dict}
171
  else:
 
222
  message = {
223
  "user_id": user_id,
224
  "room_id": room_id,
225
+ "message": f"Student {user_id} connected to the classroom"
226
  }
227
 
228
  await socket_manager.broadcast_to_room(room_id, json.dumps(message))
 
237
  "message": data
238
  }
239
 
240
+ prompt_response_dict = sendPromptChain(QA, data)
 
 
 
 
 
 
 
 
 
241
 
242
+ await socket_manager.broadcast_to_room(room_id, prompt_response_dict)
 
 
 
 
 
243
 
244
  except WebSocketDisconnect:
245
  await socket_manager.remove_user_from_room(room_id, websocket)