Kevin Hu commited on
Commit
4e421c5
·
1 Parent(s): 1660911

Fix open AI compatible rerank issue. (#3866)

Browse files

### What problem does this PR solve?
#3700
### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (2) hide show
  1. api/apps/sdk/session.py +26 -24
  2. rag/llm/rerank_model.py +1 -1
api/apps/sdk/session.py CHANGED
@@ -35,7 +35,7 @@ from api.db.services.llm_service import LLMBundle
35
 
36
  @manager.route('/chats/<chat_id>/sessions', methods=['POST'])
37
  @token_required
38
- def create(tenant_id,chat_id):
39
  req = request.json
40
  req["dialog_id"] = chat_id
41
  dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
@@ -79,7 +79,7 @@ def create_agent_session(tenant_id, agent_id):
79
  "user_id": tenant_id,
80
  "message": [{"role": "assistant", "content": canvas.get_prologue()}],
81
  "source": "agent",
82
- "dsl":json.loads(cvs.dsl)
83
  }
84
  API4ConversationService.save(**conv)
85
  conv["agent_id"] = conv.pop("dialog_id")
@@ -88,11 +88,11 @@ def create_agent_session(tenant_id, agent_id):
88
 
89
  @manager.route('/chats/<chat_id>/sessions/<session_id>', methods=['PUT'])
90
  @token_required
91
- def update(tenant_id,chat_id,session_id):
92
  req = request.json
93
  req["dialog_id"] = chat_id
94
  conv_id = session_id
95
- conv = ConversationService.query(id=conv_id,dialog_id=chat_id)
96
  if not conv:
97
  return get_error_data_result(message="Session does not exist")
98
  if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
@@ -111,7 +111,7 @@ def update(tenant_id,chat_id,session_id):
111
  @manager.route('/chats/<chat_id>/completions', methods=['POST'])
112
  @token_required
113
  def completion(tenant_id, chat_id):
114
- dia= DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value)
115
  if not dia:
116
  return get_error_data_result(message="You do not own the chat")
117
  req = request.json
@@ -120,18 +120,18 @@ def completion(tenant_id, chat_id):
120
  "id": get_uuid(),
121
  "dialog_id": chat_id,
122
  "name": req.get("name", "New session"),
123
- "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue")}]
124
  }
125
  if not conv.get("name"):
126
  return get_error_data_result(message="`name` can not be empty.")
127
  ConversationService.save(**conv)
128
  e, conv = ConversationService.get_by_id(conv["id"])
129
- session_id=conv.id
130
  else:
131
  session_id = req.get("session_id")
132
  if not req.get("question"):
133
  return get_error_data_result(message="Please input your question.")
134
- conv = ConversationService.query(id=session_id,dialog_id=chat_id)
135
  if not conv:
136
  return get_error_data_result(message="Session does not exist")
137
  conv = conv[0]
@@ -183,18 +183,18 @@ def completion(tenant_id, chat_id):
183
  chunk_list.append(new_chunk)
184
  reference["chunks"] = chunk_list
185
  ans["id"] = message_id
186
- ans["session_id"]=session_id
187
 
188
  def stream():
189
  nonlocal dia, msg, req, conv
190
  try:
191
  for ans in chat(dia, msg, **req):
192
  fillin_conv(ans)
193
- yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
194
  ConversationService.update_by_id(conv.id, conv.to_dict())
195
  except Exception as e:
196
  yield "data:" + json.dumps({"code": 500, "message": str(e),
197
- "data": {"answer": "**ERROR**: " + str(e),"reference": []}},
198
  ensure_ascii=False) + "\n\n"
199
  yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
200
 
@@ -375,10 +375,9 @@ def agent_completion(tenant_id, agent_id):
375
  return get_result(data=result)
376
 
377
 
378
-
379
  @manager.route('/chats/<chat_id>/sessions', methods=['GET'])
380
  @token_required
381
- def list_session(chat_id,tenant_id):
382
  if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
383
  return get_error_data_result(message=f"You don't own the assistant {chat_id}.")
384
  id = request.args.get("id")
@@ -390,7 +389,7 @@ def list_session(chat_id,tenant_id):
390
  desc = False
391
  else:
392
  desc = True
393
- convs = ConversationService.get_list(chat_id,page_number,items_per_page,orderby,desc,id,name)
394
  if not convs:
395
  return get_result(data=[])
396
  for conv in convs:
@@ -429,13 +428,14 @@ def list_session(chat_id,tenant_id):
429
  del conv["reference"]
430
  return get_result(data=convs)
431
 
 
432
  @manager.route('/agents/<agent_id>/sessions', methods=['GET'])
433
  @token_required
434
- def list_agent_session(agent_id,tenant_id):
435
  if not UserCanvasService.query(user_id=tenant_id, id=agent_id):
436
  return get_error_data_result(message=f"You don't own the agent {agent_id}.")
