初始化

src/chat/knowledge/qa_manager.py

@@ -0,0 +1,124 @@
+import time
+from typing import Tuple, List, Dict, Optional
+
+from .global_logger import logger
+from .embedding_store import EmbeddingManager
+from .kg_manager import KGManager
+
+# from .lpmmconfig import global_config
+from .utils.dyn_topk import dyn_select_top_k
+from src.llm_models.utils_model import LLMRequest
+from src.chat.utils.utils import get_embedding
+from src.config.config import global_config, model_config
+
+MAX_KNOWLEDGE_LENGTH = 10000  # 最大知识长度
+
+
+class QAManager:
+    def __init__(
+        self,
+        embed_manager: EmbeddingManager,
+        kg_manager: KGManager,
+    ):
+        self.embed_manager = embed_manager
+        self.kg_manager = kg_manager
+        self.qa_model = LLMRequest(model_set=model_config.model_task_config.lpmm_qa, request_type="lpmm.qa")
+
+    async def process_query(self, question: str) -> Optional[Tuple[List[Tuple[str, float, float]], Optional[Dict[str, float]]]]:
+        """处理查询"""
+
+        # 生成问题的Embedding
+        part_start_time = time.perf_counter()
+        question_embedding = await get_embedding(question)
+        if question_embedding is None:
+            logger.error("生成问题Embedding失败")
+            return None
+        part_end_time = time.perf_counter()
+        logger.debug(f"Embedding用时：{part_end_time - part_start_time:.5f}s")
+
+        # 根据问题Embedding查询Relation Embedding库
+        part_start_time = time.perf_counter()
+        relation_search_res = self.embed_manager.relation_embedding_store.search_top_k(
+            question_embedding,
+            global_config.lpmm_knowledge.qa_relation_search_top_k,
+        )
+        if relation_search_res is None:
+            return None
+        # 过滤阈值
+        # 考虑动态阈值：当存在显著数值差异的结果时，保留显著结果；否则，保留所有结果
+        relation_search_res = dyn_select_top_k(relation_search_res, 0.5, 1.0)
+        if not relation_search_res or relation_search_res[0][1] < global_config.lpmm_knowledge.qa_relation_threshold:
+            # 未找到相关关系
+            logger.debug("未找到相关关系，跳过关系检索")
+            relation_search_res = []
+
+        part_end_time = time.perf_counter()
+        logger.debug(f"关系检索用时：{part_end_time - part_start_time:.5f}s")
+
+        for res in relation_search_res:
+            rel_str = self.embed_manager.relation_embedding_store.store.get(res[0]).str
+            print(f"找到相关关系，相似度：{(res[1] * 100):.2f}%  -  {rel_str}")
+
+        # TODO: 使用LLM过滤三元组结果
+        # logger.info(f"LLM过滤三元组用时：{time.time() - part_start_time:.2f}s")
+        # part_start_time = time.time()
+
+        # 根据问题Embedding查询Paragraph Embedding库
+        part_start_time = time.perf_counter()
+        paragraph_search_res = self.embed_manager.paragraphs_embedding_store.search_top_k(
+            question_embedding,
+            global_config.lpmm_knowledge.qa_paragraph_search_top_k,
+        )
+        part_end_time = time.perf_counter()
+        logger.debug(f"文段检索用时：{part_end_time - part_start_time:.5f}s")
+
+        if len(relation_search_res) != 0:
+            logger.info("找到相关关系，将使用RAG进行检索")
+            # 使用KG检索
+            part_start_time = time.perf_counter()
+            result, ppr_node_weights = self.kg_manager.kg_search(
+                relation_search_res, paragraph_search_res, self.embed_manager
+            )
+            part_end_time = time.perf_counter()
+            logger.info(f"RAG检索用时：{part_end_time - part_start_time:.5f}s")
+        else:
+            logger.info("未找到相关关系，将使用文段检索结果")
+            result = paragraph_search_res
+            ppr_node_weights = None
+
+        # 过滤阈值
+        result = dyn_select_top_k(result, 0.5, 1.0)
+
+        for res in result:
+            raw_paragraph = self.embed_manager.paragraphs_embedding_store.store[res[0]].str
+            print(f"找到相关文段，相关系数：{res[1]:.8f}\n{raw_paragraph}\n\n")
+
+        return result, ppr_node_weights
+
+    async def get_knowledge(self, question: str) -> Optional[str]:
+        """获取知识"""
+        # 处理查询
+        processed_result = await self.process_query(question)
+        if processed_result is not None:
+            query_res = processed_result[0]
+            # 检查查询结果是否为空
+            if not query_res:
+                logger.debug("知识库查询结果为空，可能是知识库中没有相关内容")
+                return None
+                
+            knowledge = [
+                (
+                    self.embed_manager.paragraphs_embedding_store.store[res[0]].str,
+                    res[1],
+                )
+                for res in query_res
+            ]
+            found_knowledge = "\n".join(
+                [f"第{i + 1}条知识：{k[0]}\n 该条知识对于问题的相关性：{k[1]}" for i, k in enumerate(knowledge)]
+            )
+            if len(found_knowledge) > MAX_KNOWLEDGE_LENGTH:
+                found_knowledge = found_knowledge[:MAX_KNOWLEDGE_LENGTH] + "\n"
+            return found_knowledge
+        else:
+            logger.debug("LPMM知识库并未初始化，可能是从未导入过知识...")
+            return None