Update vectorize_knowledge_base.py via admin tool

vectorize_knowledge_base.py

@@ -9,21 +9,28 @@ from datetime import datetime
 import csv
 
 class KnowledgeBaseVectorizer:
-    def __init__(self, api_key: str, data_path: str = ""):
+    def __init__(self, api_key: str, data_path: str = "", vector_db_dir: str = ""):
         """
-        初始化向量化器
+        初始化向量化器（适配学生Space）
         
         Args:
             api_key: OpenAI API密钥
-            data_path: knowledge_base.md文件的路径
+            data_path: knowledge_base.md文件的路径（如果为空，使用vector_db_dir中的文件）
+            vector_db_dir: 向量数据库所在目录（通常是数据存储仓库的本地目录）
         """
         self.client = OpenAI(api_key=api_key)
-        self.data_path = data_path
         self.embedding_model = "text-embedding-3-small"
-        #self.vector_db_path = os.path.join(os.path.dirname(data_path), "vector_database.csv")
-        #self.metadata_path = os.path.join(os.path.dirname(data_path), "vector_metadata.json")
-        self.vector_db_path = "vector_database.csv"
-        self.metadata_path = "vector_metadata.json"
+        
+        # 如果指定了vector_db_dir，优先使用该目录中的文件
+        if vector_db_dir:
+            self.data_path = os.path.join(vector_db_dir, "knowledge_base.md")
+            self.vector_db_path = os.path.join(vector_db_dir, "vector_database.csv")
+            self.metadata_path = os.path.join(vector_db_dir, "vector_metadata.json")
+        else:
+            # 保持原有逻辑用于向后兼容
+            self.data_path = data_path if data_path else "knowledge_base.md"
+            self.vector_db_path = "vector_database.csv"
+            self.metadata_path = "vector_metadata.json"
         
         # 缓存相关属性
         self._cached_df = None
@@ -31,6 +38,11 @@ class KnowledgeBaseVectorizer:
         self._cached_embeddings = {}  # 缓存不同类型的向量矩阵
         self._last_load_time = None
         
+        print(f"[KnowledgeBaseVectorizer] Initialized with:")
+        print(f"  - Knowledge base: {self.data_path}")
+        print(f"  - Vector database: {self.vector_db_path}")
+        print(f"  - Metadata: {self.metadata_path}")
+        
     def parse_knowledge_base(self) -> List[Dict]:
         """
         解析knowledge_base.md文件，提取所有数据条目
@@ -44,8 +56,12 @@ class KnowledgeBaseVectorizer:
         try:
             with open(self.data_path, 'r', encoding='utf-8') as f:
                 content = f.read()
+            print(f"[parse_knowledge_base] Successfully read file: {self.data_path}")
         except FileNotFoundError:
-            print(f"错误：找不到文件 {self.data_path}")
+            print(f"[parse_knowledge_base] Error: File not found - {self.data_path}")
+            return entries
+        except Exception as e:
+            print(f"[parse_knowledge_base] Error reading file: {e}")
             return entries
         
         # 改进的匹配策略：使用更精确的正则表达式
@@ -79,15 +95,15 @@ class KnowledgeBaseVectorizer:
             }
             entries.append(entry)
             
-        print(f"成功解析 {len(entries)} 个数据条目")
+        print(f"[parse_knowledge_base] Successfully parsed {len(entries)} entries")
         
         # 打印一些调试信息
         if entries:
-            print("前3个条目的内容长度:")
+            print("[parse_knowledge_base] First 3 entries info:")
             for i, entry in enumerate(entries[:3]):
                 content_lines = entry['content'].count('\n') + 1
                 has_table = '|' in entry['content']
-                print(f"  条目 {entry['id']}: {len(entry['content'])} 字符, {content_lines} 行, 包含表格: {has_table}")
+                print(f"  Entry {entry['id']}: {len(entry['content'])} chars, {content_lines} lines, has table: {has_table}")
         
         return entries
     
@@ -108,7 +124,7 @@ class KnowledgeBaseVectorizer:
             )
             return response.data[0].embedding
         except Exception as e:
-            print(f"获取向量时出错: {e}")
+            print(f"[get_embedding] Error: {e}")
             return []
     
     def batch_get_embeddings(self, texts: List[str], batch_size: int = 10) -> List[List[float]]:
@@ -126,7 +142,7 @@ class KnowledgeBaseVectorizer:
         
         for i in range(0, len(texts), batch_size):
             batch = texts[i:i + batch_size]
-            print(f"处理批次 {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")
+            print(f"[batch_get_embeddings] Processing batch {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")
             
             try:
                 response = self.client.embeddings.create(
@@ -136,7 +152,7 @@ class KnowledgeBaseVectorizer:
                 batch_embeddings = [item.embedding for item in response.data]
                 embeddings.extend(batch_embeddings)
             except Exception as e:
-                print(f"批次处理出错: {e}")
+                print(f"[batch_get_embeddings] Batch error: {e}")
                 # 如果批处理失败，尝试单个处理
                 for text in batch:
                     embedding = self.get_embedding(text)
@@ -149,12 +165,12 @@ class KnowledgeBaseVectorizer:
         创建向量数据库并保存为CSV文件
         支持标题和内容的分别向量化
         """
