feat: 同步知识库与工作台相关改动
This commit is contained in:
Jackzhou
@@ -0,0 +1,385 @@
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
|
||||
app = FastAPI(title="Bosun Knowledge Service", version="1.0.0")
|
||||
|
||||
|
||||
class ClassifyRequest(BaseModel):
|
||||
query: str
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
query: str
|
||||
space_key: str = "general"
|
||||
top_k: int = 5
|
||||
|
||||
|
||||
class IndexItem(BaseModel):
|
||||
id: int
|
||||
title: str
|
||||
content: str
|
||||
chunk_index: int = 0
|
||||
source_type: str = ""
|
||||
source_id: str = ""
|
||||
knowledge_space_key: str = "general"
|
||||
knowledge_space_name: str = "通用知识库"
|
||||
|
||||
|
||||
class RebuildIndexRequest(BaseModel):
|
||||
index_dir: str
|
||||
items: list[IndexItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
class KnowledgeEngine:
|
||||
def __init__(self) -> None:
|
||||
self.device = "cuda" if os.getenv("KNOWLEDGE_DEVICE", "cpu") == "cuda" and torch.cuda.is_available() else "cpu"
|
||||
self.classifier_encoder_name = os.getenv("KNOWLEDGE_BERT_MODEL", "bert-base-chinese")
|
||||
self.classifier_model_name = os.getenv("KNOWLEDGE_BERT_CLASSIFIER_MODEL", "")
|
||||
self.embedding_model_name = os.getenv("KNOWLEDGE_BGE_MODEL", "BAAI/bge-small-zh-v1.5")
|
||||
self.lock = threading.RLock()
|
||||
self.classifier_tokenizer = None
|
||||
self.classifier_encoder = None
|
||||
self.classifier_head_tokenizer = None
|
||||
self.classifier_head_model = None
|
||||
self.embedding_tokenizer = None
|
||||
self.embedding_model = None
|
||||
self.index = None
|
||||
self.metadata: list[dict[str, Any]] = []
|
||||
self.index_dir = Path(os.getenv("KNOWLEDGE_INDEX_DIR", Path(__file__).resolve().parent / "data" / "faiss"))
|
||||
self.index_path = self.index_dir / "knowledge.index"
|
||||
self.meta_path = self.index_dir / "knowledge_meta.json"
|
||||
self.label_descriptions = {
|
||||
"faq": "标准业务问答、高频固定问题、制度规则查询、功能操作说明",
|
||||
"document": "需要查阅资料、归纳文档、分析内容、总结合同与报告的复杂问题",
|
||||
"smalltalk": "闲聊、故事、笑话、天气、娱乐、无业务价值的聊天问题",
|
||||
"out_of_scope": "股票、医疗、法律、旅游、影视等超出企业知识库范围的问题",
|
||||
"invalid": "信息过少、语义不完整、无法判断意图的无效提问",
|
||||
}
|
||||
self._load_index_if_exists()
|
||||
|
||||
def health(self) -> dict[str, Any]:
|
||||
return {
|
||||
"device": self.device,
|
||||
"classifier_encoder": self.classifier_encoder_name,
|
||||
"classifier_model": self.classifier_model_name or "prototype-bert",
|
||||
"embedding_model": self.embedding_model_name,
|
||||
"index_dir": str(self.index_dir),
|
||||
"indexed_items": len(self.metadata),
|
||||
}
|
||||
|
||||
def classify(self, query: str) -> dict[str, Any]:
|
||||
query = (query or "").strip()
|
||||
if len(query) < 2:
|
||||
return {"intent": "invalid", "score": 1.0, "reason": "query_too_short"}
|
||||
|
||||
if self.classifier_model_name:
|
||||
try:
|
||||
return self._classify_with_sequence_model(query)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
return self._classify_with_prototypes(query)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
rule_result = self._rule_classify(query)
|
||||
if rule_result:
|
||||
return rule_result
|
||||
|
||||
return {"intent": "document", "score": 0.0, "reason": "bert_classifier_unavailable"}
|
||||
|
||||
def _rule_classify(self, query: str) -> dict[str, Any] | None:
|
||||
text = query.strip().lower()
|
||||
if not text:
|
||||
return {"intent": "invalid", "score": 1.0, "reason": "empty_query"}
|
||||
|
||||
smalltalk_keywords = [
|
||||
"讲个故事",
|
||||
"小故事",
|
||||
"笑话",
|
||||
"闲聊",
|
||||
"聊天",
|
||||
"夸夸我",
|
||||
"今天天气",
|
||||
"天气",
|
||||
"星座",
|
||||
"你好",
|
||||
"hi",
|
||||
"hello",
|
||||
"你是谁",
|
||||
"你能做什么",
|
||||
"你是怎么工作的",
|
||||
]
|
||||
if any(keyword in text for keyword in smalltalk_keywords):
|
||||
return {"intent": "smalltalk", "score": 0.98, "reason": "rule_smalltalk"}
|
||||
|
||||
out_of_scope_keywords = [
|
||||
"股票",
|
||||
"彩票",
|
||||
"电影",
|
||||
"明星",
|
||||
"旅游",
|
||||
"菜谱",
|
||||
"医学",
|
||||
"法律咨询",
|
||||
"写诗",
|
||||
"翻译成英文",
|
||||
]
|
||||
if any(keyword in text for keyword in out_of_scope_keywords):
|
||||
return {"intent": "out_of_scope", "score": 0.98, "reason": "rule_out_of_scope"}
|
||||
|
||||
document_keywords = [
|
||||
"总结",
|
||||
"梳理",
|
||||
"分析",
|
||||
"对比",
|
||||
"归纳",
|
||||
"提纲",
|
||||
"解读",
|
||||
"起草",
|
||||
"生成",
|
||||
"合同",
|
||||
"报告",
|
||||
"方案",
|
||||
"条款",
|
||||
"根据资料",
|
||||
"根据文档",
|
||||
"根据知识库",
|
||||
"整理",
|
||||
"审批要点",
|
||||
]
|
||||
if any(keyword in text for keyword in document_keywords):
|
||||
return {"intent": "document", "score": 0.97, "reason": "rule_document"}
|
||||
|
||||
faq_keywords = [
|
||||
"如何",
|
||||
"怎么",
|
||||
"哪里",
|
||||
"在哪",
|
||||
"是否",
|
||||
"有没有",
|
||||
"可以",
|
||||
"支持",
|
||||
"密码",
|
||||
"登录",
|
||||
"佣金",
|
||||
"规则",
|
||||
"流程",
|
||||
"审批",
|
||||
"上传",
|
||||
"删除",
|
||||
"新建",
|
||||
]
|
||||
if any(keyword in text for keyword in faq_keywords):
|
||||
return {"intent": "faq", "score": 0.96, "reason": "rule_faq"}
|
||||
|
||||
if len(query) <= 18:
|
||||
return {"intent": "faq", "score": 0.75, "reason": "rule_short_query"}
|
||||
|
||||
return None
|
||||
|
||||
def search(self, query: str, space_key: str, top_k: int) -> list[dict[str, Any]]:
|
||||
query = (query or "").strip()
|
||||
if not query:
|
||||
return []
|
||||
with self.lock:
|
||||
if self.index is None or not self.metadata:
|
||||
self._load_index_if_exists()
|
||||
if self.index is None or not self.metadata:
|
||||
return []
|
||||
|
||||
query_vec = self._encode_embedding([query])[0].astype("float32")
|
||||
query_vec = np.expand_dims(query_vec, axis=0)
|
||||
fetch_k = min(len(self.metadata), max(top_k * 8, 32))
|
||||
scores, indices = self.index.search(query_vec, fetch_k)
|
||||
items: list[dict[str, Any]] = []
|
||||
for score, idx in zip(scores[0].tolist(), indices[0].tolist()):
|
||||
if idx < 0 or idx >= len(self.metadata):
|
||||
continue
|
||||
meta = self.metadata[idx]
|
||||
if not self._match_space(space_key, meta.get("knowledge_space_key", "general")):
|
||||
continue
|
||||
items.append(
|
||||
{
|
||||
**meta,
|
||||
"snippet": self._build_snippet(meta.get("content", ""), query),
|
||||
"score": float(score),
|
||||
}
|
||||
)
|
||||
if len(items) >= top_k:
|
||||
break
|
||||
return items
|
||||
|
||||
def rebuild_index(self, index_dir: str, items: list[IndexItem]) -> dict[str, Any]:
|
||||
target_dir = Path(index_dir or self.index_dir)
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
records = [item.model_dump() for item in items if item.content.strip()]
|
||||
|
||||
if not records:
|
||||
empty_index = faiss.IndexFlatIP(384)
|
||||
faiss.write_index(empty_index, str(target_dir / "knowledge.index"))
|
||||
(target_dir / "knowledge_meta.json").write_text("[]", encoding="utf-8")
|
||||
with self.lock:
|
||||
self.index_dir = target_dir
|
||||
self.index_path = target_dir / "knowledge.index"
|
||||
self.meta_path = target_dir / "knowledge_meta.json"
|
||||
self.index = empty_index
|
||||
self.metadata = []
|
||||
return {"indexed_items": 0}
|
||||
|
||||
embeddings = self._encode_embedding([item["content"] for item in records]).astype("float32")
|
||||
dim = int(embeddings.shape[1])
|
||||
index = faiss.IndexFlatIP(dim)
|
||||
index.add(embeddings)
|
||||
faiss.write_index(index, str(target_dir / "knowledge.index"))
|
||||
(target_dir / "knowledge_meta.json").write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
with self.lock:
|
||||
self.index_dir = target_dir
|
||||
self.index_path = target_dir / "knowledge.index"
|
||||
self.meta_path = target_dir / "knowledge_meta.json"
|
||||
self.index = index
|
||||
self.metadata = records
|
||||
return {"indexed_items": len(records)}
|
||||
|
||||
def _load_index_if_exists(self) -> None:
|
||||
with self.lock:
|
||||
if self.index_path.exists() and self.meta_path.exists():
|
||||
self.index = faiss.read_index(str(self.index_path))
|
||||
self.metadata = json.loads(self.meta_path.read_text(encoding="utf-8"))
|
||||
|
||||
def _classify_with_sequence_model(self, query: str) -> dict[str, Any]:
|
||||
if self.classifier_head_tokenizer is None or self.classifier_head_model is None:
|
||||
self.classifier_head_tokenizer = AutoTokenizer.from_pretrained(self.classifier_model_name)
|
||||
self.classifier_head_model = AutoModelForSequenceClassification.from_pretrained(self.classifier_model_name).to(self.device)
|
||||
self.classifier_head_model.eval()
|
||||
encoded = self.classifier_head_tokenizer(query, return_tensors="pt", truncation=True, max_length=256).to(self.device)
|
||||
with torch.no_grad():
|
||||
logits = self.classifier_head_model(**encoded).logits
|
||||
probs = F.softmax(logits, dim=-1)[0].detach().cpu().numpy()
|
||||
labels = self.classifier_head_model.config.id2label or {}
|
||||
best_idx = int(np.argmax(probs))
|
||||
intent = str(labels.get(best_idx, "document")).lower()
|
||||
if intent not in self.label_descriptions:
|
||||
intent = "document"
|
||||
return {"intent": intent, "score": float(probs[best_idx]), "reason": "sequence_classifier"}
|
||||
|
||||
def _classify_with_prototypes(self, query: str) -> dict[str, Any]:
|
||||
texts = [query] + list(self.label_descriptions.values())
|
||||
vectors = self._encode_classifier(texts)
|
||||
query_vec = vectors[0]
|
||||
label_vecs = vectors[1:]
|
||||
scores = np.dot(label_vecs, query_vec)
|
||||
labels = list(self.label_descriptions.keys())
|
||||
best_idx = int(np.argmax(scores))
|
||||
return {
|
||||
"intent": labels[best_idx],
|
||||
"score": float(scores[best_idx]),
|
||||
"reason": "prototype_bert",
|
||||
}
|
||||
|
||||
def _encode_classifier(self, texts: list[str]) -> np.ndarray:
|
||||
if self.classifier_tokenizer is None or self.classifier_encoder is None:
|
||||
self.classifier_tokenizer = AutoTokenizer.from_pretrained(self.classifier_encoder_name)
|
||||
self.classifier_encoder = AutoModel.from_pretrained(self.classifier_encoder_name).to(self.device)
|
||||
self.classifier_encoder.eval()
|
||||
return self._encode(texts, self.classifier_tokenizer, self.classifier_encoder)
|
||||
|
||||
def _encode_embedding(self, texts: list[str]) -> np.ndarray:
|
||||
if self.embedding_tokenizer is None or self.embedding_model is None:
|
||||
self.embedding_tokenizer = AutoTokenizer.from_pretrained(self.embedding_model_name)
|
||||
self.embedding_model = AutoModel.from_pretrained(self.embedding_model_name).to(self.device)
|
||||
self.embedding_model.eval()
|
||||
return self._encode(texts, self.embedding_tokenizer, self.embedding_model)
|
||||
|
||||
def _encode(self, texts: list[str], tokenizer: Any, model: Any) -> np.ndarray:
|
||||
batches: list[np.ndarray] = []
|
||||
batch_size = 16
|
||||
for start in range(0, len(texts), batch_size):
|
||||
batch = texts[start : start + batch_size]
|
||||
encoded = tokenizer(
|
||||
batch,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=512,
|
||||
return_tensors="pt",
|
||||
).to(self.device)
|
||||
with torch.no_grad():
|
||||
output = model(**encoded)
|
||||
pooled = self._mean_pool(output.last_hidden_state, encoded["attention_mask"])
|
||||
pooled = F.normalize(pooled, p=2, dim=1)
|
||||
batches.append(pooled.detach().cpu().numpy())
|
||||
return np.vstack(batches)
|
||||
|
||||
@staticmethod
|
||||
def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
|
||||
mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
|
||||
summed = torch.sum(last_hidden_state * mask, dim=1)
|
||||
counts = torch.clamp(mask.sum(dim=1), min=1e-9)
|
||||
return summed / counts
|
||||
|
||||
@staticmethod
|
||||
def _match_space(selected: str, current: str) -> bool:
|
||||
selected = (selected or "general").strip().lower()
|
||||
current = (current or "general").strip().lower()
|
||||
if selected in {"", "general", "all"}:
|
||||
return True
|
||||
return current in {selected, "general"}
|
||||
|
||||
@staticmethod
|
||||
def _build_snippet(content: str, query: str) -> str:
|
||||
text = " ".join(content.split())
|
||||
if not text:
|
||||
return ""
|
||||
terms = [term for term in query.replace("，", " ").replace("。", " ").split() if term]
|
||||
hit = next((term for term in terms if term in text), "")
|
||||
if not hit:
|
||||
return text[:180]
|
||||
pos = text.find(hit)
|
||||
start = max(0, pos - 60)
|
||||
end = min(len(text), pos + 120)
|
||||
return text[start:end]
|
||||
|
||||
|
||||
engine = KnowledgeEngine()
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> dict[str, Any]:
|
||||
return {"status": "ok", "data": engine.health()}
|
||||
|
||||
|
||||
@app.post("/classify")
|
||||
def classify(req: ClassifyRequest) -> dict[str, Any]:
|
||||
try:
|
||||
return {"data": engine.classify(req.query)}
|
||||
except Exception as exc: # pragma: no cover - runtime guard
|
||||
raise HTTPException(status_code=500, detail=f"classify failed: {exc}") from exc
|
||||
|
||||
|
||||
@app.post("/search")
|
||||
def search(req: SearchRequest) -> dict[str, Any]:
|
||||
try:
|
||||
return {"data": {"items": engine.search(req.query, req.space_key, max(1, min(req.top_k, 20))) }}
|
||||
except Exception as exc: # pragma: no cover - runtime guard
|
||||
raise HTTPException(status_code=500, detail=f"search failed: {exc}") from exc
|
||||
|
||||
|
||||
@app.post("/index/rebuild")
|
||||
def rebuild_index(req: RebuildIndexRequest) -> dict[str, Any]:
|
||||
try:
|
||||
return {"data": engine.rebuild_index(req.index_dir, req.items)}
|
||||
except Exception as exc: # pragma: no cover - runtime guard
|
||||
raise HTTPException(status_code=500, detail=f"rebuild failed: {exc}") from exc
|
||||
Reference in New Issue
Block a user