Enhance dialog context handling for Text2SQL queries by integrating session history. Introduce new methods for summarizing previous assistant messages and determining if the last interaction was a data query. Update environment configuration for embedding options and improve error handling in SQL generation. Add user-facing delivery messages for successful SQL execution. This update supports more coherent follow-up questions and improves user experience in conversational interactions.

2026-04-14 18:02:12 +08:00
parent 026d3bf88c
commit ca8bc5e7de
87 changed files with 5808 additions and 6351 deletions
@@ -0,0 +1,214 @@
+"""
+Few-shot 经验样本的 Chroma 向量库存储与检索。
+
+与 SchemaIndexer 一致：使用项目统一 ``get_embedder()``，持久化目录默认
+``./data/embeddings/chroma_fewshot``，集合名 ``fewshot_samples``。
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import chromadb
+from chromadb.config import Settings as ChromaSettings
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_FEWSHOT_CHROMA_DIR = "./data/embeddings/chroma_fewshot"
+COLLECTION_NAME = "fewshot_samples"
+# Chroma metadata 单值不宜过大，SQL/说明超长时截断
+_MAX_SQL_META = 16000
+_MAX_EXPLAIN_META = 6000
+_MAX_QEN_META = 2000
+
+
+class FewShotChromaStore:
+    """Few-shot JSONL → Chroma 写入与按向量检索。"""
+
+    def __init__(
+        self,
+        embedder: Any,
+        persist_dir: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ):
+        self.embedder = embedder
+        # Chroma 的 collection 名；显式参数优先，否则读 FEWSHOT_CHROMA_COLLECTION，再回退默认
+        explicit = (collection_name or "").strip()
+        from_env = (os.getenv("FEWSHOT_CHROMA_COLLECTION") or "").strip()
+        self.collection_name = explicit or from_env or COLLECTION_NAME
+        self.persist_dir = Path(
+            (persist_dir or os.getenv("FEWSHOT_CHROMA_PATH") or DEFAULT_FEWSHOT_CHROMA_DIR).strip()
+        )
+        self.persist_dir.mkdir(parents=True, exist_ok=True)
+        self.client = chromadb.PersistentClient(
+            path=str(self.persist_dir),
+            settings=ChromaSettings(anonymized_telemetry=False),
+        )
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+        logger.info(
+            "[OK] FewShotChromaStore: path=%s collection=%s count=%s",
+            self.persist_dir,
+            self.collection_name,
+            self.collection.count(),
+        )
+
+    def count(self) -> int:
+        return int(self.collection.count())
+
+    def clear(self) -> None:
+        """删除集合内全部文档（用于 force rebuild）。"""
+        if self.collection.count() > 0:
+            self.client.delete_collection(self.collection_name)
+            self.collection = self.client.get_or_create_collection(
+                name=self.collection_name,
+                metadata={"hnsw:space": "cosine"},
+            )
+            logger.info("[OK] Few-shot Chroma 集合已清空并重建")
+
+    def build_from_samples(
+        self,
+        samples: List["ExperienceSample"],
+        *,
+        force_rebuild: bool = False,
+        batch_size: int = 32,
+    ) -> int:
+        """
+        将已解析的样本列表写入 Chroma（向量由 question_zh 编码）。
+
+        Returns:
+            写入条数
+        """
+        if not samples:
+            logger.warning("Few-shot Chroma: 无样本，跳过构建")
+            return 0
+
+        if force_rebuild and self.count() > 0:
+            self.clear()
+
+        if not force_rebuild and self.count() > 0:
+            logger.info("Few-shot Chroma 已有 %s 条，跳过构建（加 --force 可重建）", self.count())
+            return self.count()
+
+        texts = [(s.question_zh or "").strip() or " " for s in samples]
+        # Chroma 要求 ids 为唯一字符串
+        ids = [str(s.qid) for s in samples]
+        metadatas: List[Dict[str, Any]] = []
+        for s in samples:
+            metadatas.append(
+                {
+                    "qid": str(s.qid),
+                    "question_en": ((s.question_en or "")[:_MAX_QEN_META]),
+                    "sql": (s.sql or "")[:_MAX_SQL_META],
+                    "explanation": (s.explanation or "")[:_MAX_EXPLAIN_META],
+                    "rating": int(s.rating) if s.rating is not None else -1,
+                    "difficulty": (s.difficulty or "medium")[:32],
+                    "tags_json": json.dumps(s.tags or [], ensure_ascii=False)[:4000],
+                }
+            )
+
+        logger.info("Few-shot Chroma: 计算 %s 条 embedding...", len(texts))
+        embeddings = self.embedder.encode(
+            texts,
+            batch_size=min(batch_size, len(texts)),
+            normalize=True,
+            show_progress=True,
+        )
+
+        self.collection.add(
+            embeddings=embeddings.tolist(),
+            documents=texts,
+            metadatas=metadatas,
+            ids=ids,
+        )
+        logger.info("[OK] Few-shot Chroma 索引完成: %s 条", len(ids))
+        return len(ids)
+
+    def search_raw(
+        self,
+        query: str,
+        top_k: int = 20,
+    ) -> List[Tuple[float, Dict[str, Any], str]]:
+        """
+        Returns:
+            (score, metadata_dict, document) 列表，score 同 SchemaIndexer 为 1-distance
+        """
+        n = self.count()
+        if n == 0:
+            return []
+        q_emb = self.embedder.encode(
+            [(query or "").strip() or " "],
+            batch_size=1,
+            normalize=True,
+            show_progress=False,
+        )
+        k = min(max(top_k, 1), n)
+        results = self.collection.query(
+            query_embeddings=q_emb.tolist(),
+            n_results=k,
+        )
+        out: List[Tuple[float, Dict[str, Any], str]] = []
+        if not results["ids"] or not results["ids"][0]:
+            return out
+        for tid, dist, meta, doc in zip(
+            results["ids"][0],
+            results["distances"][0],
+            results["metadatas"][0],
+            results["documents"][0],
+        ):
+            score = 1.0 - float(dist)
+            m = dict(meta or {})
+            m["qid"] = tid
+            out.append((score, m, doc or ""))
+        return out
+
+    def get_all_as_samples(self) -> List[Any]:
+        """导出集合中全部样本（用于统计/按标签列举；条数大时慎用）。"""
+        if self.count() == 0:
+            return []
+        data = self.collection.get(include=["metadatas", "documents"])
+        ids = data.get("ids") or []
+        metas = data.get("metadatas") or []
+        docs = data.get("documents") or []
+        out: List[Any] = []
+        for i, qid in enumerate(ids):
+            meta = dict(metas[i] or {}) if i < len(metas) else {}
+            meta["qid"] = qid
+            doc = docs[i] if i < len(docs) else ""
+            out.append(sample_from_chroma_metadata(meta, doc))
+        return out
+
+
+def sample_from_chroma_metadata(meta: Dict[str, Any], document: str) -> "ExperienceSample":
+    from utils.fewshot_selector import ExperienceSample
+
+    tags_raw = meta.get("tags_json") or "[]"
+    try:
+        tags = json.loads(tags_raw) if isinstance(tags_raw, str) else []
+    except json.JSONDecodeError:
+        tags = []
+    if not isinstance(tags, list):
+        tags = []
+    r = meta.get("rating")
+    try:
+        ri = int(r) if r is not None else -1
+    except (TypeError, ValueError):
+        ri = -1
+    rating = ri if ri >= 0 else None
+    qen = meta.get("question_en")
+    return ExperienceSample(
+        qid=str(meta.get("qid", "")),
+        question_zh=document or str(meta.get("question_zh", "")),
+        question_en=qen if qen else None,
+        sql=str(meta.get("sql", "")),
+        explanation=str(meta.get("explanation", "")),
+        rating=rating,
+        tags=tags,
+        difficulty=str(meta.get("difficulty", "medium") or "medium"),
+    )