-        print("开始创建向量数据库...")
+        print("[create_vector_database] Starting to create vector database...")
         
         # 1. 解析知识库
         entries = self.parse_knowledge_base()
         if not entries:
-            print("没有找到任何数据条目")
+            print("[create_vector_database] No entries found")
             return
         
         # 2. 准备要向量化的文本
@@ -163,17 +179,17 @@ class KnowledgeBaseVectorizer:
         full_texts = [entry['full_text'] for entry in entries]
         
         # 3. 批量获取向量
-        print("开始向量化标题...")
+        print("[create_vector_database] Vectorizing titles...")
         title_embeddings = self.batch_get_embeddings(titles)
         
-        print("开始向量化内容...")
+        print("[create_vector_database] Vectorizing contents...")
         content_embeddings = self.batch_get_embeddings(contents)
         
-        print("开始向量化完整文本...")
+        print("[create_vector_database] Vectorizing full texts...")
         full_embeddings = self.batch_get_embeddings(full_texts)
         
         # 4. 创建DataFrame来存储数据
-        print("创建向量数据库DataFrame...")
+        print("[create_vector_database] Creating DataFrame...")
         
         # 准备数据行
         rows = []
@@ -205,7 +221,7 @@ class KnowledgeBaseVectorizer:
         df = pd.DataFrame(rows)
         
         # 5. 保存为CSV文件
-        print("保存向量数据库到CSV...")
+        print(f"[create_vector_database] Saving to {self.vector_db_path}...")
         df.to_csv(self.vector_db_path, index=False, encoding='utf-8')
         
         # 6. 保存元数据（JSON格式，便于查看）
@@ -228,11 +244,11 @@ class KnowledgeBaseVectorizer:
         with open(self.metadata_path, 'w', encoding='utf-8') as f:
             json.dump(metadata, f, ensure_ascii=False, indent=2)
         
-        print(f"向量数据库创建完成！")
-        print(f"向量数据库保存在: {self.vector_db_path}")
-        print(f"元数据保存在: {self.metadata_path}")
-        print(f"总共处理了 {len(entries)} 个条目")
-        print(f"每个向量的维度: {len(title_embeddings[0]) if title_embeddings else 0}")
+        print(f"[create_vector_database] Vector database created successfully!")
+        print(f"  - Vector database saved to: {self.vector_db_path}")
+        print(f"  - Metadata saved to: {self.metadata_path}")
+        print(f"  - Processed {len(entries)} entries")
+        print(f"  - Vector dimensions: {len(title_embeddings[0]) if title_embeddings else 0}")
         
         # 清除缓存以便重新加载
         self.clear_cache()
