feat: add knowledge retrieval service
This commit is contained in:
@@ -9,8 +9,10 @@ from apps.knowledge.schemas import (
|
||||
KnowledgeDocumentOut,
|
||||
KnowledgeItemIn,
|
||||
KnowledgeItemOut,
|
||||
KnowledgeSearchIn,
|
||||
KnowledgeSearchResultOut,
|
||||
)
|
||||
from apps.knowledge.services import ingest_document_text, rebuild_item_chunks
|
||||
from apps.knowledge.services import ingest_document_text, rebuild_item_chunks, retrieve_knowledge
|
||||
|
||||
router = Router(tags=["knowledge"])
|
||||
|
||||
@@ -81,3 +83,13 @@ def create_knowledge_document(request, payload: KnowledgeDocumentIn):
|
||||
content=payload.content,
|
||||
operator=request.user if request.user.is_authenticated else None,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/search", response=list[KnowledgeSearchResultOut])
|
||||
def search_knowledge(request, payload: KnowledgeSearchIn):
|
||||
require_permission(request, "knowledge.view")
|
||||
return retrieve_knowledge(
|
||||
query=payload.query,
|
||||
knowledge_base_id=payload.knowledge_base_id,
|
||||
limit=min(max(payload.limit, 1), 20),
|
||||
)
|
||||
|
||||
@@ -65,3 +65,19 @@ class KnowledgeItemOut(Schema):
|
||||
source_type: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class KnowledgeSearchIn(Schema):
|
||||
query: str
|
||||
knowledge_base_id: int | None = None
|
||||
limit: int = 5
|
||||
|
||||
|
||||
class KnowledgeSearchResultOut(Schema):
|
||||
id: int
|
||||
knowledge_base_id: int
|
||||
title: str
|
||||
source_type: str
|
||||
content: str
|
||||
score: float
|
||||
metadata: dict
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import hashlib
|
||||
import math
|
||||
|
||||
from apps.knowledge.models import KnowledgeChunk, KnowledgeDocument, KnowledgeItem
|
||||
from apps.knowledge.models import KnowledgeBase, KnowledgeChunk, KnowledgeDocument, KnowledgeItem
|
||||
from apps.solutions.models import Solution
|
||||
|
||||
DEFAULT_CHUNK_SIZE = 600
|
||||
|
||||
@@ -23,6 +24,18 @@ def embed_text(content: str, dimensions: int = 16) -> list[float]:
|
||||
return [round(value / length, 6) for value in values]
|
||||
|
||||
|
||||
def keyword_score(query: str, content: str) -> float:
|
||||
normalized_query = query.lower()
|
||||
terms = [term for term in normalized_query.split() if term]
|
||||
terms.extend(character for character in normalized_query if not character.isspace())
|
||||
lowered = content.lower()
|
||||
return sum(1 for term in terms if term in lowered)
|
||||
|
||||
|
||||
def vector_score(left: list[float], right: list[float]) -> float:
|
||||
return sum(a * b for a, b in zip(left, right, strict=False))
|
||||
|
||||
|
||||
def rebuild_item_chunks(item: KnowledgeItem) -> list[KnowledgeChunk]:
|
||||
item.chunks.all().delete()
|
||||
content = "\n".join(
|
||||
@@ -85,3 +98,70 @@ def ingest_document_text(
|
||||
},
|
||||
)
|
||||
return document
|
||||
|
||||
|
||||
def index_published_solution(solution: Solution) -> list[KnowledgeChunk]:
|
||||
base, _ = KnowledgeBase.objects.get_or_create(
|
||||
name="解决方案知识库",
|
||||
defaults={"description": "从已发布解决方案生成的检索知识", "scope": "solutions"},
|
||||
)
|
||||
KnowledgeChunk.objects.filter(metadata__solution_id=solution.id).delete()
|
||||
content = "\n".join(
|
||||
[
|
||||
solution.title,
|
||||
solution.problem_signature,
|
||||
solution.related_system,
|
||||
solution.environment_scope,
|
||||
solution.root_cause,
|
||||
solution.solution,
|
||||
solution.risk_note,
|
||||
]
|
||||
)
|
||||
chunks = []
|
||||
for index, chunk_text in enumerate(split_text(content)):
|
||||
chunks.append(
|
||||
KnowledgeChunk.objects.create(
|
||||
knowledge_base=base,
|
||||
chunk_index=index,
|
||||
content=chunk_text,
|
||||
embedding=embed_text(chunk_text),
|
||||
metadata={
|
||||
"source_type": "solution",
|
||||
"title": solution.title,
|
||||
"solution_id": solution.id,
|
||||
},
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def retrieve_knowledge(
|
||||
query: str,
|
||||
knowledge_base_id: int | None = None,
|
||||
limit: int = 5,
|
||||
) -> list[dict]:
|
||||
query_embedding = embed_text(query)
|
||||
queryset = KnowledgeChunk.objects.select_related("knowledge_base", "item", "document")
|
||||
if knowledge_base_id:
|
||||
queryset = queryset.filter(knowledge_base_id=knowledge_base_id)
|
||||
scored = []
|
||||
for chunk in queryset[:1000]:
|
||||
score = keyword_score(query, chunk.content) * 10 + vector_score(
|
||||
query_embedding,
|
||||
chunk.embedding,
|
||||
)
|
||||
if score > 0:
|
||||
scored.append((score, chunk))
|
||||
scored.sort(key=lambda pair: pair[0], reverse=True)
|
||||
return [
|
||||
{
|
||||
"id": chunk.id,
|
||||
"knowledge_base_id": chunk.knowledge_base_id,
|
||||
"title": chunk.metadata.get("title", chunk.knowledge_base.name),
|
||||
"source_type": chunk.metadata.get("source_type", "chunk"),
|
||||
"content": chunk.content,
|
||||
"score": round(score, 4),
|
||||
"metadata": chunk.metadata,
|
||||
}
|
||||
for score, chunk in scored[:limit]
|
||||
]
|
||||
|
||||
53
backend/apps/knowledge/tests/test_retrieval.py
Normal file
53
backend/apps/knowledge/tests/test_retrieval.py
Normal file
@@ -0,0 +1,53 @@
|
||||
from apps.knowledge.models import KnowledgeBase, KnowledgeChunk, KnowledgeItem
|
||||
from apps.knowledge.services import retrieve_knowledge
|
||||
from apps.solutions.models import Solution
|
||||
|
||||
|
||||
def test_retrieve_knowledge_prefers_relevant_chunk(db):
|
||||
base = KnowledgeBase.objects.create(name="采集知识库", scope="数据采集")
|
||||
item = KnowledgeItem.objects.create(
|
||||
knowledge_base=base,
|
||||
title="税局采集进度",
|
||||
problem="采集进度不更新",
|
||||
solution="查询 Python 日志,确认税局状态。",
|
||||
status="published",
|
||||
)
|
||||
KnowledgeChunk.objects.create(
|
||||
knowledge_base=base,
|
||||
item=item,
|
||||
chunk_index=0,
|
||||
content="税局采集进度不更新时,需要查询 Python 日志。",
|
||||
embedding=[0.1, 0.2],
|
||||
metadata={"source_type": "knowledge_item", "title": item.title},
|
||||
)
|
||||
KnowledgeChunk.objects.create(
|
||||
knowledge_base=base,
|
||||
chunk_index=1,
|
||||
content="登录失败通常需要检查账号密码。",
|
||||
embedding=[0.9, 0.1],
|
||||
metadata={"source_type": "document", "title": "登录问题"},
|
||||
)
|
||||
|
||||
results = retrieve_knowledge("税局采集进度怎么查询", limit=1)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["title"] == "税局采集进度"
|
||||
assert "Python 日志" in results[0]["content"]
|
||||
|
||||
|
||||
def test_published_solution_can_be_indexed_for_retrieval(db):
|
||||
solution = Solution.objects.create(
|
||||
title="税局不稳定采集处理中",
|
||||
problem_signature="采集进度长时间不更新",
|
||||
root_cause="税局服务不稳定",
|
||||
solution="回复客户税局不稳定,任务仍在采集中。",
|
||||
status="published",
|
||||
)
|
||||
|
||||
from apps.knowledge.services import index_published_solution
|
||||
|
||||
chunks = index_published_solution(solution)
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].metadata["source_type"] == "solution"
|
||||
assert chunks[0].metadata["solution_id"] == solution.id
|
||||
@@ -20,7 +20,7 @@ router = Router(tags=["solutions"])
|
||||
@router.post("", response=SolutionOut)
|
||||
def create_solution(request, payload: SolutionIn):
|
||||
require_permission(request, "solution.manage")
|
||||
return Solution.objects.create(
|
||||
solution = Solution.objects.create(
|
||||
title=payload.title,
|
||||
problem_signature=payload.problem_signature,
|
||||
related_system=payload.related_system,
|
||||
@@ -31,6 +31,11 @@ def create_solution(request, payload: SolutionIn):
|
||||
status=payload.status,
|
||||
created_by=request.user if request.user.is_authenticated else None,
|
||||
)
|
||||
if solution.status == "published":
|
||||
from apps.knowledge.services import index_published_solution
|
||||
|
||||
index_published_solution(solution)
|
||||
return solution
|
||||
|
||||
|
||||
@router.get("", response=list[SolutionOut])
|
||||
|
||||
Reference in New Issue
Block a user