230 lines
8.0 KiB
Python
230 lines
8.0 KiB
Python
"""
|
||
Few-shot 经验样本的 Chroma 向量库存储与检索。
|
||
|
||
与 SchemaIndexer 一致：使用项目统一 ``get_embedder()``；默认 Chroma 为内存
|
||
（``EphemeralClient``）。离线灌库脚本可设 ``persist_to_disk=True`` 写入磁盘。
|
||
集合名默认 ``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,
|
||
*,
|
||
persist_to_disk: bool = False,
|
||
):
|
||
"""
|
||
Args:
|
||
embedder: 编码器实例。
|
||
persist_dir: 磁盘模式下的持久化目录；内存模式下仍解析为配置引用路径。
|
||
collection_name: 集合名；可空，空则读环境变量或默认。
|
||
persist_to_disk: 为 True 时使用 ``PersistentClient`` 落盘（如灌库脚本）。
|
||
"""
|
||
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_to_disk = persist_to_disk
|
||
if persist_to_disk:
|
||
self.persist_dir.mkdir(parents=True, exist_ok=True)
|
||
self.client = chromadb.PersistentClient(
|
||
path=str(self.persist_dir),
|
||
settings=ChromaSettings(anonymized_telemetry=False),
|
||
)
|
||
else:
|
||
self.client = chromadb.EphemeralClient(
|
||
settings=ChromaSettings(anonymized_telemetry=False),
|
||
)
|
||
self.collection = self.client.get_or_create_collection(
|
||
name=self.collection_name,
|
||
metadata={"hnsw:space": "cosine"},
|
||
)
|
||
backend = f"磁盘 {self.persist_dir}" if persist_to_disk else f"内存(配置路径={self.persist_dir})"
|
||
logger.info(
|
||
f"[OK] FewShotChromaStore: {backend} collection={self.collection_name} "
|
||
f"count={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"),
|
||
)
|