@@ -243,7 +259,7 @@ class KnowledgeBaseVectorizer:
         self._cached_metadata = None
         self._cached_embeddings = {}
         self._last_load_time = None
-        print("向量数据库缓存已清除")
+        print("[clear_cache] Vector database cache cleared")
     
     def load_vector_database(self, force_reload: bool = False) -> Tuple[Optional[pd.DataFrame], Optional[Dict]]:
         """
@@ -261,9 +277,11 @@ class KnowledgeBaseVectorizer:
         
         try:
             # 加载CSV文件
+            print(f"[load_vector_database] Loading from {self.vector_db_path}")
             df = pd.read_csv(self.vector_db_path, encoding='utf-8')
             
             # 加载元数据
+            print(f"[load_vector_database] Loading metadata from {self.metadata_path}")
             with open(self.metadata_path, 'r', encoding='utf-8') as f:
                 metadata = json.load(f)
             
@@ -275,13 +293,13 @@ class KnowledgeBaseVectorizer:
             # 预加载向量矩阵到缓存
             self._preload_embeddings()
             
-            print(f"成功加载向量数据库，包含 {len(df)} 个条目")
+            print(f"[load_vector_database] Successfully loaded vector database with {len(df)} entries")
             return df, metadata
         except FileNotFoundError as e:
-            print(f"错误：找不到文件 - {e}")
+            print(f"[load_vector_database] Error: File not found - {e}")
             return None, None
         except Exception as e:
-            print(f"加载向量数据库时出错: {e}")
+            print(f"[load_vector_database] Error loading vector database: {e}")
             return None, None
     
     def _preload_embeddings(self):
@@ -300,7 +318,7 @@ class KnowledgeBaseVectorizer:
                     'normalized': embeddings_norm
                 }
         
-        print(f"预加载了 {len(vector_types)} 种类型的向量矩阵")
+        print(f"[_preload_embeddings] Preloaded {len(vector_types)} types of vector matrices")
     
     def get_embeddings_from_df(self, df: pd.DataFrame, vector_type: str = 'full') -> np.ndarray:
         """
@@ -353,11 +371,11 @@ class KnowledgeBaseVectorizer:
             return [[] for _ in queries]
         
         # 批量获取查询向量
-        print(f"批量生成 {len(queries)} 个查询的向量...")
+        print(f"[batch_search_similar] Generating vectors for {len(queries)} queries...")
         query_embeddings = self.batch_get_embeddings(queries, batch_size=min(10, len(queries)))
         
         if len(query_embeddings) != len(queries):
-            print("查询向量生成失败")
+            print("[batch_search_similar] Query vector generation failed")
             return [[] for _ in queries]
         
         # 获取缓存的归一化向量矩阵
@@ -413,7 +431,7 @@ class KnowledgeBaseVectorizer:
                 query_results.append((entry, float(combined_similarities[idx]), similarity_details))
             
             all_results.append(query_results)
-            print(f"完成查询 {i+1}/{len(queries)}: '{query[:50]}...'")
+            print(f"[batch_search_similar] Completed query {i+1}/{len(queries)}: '{query[:50]}...'")
         
         return all_results
     
@@ -439,7 +457,7 @@ class KnowledgeBaseVectorizer:
         results = self.batch_search_similar([query], top_k, title_weight, content_weight, full_weight)
         return results[0] if results else []
     
