Create app_chromadb.py

app_chromadb.py

@@ -0,0 +1,176 @@
+# app_chromadb.py
+# 这个文件只定义类和方法，它是一个被 app.py 调用的“模块”。
+
+import os
+import requests
+import hashlib
+from pathlib import Path
+from typing import List, Dict
+import time
+from datetime import datetime
+import uuid
+
+class MarkdownKnowledgeBase:
+    """
+    负责处理 Markdown 文件、与 SiliconFlow API 交互以获取向量，
+    并将最终数据存入 ChromaDB 的核心类。
+    """
+    def __init__(self, api_token: str, chroma_collection, base_url: str = "https://api.siliconflow.cn/v1"):
+        self.api_token = api_token
+        self.base_url = base_url
+        self.headers = {
+            "Authorization": f"Bearer {api_token}",
+            "Content-Type": "application/json"
+        }
+        self.collection = chroma_collection
+        
+    def get_embeddings(self, texts: List[str], model: str = "BAAI/bge-m3") -> List[List[float]]:
+        """
+        调用 SiliconFlow API 获取文本的嵌入向量。
+        """
+        url = f"{self.base_url}/embeddings"
+        embeddings = []
+        batch_size = 32
+        total_batches = (len(texts) + batch_size - 1) // batch_size
+        
+        for batch_idx in range(0, len(texts), batch_size):
+            batch = texts[batch_idx:batch_idx + batch_size]
+            current_batch = batch_idx // batch_size + 1
+            print(f"处理批次 {current_batch}/{total_batches} ({len(batch)} 个文本)")
+            payload = {"model": model, "input": batch, "encoding_format": "float"}
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    response = requests.post(url, json=payload, headers=self.headers, timeout=60)
+                    response.raise_for_status()
+                    result = response.json()
+                    if 'data' in result:
+                        embeddings.extend([item['embedding'] for item in result['data']])
+                        break 
+                    else:
+                        if attempt == max_retries - 1: embeddings.extend([[] for _ in batch])
+                except requests.exceptions.RequestException as e:
+                    print(f"  ✗ 请求失败 (尝试 {attempt + 1}/{max_retries}): {e}")
+                    if attempt == max_retries - 1: embeddings.extend([[] for _ in batch])
+            time.sleep(0.1)
+        return embeddings
+
+    def build_knowledge_base(self, folder_path: str, chunk_size: int = 4096, overlap: int = 400, 
+                           max_files: int = None, sample_mode: str = "random"):
+        """
+        扫描、分块、向量化并最终将数据存入 ChromaDB。
+        """
+        print("扫描文件并生成文本块...")
+        md_files = self._scan_files(folder_path)
+        if max_files and len(md_files) > max_files:
+            md_files = self._sample_files(md_files, max_files, sample_mode)
+        
+        all_chunks, all_metadatas = [], []
+        for file_path in md_files:
+            file_info = self._read_content(file_path)
+            if not file_info or len(file_info['content'].strip()) < 50:
+                continue
+            chunks = self._chunk_text(file_info['content'], chunk_size, overlap)
+            for j, chunk in enumerate(chunks):
+                if len(chunk.strip()) > 20:
+                    all_chunks.append(chunk)
+                    all_metadatas.append({'file_name': file_info['file_name'], 'source': file_info['file_path']})
+        
+        if not all_chunks:
+            print("没有有效的文本块可供处理。")
+            return
+            
+        print(f"总共生成 {len(all_chunks)} 个文本块，开始获取向量...")
+        embeddings = self.get_embeddings(all_chunks)
+        
+        valid_indices = [i for i, emb in enumerate(embeddings) if emb]
+        if not valid_indices:
+             print("未能成功获取任何向量，无法添加到知识库。")
+             return
+
+        valid_embeddings = [embeddings[i] for i in valid_indices]
+        valid_chunks = [all_chunks[i] for i in valid_indices]
+        valid_metadatas = [all_metadatas[i] for i in valid_indices]
+        ids = [str(uuid.uuid4()) for _ in valid_chunks]
+        
+        print(f"获取向量完成，正在将 {len(ids)} 个有效条目批量写入 ChromaDB...")
+        
+        if ids: # 确保有内容可以添加
+            self.collection.add(
+                embeddings=valid_embeddings,
+                documents=valid_chunks,
+                metadatas=valid_metadatas,
+                ids=ids
+            )
+        
+        print("知识库构建并存入 ChromaDB 成功！")
+
+    def search(self, query: str, top_k: int = 5) -> List[Dict]:
+        """
+        在 ChromaDB 中执行向量搜索。
+        """
+        print(f"在 ChromaDB 中搜索: '{query}'")
+        query_embedding = self.get_embeddings([query])[0]
+        if not query_embedding:
+            return []
+
+        results = self.collection.query(
+            query_embeddings=[query_embedding],
+            n_results=top_k
+        )
+        
+        formatted_results = []
+        if results and results['ids'][0]:
+            for i in range(len(results['ids'][0])):
+                formatted_results.append({
+                    "id": results['ids'][0][i],
+                    "content": results['documents'][0][i],
+                    "metadata": results['metadatas'][0][i],
+                    "distance": results['distances'][0][i] 
+                })
+        return formatted_results
+
+    # --- 私有辅助方法 ---
+    def _scan_files(self, folder_path: str) -> List[str]:
+        md_files = []
+        folder = Path(folder_path)
+        if not folder.exists(): return []
+        for md_file in folder.rglob("*.md"):
+            if md_file.is_file(): md_files.append(str(md_file.resolve()))
+        return md_files
+
+    def _read_content(self, file_path: str) -> Dict:
+        try:
+            encodings = ['utf-8', 'utf-8-sig', 'gbk', 'cp1252', 'latin1']
+            content = None
+            for encoding in encodings:
+                try:
+                    with open(file_path, 'r', encoding=encoding) as file:
+                        content = file.read()
+                        break
+                except UnicodeDecodeError: continue
+            if content is None: return None
+            return {'file_name': os.path.basename(file_path), 'content': content, 'file_path': file_path}
+        except Exception:
+            return None
+        
+    def _sample_files(self, md_files: List[str], max_files: int, mode: str) -> List[str]:
+         if mode == "random":
+             import random
+             return random.sample(md_files, min(len(md_files), max_files))
+         elif mode == "largest":
+             return sorted(md_files, key=lambda f: os.path.getsize(f) if os.path.exists(f) else 0, reverse=True)[:max_files]
+         elif mode == "recent":
+             return sorted(md_files, key=lambda f: os.path.getmtime(f) if os.path.exists(f) else 0, reverse=True)[:max_files]
+         return md_files[:max_files]
+
+    def _chunk_text(self, text: str, chunk_size: int, overlap: int) -> List[str]:
+        if len(text) <= chunk_size: return [text]
+        chunks = []
+        start = 0
+        while start < len(text):
+            end = start + chunk_size
+            chunk = text[start:end]
+            chunks.append(chunk)
+            start += chunk_size - overlap
+        return chunks