968 lines
36 KiB
Python
968 lines
36 KiB
Python
import asyncio
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from config import WIKI_CONFIG
|
||
from wiki_engine import db
|
||
from wiki_engine.chunk_enhancer import enhance_slices
|
||
from wiki_engine.htknow_client import HTKnowClient
|
||
from wiki_engine.wiki_builder import build_pages_for_file
|
||
|
||
|
||
_FILE_SLICE_TEXT_CACHE: Dict[str, Dict[str, str]] = {}
|
||
|
||
|
||
async def sync_file(
|
||
file_id: str,
|
||
kb_id: Optional[str] = None,
|
||
file_info: Optional[Dict[str, Any]] = None,
|
||
use_llm: bool = False,
|
||
job_id: Optional[str] = None,
|
||
force: bool = False,
|
||
) -> Dict[str, Any]:
|
||
client = HTKnowClient()
|
||
file_info = dict(file_info or {})
|
||
filename = str(file_info.get("filename") or file_info.get("name") or "")
|
||
if job_id:
|
||
await db.update_job(
|
||
job_id,
|
||
stage="fetch_file",
|
||
current_file_id=str(file_id),
|
||
current_filename=filename,
|
||
)
|
||
if not file_info or not (file_info.get("filename") or file_info.get("name")):
|
||
try:
|
||
file_info.update(await client.get_file(str(file_id)))
|
||
except Exception as exc:
|
||
print(f"[wiki_engine] get file detail failed for {file_id}: {exc}")
|
||
file_info.setdefault("id", str(file_id))
|
||
file_info.setdefault("file_id", str(file_id))
|
||
if kb_id:
|
||
file_info.setdefault("kb_id", str(kb_id))
|
||
filename = str(file_info.get("filename") or file_info.get("name") or "")
|
||
if job_id:
|
||
await db.update_job(job_id, stage="fetch_slices", current_filename=filename)
|
||
|
||
slices = await client.get_file_slices(str(file_id), limit=WIKI_CONFIG.get("max_slices_per_file"))
|
||
if job_id:
|
||
await db.update_job(job_id, stage="enhance_slices")
|
||
refs = enhance_slices(
|
||
slices,
|
||
file_info=file_info,
|
||
max_slices=int(WIKI_CONFIG.get("max_slices_per_file") or 300),
|
||
)
|
||
if not refs:
|
||
raise ValueError(f"no slices returned for file {file_id}; keep existing wiki data unchanged")
|
||
|
||
content_signature = _file_content_signature(refs)
|
||
existing_state = await db.get_file_state(str(file_id))
|
||
existing_page_count = await db.count_file_pages(str(file_id)) if existing_state else 0
|
||
if not force and _is_unchanged_file(existing_state, refs, content_signature) and existing_page_count > 0:
|
||
if job_id:
|
||
await db.update_job(job_id, stage="skipped_unchanged")
|
||
return {
|
||
"success": True,
|
||
"skipped": True,
|
||
"reason": "unchanged",
|
||
"file_id": str(file_id),
|
||
"kb_id": str(kb_id or file_info.get("kb_id") or ""),
|
||
"slice_count": len(refs),
|
||
"embedding_count": 0,
|
||
"embedding_error": "",
|
||
"page_count": existing_page_count,
|
||
"page_ids": [],
|
||
}
|
||
|
||
source_hash = str(file_info.get("hash") or "")
|
||
metadata = file_info.get("metadata") if isinstance(file_info.get("metadata"), dict) else {}
|
||
metadata = dict(metadata)
|
||
metadata["content_signature"] = content_signature
|
||
metadata["source_hash"] = source_hash
|
||
metadata["wiki_pipeline_version"] = "wiki-sync-v3"
|
||
file_info["metadata"] = metadata
|
||
file_info["content_signature"] = content_signature
|
||
if not source_hash:
|
||
file_info["hash"] = content_signature
|
||
|
||
await db.upsert_file_state(file_info, slice_count=len(refs))
|
||
if job_id:
|
||
await db.update_job(job_id, stage="write_slices")
|
||
await db.replace_file_slice_refs(str(file_id), refs)
|
||
if job_id:
|
||
await db.update_job(job_id, stage="build_pages")
|
||
async def progress(stage: str) -> None:
|
||
if job_id and WIKI_CONFIG.get("stage_progress_enabled", True):
|
||
await db.update_job(job_id, stage=stage)
|
||
|
||
pages = await build_pages_for_file(file_info, refs, use_llm=use_llm, progress=progress if job_id else None)
|
||
if job_id:
|
||
await db.update_job(job_id, stage="write_pages")
|
||
page_ids = []
|
||
for page in pages:
|
||
page_id = await _upsert_page_with_slug_lock(page, owner=job_id or f"sync:{file_id}")
|
||
page_ids.append(page_id)
|
||
await db.replace_page_slices(page_id, page.get("links", []))
|
||
if job_id:
|
||
await db.update_job(job_id, stage="file_done")
|
||
|
||
return {
|
||
"success": True,
|
||
"skipped": False,
|
||
"file_id": str(file_id),
|
||
"kb_id": str(kb_id or file_info.get("kb_id") or ""),
|
||
"slice_count": len(refs),
|
||
"embedding_count": 0,
|
||
"embedding_error": "",
|
||
"page_count": len(page_ids),
|
||
"page_ids": page_ids,
|
||
}
|
||
|
||
|
||
async def sync_kb(
|
||
kb_id: str,
|
||
limit: Optional[int] = None,
|
||
use_llm: bool = False,
|
||
job_id: Optional[str] = None,
|
||
force: bool = False,
|
||
) -> Dict[str, Any]:
|
||
limit = int(limit or WIKI_CONFIG.get("default_sync_limit") or 20)
|
||
client = HTKnowClient()
|
||
files = await client.list_files(kb_id=kb_id, limit=limit)
|
||
if job_id:
|
||
await db.update_job(job_id, total_files=len(files))
|
||
|
||
processed = 0
|
||
total_slices = 0
|
||
skipped = 0
|
||
page_count = 0
|
||
errors: List[Dict[str, str]] = []
|
||
for file_info in files:
|
||
file_id = str(file_info.get("id") or file_info.get("file_id") or "")
|
||
if not file_id:
|
||
continue
|
||
try:
|
||
result = await sync_file(file_id, kb_id=kb_id, file_info=file_info, use_llm=use_llm, job_id=job_id, force=force)
|
||
processed += 1
|
||
total_slices += int(result.get("slice_count") or 0)
|
||
skipped += 1 if result.get("skipped") else 0
|
||
if not result.get("skipped"):
|
||
page_count += int(result.get("page_count") or 0)
|
||
except Exception as exc:
|
||
errors.append({"file_id": file_id, "error": str(exc)})
|
||
if job_id:
|
||
await db.update_job(
|
||
job_id,
|
||
processed_files=processed,
|
||
total_slices=total_slices,
|
||
skipped_files=skipped,
|
||
page_count=page_count,
|
||
)
|
||
await asyncio.sleep(0)
|
||
|
||
if job_id and WIKI_CONFIG.get("finalize_enabled", True):
|
||
await enqueue_finalize_job(kb_id, job_id)
|
||
|
||
return {
|
||
"success": not errors,
|
||
"kb_id": str(kb_id),
|
||
"file_count": len(files),
|
||
"processed_files": processed,
|
||
"total_slices": total_slices,
|
||
"skipped_files": skipped,
|
||
"page_count": page_count,
|
||
"errors": errors[:20],
|
||
}
|
||
|
||
|
||
async def start_sync_job(
|
||
target_type: str,
|
||
target_id: str,
|
||
kb_id: Optional[str] = None,
|
||
limit: Optional[int] = None,
|
||
use_llm: bool = False,
|
||
force: bool = False,
|
||
) -> str:
|
||
job_id = str(uuid.uuid4())
|
||
await db.create_job(job_id, target_type=target_type, target_id=target_id, kb_id=kb_id or "")
|
||
if WIKI_CONFIG.get("queue_enabled", True):
|
||
if target_type == "file":
|
||
await db.set_job_totals(job_id, 1)
|
||
await db.enqueue_wiki_file_op(
|
||
job_id=job_id,
|
||
kb_id=kb_id or "",
|
||
file_id=target_id,
|
||
file_info={},
|
||
use_llm=use_llm,
|
||
force=force,
|
||
)
|
||
return job_id
|
||
if target_type == "kb":
|
||
client = HTKnowClient()
|
||
files = await client.list_files(kb_id=target_id, limit=limit or WIKI_CONFIG.get("default_sync_limit"))
|
||
valid_files = [
|
||
(str(file_info.get("id") or file_info.get("file_id") or ""), file_info)
|
||
for file_info in files
|
||
if str(file_info.get("id") or file_info.get("file_id") or "")
|
||
]
|
||
await db.set_job_totals(job_id, len(valid_files))
|
||
if not valid_files:
|
||
await db.update_job(job_id, status="success")
|
||
return job_id
|
||
for file_id, file_info in valid_files:
|
||
await db.enqueue_wiki_file_op(
|
||
job_id=job_id,
|
||
kb_id=target_id,
|
||
file_id=file_id,
|
||
file_info=file_info,
|
||
use_llm=use_llm,
|
||
force=force,
|
||
)
|
||
return job_id
|
||
|
||
async def _runner() -> None:
|
||
try:
|
||
if target_type == "file":
|
||
await sync_file(target_id, kb_id=kb_id, use_llm=use_llm, job_id=job_id, force=force)
|
||
elif target_type == "kb":
|
||
await sync_kb(target_id, limit=limit, use_llm=use_llm, job_id=job_id, force=force)
|
||
else:
|
||
raise ValueError(f"unsupported target_type: {target_type}")
|
||
await db.update_job(job_id, status="success")
|
||
except Exception as exc:
|
||
await db.update_job(job_id, status="failed", error=str(exc))
|
||
|
||
asyncio.create_task(_runner())
|
||
return job_id
|
||
|
||
|
||
async def get_queue_stats() -> Dict[str, Any]:
|
||
return await db.get_wiki_queue_stats()
|
||
|
||
|
||
async def get_sync_job(job_id: str) -> Optional[Dict[str, Any]]:
|
||
return await db.get_job(job_id)
|
||
|
||
|
||
async def enqueue_finalize_job(kb_id: str, job_id: str = "") -> Dict[str, Any]:
|
||
if not WIKI_CONFIG.get("finalize_enabled", True):
|
||
return {"success": True, "enqueued": False, "reason": "disabled"}
|
||
op_id = await db.enqueue_wiki_finalize_op(
|
||
job_id=job_id,
|
||
kb_id=kb_id,
|
||
delay_seconds=int(WIKI_CONFIG.get("finalize_debounce_seconds") or 20),
|
||
)
|
||
return {"success": True, "enqueued": bool(op_id), "op_id": op_id}
|
||
|
||
|
||
async def finalize_kb(kb_id: str, job_id: str = "") -> Dict[str, Any]:
|
||
if job_id:
|
||
await db.update_job(job_id, stage="finalize")
|
||
page_count = await db.count_pages_for_kb(kb_id)
|
||
if job_id:
|
||
await db.update_job(job_id, stage="completed", page_count=page_count, current_file_id="", current_filename="")
|
||
return {"success": True, "kb_id": str(kb_id or ""), "page_count": page_count}
|
||
|
||
|
||
async def search_wiki(
|
||
query: str,
|
||
kb_id: Optional[str] = None,
|
||
limit: int = 8,
|
||
include_slices: bool = True,
|
||
) -> Dict[str, Any]:
|
||
page_limit = max(int(WIKI_CONFIG.get("page_search_top_k") or limit or 8), limit)
|
||
pages = await db.search_pages(query=query, kb_id=kb_id, limit=page_limit)
|
||
|
||
slice_map = {}
|
||
if include_slices and pages:
|
||
slice_map = await db.get_page_slices(
|
||
[int(p["id"]) for p in pages],
|
||
max_slices_per_page=int(WIKI_CONFIG.get("evidence_slices_per_page") or 6),
|
||
)
|
||
items = []
|
||
for rank, page in enumerate(pages, start=1):
|
||
page_id = int(page["id"])
|
||
slices = slice_map.get(page_id, [])
|
||
related_slice_ids = [str(s.get("slice_id")) for s in slices if s.get("slice_id")]
|
||
related_file_ids = list({
|
||
str(s.get("file_id"))
|
||
for s in slices
|
||
if s.get("file_id")
|
||
} | set([str(x) for x in (page.get("related_file_ids") or []) if x]))
|
||
items.append({
|
||
"source_type": "wiki",
|
||
"source_id": str(page_id),
|
||
"kb_id": str(page.get("kb_id") or ""),
|
||
"rank_sources": {"page": rank},
|
||
"slug": page.get("slug", ""),
|
||
"title": page.get("title", ""),
|
||
"page_type": page.get("page_type", ""),
|
||
"summary": page.get("summary", ""),
|
||
"content": page.get("content", ""),
|
||
"keywords": page.get("keywords", []),
|
||
"aliases": page.get("aliases", []),
|
||
"score": page.get("score", 0.0),
|
||
"related_file_ids": related_file_ids,
|
||
"related_slice_ids": related_slice_ids,
|
||
"slices": slices,
|
||
})
|
||
|
||
items = _fuse_wiki_items(items, limit=max(limit, 1))
|
||
await db.log_query(query, kb_id, len(items))
|
||
return {
|
||
"success": True,
|
||
"query": query,
|
||
"kb_id": kb_id or "",
|
||
"items": items,
|
||
"debug": {"page_hits": len(pages), "vector_hits": 0, "text_hits": 0},
|
||
}
|
||
|
||
|
||
async def evidence_search(query: str, kb_id: Optional[str] = None, limit: int = 5) -> Dict[str, Any]:
|
||
result = await search_wiki(
|
||
query=query,
|
||
kb_id=kb_id,
|
||
limit=max(limit * 2, limit),
|
||
include_slices=True,
|
||
)
|
||
evidence_items = []
|
||
for item in result.get("items", []):
|
||
slices = await _hydrate_wiki_slices(item.get("slices", []))
|
||
slice_texts = []
|
||
for s in slices[:6]:
|
||
text = str(s.get("content") or s.get("brief") or "").strip()
|
||
if not text:
|
||
continue
|
||
filename = str(s.get("filename") or "").strip()
|
||
title_path = " / ".join(str(x) for x in (s.get("title_path") or []) if x)
|
||
prefix_parts = [x for x in (filename, title_path) if x]
|
||
prefix = f"{' - '.join(prefix_parts)}\n" if prefix_parts else ""
|
||
slice_texts.append(
|
||
f"{prefix}{text}\n"
|
||
f"(file_id={s.get('file_id') or ''}, slice_id={s.get('slice_id') or ''})"
|
||
)
|
||
deep_text = "\n\n".join(slice_texts)
|
||
page_text = "\n".join(x for x in [
|
||
str(item.get("summary") or "").strip(),
|
||
str(item.get("content") or "").strip(),
|
||
] if x).strip()
|
||
text = "\n\n".join(x for x in [page_text, deep_text] if x).strip()
|
||
if not text:
|
||
continue
|
||
evidence_items.append({
|
||
"source_type": "wiki",
|
||
"source_id": item.get("source_id"),
|
||
"title": item.get("title"),
|
||
"text": text,
|
||
"score": item.get("score", 0.0),
|
||
"kb_id": item.get("kb_id", ""),
|
||
"metadata": {
|
||
"kb_id": item.get("kb_id", ""),
|
||
"page_type": item.get("page_type"),
|
||
"keywords": item.get("keywords", []),
|
||
"related_file_ids": item.get("related_file_ids", []),
|
||
"related_slice_ids": item.get("related_slice_ids", []),
|
||
"slices": slices,
|
||
"rank_sources": item.get("rank_sources", {}),
|
||
},
|
||
})
|
||
# Keep the page-search order here. The baike workflow does a later cross-source
|
||
# rerank, and preserving the Wiki top hit makes the displayed Wiki evidence
|
||
# match `/api/wiki/search`.
|
||
evidence_items.sort(key=lambda x: float(x.get("score") or 0.0), reverse=True)
|
||
return {"success": True, "query": query, "items": evidence_items[:limit], "debug": result.get("debug", {})}
|
||
|
||
|
||
def _fuse_wiki_items(items: List[Dict[str, Any]], limit: int) -> List[Dict[str, Any]]:
|
||
rrf_k = max(int(WIKI_CONFIG.get("rrf_k") or 60), 1)
|
||
by_key: Dict[str, Dict[str, Any]] = {}
|
||
for item in items:
|
||
key = str(item.get("source_id") or item.get("title") or "")
|
||
if not key:
|
||
continue
|
||
fused_score = float(item.get("score") or 0.0)
|
||
for rank in (item.get("rank_sources") or {}).values():
|
||
try:
|
||
fused_score += 1.0 / (rrf_k + int(rank))
|
||
except Exception:
|
||
continue
|
||
item = dict(item)
|
||
item["score"] = fused_score
|
||
existing = by_key.get(key)
|
||
if existing:
|
||
merged_sources = dict(existing.get("rank_sources") or {})
|
||
for source, rank in (item.get("rank_sources") or {}).items():
|
||
if source not in merged_sources or int(rank) < int(merged_sources.get(source) or 999999):
|
||
merged_sources[source] = rank
|
||
item["rank_sources"] = merged_sources
|
||
fused_score = max(fused_score, float(existing.get("score") or 0.0))
|
||
for rank in merged_sources.values():
|
||
try:
|
||
fused_score += 1.0 / (rrf_k + int(rank))
|
||
except Exception:
|
||
continue
|
||
item["score"] = fused_score
|
||
if not existing or fused_score > float(existing.get("score") or 0.0):
|
||
by_key[key] = item
|
||
return sorted(by_key.values(), key=lambda x: float(x.get("score") or 0.0), reverse=True)[:limit]
|
||
|
||
|
||
async def _rerank_wiki_evidence(query: str, items: List[Dict[str, Any]], top_k: int) -> List[Dict[str, Any]]:
|
||
if not query.strip() or not items:
|
||
return items[:top_k]
|
||
try:
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
|
||
selected = items[:24]
|
||
documents = [
|
||
"\n".join([
|
||
str(item.get("title") or ""),
|
||
str(item.get("text") or ""),
|
||
])[:4000]
|
||
for item in selected
|
||
]
|
||
rows = await OpenaiAPI.rerank(query=query, documents=documents, top_k=top_k)
|
||
if not rows:
|
||
return items[:top_k]
|
||
ranked: List[Dict[str, Any]] = []
|
||
used = set()
|
||
for row in rows:
|
||
idx = int(row.get("index"))
|
||
if idx < 0 or idx >= len(selected) or idx in used:
|
||
continue
|
||
item = dict(selected[idx])
|
||
item["rerank_score"] = float(row.get("relevance_score") or 0.0)
|
||
item["score"] = max(float(item.get("score") or 0.0), item["rerank_score"])
|
||
ranked.append(item)
|
||
used.add(idx)
|
||
for idx, item in enumerate(selected):
|
||
if idx not in used and len(ranked) < top_k:
|
||
ranked.append(item)
|
||
return ranked[:top_k]
|
||
except Exception as exc:
|
||
print(f"[wiki_engine] wiki evidence rerank failed: {exc}")
|
||
return items[:top_k]
|
||
|
||
|
||
async def _hydrate_wiki_slices(slices: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""Fill missing slice content from 11000 for old Wiki rows."""
|
||
out = [dict(s or {}) for s in slices]
|
||
missing_by_file: Dict[str, set[str]] = {}
|
||
for s in out:
|
||
if str(s.get("content") or "").strip():
|
||
continue
|
||
file_id = str(s.get("file_id") or "").strip()
|
||
slice_id = str(s.get("slice_id") or "").strip()
|
||
if file_id and slice_id:
|
||
missing_by_file.setdefault(file_id, set()).add(slice_id)
|
||
if not missing_by_file:
|
||
return out
|
||
|
||
client = HTKnowClient()
|
||
fetched: Dict[str, str] = {}
|
||
for file_id, wanted_ids in missing_by_file.items():
|
||
cached = _FILE_SLICE_TEXT_CACHE.get(file_id)
|
||
if cached is not None:
|
||
fetched.update({sid: cached.get(sid, "") for sid in wanted_ids})
|
||
continue
|
||
try:
|
||
raw_slices = await client.get_file_slices(file_id, limit=WIKI_CONFIG.get("max_slices_per_file"))
|
||
except Exception as exc:
|
||
print(f"[wiki_engine] hydrate slices failed for file {file_id}: {exc}")
|
||
continue
|
||
file_texts: Dict[str, str] = {}
|
||
for idx, raw in enumerate(raw_slices):
|
||
sid = str(raw.get("id") or raw.get("slice_id") or raw.get("uuid") or f"{file_id}:{idx}")
|
||
text = _raw_slice_text(raw)
|
||
if text:
|
||
file_texts[sid] = text
|
||
if sid in wanted_ids:
|
||
fetched[sid] = text
|
||
_FILE_SLICE_TEXT_CACHE[file_id] = file_texts
|
||
|
||
for s in out:
|
||
sid = str(s.get("slice_id") or "")
|
||
if sid in fetched and fetched[sid]:
|
||
s["content"] = fetched[sid]
|
||
return out
|
||
|
||
|
||
def _file_content_signature(refs: List[Dict[str, Any]]) -> str:
|
||
rows = []
|
||
for ref in sorted(refs, key=lambda x: (int(x.get("ordinal") or 0), str(x.get("slice_id") or ""))):
|
||
rows.append({
|
||
"slice_id": str(ref.get("slice_id") or ""),
|
||
"content_hash": str(ref.get("content_hash") or ""),
|
||
"title_path": ref.get("title_path") or [],
|
||
"keywords": ref.get("keywords") or [],
|
||
})
|
||
raw = json.dumps(rows, ensure_ascii=False, sort_keys=True)
|
||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def _is_unchanged_file(
|
||
state: Optional[Dict[str, Any]],
|
||
refs: List[Dict[str, Any]],
|
||
content_signature: str,
|
||
) -> bool:
|
||
if not state:
|
||
return False
|
||
if int(state.get("slice_count") or 0) != len(refs):
|
||
return False
|
||
metadata = state.get("metadata") if isinstance(state.get("metadata"), dict) else {}
|
||
stored_signature = str(metadata.get("content_signature") or "")
|
||
if stored_signature and stored_signature == content_signature:
|
||
return True
|
||
return str(state.get("file_hash") or "") == content_signature
|
||
|
||
|
||
async def _upsert_page_with_slug_lock(page: Dict[str, Any], owner: str) -> int:
|
||
kb_id = str(page.get("kb_id") or "")
|
||
slug = str(page.get("slug") or "")
|
||
if not kb_id or not slug:
|
||
return await db.upsert_wiki_page(page)
|
||
|
||
ttl = int(WIKI_CONFIG.get("slug_lock_ttl_seconds") or 600)
|
||
acquired = False
|
||
for _ in range(600):
|
||
acquired = await db.acquire_wiki_page_lock(kb_id, slug, owner=owner, ttl_seconds=ttl)
|
||
if acquired:
|
||
break
|
||
await asyncio.sleep(0.1)
|
||
if not acquired:
|
||
raise TimeoutError(f"wiki page lock timeout: {kb_id}/{slug}")
|
||
|
||
try:
|
||
return await db.upsert_wiki_page(page)
|
||
finally:
|
||
await db.release_wiki_page_lock(kb_id, slug, owner=owner)
|
||
|
||
|
||
def _raw_slice_text(raw: Dict[str, Any]) -> str:
|
||
for key in ("content", "text", "page_content", "slice_content", "markdown"):
|
||
value = raw.get(key)
|
||
if isinstance(value, str) and value.strip():
|
||
return value
|
||
return ""
|
||
|
||
|
||
async def get_page(page_id: int) -> Optional[Dict[str, Any]]:
|
||
return await db.get_page(page_id)
|
||
|
||
|
||
async def get_stats() -> Dict[str, Any]:
|
||
return await db.get_stats()
|
||
|
||
|
||
async def get_kb_tree() -> Dict[str, Any]:
|
||
client = HTKnowClient()
|
||
tree = await client.get_kb_tree()
|
||
return {"success": True, "tree": tree}
|
||
|
||
|
||
async def resolve_kb_reference(reference: str) -> Dict[str, Any]:
|
||
"""Resolve a user-facing knowledge-base name or admin id to a kb_id."""
|
||
ref = str(reference or "").strip()
|
||
if not ref:
|
||
return {"success": False, "error": "knowledge base name is empty", "matches": []}
|
||
|
||
client = HTKnowClient()
|
||
tree = await client.get_kb_tree()
|
||
nodes = _flatten_kb_tree(tree)
|
||
ref_norm = _normalize_kb_text(ref)
|
||
|
||
scored: List[Dict[str, Any]] = []
|
||
for node in nodes:
|
||
node_id = str(node.get("id") or "")
|
||
names = _node_names(node)
|
||
best_score = 0
|
||
best_name = names[0] if names else ""
|
||
if node_id and node_id == ref:
|
||
best_score = 10000
|
||
for name in names:
|
||
name_norm = _normalize_kb_text(name)
|
||
if not name_norm:
|
||
continue
|
||
score = 0
|
||
if name_norm == ref_norm:
|
||
score = 9000 + len(name_norm)
|
||
elif name_norm in ref_norm:
|
||
score = 7000 + len(name_norm)
|
||
elif ref_norm in name_norm:
|
||
score = 5000 + len(ref_norm)
|
||
if score > best_score:
|
||
best_score = score
|
||
best_name = name
|
||
if best_score:
|
||
scored.append({
|
||
"kb_id": node_id,
|
||
"name": best_name,
|
||
"path": node.get("_path", best_name),
|
||
"score": best_score,
|
||
})
|
||
|
||
scored.sort(key=lambda item: (-int(item.get("score") or 0), -len(str(item.get("path") or ""))))
|
||
if not scored:
|
||
return {
|
||
"success": False,
|
||
"error": f"未找到名为“{ref}”的知识库",
|
||
"matches": [],
|
||
"candidates": [
|
||
{"kb_id": str(node.get("id") or ""), "name": (_node_names(node) or [""])[0], "path": node.get("_path", "")}
|
||
for node in nodes[:20]
|
||
],
|
||
}
|
||
|
||
top_score = int(scored[0].get("score") or 0)
|
||
top_matches = [item for item in scored if int(item.get("score") or 0) == top_score]
|
||
if len({item.get("kb_id") for item in top_matches}) > 1 and top_score < 10000:
|
||
return {
|
||
"success": False,
|
||
"error": f"知识库名称“{ref}”匹配到多个结果，请说得更完整一点",
|
||
"matches": top_matches[:10],
|
||
}
|
||
|
||
best = scored[0]
|
||
return {
|
||
"success": True,
|
||
"kb_id": str(best.get("kb_id") or ""),
|
||
"kb_name": str(best.get("name") or ""),
|
||
"kb_path": str(best.get("path") or ""),
|
||
"matches": scored[:10],
|
||
}
|
||
|
||
|
||
async def diff_kb(
|
||
kb_id: str,
|
||
include_children: bool = False,
|
||
limit: int = 200,
|
||
parsed_only: bool = True,
|
||
) -> Dict[str, Any]:
|
||
client = HTKnowClient()
|
||
kb_ids = await _resolve_kb_ids(client, kb_id, include_children=include_children)
|
||
ht_files: Dict[str, Dict[str, Any]] = {}
|
||
skipped_files: List[Dict[str, Any]] = []
|
||
|
||
for kid in kb_ids:
|
||
files = await client.list_files(kb_id=kid, limit=max(int(limit or 200), 1))
|
||
for file_info in files:
|
||
fid = str(file_info.get("id") or file_info.get("file_id") or "")
|
||
if not fid:
|
||
continue
|
||
if parsed_only and str(file_info.get("status")) != "1":
|
||
skipped_files.append(_compact_file(file_info))
|
||
continue
|
||
file_info["file_id"] = fid
|
||
file_info["kb_id"] = str(file_info.get("kb_id") or kid)
|
||
ht_files[fid] = file_info
|
||
|
||
wiki_states = await db.list_file_states(kb_ids)
|
||
wiki_by_file = {str(item.get("file_id")): item for item in wiki_states}
|
||
|
||
missing_in_wiki = []
|
||
stale_in_wiki = []
|
||
for fid, file_info in ht_files.items():
|
||
state = wiki_by_file.get(fid)
|
||
if not state:
|
||
missing_in_wiki.append(_compact_file(file_info))
|
||
elif _is_stale(file_info, state):
|
||
stale_in_wiki.append({
|
||
"file_id": fid,
|
||
"kb_id": str(file_info.get("kb_id") or ""),
|
||
"filename": str(file_info.get("filename") or ""),
|
||
"htknow_hash": str(file_info.get("hash") or ""),
|
||
"wiki_hash": str(state.get("file_hash") or ""),
|
||
"htknow_updated_at": str(file_info.get("updated_at") or ""),
|
||
"wiki_source_updated_at": str(state.get("source_updated_at") or ""),
|
||
"last_synced_at": state.get("last_synced_at") or "",
|
||
})
|
||
|
||
orphan_in_wiki = []
|
||
for fid, state in wiki_by_file.items():
|
||
if fid not in ht_files:
|
||
orphan_in_wiki.append({
|
||
"file_id": fid,
|
||
"kb_id": str(state.get("kb_id") or ""),
|
||
"filename": str(state.get("filename") or ""),
|
||
"last_synced_at": state.get("last_synced_at") or "",
|
||
})
|
||
|
||
return {
|
||
"success": True,
|
||
"kb_id": str(kb_id),
|
||
"include_children": include_children,
|
||
"kb_ids": kb_ids,
|
||
"summary": {
|
||
"htknow_files": len(ht_files),
|
||
"wiki_files": len(wiki_by_file),
|
||
"missing_in_wiki": len(missing_in_wiki),
|
||
"stale_in_wiki": len(stale_in_wiki),
|
||
"orphan_in_wiki": len(orphan_in_wiki),
|
||
"skipped_unparsed": len(skipped_files),
|
||
},
|
||
"missing_in_wiki": missing_in_wiki,
|
||
"stale_in_wiki": stale_in_wiki,
|
||
"orphan_in_wiki": orphan_in_wiki,
|
||
"skipped_unparsed": skipped_files[:50],
|
||
}
|
||
|
||
|
||
async def reconcile_kb(
|
||
kb_id: str,
|
||
include_children: bool = False,
|
||
limit: int = 200,
|
||
sync_missing: bool = True,
|
||
rebuild_stale: bool = True,
|
||
delete_orphan: bool = False,
|
||
use_llm: bool = False,
|
||
dry_run: bool = True,
|
||
) -> Dict[str, Any]:
|
||
diff = await diff_kb(kb_id=kb_id, include_children=include_children, limit=limit, parsed_only=True)
|
||
actions: List[Dict[str, Any]] = []
|
||
if sync_missing:
|
||
actions.extend({"action": "sync_file", "file_id": x["file_id"], "kb_id": x.get("kb_id", "")} for x in diff["missing_in_wiki"])
|
||
if rebuild_stale:
|
||
actions.extend({"action": "rebuild_file", "file_id": x["file_id"], "kb_id": x.get("kb_id", "")} for x in diff["stale_in_wiki"])
|
||
if delete_orphan:
|
||
actions.extend({"action": "delete_wiki_file", "file_id": x["file_id"], "kb_id": x.get("kb_id", "")} for x in diff["orphan_in_wiki"])
|
||
|
||
executed: List[Dict[str, Any]] = []
|
||
if not dry_run:
|
||
client = HTKnowClient()
|
||
for action in actions:
|
||
try:
|
||
if action["action"] in {"sync_file", "rebuild_file"}:
|
||
file_info = await client.get_file(action["file_id"])
|
||
if action.get("kb_id"):
|
||
file_info["kb_id"] = action["kb_id"]
|
||
result = await sync_file(
|
||
action["file_id"],
|
||
kb_id=action.get("kb_id"),
|
||
file_info=file_info,
|
||
use_llm=use_llm,
|
||
force=action["action"] == "rebuild_file",
|
||
)
|
||
elif action["action"] == "delete_wiki_file":
|
||
result = await delete_by_file(action["file_id"])
|
||
else:
|
||
result = {"success": False, "error": "unknown action"}
|
||
executed.append({"action": action, "result": result})
|
||
except Exception as exc:
|
||
executed.append({"action": action, "result": {"success": False, "error": str(exc)}})
|
||
|
||
return {
|
||
"success": True,
|
||
"dry_run": dry_run,
|
||
"kb_id": str(kb_id),
|
||
"include_children": include_children,
|
||
"summary": diff.get("summary", {}),
|
||
"planned_actions": actions,
|
||
"executed": executed,
|
||
}
|
||
|
||
|
||
async def create_admin_plan(
|
||
chat_id: str,
|
||
intent: str,
|
||
scope: str,
|
||
actions: List[Dict[str, Any]],
|
||
kb_id: str = "",
|
||
kb_name: str = "",
|
||
summary: Optional[Dict[str, Any]] = None,
|
||
metadata: Optional[Dict[str, Any]] = None,
|
||
) -> Dict[str, Any]:
|
||
return await db.create_admin_plan(
|
||
plan_id=str(uuid.uuid4()),
|
||
chat_id=chat_id,
|
||
intent=intent,
|
||
scope=scope,
|
||
kb_id=kb_id,
|
||
kb_name=kb_name,
|
||
summary=summary or {},
|
||
actions=actions,
|
||
metadata=metadata or {},
|
||
)
|
||
|
||
|
||
async def get_pending_admin_plan(chat_id: str) -> Optional[Dict[str, Any]]:
|
||
return await db.get_latest_pending_admin_plan(chat_id)
|
||
|
||
|
||
async def cancel_admin_plan(plan_id: str) -> Optional[Dict[str, Any]]:
|
||
return await db.update_admin_plan_status(plan_id, "cancelled")
|
||
|
||
|
||
async def execute_admin_plan(
|
||
plan: Dict[str, Any],
|
||
selected_indexes: Optional[List[int]] = None,
|
||
use_llm: bool = False,
|
||
) -> Dict[str, Any]:
|
||
actions = plan.get("actions") or []
|
||
if selected_indexes:
|
||
selected = {idx for idx in selected_indexes if idx >= 1}
|
||
actions = [action for idx, action in enumerate(actions, start=1) if idx in selected]
|
||
executed: List[Dict[str, Any]] = []
|
||
client = HTKnowClient()
|
||
for action in actions:
|
||
try:
|
||
action_type = action.get("action")
|
||
if action_type in {"sync_file", "rebuild_file"}:
|
||
job_id = await start_sync_job(
|
||
target_type="file",
|
||
target_id=str(action.get("file_id") or ""),
|
||
kb_id=str(action.get("kb_id") or ""),
|
||
use_llm=bool(action.get("use_llm") or use_llm),
|
||
force=action_type == "rebuild_file",
|
||
)
|
||
job = await db.get_job(job_id)
|
||
result = {"success": True, "status": "running", "job_id": job_id, "job": job}
|
||
elif action_type == "delete_wiki_file":
|
||
result = await delete_by_file(str(action.get("file_id") or ""))
|
||
elif action_type == "delete_wiki_kb":
|
||
result = await delete_by_kb(str(action.get("kb_id") or plan.get("kb_id") or ""))
|
||
else:
|
||
result = {"success": False, "error": f"unknown action: {action_type}"}
|
||
executed.append({"action": action, "result": result})
|
||
except Exception as exc:
|
||
executed.append({"action": action, "result": {"success": False, "error": str(exc)}})
|
||
await asyncio.sleep(0)
|
||
|
||
failed = [item for item in executed if not item.get("result", {}).get("success")]
|
||
status = "failed" if failed else "executed"
|
||
await db.update_admin_plan_status(str(plan.get("id") or ""), status, executed=executed)
|
||
return {
|
||
"success": not failed,
|
||
"plan_id": plan.get("id", ""),
|
||
"status": status,
|
||
"executed_count": len(executed),
|
||
"failed_count": len(failed),
|
||
"executed": executed,
|
||
}
|
||
|
||
|
||
async def delete_by_file(file_id: str) -> Dict[str, Any]:
|
||
counts = await db.delete_wiki_by_file(file_id)
|
||
return {"success": True, "scope": "file", "file_id": str(file_id), "deleted": counts}
|
||
|
||
|
||
async def delete_by_kb(kb_id: str) -> Dict[str, Any]:
|
||
counts = await db.delete_wiki_by_kb(kb_id)
|
||
return {"success": True, "scope": "kb", "kb_id": str(kb_id), "deleted": counts}
|
||
|
||
|
||
async def delete_by_time(
|
||
before: Optional[datetime] = None,
|
||
after: Optional[datetime] = None,
|
||
kb_id: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
counts = await db.delete_wiki_by_time(before=before, after=after, kb_id=kb_id)
|
||
return {
|
||
"success": True,
|
||
"scope": "time",
|
||
"kb_id": str(kb_id or ""),
|
||
"before": before.isoformat() if before else "",
|
||
"after": after.isoformat() if after else "",
|
||
"deleted": counts,
|
||
}
|
||
|
||
|
||
async def _resolve_kb_ids(client: HTKnowClient, kb_id: str, include_children: bool = False) -> List[str]:
|
||
root = str(kb_id)
|
||
if not include_children:
|
||
return [root]
|
||
tree = await client.get_kb_tree()
|
||
found = _find_kb_node(tree, root)
|
||
if not found:
|
||
return [root]
|
||
ids: List[str] = []
|
||
_collect_kb_ids(found, ids)
|
||
return ids or [root]
|
||
|
||
|
||
def _find_kb_node(nodes: List[Dict[str, Any]], kb_id: str) -> Optional[Dict[str, Any]]:
|
||
for node in nodes:
|
||
if str(node.get("id")) == str(kb_id):
|
||
return node
|
||
children = node.get("children") or []
|
||
if isinstance(children, list):
|
||
found = _find_kb_node(children, kb_id)
|
||
if found:
|
||
return found
|
||
return None
|
||
|
||
|
||
def _collect_kb_ids(node: Dict[str, Any], out: List[str]) -> None:
|
||
kid = node.get("id")
|
||
if kid is not None:
|
||
out.append(str(kid))
|
||
children = node.get("children") or []
|
||
if isinstance(children, list):
|
||
for child in children:
|
||
if isinstance(child, dict):
|
||
_collect_kb_ids(child, out)
|
||
|
||
|
||
def _flatten_kb_tree(nodes: List[Dict[str, Any]], parent_path: str = "") -> List[Dict[str, Any]]:
|
||
flattened: List[Dict[str, Any]] = []
|
||
for node in nodes:
|
||
if not isinstance(node, dict):
|
||
continue
|
||
item = dict(node)
|
||
names = _node_names(item)
|
||
display_name = names[0] if names else str(item.get("id") or "")
|
||
item["_path"] = f"{parent_path}/{display_name}" if parent_path else display_name
|
||
flattened.append(item)
|
||
children = item.get("children") or []
|
||
if isinstance(children, list):
|
||
flattened.extend(_flatten_kb_tree(children, item["_path"]))
|
||
return flattened
|
||
|
||
|
||
def _node_names(node: Dict[str, Any]) -> List[str]:
|
||
keys = (
|
||
"name",
|
||
"title",
|
||
"label",
|
||
"kb_name",
|
||
"knowledge_base_name",
|
||
"knowledgeBaseName",
|
||
)
|
||
names: List[str] = []
|
||
for key in keys:
|
||
value = node.get(key)
|
||
if value is not None:
|
||
text = str(value).strip()
|
||
if text and text not in names:
|
||
names.append(text)
|
||
return names
|
||
|
||
|
||
def _normalize_kb_text(text: str) -> str:
|
||
return re.sub(r"[\s:：,，.。;；\"'“”‘’《》<>（）()【】\\[\\]_-]+", "", str(text or "").lower())
|
||
|
||
|
||
def _compact_file(file_info: Dict[str, Any]) -> Dict[str, Any]:
|
||
return {
|
||
"file_id": str(file_info.get("file_id") or file_info.get("id") or ""),
|
||
"kb_id": str(file_info.get("kb_id") or ""),
|
||
"filename": str(file_info.get("filename") or file_info.get("name") or ""),
|
||
"hash": str(file_info.get("hash") or ""),
|
||
"status": str(file_info.get("status") or ""),
|
||
"updated_at": str(file_info.get("updated_at") or ""),
|
||
}
|
||
|
||
|
||
def _is_stale(file_info: Dict[str, Any], state: Dict[str, Any]) -> bool:
|
||
source_hash = str(file_info.get("hash") or "")
|
||
wiki_hash = str(state.get("file_hash") or "")
|
||
if source_hash and wiki_hash and source_hash != wiki_hash:
|
||
return True
|
||
source_updated = str(file_info.get("updated_at") or "")
|
||
wiki_source_updated = str(state.get("source_updated_at") or "")
|
||
if source_updated and wiki_source_updated and source_updated != wiki_source_updated:
|
||
return True
|
||
return False
|