-    def search_with_entities_optimized(self, entities: List[str], top_k: int = 3) -> List[Tuple[Dict, float, Dict]]:
+    def search_with_entities_optimized(self, entities: List[str], top_k: int = 5) -> List[Tuple[Dict, float, Dict]]:
         """
         优化版本：使用实体列表搜索知识库，只加载一次向量数据库
         
@@ -457,8 +475,8 @@ class KnowledgeBaseVectorizer:
         batch_results = self.batch_search_similar(
             entities, 
             top_k=top_k,
-            title_weight=0.5,    # 对于实体搜索，标题权重更高
-            content_weight=0.3,
+            title_weight=0.3,    # 对于实体搜索，标题权重更高
+            content_weight=0.5,
             full_weight=0.2
         )
         
@@ -477,98 +495,6 @@ class KnowledgeBaseVectorizer:
         sorted_results = sorted(all_results, key=lambda x: x[1], reverse=True)
         return sorted_results
     
-    def add_new_entry(self, id: str, title: str, source: str, content: str):
-        """
-        添加新条目到向量数据库
-        
-        Args:
-            id: 条目ID
-            title: 标题
-            source: 来源
-            content: 内容
-        """
-        # 加载现有数据库
-        df, metadata = self.load_vector_database()
-        
-        if df is None:
-            print("向量数据库不存在，将创建新的数据库")
-            df = pd.DataFrame()
-        
-        # 创建新条目
-        full_text = f"{title} {content}"
-        
-        # 获取三种类型的向量
-        print(f"正在为新条目 {id} 生成向量...")
-        title_embedding = self.get_embedding(title)
-        content_embedding = self.get_embedding(content)
-        full_embedding = self.get_embedding(full_text)
-        
-        if not all([title_embedding, content_embedding, full_embedding]):
-            print("无法生成向量")
-            return
-        
-        # 创建新条目
-        new_entry = {
-            'index': len(df),
-            'id': id,
-            'title': title,
-            'source': source,
-            'content': content,
-            'full_text': full_text
-        }
-        
-        # 添加向量维度
-        for j, val in enumerate(title_embedding):
-            new_entry[f'title_dim_{j}'] = val
-        for j, val in enumerate(content_embedding):
-            new_entry[f'content_dim_{j}'] = val
-        for j, val in enumerate(full_embedding):
-            new_entry[f'full_dim_{j}'] = val
-        
-        # 添加到DataFrame
-        new_df = pd.DataFrame([new_entry])
-        df = pd.concat([df, new_df], ignore_index=True)
-        
-        # 保存更新后的数据库
-        df.to_csv(self.vector_db_path, index=False, encoding='utf-8')
-        
-        # 更新元数据
-        if metadata:
-            metadata['num_entries'] = len(df)
-            metadata['updated_at'] = datetime.now().isoformat()
-            with open(self.metadata_path, 'w', encoding='utf-8') as f:
-                json.dump(metadata, f, ensure_ascii=False, indent=2)
-        
-        # 清除缓存以便重新加载
-        self.clear_cache()
-        
-        print(f"成功添加新条目 {id}")
-    
-    def export_to_readable_format(self, output_path: str = None):
-        """
-        导出向量数据库到更易读的格式（不包含向量维度）
-        
-        Args:
-            output_path: 输出文件路径
-        """
-        df, _ = self.load_vector_database()
-        if df is None:
-            return
-        
-        if output_path is None:
-            output_path = os.path.join(
-                os.path.dirname(self.data_path), 
-                "vector_database_readable.csv"
-            )
-        
-        # 只保留非向量列
-        non_vector_cols = [col for col in df.columns if not any(col.startswith(prefix) for prefix in ['title_dim_', 'content_dim_', 'full_dim_'])]
-        readable_df = df[non_vector_cols]
-        
-        # 保存
-        readable_df.to_csv(output_path, index=False, encoding='utf-8')
-        print(f"可读格式的数据库已保存到: {output_path}")
-    
     def get_cache_info(self) -> Dict:
         """
         获取缓存状态信息
@@ -580,5 +506,10 @@ class KnowledgeBaseVectorizer:
             'is_cached': self._cached_df is not None,
             'cache_size': len(self._cached_df) if self._cached_df is not None else 0,
             'cached_embeddings': list(self._cached_embeddings.keys()),
-            'last_load_time': self._last_load_time.isoformat() if self._last_load_time else None
-        }
+            'last_load_time': self._last_load_time.isoformat() if self._last_load_time else None,
+            'data_paths': {
+                'knowledge_base': self.data_path,
+                'vector_database': self.vector_db_path,
+                'metadata': self.metadata_path
+            }
+        }