Update core/rag_engine.py

core/rag_engine.py

@@ -1,43 +1,62 @@
 # -*- coding: utf-8 -*-
-"""轻量级 RAG 引擎 - 解决 20 篇文献的上下文溢出与精准召回"""
-
 from typing import List
 import numpy as np
+from config import APP_CONFIG
+
+try:
+    from sklearn.feature_extraction.text import TfidfVectorizer
+    from sklearn.metrics.pairwise import cosine_similarity
+    SKLEARN_AVAILABLE = True
+except ImportError:
+    SKLEARN_AVAILABLE = False
 
 class ResearchRAGEngine:
     def __init__(self, raw_texts: List[str]):
         self.chunks = []
+        self.vectorizer = None
+        self.tfidf_matrix = None
+        self.chunk_size = APP_CONFIG["rag_chunk_size"]
         self._initialize_chunks(raw_texts)
 
     def _initialize_chunks(self, raw_texts: List[str]):
-        """将文献全文切割为 800-1000 字左右的语义块"""
         for text in raw_texts:
-            # 简单的按长度切分，实际可按段落 \n\n 切分
-            size = 3000
-            for i in range(0, len(text), size):
-                self.chunks.append(text[i:i + size])
+            for i in range(0, len(text), self.chunk_size):
+                chunk = text[i:i + self.chunk_size].strip()
+                if len(chunk) > 50:
+                    self.chunks.append(chunk)
+                    
+        if SKLEARN_AVAILABLE and self.chunks:
+            self.vectorizer = TfidfVectorizer(stop_words='english', max_df=0.85)
+            try:
+                self.tfidf_matrix = self.vectorizer.fit_transform(self.chunks)
+            except ValueError:
+                self.vectorizer = None
 
-    def get_relevant_context(self, query: str, top_k: int = 3) -> str:
-        """
-        根据当前章节主题，通过简单的关键词加权检索相关文本块。
-        (为了保持低依赖，此处使用词频匹配，如有需要可升级为 Embedding)
-        """
-        if not self.chunks: return ""
+    def get_relevant_context(self, query: str, top_k: int = None) -> str:
+        if not self.chunks: 
+            return ""
+        if top_k is None: 
+            top_k = APP_CONFIG["rag_top_k"]
         
-        scores = []
-        query_words = set(query.lower().split())
-        
-        for chunk in self.chunks:
-            chunk_lower = chunk.lower()
-            # 简单评分：计算查询词在块中的覆盖度
-            score = sum(1 for word in query_words if word in chunk_lower)
-            scores.append(score)
+        relevant_chunks = []
+        if self.vectorizer and self.tfidf_matrix is not None:
+            query_vec = self.vectorizer.transform([query])
+            similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
+            top_indices = similarities.argsort()[-top_k:][::-1]
+            relevant_chunks = [self.chunks[i] for i in top_indices if similarities[i] > 0.05]
+        else:
+            scores = []
+            query_words = set(query.lower().split())
+            for chunk in self.chunks:
+                chunk_lower = chunk.lower()
+                score = sum(1 for word in query_words if word in chunk_lower)
+                scores.append(score)
+            top_indices = np.argsort(scores)[-top_k:]
+            relevant_chunks = [self.chunks[i] for i in top_indices if scores[i] > 0]
             
-        # 获取评分最高的前 k 个块
-        top_indices = np.argsort(scores)[-top_k:]
-        relevant_chunks = [self.chunks[i] for i in top_indices if scores[i] > 0]
-        
-        return "\n--- Context Snippet ---\n".join(relevant_chunks)
+        if not relevant_chunks: 
+            return ""
+        return "\n--- Target Context Snippet ---\n" + "\n\n".join(relevant_chunks)
 
 def initialize_rag(texts: List[str]):
     return ResearchRAGEngine(texts)