feat: knowledge base first OK
This commit is contained in:
qixinbo
@@ -0,0 +1,250 @@
|
||||
import math
|
||||
import re
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from app.services.knowledge_base_store import knowledge_base_store
|
||||
from app.services.knowledge_global_config_store import knowledge_global_config_store
|
||||
from app.services.openai_compat import normalize_openai_base_url
|
||||
|
||||
try:
|
||||
from llama_index.core import Document, VectorStoreIndex
|
||||
from llama_index.core.node_parser import SentenceSplitter
|
||||
|
||||
LLAMAINDEX_AVAILABLE = True
|
||||
except Exception:
|
||||
Document = Any
|
||||
VectorStoreIndex = Any
|
||||
SentenceSplitter = Any
|
||||
LLAMAINDEX_AVAILABLE = False
|
||||
|
||||
|
||||
def _tokenize(text: str) -> List[str]:
|
||||
return re.findall(r"[a-zA-Z0-9]+|[\u4e00-\u9fff]", (text or "").lower())
|
||||
|
||||
|
||||
def _normalize_embedding_api_base(api_base: str) -> str:
|
||||
return normalize_openai_base_url(api_base)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchHit:
|
||||
doc_id: str
|
||||
title: str
|
||||
chunk: str
|
||||
score: float
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
|
||||
class KnowledgeIndexService:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.RLock()
|
||||
self._cache: Dict[str, Tuple[str, Any, List[Dict[str, Any]]]] = {}
|
||||
|
||||
@staticmethod
|
||||
def _signature(kb: Dict[str, Any]) -> str:
|
||||
doc_parts = []
|
||||
for doc in kb.get("documents", []):
|
||||
doc_parts.append(f"{doc.get('id')}:{doc.get('updated_at')}:{len(doc.get('content', ''))}")
|
||||
return "|".join(
|
||||
[
|
||||
str(kb.get("updated_at")),
|
||||
str(kb.get("chunk_size")),
|
||||
str(kb.get("chunk_overlap")),
|
||||
*doc_parts,
|
||||
]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fallback_chunks(kb: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
chunks: List[Dict[str, Any]] = []
|
||||
chunk_size = int(kb.get("chunk_size") or 512)
|
||||
overlap = int(kb.get("chunk_overlap") or 50)
|
||||
step = max(1, chunk_size - overlap)
|
||||
for doc in kb.get("documents", []):
|
||||
text = doc.get("content") or ""
|
||||
if not text:
|
||||
continue
|
||||
if len(text) <= chunk_size:
|
||||
chunks.append(
|
||||
{
|
||||
"doc_id": doc.get("id", ""),
|
||||
"title": doc.get("title", ""),
|
||||
"chunk": text,
|
||||
"metadata": doc.get("metadata") or {},
|
||||
}
|
||||
)
|
||||
continue
|
||||
for start in range(0, len(text), step):
|
||||
piece = text[start : start + chunk_size]
|
||||
if not piece:
|
||||
continue
|
||||
chunks.append(
|
||||
{
|
||||
"doc_id": doc.get("id", ""),
|
||||
"title": doc.get("title", ""),
|
||||
"chunk": piece,
|
||||
"metadata": doc.get("metadata") or {},
|
||||
}
|
||||
)
|
||||
return chunks
|
||||
|
||||
def _build_index(self, kb: Dict[str, Any]) -> Tuple[Any, List[Dict[str, Any]]]:
|
||||
fallback_chunks = self._fallback_chunks(kb)
|
||||
if not LLAMAINDEX_AVAILABLE:
|
||||
return None, fallback_chunks
|
||||
chunk_size = int(kb.get("chunk_size") or 512)
|
||||
overlap = int(kb.get("chunk_overlap") or 50)
|
||||
splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
|
||||
docs = [
|
||||
Document(
|
||||
text=(doc.get("content") or ""),
|
||||
metadata={
|
||||
"doc_id": doc.get("id", ""),
|
||||
"title": doc.get("title", ""),
|
||||
**(doc.get("metadata") or {}),
|
||||
},
|
||||
)
|
||||
for doc in kb.get("documents", [])
|
||||
if (doc.get("content") or "").strip()
|
||||
]
|
||||
if not docs:
|
||||
return None, fallback_chunks
|
||||
embed_model = self._build_embed_model(kb)
|
||||
if embed_model is not None:
|
||||
index = VectorStoreIndex.from_documents(
|
||||
docs,
|
||||
transformations=[splitter],
|
||||
embed_model=embed_model,
|
||||
)
|
||||
else:
|
||||
index = VectorStoreIndex.from_documents(docs, transformations=[splitter])
|
||||
return index, fallback_chunks
|
||||
|
||||
@staticmethod
|
||||
def _build_embed_model(kb: Dict[str, Any]) -> Any:
|
||||
global_config = knowledge_global_config_store.get()
|
||||
api_base = global_config.get("api_base")
|
||||
api_key = global_config.get("api_key")
|
||||
model_name = kb.get("embedding_model") or global_config.get("default_embedding_model")
|
||||
if not api_base or not api_key or not model_name:
|
||||
return None
|
||||
api_base = _normalize_embedding_api_base(api_base)
|
||||
try:
|
||||
from llama_index.embeddings.openai_like import OpenAILikeEmbedding
|
||||
|
||||
return OpenAILikeEmbedding(
|
||||
model_name=model_name,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
embed_batch_size=10,
|
||||
)
|
||||
except Exception:
|
||||
try:
|
||||
from llama_index.embeddings.openai import OpenAIEmbedding
|
||||
|
||||
return OpenAIEmbedding(
|
||||
model_name=model_name,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
embed_batch_size=10,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def reindex(self, kb_id: str) -> Dict[str, Any]:
|
||||
kb = knowledge_base_store.get(kb_id)
|
||||
if not kb:
|
||||
raise ValueError("Knowledge base not found")
|
||||
with self._lock:
|
||||
signature = self._signature(kb)
|
||||
index, fallback_chunks = self._build_index(kb)
|
||||
self._cache[kb_id] = (signature, index, fallback_chunks)
|
||||
return {
|
||||
"kb_id": kb_id,
|
||||
"status": "ok",
|
||||
"documents": len(kb.get("documents", [])),
|
||||
"engine": "llamaindex" if LLAMAINDEX_AVAILABLE and index is not None else "fallback",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _fallback_search(query: str, chunks: List[Dict[str, Any]], top_k: int) -> List[SearchHit]:
|
||||
q_tokens = _tokenize(query)
|
||||
if not q_tokens:
|
||||
return []
|
||||
q_set = set(q_tokens)
|
||||
scored: List[SearchHit] = []
|
||||
for chunk_item in chunks:
|
||||
c_tokens = _tokenize(chunk_item.get("chunk", ""))
|
||||
if not c_tokens:
|
||||
continue
|
||||
overlap = sum(1 for t in c_tokens if t in q_set)
|
||||
if overlap == 0:
|
||||
continue
|
||||
score = overlap / math.sqrt(len(c_tokens))
|
||||
scored.append(
|
||||
SearchHit(
|
||||
doc_id=chunk_item.get("doc_id", ""),
|
||||
title=chunk_item.get("title", ""),
|
||||
chunk=chunk_item.get("chunk", ""),
|
||||
score=float(score),
|
||||
metadata=chunk_item.get("metadata") or {},
|
||||
)
|
||||
)
|
||||
scored.sort(key=lambda x: x.score, reverse=True)
|
||||
return scored[:top_k]
|
||||
|
||||
def search(self, kb_id: str, query: str, top_k: int | None = None) -> Dict[str, Any]:
|
||||
kb = knowledge_base_store.get(kb_id)
|
||||
if not kb:
|
||||
raise ValueError("Knowledge base not found")
|
||||
if not kb.get("documents"):
|
||||
return {"answer": "", "hits": []}
|
||||
effective_top_k = int(top_k or kb.get("top_k") or 3)
|
||||
with self._lock:
|
||||
signature = self._signature(kb)
|
||||
cached = self._cache.get(kb_id)
|
||||
if not cached or cached[0] != signature:
|
||||
index, fallback_chunks = self._build_index(kb)
|
||||
cached = (signature, index, fallback_chunks)
|
||||
self._cache[kb_id] = cached
|
||||
_, index, fallback_chunks = cached
|
||||
if index is None:
|
||||
hits = self._fallback_search(query=query, chunks=fallback_chunks, top_k=effective_top_k)
|
||||
answer = "\n\n".join(hit.chunk for hit in hits)
|
||||
return {
|
||||
"answer": answer,
|
||||
"hits": [hit.__dict__ for hit in hits],
|
||||
}
|
||||
retriever = index.as_retriever(similarity_top_k=effective_top_k)
|
||||
response_nodes = retriever.retrieve(query)
|
||||
hits: List[Dict[str, Any]] = []
|
||||
for node_with_score in response_nodes:
|
||||
node = getattr(node_with_score, "node", None)
|
||||
metadata = getattr(node, "metadata", {}) if node is not None else {}
|
||||
chunk_text = ""
|
||||
if node is not None and hasattr(node, "get_content"):
|
||||
chunk_text = node.get_content()
|
||||
elif node is not None:
|
||||
chunk_text = str(getattr(node, "text", ""))
|
||||
hits.append(
|
||||
{
|
||||
"doc_id": metadata.get("doc_id", ""),
|
||||
"title": metadata.get("title", ""),
|
||||
"chunk": chunk_text,
|
||||
"score": float(getattr(node_with_score, "score", 0.0) or 0.0),
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
if not hits:
|
||||
fallback_hits = self._fallback_search(query=query, chunks=fallback_chunks, top_k=effective_top_k)
|
||||
return {
|
||||
"answer": "\n\n".join(hit.chunk for hit in fallback_hits),
|
||||
"hits": [hit.__dict__ for hit in fallback_hits],
|
||||
}
|
||||
answer = "\n\n".join(item.get("chunk", "") for item in hits if item.get("chunk"))
|
||||
return {"answer": answer, "hits": hits}
|
||||
|
||||
|
||||
knowledge_index_service = KnowledgeIndexService()
|
||||
Reference in New Issue
Block a user