545 lines
19 KiB
Python
545 lines
19 KiB
Python
"""知识索引与混合检索服务。
|
||
|
||
生产环境组合 PostgreSQL 全文、pgvector、结构化范围和可选重排;非 PostgreSQL 环境仅用于
|
||
测试与本地兼容验证。
|
||
"""
|
||
|
||
import math
|
||
import re
|
||
|
||
from django.conf import settings
|
||
from django.contrib.postgres.search import SearchQuery, SearchRank, SearchVector
|
||
from django.db import connection, transaction
|
||
from django.db.models import Q
|
||
from django.utils import timezone
|
||
from pgvector.django import CosineDistance
|
||
|
||
from apps.knowledge.embeddings import embed_contents, hash_embedding
|
||
from apps.knowledge.models import (
|
||
GovernedAssetType,
|
||
KnowledgeBase,
|
||
KnowledgeChunk,
|
||
KnowledgeDocument,
|
||
KnowledgeItem,
|
||
TroubleshootingPlan,
|
||
)
|
||
from apps.knowledge.reranking import rerank_candidates
|
||
from apps.solutions.models import Solution
|
||
|
||
DEFAULT_CHUNK_SIZE = 600
|
||
|
||
|
||
def split_text(content: str, chunk_size: int = DEFAULT_CHUNK_SIZE) -> list[str]:
|
||
normalized = "\n".join(line.strip() for line in content.splitlines() if line.strip())
|
||
if not normalized:
|
||
return []
|
||
return [
|
||
normalized[index : index + chunk_size]
|
||
for index in range(0, len(normalized), chunk_size)
|
||
]
|
||
|
||
|
||
def embed_text(content: str, dimensions: int | None = None) -> list[float]:
|
||
return hash_embedding(content, dimensions)
|
||
|
||
|
||
def _terms(value: str) -> set[str]:
|
||
normalized = value.casefold()
|
||
words = set(re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}|[\u4e00-\u9fff]{2,}", normalized))
|
||
chinese = "".join(re.findall(r"[\u4e00-\u9fff]", normalized))
|
||
words.update(chinese[index : index + 2] for index in range(max(len(chinese) - 1, 0)))
|
||
return {term for term in words if term}
|
||
|
||
|
||
def keyword_score(query: str, content: str) -> float:
|
||
query_terms = _terms(query)
|
||
content_terms = _terms(content)
|
||
if not query_terms or not content_terms:
|
||
return 0.0
|
||
return len(query_terms & content_terms) / len(query_terms)
|
||
|
||
|
||
def vector_score(left: list[float], right: list[float]) -> float:
|
||
if not left or not right:
|
||
return 0.0
|
||
left_length = math.sqrt(sum(value * value for value in left)) or 1.0
|
||
right_length = math.sqrt(sum(value * value for value in right)) or 1.0
|
||
return float(
|
||
sum(a * b for a, b in zip(left, right, strict=False))
|
||
/ (left_length * right_length)
|
||
)
|
||
|
||
|
||
def _create_chunks(
|
||
*,
|
||
knowledge_base: KnowledgeBase,
|
||
texts: list[str],
|
||
metadata: dict,
|
||
item: KnowledgeItem | None = None,
|
||
troubleshooting_plan: TroubleshootingPlan | None = None,
|
||
document: KnowledgeDocument | None = None,
|
||
historical_ticket=None,
|
||
replace_queryset=None,
|
||
) -> list[KnowledgeChunk]:
|
||
batch = embed_contents(texts)
|
||
with transaction.atomic():
|
||
if replace_queryset is not None:
|
||
replace_queryset.delete()
|
||
return [
|
||
KnowledgeChunk.objects.create(
|
||
knowledge_base=knowledge_base,
|
||
item=item,
|
||
troubleshooting_plan=troubleshooting_plan,
|
||
document=document,
|
||
historical_ticket=historical_ticket,
|
||
chunk_index=index,
|
||
content=chunk_text,
|
||
embedding=vector,
|
||
embedding_vector=vector,
|
||
embedding_model=batch.model_name,
|
||
metadata={
|
||
**metadata,
|
||
"embedding_deployment_id": batch.deployment_id,
|
||
"embedding_fallback": batch.fallback,
|
||
},
|
||
)
|
||
for index, (chunk_text, vector) in enumerate(
|
||
zip(texts, batch.vectors, strict=True)
|
||
)
|
||
]
|
||
|
||
|
||
def rebuild_item_chunks(item: KnowledgeItem) -> list[KnowledgeChunk]:
|
||
content = "\n".join(
|
||
[
|
||
item.title,
|
||
item.problem,
|
||
item.applicable_scope,
|
||
item.category_code,
|
||
item.related_system,
|
||
item.environment_scope,
|
||
item.error_code,
|
||
item.root_cause,
|
||
item.troubleshooting_steps,
|
||
item.solution,
|
||
item.risk_note,
|
||
]
|
||
)
|
||
return _create_chunks(
|
||
knowledge_base=item.knowledge_base,
|
||
item=item,
|
||
texts=split_text(content),
|
||
metadata={
|
||
"source_type": "knowledge_item",
|
||
"title": item.title,
|
||
"item_id": item.id,
|
||
"category_code": item.category_code,
|
||
"related_system": item.related_system,
|
||
"environment": item.environment_scope,
|
||
},
|
||
replace_queryset=item.chunks.all(),
|
||
)
|
||
|
||
|
||
def rebuild_troubleshooting_plan_chunks(
|
||
plan: TroubleshootingPlan,
|
||
) -> list[KnowledgeChunk]:
|
||
content = "\n".join(
|
||
[
|
||
plan.title,
|
||
plan.problem_pattern,
|
||
plan.applicable_scope,
|
||
plan.category_code,
|
||
plan.related_system,
|
||
plan.environment_scope,
|
||
plan.decision_rules,
|
||
plan.troubleshooting_steps,
|
||
plan.risk_note,
|
||
]
|
||
)
|
||
return _create_chunks(
|
||
knowledge_base=plan.knowledge_base,
|
||
troubleshooting_plan=plan,
|
||
texts=split_text(content),
|
||
metadata={
|
||
"source_type": "troubleshooting_plan",
|
||
"title": plan.title,
|
||
"troubleshooting_plan_id": plan.id,
|
||
"category_code": plan.category_code,
|
||
"related_system": plan.related_system,
|
||
"environment": plan.environment_scope,
|
||
},
|
||
replace_queryset=plan.chunks.all(),
|
||
)
|
||
|
||
|
||
def ingest_document_text(
|
||
knowledge_base_id: int,
|
||
file_name: str,
|
||
file_type: str,
|
||
content: str,
|
||
operator=None,
|
||
) -> KnowledgeDocument:
|
||
document = KnowledgeDocument.objects.create(
|
||
knowledge_base_id=knowledge_base_id,
|
||
file_name=file_name,
|
||
file_type=file_type,
|
||
content_text=content,
|
||
scan_status="not_applicable",
|
||
parse_status="success",
|
||
created_by=operator,
|
||
)
|
||
rebuild_document_chunks(document)
|
||
document.indexed_at = timezone.now()
|
||
document.save(update_fields=["indexed_at", "updated_at"])
|
||
return document
|
||
|
||
|
||
def rebuild_document_chunks(document: KnowledgeDocument) -> list[KnowledgeChunk]:
|
||
return _create_chunks(
|
||
knowledge_base=document.knowledge_base,
|
||
document=document,
|
||
texts=split_text(document.content_text),
|
||
metadata={
|
||
"source_type": "document",
|
||
"title": document.file_name,
|
||
"document_id": document.id,
|
||
},
|
||
replace_queryset=document.chunks.all(),
|
||
)
|
||
|
||
|
||
def index_published_solution(solution: Solution) -> list[KnowledgeChunk]:
|
||
base, _ = KnowledgeBase.objects.get_or_create(
|
||
name="解决方案知识库",
|
||
defaults={"description": "从已发布解决方案生成的检索知识", "scope": "solutions"},
|
||
)
|
||
content = "\n".join(
|
||
[
|
||
solution.title,
|
||
solution.problem_signature,
|
||
solution.related_system,
|
||
solution.environment_scope,
|
||
solution.root_cause,
|
||
solution.solution,
|
||
solution.risk_note,
|
||
]
|
||
)
|
||
return _create_chunks(
|
||
knowledge_base=base,
|
||
texts=split_text(content),
|
||
metadata={
|
||
"source_type": "solution",
|
||
"title": solution.title,
|
||
"solution_id": solution.id,
|
||
"solution_version": solution.updated_at.isoformat(),
|
||
"category_code": solution.category.code if solution.category_id else "",
|
||
"related_system": solution.related_system,
|
||
"environment": solution.environment_scope,
|
||
},
|
||
replace_queryset=KnowledgeChunk.objects.filter(metadata__solution_id=solution.id),
|
||
)
|
||
|
||
|
||
def index_resolved_ticket(ticket) -> list[KnowledgeChunk]:
|
||
if ticket.status not in {"resolved", "closed"} or ticket.archived_at is not None:
|
||
KnowledgeChunk.objects.filter(metadata__historical_ticket_id=ticket.id).delete()
|
||
return []
|
||
base, _ = KnowledgeBase.objects.get_or_create(
|
||
name="已解决工单案例库",
|
||
defaults={
|
||
"description": "从已解决或已关闭工单生成的检索案例",
|
||
"scope": "historical_tickets",
|
||
},
|
||
)
|
||
resolution = getattr(ticket, "resolution", None)
|
||
content = "\n".join(
|
||
[
|
||
ticket.title,
|
||
ticket.description,
|
||
ticket.problem_performance,
|
||
ticket.error_information,
|
||
resolution.root_cause if resolution else "",
|
||
resolution.solution_text if resolution else "",
|
||
]
|
||
)
|
||
return _create_chunks(
|
||
knowledge_base=base,
|
||
historical_ticket=ticket,
|
||
texts=split_text(content),
|
||
metadata={
|
||
"source_type": "historical_ticket",
|
||
"title": ticket.title,
|
||
"historical_ticket_id": ticket.id,
|
||
"ticket_version": ticket.updated_at.isoformat(),
|
||
"category_code": ticket.category.code if ticket.category_id else "",
|
||
"related_system": ticket.related_system,
|
||
"environment": ticket.environment,
|
||
},
|
||
replace_queryset=KnowledgeChunk.objects.filter(
|
||
metadata__historical_ticket_id=ticket.id
|
||
),
|
||
)
|
||
|
||
|
||
def sync_ticket_search_index(ticket_id: int) -> int:
|
||
from apps.tickets.models import Ticket
|
||
|
||
ticket = (
|
||
Ticket.objects.select_related("category", "resolution")
|
||
.filter(id=ticket_id)
|
||
.first()
|
||
)
|
||
if ticket is None:
|
||
KnowledgeChunk.objects.filter(metadata__historical_ticket_id=ticket_id).delete()
|
||
return 0
|
||
return len(index_resolved_ticket(ticket))
|
||
|
||
|
||
def sync_governed_asset_index(asset_type: str, object_id: int) -> int:
|
||
if asset_type == GovernedAssetType.KNOWLEDGE_ITEM:
|
||
item = KnowledgeItem.objects.filter(id=object_id).first()
|
||
if item is None:
|
||
return 0
|
||
if item.status != "published":
|
||
item.chunks.all().delete()
|
||
return 0
|
||
return len(rebuild_item_chunks(item))
|
||
if asset_type == GovernedAssetType.TROUBLESHOOTING_PLAN:
|
||
plan = TroubleshootingPlan.objects.filter(id=object_id).first()
|
||
if plan is None:
|
||
return 0
|
||
if plan.status != "published":
|
||
plan.chunks.all().delete()
|
||
return 0
|
||
return len(rebuild_troubleshooting_plan_chunks(plan))
|
||
if asset_type == GovernedAssetType.SOLUTION:
|
||
solution = Solution.objects.filter(id=object_id).first()
|
||
if solution is None or solution.status != "published":
|
||
KnowledgeChunk.objects.filter(metadata__solution_id=object_id).delete()
|
||
return 0
|
||
return len(index_published_solution(solution))
|
||
raise ValueError("Unknown governed asset type")
|
||
|
||
|
||
def _base_retrieval_queryset(
|
||
knowledge_base_id: int | None,
|
||
published_only: bool,
|
||
scope: dict[str, str] | None,
|
||
):
|
||
queryset = KnowledgeChunk.objects.select_related(
|
||
"knowledge_base",
|
||
"item",
|
||
"troubleshooting_plan",
|
||
"document",
|
||
"historical_ticket",
|
||
).filter(knowledge_base__enabled=True)
|
||
if knowledge_base_id:
|
||
queryset = queryset.filter(knowledge_base_id=knowledge_base_id)
|
||
if published_only:
|
||
published_solution_ids = Solution.objects.filter(status="published").values_list(
|
||
"id",
|
||
flat=True,
|
||
)
|
||
queryset = queryset.filter(
|
||
Q(item__status="published")
|
||
| Q(troubleshooting_plan__status="published")
|
||
| Q(metadata__solution_id__in=list(published_solution_ids))
|
||
| Q(
|
||
historical_ticket__status__in=["resolved", "closed"],
|
||
historical_ticket__archived_at__isnull=True,
|
||
)
|
||
)
|
||
for key, value in (scope or {}).items():
|
||
normalized = value.strip()
|
||
if not normalized:
|
||
continue
|
||
queryset = queryset.filter(
|
||
Q(**{f"metadata__{key}__isnull": True})
|
||
| Q(**{f"metadata__{key}": ""})
|
||
| Q(**{f"metadata__{key}__iexact": normalized})
|
||
)
|
||
return queryset
|
||
|
||
|
||
def _structured_scope_score(metadata: dict, scope: dict[str, str] | None) -> float:
|
||
expected = {
|
||
key: value.strip().casefold()
|
||
for key, value in (scope or {}).items()
|
||
if value.strip()
|
||
}
|
||
if not expected:
|
||
return 1.0
|
||
scores = []
|
||
for key, value in expected.items():
|
||
candidate = str(metadata.get(key, "") or "").strip().casefold()
|
||
scores.append(1.0 if candidate == value else 0.5 if not candidate else 0.0)
|
||
return sum(scores) / len(scores)
|
||
|
||
|
||
def _postgres_candidates(queryset, query: str, query_vector: list[float], candidate_limit: int):
|
||
text_query = SearchQuery(query, search_type="websearch", config="simple")
|
||
text_rows = list(
|
||
queryset.annotate(
|
||
text_rank=SearchRank(SearchVector("content", config="simple"), text_query),
|
||
)
|
||
.filter(text_rank__gt=0)
|
||
.order_by("-text_rank")[:candidate_limit]
|
||
)
|
||
vector_rows = list(
|
||
queryset.filter(embedding_vector__isnull=False)
|
||
.annotate(vector_distance=CosineDistance("embedding_vector", query_vector))
|
||
.order_by("vector_distance")[:candidate_limit]
|
||
)
|
||
candidates: dict[int, dict] = {}
|
||
for chunk in text_rows:
|
||
candidates.setdefault(chunk.id, {"chunk": chunk, "text_score": 0.0, "vector_score": 0.0})
|
||
candidates[chunk.id]["text_score"] = max(min(float(chunk.text_rank), 1.0), 0.0)
|
||
for chunk in vector_rows:
|
||
candidates.setdefault(chunk.id, {"chunk": chunk, "text_score": 0.0, "vector_score": 0.0})
|
||
candidates[chunk.id]["vector_score"] = max(
|
||
min(1.0 - float(chunk.vector_distance), 1.0),
|
||
0.0,
|
||
)
|
||
return list(candidates.values())
|
||
|
||
|
||
def _portable_candidates(queryset, query: str, query_vector: list[float], candidate_limit: int):
|
||
candidates = []
|
||
for chunk in queryset[:1000]:
|
||
stored_vector = (
|
||
list(chunk.embedding_vector)
|
||
if chunk.embedding_vector is not None
|
||
else list(chunk.embedding or [])
|
||
)
|
||
text_value = keyword_score(query, chunk.content)
|
||
vector_value = max(vector_score(query_vector, stored_vector), 0.0)
|
||
if text_value > 0 or vector_value > 0:
|
||
candidates.append(
|
||
{
|
||
"chunk": chunk,
|
||
"text_score": text_value,
|
||
"vector_score": vector_value,
|
||
}
|
||
)
|
||
candidates.sort(
|
||
key=lambda item: (item["text_score"], item["vector_score"]),
|
||
reverse=True,
|
||
)
|
||
return candidates[: candidate_limit * 2]
|
||
|
||
|
||
def retrieve_knowledge(
|
||
query: str,
|
||
knowledge_base_id: int | None = None,
|
||
limit: int = 5,
|
||
*,
|
||
published_only: bool = False,
|
||
scope: dict[str, str] | None = None,
|
||
) -> list[dict]:
|
||
"""执行全文、向量、范围加权和可选模型重排,并返回完整检索审计元数据。"""
|
||
|
||
batch = embed_contents([query])
|
||
query_vector = batch.vectors[0]
|
||
queryset = _base_retrieval_queryset(knowledge_base_id, published_only, scope)
|
||
candidate_limit = max(limit * 10, 50)
|
||
if connection.vendor == "postgresql":
|
||
candidates = _postgres_candidates(queryset, query, query_vector, candidate_limit)
|
||
channels = ["postgresql_full_text", "pgvector_cosine"]
|
||
else:
|
||
# SQLite 路径只服务测试和本地兼容,不作为生产召回质量的验收证据。
|
||
candidates = _portable_candidates(queryset, query, query_vector, candidate_limit)
|
||
channels = ["portable_text", "portable_vector"]
|
||
|
||
text_weight = settings.KNOWLEDGE_RETRIEVAL_TEXT_WEIGHT
|
||
vector_weight = settings.KNOWLEDGE_RETRIEVAL_VECTOR_WEIGHT
|
||
scope_weight = settings.KNOWLEDGE_RETRIEVAL_SCOPE_WEIGHT
|
||
for candidate in candidates:
|
||
candidate["scope_score"] = _structured_scope_score(
|
||
candidate["chunk"].metadata,
|
||
scope,
|
||
)
|
||
candidate["score"] = min(
|
||
candidate["text_score"] * text_weight
|
||
+ candidate["vector_score"] * vector_weight
|
||
+ candidate["scope_score"] * scope_weight,
|
||
1.0,
|
||
)
|
||
candidates.sort(
|
||
key=lambda item: (item["score"], item["text_score"], item["vector_score"]),
|
||
reverse=True,
|
||
)
|
||
rerank_limit = min(
|
||
max(settings.KNOWLEDGE_RERANK_CANDIDATE_LIMIT, limit),
|
||
len(candidates),
|
||
)
|
||
rerank_input = [
|
||
{
|
||
"id": item["chunk"].id,
|
||
"title": item["chunk"].metadata.get(
|
||
"title",
|
||
item["chunk"].knowledge_base.name,
|
||
),
|
||
"content": item["chunk"].content,
|
||
"base_score": item["score"],
|
||
}
|
||
for item in candidates[:rerank_limit]
|
||
]
|
||
rerank = rerank_candidates(query, rerank_input)
|
||
model_weight = (
|
||
min(max(settings.KNOWLEDGE_RERANK_MODEL_WEIGHT, 0.0), 1.0)
|
||
if rerank.status == "succeeded"
|
||
else 0.0
|
||
)
|
||
for candidate in candidates:
|
||
candidate_id = candidate["chunk"].id
|
||
model_score = rerank.scores.get(candidate_id, candidate["score"])
|
||
candidate["rerank_score"] = model_score
|
||
candidate["score"] = (
|
||
candidate["score"] * (1.0 - model_weight) + model_score * model_weight
|
||
)
|
||
candidates.sort(
|
||
key=lambda item: (item["score"], item["text_score"], item["vector_score"]),
|
||
reverse=True,
|
||
)
|
||
return [
|
||
{
|
||
"id": item["chunk"].id,
|
||
"knowledge_base_id": item["chunk"].knowledge_base_id,
|
||
"title": item["chunk"].metadata.get(
|
||
"title",
|
||
item["chunk"].knowledge_base.name,
|
||
),
|
||
"source_type": item["chunk"].metadata.get("source_type", "chunk"),
|
||
"content": item["chunk"].content,
|
||
"score": round(item["score"], 4),
|
||
"metadata": {
|
||
**item["chunk"].metadata,
|
||
"retrieval": {
|
||
"channels": channels,
|
||
"text_score": round(item["text_score"], 4),
|
||
"vector_score": round(item["vector_score"], 4),
|
||
"scope_score": round(item["scope_score"], 4),
|
||
"weights": {
|
||
"text": text_weight,
|
||
"vector": vector_weight,
|
||
"scope": scope_weight,
|
||
},
|
||
"query_embedding_model": batch.model_name,
|
||
"query_embedding_deployment_id": batch.deployment_id,
|
||
"query_embedding_fallback": batch.fallback,
|
||
"chunk_embedding_model": item["chunk"].embedding_model,
|
||
"rerank": {
|
||
"status": rerank.status,
|
||
"model_name": rerank.model_name,
|
||
"deployment_id": rerank.deployment_id,
|
||
"model_score": round(item["rerank_score"], 4),
|
||
"model_weight": model_weight,
|
||
"candidate_limit": rerank_limit,
|
||
},
|
||
},
|
||
},
|
||
}
|
||
for item in candidates[:limit]
|
||
]
|