437
  id = request.args.get("id")
438
- if not API4ConversationService.query(id=id,user_id=tenant_id):
439
  return get_error_data_result(f"You don't own the session {id}")
440
  page_number = int(request.args.get("page", 1))
441
  items_per_page = int(request.args.get("page_size", 30))
@@ -444,7 +444,7 @@ def list_agent_session(agent_id,tenant_id):
444
  desc = False
445
  else:
446
  desc = True
447
- convs = API4ConversationService.get_list(agent_id,tenant_id,page_number,items_per_page,orderby,desc,id)
448
  if not convs:
449
  return get_result(data=[])
450
  for conv in convs:
@@ -486,7 +486,7 @@ def list_agent_session(agent_id,tenant_id):
486
 
487
  @manager.route('/chats/<chat_id>/sessions', methods=["DELETE"])
488
  @token_required
489
- def delete(tenant_id,chat_id):
490
  if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
491
  return get_error_data_result(message="You don't own the chat")
492
  req = request.json
@@ -494,21 +494,22 @@ def delete(tenant_id,chat_id):
494
  if not req:
495
  ids = None
496
  else:
497
- ids=req.get("ids")
498
 
499
  if not ids:
500
  conv_list = []
501
  for conv in convs:
502
  conv_list.append(conv.id)
503
  else:
504
- conv_list=ids
505
  for id in conv_list:
506
- conv = ConversationService.query(id=id,dialog_id=chat_id)
507
  if not conv:
508
  return get_error_data_result(message="The chat doesn't own the session")
509
  ConversationService.delete_by_id(id)
510
  return get_result()
511
 
 
512
  @manager.route('/sessions/ask', methods=['POST'])
513
  @token_required
514
  def ask_about(tenant_id):
@@ -517,17 +518,18 @@ def ask_about(tenant_id):
517
  return get_error_data_result("`question` is required.")
518
  if not req.get("dataset_ids"):
519
  return get_error_data_result("`dataset_ids` is required.")
520
- if not isinstance(req.get("dataset_ids"),list):
521
  return get_error_data_result("`dataset_ids` should be a list.")
522
- req["kb_ids"]=req.pop("dataset_ids")
523
  for kb_id in req["kb_ids"]:
524
- if not KnowledgebaseService.accessible(kb_id,tenant_id):
525
  return get_error_data_result(f"You don't own the dataset {kb_id}.")
526
  kbs = KnowledgebaseService.query(id=kb_id)
527
  kb = kbs[0]
528
  if kb.chunk_num == 0:
529
  return get_error_data_result(f"The dataset {kb_id} doesn't own parsed file")
530
  uid = tenant_id
 
531
  def stream():
532
  nonlocal req, uid
533
  try:
 
35
 
36
  @manager.route('/chats/<chat_id>/sessions', methods=['POST'])
37
  @token_required
38
+ def create(tenant_id, chat_id):
39
  req = request.json
40
  req["dialog_id"] = chat_id
41
  dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
 
79
  "user_id": tenant_id,
80
  "message": [{"role": "assistant", "content": canvas.get_prologue()}],
81
  "source": "agent",
82
+ "dsl": json.loads(cvs.dsl)
83
  }
84
  API4ConversationService.save(**conv)
85
  conv["agent_id"] = conv.pop("dialog_id")
 
88
 
89
  @manager.route('/chats/<chat_id>/sessions/<session_id>', methods=['PUT'])
90
  @token_required
91
+ def update(tenant_id, chat_id, session_id):
92
  req = request.json
93
  req["dialog_id"] = chat_id
94
  conv_id = session_id
95
+ conv = ConversationService.query(id=conv_id, dialog_id=chat_id)
96
  if not conv:
97
  return get_error_data_result(message="Session does not exist")
98
  if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
 
111
  @manager.route('/chats/<chat_id>/completions', methods=['POST'])
112
  @token_required
113
  def completion(tenant_id, chat_id):
114
+ dia = DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value)
115
  if not dia:
116
  return get_error_data_result(message="You do not own the chat")
117
  req = request.json
 
120
  "id": get_uuid(),
121
  "dialog_id": chat_id,
122
  "name": req.get("name", "New session"),
123
+ "message": [{"role": "assistant", "content": dia[0].prompt_config.get("prologue")}]
124
  }
125
  if not conv.get("name"):
126
  return get_error_data_result(message="`name` can not be empty.")
127
  ConversationService.save(**conv)
128
  e, conv = ConversationService.get_by_id(conv["id"])
129
+ session_id = conv.id
130
  else:
131
  session_id = req.get("session_id")
132
  if not req.get("question"):
133
  return get_error_data_result(message="Please input your question.")
134
+ conv = ConversationService.query(id=session_id, dialog_id=chat_id)
135
  if not conv:
136
  return get_error_data_result(message="Session does not exist")
