Update vectorize_knowledge_base.py from CIV3283/CIV3283_admin

vectorize_knowledge_base.py

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