0.1.1 暂存

2026-04-14 10:28:22 +08:00
parent 4a0638ba2e
commit cc83fe963a
83 changed files with 4005 additions and 179 deletions
@@ -0,0 +1,229 @@
+"""
+Schema向量索引构建器 - 基于ChromaDB
+"""
+
+import chromadb
+from chromadb.config import Settings as ChromaSettings
+from typing import Any, List, Dict, Optional
+import logging
+from pathlib import Path
+
+from schema.manager import SchemaManager
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaIndexer:
+    """
+    Schema向量索引器
+
+    功能：
+    1. 为表结构构建向量索引
+    2. 基于问题的表检索
+    3. 持久化存储和加载
+    """
+
+    def __init__(
+        self,
+        embedder: Any,
+        persist_dir: str = "./data/embeddings",
+        collection_name: str = "schema_tables",
+    ):
+        """
+        初始化索引器
+
+        Args:
+            embedder: Embedding模型实例
+            persist_dir: 向量数据库持久化目录
+            collection_name: 集合名称
+        """
+        self.embedder = embedder
+        self.persist_dir = Path(persist_dir)
+        self.persist_dir.mkdir(parents=True, exist_ok=True)
+
+        # 初始化ChromaDB客户端
+        self.client = chromadb.PersistentClient(
+            path=str(self.persist_dir),
+            settings=ChromaSettings(anonymized_telemetry=False),
+        )
+
+        # 获取或创建集合
+        self.collection = self.client.get_or_create_collection(
+            name=collection_name,
+            metadata={"hnsw:space": "cosine"},  # 使用余弦相似度
+        )
+
+        logger.info(f"[OK] 初始化SchemaIndexer: collection={collection_name}")
+
+    def build_index(
+        self,
+        schema_manager: SchemaManager,
+        batch_size: int = 32,
+        force_rebuild: bool = False,
+    ) -> bool:
+        """
+        为Schema构建向量索引
+
+        Args:
+            schema_manager: SchemaManager实例
+            batch_size: 批处理大小
+            force_rebuild: 是否强制重建（默认为False，增量更新）
+
+        Returns:
+            True=成功，False=已存在且未强制重建
+        """
+        existing_ids = set(self.collection.get()["ids"]) if self.collection.count() > 0 else set()
+
+        if not force_rebuild and existing_ids:
+            logger.info(f"索引已存在（{len(existing_ids)}条记录），跳过构建")
+            return False
+
+        if force_rebuild and existing_ids:
+            logger.info(f"强制重建索引，删除{len(existing_ids)}条旧记录")
+            self.collection.delete()
+
+        # 准备数据
+        table_texts = []
+        table_ids = []
+        table_metadatas = []
+
+        for table_info in schema_manager.get_table_for_embedding():
+            table_texts.append(table_info["text"])
+            table_ids.append(table_info["name"])
+            table_metadatas.append({
+                "table_name": table_info["name"],
+                "column_count": len(schema_manager.get_table(table_info["name"]).columns),
+            })
+
+        # 批量计算embedding
+        logger.info(f"计算{len(table_texts)}张表的embedding...")
+        embeddings = self.embedder.encode(table_texts, batch_size=batch_size)
+
+        # 存入ChromaDB
+        self.collection.add(
+            embeddings=embeddings.tolist(),
+            documents=table_texts,
+            metadatas=table_metadatas,
+            ids=table_ids,
+        )
+
+        logger.info(f"[OK] 索引构建完成，共{len(table_ids)}张表")
+        return True
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 20,
+        score_threshold: float = 0.3,
+    ) -> List[Dict]:
+        """
+        检索相关表
+
+        Args:
+            query: 查询文本（用户问题）
+            top_k: 返回前K个结果
+            score_threshold: 相似度阈值（0-1），低于此值的结果会被过滤
+
+        Returns:
+            检索结果列表，按相似度降序排列
+            [
+                {
+                    "table_name": "表名",
+                    "score": 0.95,
+                    "document": "表描述文本",
+                    "metadata": {...}
+                },
+                ...
+            ]
+        """
+        # 编码查询文本
+        query_embedding = self.embedder.encode([query])
+
+        # 执行检索
+        results = self.collection.query(
+            query_embeddings=query_embedding.tolist(),
+            n_results=min(top_k, self.collection.count()),
+        )
+
+        # 格式化结果
+        formatted = []
+        if results["ids"] and len(results["ids"][0]) > 0:
+            for idx, (table_id, distance, metadata, document) in enumerate(zip(
+                results["ids"][0],
+                results["distances"][0],
+                results["metadatas"][0],
+                results["documents"][0],
+            )):
+                score = 1 - distance  # 余弦相似度（已归一化，转换为0-1，越大越相似）
+                if score >= score_threshold:
+                    formatted.append({
+                        "table_name": table_id,
+                        "score": float(score),
+                        "document": document,
+                        "metadata": metadata,
+                        "rank": idx + 1,
+                    })
+
+        logger.debug(f"检索 '{query[:50]}...' -> 找到{len(formatted)}个相关表（阈值={score_threshold}）")
+        return formatted
+
+    def search_by_table_names(self, table_names: List[str]) -> List[Dict]:
+        """
+        直接通过表名获取表信息
+
+        Args:
+            table_names: 表名列表
+
+        Returns:
+            表信息列表
+        """
+        results = []
+        for name in table_names:
+            try:
+                result = self.collection.get(ids=[name], include=["metadatas", "documents"])
+                if result["ids"]:
+                    results.append({
+                        "table_name": name,
+                        "score": 1.0,
+                        "document": result["documents"][0],
+                        "metadata": result["metadatas"][0],
+                    })
+            except Exception as e:
+                logger.warning(f"获取表信息失败 {name}: {e}")
+
+        return results
+
+    def get_all_tables(self) -> List[str]:
+        """获取索引中的所有表名"""
+        result = self.collection.get()
+        return result["ids"] if result["ids"] else []
+
+    def delete_table(self, table_name: str):
+        """从索引中删除表"""
+        self.collection.delete(ids=[table_name])
+        logger.info(f"已删除表索引: {table_name}")
+
+    def clear(self):
+        """清空索引"""
+        self.collection.delete()
+        logger.info("索引已清空")
+
+    def count(self) -> int:
+        """获取索引中的表数量"""
+        return self.collection.count()
+
+    # ========== 统计与调试 ==========
+
+    def get_statistics(self) -> Dict:
+        """获取索引统计信息"""
+        count = self.collection.count()
+        result = self.collection.get(include=["metadatas"])
+
+        total_columns = sum(m.get("column_count", 0) for m in result["metadatas"])
+
+        return {
+            "indexed_tables": count,
+            "total_columns": total_columns,
+            "avg_columns": total_columns / count if count > 0 else 0,
+            "persist_dir": str(self.persist_dir),
+        }