Fix: renrank_model and pdf_parser bugs | Update: session API (#2601)

### What problem does this PR solve?

Fix: renrank_model and pdf_parser bugs | Update: session API
#2575
#2559
### Type of change

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

---------

Co-authored-by: liuhua <[email protected]>

api/apps/sdk/session.py

@@ -87,9 +87,9 @@ def completion(tenant_id):
     # req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
     #    {"role": "user", "content": "上海有吗？"}
     # ]}
-    if "id" not in req:
-        return get_data_error_result(retmsg="id is required")
-    conv = ConversationService.query(id=req["id"])
+    if "session_id" not in req:
+        return get_data_error_result(retmsg="session_id is required")
+    conv = ConversationService.query(id=req["session_id"])
     if not conv:
         return get_data_error_result(retmsg="Session does not exist")
     conv = conv[0]
@@ -108,7 +108,7 @@ def completion(tenant_id):
         msg.append(m)
     message_id = msg[-1].get("id")
     e, dia = DialogService.get_by_id(conv.dialog_id)
-    del req["id"]
+    del req["session_id"]
 
     if not conv.reference:
         conv.reference = []
@@ -168,6 +168,9 @@ def get(tenant_id):
         return get_data_error_result(retmsg="Session does not exist")
     if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
         return get_data_error_result(retmsg="You do not own the session")
+    if "assistant_id" in req:
+        if req["assistant_id"] != conv[0].dialog_id:
+            return get_data_error_result(retmsg="The session doesn't belong to the assistant")
     conv = conv[0].to_dict()
     conv['messages'] = conv.pop("message")
     conv["assistant_id"] = conv.pop("dialog_id")
@@ -207,7 +210,7 @@ def list(tenant_id):
     assistant_id = request.args["assistant_id"]
     if not DialogService.query(tenant_id=tenant_id, id=assistant_id, status=StatusEnum.VALID.value):
         return get_json_result(
-            data=False, retmsg=f'Only owner of the assistant is authorized for this operation.',
+            data=False, retmsg=f"You don't own the assistant.",
             retcode=RetCode.OPERATING_ERROR)
     convs = ConversationService.query(
         dialog_id=assistant_id,

deepdoc/parser/pdf_parser.py

@@ -488,7 +488,7 @@ class RAGFlowPdfParser:
                         i += 1
                         continue
 
-                    if not down["text"].strip():
+                    if not down["text"].strip() or not up["text"].strip():
                         i += 1
                         continue
 

rag/llm/rerank_model.py

@@ -26,9 +26,11 @@ from api.utils.file_utils import get_home_cache_dir
 from rag.utils import num_tokens_from_string, truncate
 import json
 
+
 def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
+
 class Base(ABC):
     def __init__(self, key, model_name):
         pass
@@ -59,16 +61,19 @@ class DefaultRerank(Base):
             with DefaultRerank._model_lock:
                 if not DefaultRerank._model:
                     try:
-                        DefaultRerank._model = FlagReranker(os.path.join(get_home_cache_dir(), re.sub(r"^[a-zA-Z]+/", "", model_name)), use_fp16=torch.cuda.is_available())
+                        DefaultRerank._model = FlagReranker(
+                            os.path.join(get_home_cache_dir(), re.sub(r"^[a-zA-Z]+/", "", model_name)),
+                            use_fp16=torch.cuda.is_available())
                     except Exception as e:
-                        model_dir = snapshot_download(repo_id= model_name,
-                                                      local_dir=os.path.join(get_home_cache_dir(), re.sub(r"^[a-zA-Z]+/", "", model_name)),
+                        model_dir = snapshot_download(repo_id=model_name,
+                                                      local_dir=os.path.join(get_home_cache_dir(),
+                                                                             re.sub(r"^[a-zA-Z]+/", "", model_name)),
                                                       local_dir_use_symlinks=False)
                         DefaultRerank._model = FlagReranker(model_dir, use_fp16=torch.cuda.is_available())
         self._model = DefaultRerank._model
 
     def similarity(self, query: str, texts: list):
-        pairs = [(query,truncate(t, 2048)) for t in texts]
+        pairs = [(query, truncate(t, 2048)) for t in texts]
         token_count = 0
         for _, t in pairs:
             token_count += num_tokens_from_string(t)
@@ -77,8 +82,10 @@ class DefaultRerank(Base):
         for i in range(0, len(pairs), batch_size):
             scores = self._model.compute_score(pairs[i:i + batch_size], max_length=2048)
             scores = sigmoid(np.array(scores)).tolist()
-            if isinstance(scores, float): res.append(scores)
-            else:  res.extend(scores)
+            if isinstance(scores, float):
+                res.append(scores)
+            else:
+                res.extend(scores)
         return np.array(res), token_count
 
 
@@ -101,7 +108,10 @@ class JinaRerank(Base):
             "top_n": len(texts)
         }
         res = requests.post(self.base_url, headers=self.headers, json=data).json()
-        return np.array([d["relevance_score"] for d in res["results"]]), res["usage"]["total_tokens"]
+        rank = np.zeros(len(texts), dtype=float)
+        for d in res["results"]:
+            rank[d["index"]] = d["relevance_score"]
+        return rank, res["usage"]["total_tokens"]
 
 
 class YoudaoRerank(DefaultRerank):
@@ -124,7 +134,7 @@ class YoudaoRerank(DefaultRerank):
                                 "maidalun1020", "InfiniFlow"))
 
         self._model = YoudaoRerank._model