137
  conv = conv[0]
 
183
  chunk_list.append(new_chunk)
184
  reference["chunks"] = chunk_list
185
  ans["id"] = message_id
186
+ ans["session_id"] = session_id
187
 
188
  def stream():
189
  nonlocal dia, msg, req, conv
190
  try:
191
  for ans in chat(dia, msg, **req):
192
  fillin_conv(ans)
193
+ yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
194
  ConversationService.update_by_id(conv.id, conv.to_dict())
195
  except Exception as e:
196
  yield "data:" + json.dumps({"code": 500, "message": str(e),
197
+ "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
198
  ensure_ascii=False) + "\n\n"
199
  yield "data:" + json.dumps({"code": 0, "data": True}, ensure_ascii=False) + "\n\n"
200
 
 
375
  return get_result(data=result)
376
 
377
 
 
378
  @manager.route('/chats/<chat_id>/sessions', methods=['GET'])
379
  @token_required
380
+ def list_session(tenant_id, chat_id):
381
  if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
382
  return get_error_data_result(message=f"You don't own the assistant {chat_id}.")
383
  id = request.args.get("id")
 
389
  desc = False
390
  else:
391
  desc = True
392
+ convs = ConversationService.get_list(chat_id, page_number, items_per_page, orderby, desc, id, name)
393
  if not convs:
394
  return get_result(data=[])
395
  for conv in convs:
 
428
  del conv["reference"]
429
  return get_result(data=convs)
430
 
431
+
432
  @manager.route('/agents/<agent_id>/sessions', methods=['GET'])
433
  @token_required
434
+ def list_agent_session(tenant_id, agent_id):
435
  if not UserCanvasService.query(user_id=tenant_id, id=agent_id):
436
  return get_error_data_result(message=f"You don't own the agent {agent_id}.")
437
  id = request.args.get("id")
438
+ if not API4ConversationService.query(id=id, user_id=tenant_id):
439
  return get_error_data_result(f"You don't own the session {id}")
440
  page_number = int(request.args.get("page", 1))
441
  items_per_page = int(request.args.get("page_size", 30))
 
444
  desc = False
445
  else:
446
  desc = True
447
+ convs = API4ConversationService.get_list(agent_id, tenant_id, page_number, items_per_page, orderby, desc, id)
448
  if not convs:
449
  return get_result(data=[])
450
  for conv in convs:
 
486
 
487
  @manager.route('/chats/<chat_id>/sessions', methods=["DELETE"])
488
  @token_required
489
+ def delete(tenant_id, chat_id):
490
  if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
491
  return get_error_data_result(message="You don't own the chat")
492
  req = request.json
 
494
  if not req:
495
  ids = None
496
  else:
497
+ ids = req.get("ids")
498
 
499
  if not ids:
500
  conv_list = []
501
  for conv in convs:
502
  conv_list.append(conv.id)
503
  else:
504
+ conv_list = ids
505
  for id in conv_list:
506
+ conv = ConversationService.query(id=id, dialog_id=chat_id)
507
  if not conv:
508
  return get_error_data_result(message="The chat doesn't own the session")
509
  ConversationService.delete_by_id(id)
510
  return get_result()
511
 
512
+
513
  @manager.route('/sessions/ask', methods=['POST'])
514
  @token_required
515
  def ask_about(tenant_id):
 
518
  return get_error_data_result("`question` is required.")
519
  if not req.get("dataset_ids"):
520
  return get_error_data_result("`dataset_ids` is required.")
521
+ if not isinstance(req.get("dataset_ids"), list):
522
  return get_error_data_result("`dataset_ids` should be a list.")
523
+ req["kb_ids"] = req.pop("dataset_ids")
524
  for kb_id in req["kb_ids"]:
525
+ if not KnowledgebaseService.accessible(kb_id, tenant_id):
526
  return get_error_data_result(f"You don't own the dataset {kb_id}.")
527
  kbs = KnowledgebaseService.query(id=kb_id)
528
  kb = kbs[0]
529
  if kb.chunk_num == 0:
530
  return get_error_data_result(f"The dataset {kb_id} doesn't own parsed file")
531
  uid = tenant_id
532
+
533
  def stream():
534
  nonlocal req, uid
535
  try:
rag/llm/rerank_model.py CHANGED
@@ -286,7 +286,7 @@ class OpenAI_APIRerank(Base):
286
  "Content-Type": "application/json",
287
  "Authorization": f"Bearer {key}"
288
  }
289
- self.model_name = model_name
290
 
291
  def similarity(self, query: str, texts: list):
292
  # noway to config Ragflow , use fix setting
 
286
  "Content-Type": "application/json",
287
  "Authorization": f"Bearer {key}"
288
  }
289
+ self.model_name = model_name.split("___")[0]
290
 
291
  def similarity(self, query: str, texts: list):
292
  # noway to config Ragflow , use fix setting