eaiadminandClaude Code
0455f064ac
- 后端：新增 /api/audio/transcribe 接口，调用 Ollama whisper 进行语音识别 - 前端：新增 AudioTranscribePage.vue 页面，支持 MP3/WAV/M4A/OGG/FLAC 等格式 - 注册路由、工具卡片、智能助手欢迎语更新 Co-Authored-By: Claude Code <noreply@anthropic.com>
386 lines
15 KiB
Python
386 lines
15 KiB
Python
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
|