-    
+
     def similarity(self, query: str, texts: list):
         pairs = [(query, truncate(t, self._model.max_length)) for t in texts]
         token_count = 0
@@ -135,8 +145,10 @@ class YoudaoRerank(DefaultRerank):
         for i in range(0, len(pairs), batch_size):
             scores = self._model.compute_score(pairs[i:i + batch_size], max_length=self._model.max_length)
             scores = sigmoid(np.array(scores)).tolist()
-            if isinstance(scores, float): res.append(scores)
-            else: res.extend(scores)
+            if isinstance(scores, float):
+                res.append(scores)
+            else:
+                res.extend(scores)
         return np.array(res), token_count
 
 
@@ -162,7 +174,10 @@ class XInferenceRerank(Base):
             "documents": texts
         }
         res = requests.post(self.base_url, headers=self.headers, json=data).json()
-        return np.array([d["relevance_score"] for d in res["results"]]), res["meta"]["tokens"]["input_tokens"]+res["meta"]["tokens"]["output_tokens"]
+        rank = np.zeros(len(texts), dtype=float)
+        for d in res["results"]:
+            rank[d["index"]] = d["relevance_score"]
+        return rank, res["meta"]["tokens"]["input_tokens"] + res["meta"]["tokens"]["output_tokens"]
 
 
 class LocalAIRerank(Base):
@@ -175,7 +190,7 @@ class LocalAIRerank(Base):
 
 class NvidiaRerank(Base):
     def __init__(
-        self, key, model_name, base_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/"
+            self, key, model_name, base_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/"
     ):
         if not base_url:
             base_url = "https://ai.api.nvidia.com/v1/retrieval/nvidia/"
@@ -208,9 +223,10 @@ class NvidiaRerank(Base):
             "top_n": len(texts),
         }
         res = requests.post(self.base_url, headers=self.headers, json=data).json()
-        rank = np.array([d["logit"] for d in res["rankings"]])
-        indexs = [d["index"] for d in res["rankings"]]
-        return rank[indexs], token_count
+        rank = np.zeros(len(texts), dtype=float)
+        for d in res["rankings"]:
+            rank[d["index"]] = d["logit"]
+        return rank, token_count
 
 
 class LmStudioRerank(Base):
@@ -247,9 +263,10 @@ class CoHereRerank(Base):
             top_n=len(texts),
             return_documents=False,
         )
-        rank = np.array([d.relevance_score for d in res.results])
-        indexs = [d.index for d in res.results]
-        return rank[indexs], token_count
+        rank = np.zeros(len(texts), dtype=float)
+        for d in res.results:
+            rank[d.index] = d.relevance_score
+        return rank, token_count
 
 
 class TogetherAIRerank(Base):
