261 lines
8.3 KiB
Python
261 lines
8.3 KiB
Python
import argparse
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Any, Dict, Iterable, List, Tuple
|
||
|
||
import httpx
|
||
|
||
|
||
def _bool_env(name: str, default: bool) -> bool:
|
||
val = os.getenv(name)
|
||
if val is None:
|
||
return default
|
||
return val.strip().lower() in {"1", "true", "yes", "y", "on"}
|
||
|
||
|
||
def _cfg() -> Dict[str, Any]:
|
||
es_url = os.getenv("KB_ES_URL") or os.getenv("KB_OS_URL")
|
||
if not es_url:
|
||
raise SystemExit("Missing KB_ES_URL (or KB_OS_URL)")
|
||
|
||
return {
|
||
"es_url": es_url.rstrip("/"),
|
||
"index": os.getenv("KB_ES_INDEX", "kb_chunks_v1"),
|
||
"username": os.getenv("KB_ES_USERNAME", ""),
|
||
"password": os.getenv("KB_ES_PASSWORD", ""),
|
||
"verify_ssl": _bool_env("KB_ES_VERIFY_SSL", True),
|
||
"version": os.getenv("KB_ES_VERSION", "v1"),
|
||
}
|
||
|
||
|
||
def read_text_file(path: Path) -> str:
|
||
data = path.read_text(encoding="utf-8", errors="ignore")
|
||
return data.replace("\ufeff", "").strip()
|
||
|
||
|
||
def iter_text_files(folder: Path, exts: Tuple[str, ...]) -> Iterable[Path]:
|
||
for p in folder.rglob("*"):
|
||
if p.is_file() and p.suffix.lower() in exts:
|
||
yield p
|
||
|
||
|
||
def _slug_id(s: str) -> str:
|
||
return hashlib.sha1(s.encode("utf-8")).hexdigest()[:16]
|
||
|
||
|
||
def chunk_text(text: str, chunk_size: int, overlap: int) -> List[str]:
|
||
text = re.sub(r"\s+", " ", text).strip()
|
||
if not text:
|
||
return []
|
||
if chunk_size <= 0:
|
||
return [text]
|
||
overlap = max(0, min(overlap, chunk_size - 1))
|
||
out = []
|
||
i = 0
|
||
while i < len(text):
|
||
out.append(text[i : i + chunk_size])
|
||
if i + chunk_size >= len(text):
|
||
break
|
||
i = i + chunk_size - overlap
|
||
return out
|
||
|
||
|
||
_QA_Q_RE = re.compile(r"^\s*(\d+)\s*[\.．、)]\s*(.+?)\s*$")
|
||
_QA_A_RE = re.compile(r"^\s*[•\-*]\s*答案\s*[:：]\s*(.*)$")
|
||
|
||
# 兼容：markdown Q&A 版（常见：**Q5：...** **A：** ...）
|
||
_MD_Q_RE = re.compile(r"^\s*(?:\*\*)?Q\s*(\d+)?\s*[:：]\s*(.+?)\s*(?:\*\*)?\s*$", re.IGNORECASE)
|
||
_MD_A_RE = re.compile(r"^\s*(?:\*\*)?A\s*[:：]\s*(?:\*\*)?\s*(.*)\s*$", re.IGNORECASE)
|
||
_MD_INLINE_QA_RE = re.compile(
|
||
r"\*\*\s*Q\s*\d*\s*[:：]\s*(.+?)\s*\*\*\s*(?:\*\*)?\s*A\s*[:：]\s*(?:\*\*)?\s*(.+)\s*$",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
|
||
def parse_markdown_qa(text: str) -> List[Dict[str, str]]:
|
||
"""解析形如：
|
||
1.xxx？\n•答案：yyy
|
||
的 Q&A 列表。
|
||
"""
|
||
|
||
lines = [ln.rstrip() for ln in (text or "").splitlines()]
|
||
items: List[Dict[str, str]] = []
|
||
|
||
cur_q: str | None = None
|
||
cur_a_lines: List[str] = []
|
||
in_answer = False
|
||
|
||
def flush():
|
||
nonlocal cur_q, cur_a_lines, in_answer
|
||
if cur_q:
|
||
ans = "\n".join([a for a in cur_a_lines if a.strip()]).strip()
|
||
if ans:
|
||
items.append({"question": cur_q.strip(), "answer": ans})
|
||
cur_q = None
|
||
cur_a_lines = []
|
||
in_answer = False
|
||
|
||
for ln in lines:
|
||
if not ln.strip():
|
||
continue
|
||
|
||
m_inline = _MD_INLINE_QA_RE.match(ln)
|
||
if m_inline:
|
||
flush()
|
||
q = m_inline.group(1).strip()
|
||
a = m_inline.group(2).strip()
|
||
if q and a:
|
||
items.append({"question": q, "answer": a})
|
||
continue
|
||
|
||
m_q = _QA_Q_RE.match(ln)
|
||
if m_q:
|
||
flush()
|
||
cur_q = m_q.group(2)
|
||
continue
|
||
|
||
m_md_q = _MD_Q_RE.match(ln)
|
||
if m_md_q:
|
||
flush()
|
||
cur_q = m_md_q.group(2)
|
||
continue
|
||
|
||
m_a = _QA_A_RE.match(ln)
|
||
if m_a and cur_q:
|
||
first = m_a.group(1).strip()
|
||
if first:
|
||
cur_a_lines.append(first)
|
||
in_answer = True
|
||
continue
|
||
|
||
m_md_a = _MD_A_RE.match(ln)
|
||
if m_md_a and cur_q:
|
||
first = (m_md_a.group(1) or "").strip()
|
||
if first:
|
||
cur_a_lines.append(first)
|
||
in_answer = True
|
||
continue
|
||
if cur_q:
|
||
# 支持答案多行（没有“A：/•答案：”前缀的续行）
|
||
if in_answer:
|
||
s = ln.strip()
|
||
# 跳过与答案无关的章节标题/分隔线
|
||
if s.startswith("#") or s in {"---", "***"} or re.fullmatch(r"[-*_]{3,}", s):
|
||
continue
|
||
cur_a_lines.append(s)
|
||
|
||
flush()
|
||
return items
|
||
|
||
|
||
def build_docs(file_path: Path, root: Path, *, chunk_size: int, overlap: int) -> List[Dict[str, Any]]:
|
||
rel = str(file_path.relative_to(root)).replace("\\", "/")
|
||
content = read_text_file(file_path)
|
||
title = file_path.stem
|
||
doc_id = _slug_id(rel)
|
||
|
||
# 优先按 Q&A 结构化入库（更准、更快）；解析不到再退回 chunk
|
||
qa_items = parse_markdown_qa(content) if file_path.suffix.lower() == ".md" else []
|
||
if qa_items:
|
||
docs: List[Dict[str, Any]] = []
|
||
for i, qa in enumerate(qa_items, 1):
|
||
qa_id = f"{doc_id}_qa_{i}"
|
||
q = qa["question"].strip()
|
||
a = qa["answer"].strip()
|
||
docs.append(
|
||
{
|
||
"doc_id": doc_id,
|
||
"chunk_id": qa_id,
|
||
"qa_id": qa_id,
|
||
"title": title,
|
||
"question": q,
|
||
"answer": a,
|
||
"content": f"{q}\n{a}",
|
||
"source": rel,
|
||
}
|
||
)
|
||
return docs
|
||
|
||
parts = chunk_text(content, chunk_size=chunk_size, overlap=overlap)
|
||
docs = []
|
||
for idx, part in enumerate(parts):
|
||
docs.append(
|
||
{
|
||
"doc_id": doc_id,
|
||
"chunk_id": f"{doc_id}_{idx}",
|
||
"title": title,
|
||
"content": part,
|
||
"source": rel,
|
||
}
|
||
)
|
||
return docs
|
||
|
||
|
||
def to_bulk_ndjson(index: str, docs: List[Dict[str, Any]], *, version: str) -> str:
|
||
lines: List[str] = []
|
||
for d in docs:
|
||
d = dict(d)
|
||
d["version"] = version
|
||
_id = d.get("chunk_id")
|
||
lines.append(json.dumps({"index": {"_index": index, "_id": _id}}, ensure_ascii=False))
|
||
lines.append(json.dumps(d, ensure_ascii=False))
|
||
return "\n".join(lines) + "\n"
|
||
|
||
|
||
async def main() -> None:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--folder", required=True, help="资料目录（建议放 txt/md）")
|
||
parser.add_argument("--ext", default=".txt,.md", help="导入的扩展名，逗号分隔")
|
||
parser.add_argument("--chunk-size", type=int, default=800)
|
||
parser.add_argument("--overlap", type=int, default=120)
|
||
parser.add_argument("--batch", type=int, default=500)
|
||
args = parser.parse_args()
|
||
|
||
cfg = _cfg()
|
||
folder = Path(args.folder).resolve()
|
||
exts = tuple([e.strip().lower() for e in args.ext.split(",") if e.strip()])
|
||
files = list(iter_text_files(folder, exts=exts))
|
||
if not files:
|
||
raise SystemExit(f"No files found in {folder} for exts={exts}")
|
||
|
||
auth = None
|
||
if cfg["username"] and cfg["password"]:
|
||
auth = (cfg["username"], cfg["password"])
|
||
|
||
bulk_url = f"{cfg['es_url']}/_bulk"
|
||
headers = {"Content-Type": "application/x-ndjson"}
|
||
|
||
all_docs: List[Dict[str, Any]] = []
|
||
for f in files:
|
||
all_docs.extend(build_docs(f, folder, chunk_size=args.chunk_size, overlap=args.overlap))
|
||
|
||
total = len(all_docs)
|
||
print(f"Files={len(files)} Chunks={total} Index={cfg['index']}")
|
||
|
||
async with httpx.AsyncClient(timeout=30.0, verify=cfg["verify_ssl"]) as client:
|
||
for start in range(0, total, args.batch):
|
||
batch_docs = all_docs[start : start + args.batch]
|
||
body = to_bulk_ndjson(cfg["index"], batch_docs, version=cfg["version"])
|
||
resp = await client.post(bulk_url, content=body.encode("utf-8"), headers=headers, auth=auth)
|
||
if resp.status_code >= 300:
|
||
raise SystemExit(f"Bulk failed HTTP {resp.status_code}: {resp.text}")
|
||
payload = resp.json()
|
||
if payload.get("errors"):
|
||
first_err = None
|
||
for item in payload.get("items", []):
|
||
v = item.get("index") or item.get("create") or item.get("update")
|
||
if v and v.get("error"):
|
||
first_err = v.get("error")
|
||
break
|
||
raise SystemExit(f"Bulk had errors: {first_err}")
|
||
print(f"Ingested {min(start + len(batch_docs), total)}/{total}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
|
||
asyncio.run(main())
|