@@ -262,7 +279,7 @@ class TogetherAIRerank(Base):
 
 class SILICONFLOWRerank(Base):
     def __init__(
-        self, key, model_name, base_url="https://api.siliconflow.cn/v1/rerank"
+            self, key, model_name, base_url="https://api.siliconflow.cn/v1/rerank"
     ):
         if not base_url:
             base_url = "https://api.siliconflow.cn/v1/rerank"
@@ -287,10 +304,11 @@ class SILICONFLOWRerank(Base):
         response = requests.post(
             self.base_url, json=payload, headers=self.headers
         ).json()
-        rank = np.array([d["relevance_score"] for d in response["results"]])
-        indexs = [d["index"] for d in response["results"]]
+        rank = np.zeros(len(texts), dtype=float)
+        for d in response["results"]:
+            rank[d["index"]] = d["relevance_score"]
         return (
-            rank[indexs],
+            rank,
             response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"],
         )
 
@@ -312,9 +330,10 @@ class BaiduYiyanRerank(Base):
             documents=texts,
             top_n=len(texts),
         ).body
-        rank = np.array([d["relevance_score"] for d in res["results"]])
-        indexs = [d["index"] for d in res["results"]]
-        return rank[indexs], res["usage"]["total_tokens"]
+        rank = np.zeros(len(texts), dtype=float)
+        for d in res["results"]:
+            rank[d["index"]] = d["relevance_score"]
+        return rank, res["usage"]["total_tokens"]
 
 
 class VoyageRerank(Base):
@@ -328,6 +347,7 @@ class VoyageRerank(Base):
         res = self.client.rerank(
             query=query, documents=texts, model=self.model_name, top_k=len(texts)
         )
-        rank = np.array([r.relevance_score for r in res.results])
-        indexs = [r.index for r in res.results]
-        return rank[indexs], res.total_tokens
+        rank = np.zeros(len(texts), dtype=float)
+        for r in res.results:
+            rank[r.index] = r.relevance_score
+        return rank, res.total_tokens

sdk/python/ragflow/modules/assistant.py

@@ -76,7 +76,7 @@ class Assistant(Base):
         raise Exception(res["retmsg"])
 
     def get_session(self, id) -> Session:
-        res = self.get("/session/get", {"id": id})
+        res = self.get("/session/get", {"id": id,"assistant_id":self.id})
         res = res.json()
         if res.get("retmsg") == "success":
             return Session(self.rag, res["data"])

sdk/python/ragflow/modules/session.py

@@ -16,9 +16,12 @@ class Session(Base):
             if "reference" in message:
                 message.pop("reference")
         res = self.post("/session/completion",
-                        {"id": self.id, "question": question, "stream": stream}, stream=True)
+                        {"session_id": self.id, "question": question, "stream": True}, 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["retmsg"])
             if line.startswith("data:"):
                 json_data = json.loads(line[5:])
                 if json_data["data"] != True:
@@ -69,6 +72,7 @@ class Message(Base):
         self.reference = None
         self.role = "assistant"
         self.prompt = None
+        self.id = None
         super().__init__(rag, res_dict)
 
 
@@ -76,10 +80,10 @@ class Chunk(Base):
     def __init__(self, rag, res_dict):
         self.id = None
         self.content = None
-        self.document_id = None
-        self.document_name = None
-        self.knowledgebase_id = None
-        self.image_id = None
+        self.document_id = ""
+        self.document_name = ""
+        self.knowledgebase_id = ""
+        self.image_id = ""
         self.similarity = None
         self.vector_similarity = None
         self.term_similarity = None

sdk/python/test/t_session.py

@@ -19,7 +19,7 @@ class TestSession:
         question = "What is AI"
         for ans in session.chat(question, stream=True):
             pass
-        assert ans.content!="\n**ERROR**", "Please check this error."
+        assert not ans.content.startswith("**ERROR**"), "Please check this error."
 
     def test_delete_session_with_success(self):
         rag = RAGFlow(API_KEY, HOST_ADDRESS)