290 lines
11 KiB
Python
290 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_SOURCE_DIR = ROOT / "plugins" / "1c" / "rag" / "sources"
|
|
DEFAULT_OUTPUT = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_corpus.jsonl"
|
|
SUPPORTED_EXTENSIONS = {".md", ".txt", ".bsl", ".os"}
|
|
HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$")
|
|
FRONT_MATTER_RE = re.compile(r"^---\n(?P<body>.*?)\n---\n?", re.DOTALL)
|
|
|
|
|
|
def normalize_text(text: str) -> str:
|
|
lines = [line.rstrip() for line in text.replace("\r\n", "\n").split("\n")]
|
|
return "\n".join(lines).strip()
|
|
|
|
|
|
def parse_front_matter(text: str) -> tuple[dict[str, str], str]:
|
|
match = FRONT_MATTER_RE.match(text)
|
|
if not match:
|
|
return {}, text
|
|
metadata: dict[str, str] = {}
|
|
for line in match.group("body").splitlines():
|
|
if ":" not in line:
|
|
continue
|
|
key, value = line.split(":", 1)
|
|
value = value.strip().strip('"').strip("'")
|
|
metadata[key.strip()] = value
|
|
return metadata, text[match.end() :].strip()
|
|
|
|
|
|
def split_paragraphs(text: str) -> list[str]:
|
|
blocks: list[str] = []
|
|
current: list[str] = []
|
|
for line in text.split("\n"):
|
|
stripped = line.strip()
|
|
if not stripped:
|
|
if current:
|
|
blocks.append("\n".join(current).strip())
|
|
current = []
|
|
continue
|
|
if HEADING_RE.match(stripped) and current:
|
|
blocks.append("\n".join(current).strip())
|
|
current = [stripped]
|
|
else:
|
|
current.append(line.rstrip())
|
|
if current:
|
|
blocks.append("\n".join(current).strip())
|
|
return [block for block in blocks if block]
|
|
|
|
|
|
def hard_split_text(text: str, chunk_size: int, overlap: int) -> list[str]:
|
|
chunks: list[str] = []
|
|
start = 0
|
|
while start < len(text):
|
|
end = min(start + chunk_size, len(text))
|
|
chunk = text[start:end].strip()
|
|
if chunk:
|
|
chunks.append(chunk)
|
|
if end == len(text):
|
|
break
|
|
start = end - overlap
|
|
return chunks
|
|
|
|
|
|
def chunk_text(text: str, chunk_size: int, overlap: int) -> list[dict]:
|
|
if chunk_size <= 0:
|
|
raise ValueError("chunk_size must be positive")
|
|
if overlap < 0:
|
|
raise ValueError("overlap must not be negative")
|
|
if overlap >= chunk_size:
|
|
raise ValueError("overlap must be smaller than chunk_size")
|
|
|
|
chunks: list[dict] = []
|
|
current: list[str] = []
|
|
current_heading = ""
|
|
last_tail = ""
|
|
headings_seen: list[str] = []
|
|
|
|
def flush() -> None:
|
|
nonlocal current, last_tail
|
|
if not current:
|
|
return
|
|
content = "\n\n".join(current).strip()
|
|
if last_tail and not content.startswith(last_tail):
|
|
content = f"{last_tail}\n\n{content}".strip()
|
|
chunks.append(
|
|
{
|
|
"content": content,
|
|
"heading": current_heading,
|
|
"headings": headings_seen[-4:],
|
|
}
|
|
)
|
|
last_tail = content[-overlap:].strip() if overlap else ""
|
|
current = []
|
|
|
|
for block in split_paragraphs(text):
|
|
heading_match = HEADING_RE.match(block.split("\n", 1)[0].strip())
|
|
if heading_match:
|
|
current_heading = heading_match.group(2).strip()
|
|
headings_seen.append(current_heading)
|
|
|
|
if len(block) > chunk_size:
|
|
flush()
|
|
for part in hard_split_text(block, chunk_size, overlap):
|
|
chunks.append(
|
|
{
|
|
"content": part,
|
|
"heading": current_heading,
|
|
"headings": headings_seen[-4:],
|
|
}
|
|
)
|
|
last_tail = part[-overlap:].strip() if overlap else ""
|
|
continue
|
|
|
|
candidate = "\n\n".join([*current, block]).strip()
|
|
if current and len(candidate) > chunk_size:
|
|
flush()
|
|
current.append(block)
|
|
|
|
flush()
|
|
return chunks
|
|
|
|
|
|
def stable_id(*parts: str) -> str:
|
|
digest = hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()
|
|
return digest[:16]
|
|
|
|
|
|
def classify_source(path: Path, text: str, front_matter: dict[str, str] | None = None) -> str:
|
|
front_matter = front_matter or {}
|
|
if front_matter.get("source") == "official_1c_its":
|
|
return front_matter.get("source_type") or "official_1c_its"
|
|
name = path.name.lower()
|
|
lowered = text[:4000].lower()
|
|
if path.suffix.lower() in {".bsl", ".os"}:
|
|
return "bsl"
|
|
if "1c bsl module snapshot" in lowered or "```bsl" in lowered or "процедура " in lowered:
|
|
return "bsl"
|
|
if "1c metadata snapshot" in lowered or "metadata snapshot" in lowered:
|
|
return "metadata"
|
|
if "read-only" in lowered or "только чтение" in lowered or "выбрать" in lowered and "запрос" in lowered:
|
|
return "query"
|
|
if "production" in lowered or "резервн" in lowered or "согласован" in lowered or "опасн" in lowered:
|
|
return "safety"
|
|
if "metadata" in name:
|
|
return "metadata"
|
|
if "bsl" in name or "module" in name:
|
|
return "bsl"
|
|
return "docs"
|
|
|
|
|
|
def iter_source_files(source_dir: Path) -> list[Path]:
|
|
if not source_dir.exists():
|
|
return []
|
|
return sorted(
|
|
path
|
|
for path in source_dir.rglob("*")
|
|
if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
|
|
)
|
|
|
|
|
|
def build_records(source_dir: Path, chunk_size: int, overlap: int) -> list[dict]:
|
|
records: list[dict] = []
|
|
for path in iter_source_files(source_dir):
|
|
relative_path = path.relative_to(source_dir).as_posix()
|
|
text = normalize_text(path.read_text(encoding="utf-8"))
|
|
if not text:
|
|
continue
|
|
front_matter, body = parse_front_matter(text)
|
|
chunk_source_text = body or text
|
|
|
|
source_type = classify_source(path, chunk_source_text, front_matter)
|
|
chunks = chunk_text(chunk_source_text, chunk_size=chunk_size, overlap=overlap)
|
|
document_id = stable_id(relative_path, chunk_source_text)
|
|
base_title = front_matter.get("title") or path.stem.replace("_", " ").replace("-", " ")
|
|
|
|
for index, chunk in enumerate(chunks):
|
|
content = chunk["content"]
|
|
title = chunk.get("heading") or base_title
|
|
records.append(
|
|
{
|
|
"id": stable_id(document_id, str(index), content),
|
|
"document_id": document_id,
|
|
"source_path": relative_path,
|
|
"title": title,
|
|
"chunk_index": index,
|
|
"content": content,
|
|
"metadata": {
|
|
"domain": "1c",
|
|
"source_type": source_type,
|
|
"file_type": path.suffix.lower().lstrip("."),
|
|
"access": front_matter.get("access") or "",
|
|
"official_source": front_matter.get("source") or "",
|
|
"url": front_matter.get("url") or "",
|
|
"platform_family": front_matter.get("platform_family") or "",
|
|
"platform_version": front_matter.get("platform_version") or "",
|
|
"platform_doc_id": front_matter.get("platform_doc_id") or "",
|
|
"doc_book": front_matter.get("doc_book") or "",
|
|
"doc_bookmark": front_matter.get("doc_bookmark") or "",
|
|
"doc_content_id": front_matter.get("doc_content_id") or "",
|
|
"doc_coordinate": front_matter.get("doc_coordinate") or "",
|
|
"access_blocked": front_matter.get("access_blocked") or "",
|
|
"access_findings": front_matter.get("access_findings") or "",
|
|
"source_sha256": front_matter.get("source_sha256") or "",
|
|
"chunk_size": chunk_size,
|
|
"overlap": overlap,
|
|
"heading": chunk.get("heading") or "",
|
|
"headings": chunk.get("headings") or [],
|
|
},
|
|
}
|
|
)
|
|
return records
|
|
|
|
|
|
def build_manifest(source_dir: Path, records: list[dict], chunk_size: int, overlap: int) -> dict:
|
|
by_source: dict[str, dict] = {}
|
|
for record in records:
|
|
source_path = record["source_path"]
|
|
source = by_source.setdefault(
|
|
source_path,
|
|
{
|
|
"source_path": source_path,
|
|
"title": record.get("title"),
|
|
"source_type": (record.get("metadata") or {}).get("source_type"),
|
|
"file_type": (record.get("metadata") or {}).get("file_type"),
|
|
"document_id": record.get("document_id"),
|
|
"chunk_count": 0,
|
|
"content_hash": "",
|
|
},
|
|
)
|
|
source["chunk_count"] += 1
|
|
|
|
for relative_path, source in by_source.items():
|
|
path = source_dir / relative_path
|
|
text = normalize_text(path.read_text(encoding="utf-8")) if path.exists() else ""
|
|
source["content_hash"] = hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
return {
|
|
"schema_version": 1,
|
|
"source_dir": str(source_dir),
|
|
"chunk_size": chunk_size,
|
|
"overlap": overlap,
|
|
"source_count": len(by_source),
|
|
"chunk_count": len(records),
|
|
"sources": sorted(by_source.values(), key=lambda item: item["source_path"]),
|
|
}
|
|
|
|
|
|
def write_jsonl(path: Path, records: list[dict]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Prepare a JSONL RAG corpus for the 1C plugin.")
|
|
parser.add_argument("--source-dir", type=Path, default=DEFAULT_SOURCE_DIR)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
parser.add_argument("--manifest", type=Path)
|
|
parser.add_argument("--chunk-size", type=int, default=1800)
|
|
parser.add_argument("--overlap", type=int, default=200)
|
|
args = parser.parse_args()
|
|
|
|
records = build_records(
|
|
source_dir=args.source_dir,
|
|
chunk_size=args.chunk_size,
|
|
overlap=args.overlap,
|
|
)
|
|
write_jsonl(args.output, records)
|
|
manifest_path = args.manifest or args.output.with_suffix(".manifest.json")
|
|
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
|
manifest_path.write_text(
|
|
json.dumps(build_manifest(args.source_dir, records, args.chunk_size, args.overlap), ensure_ascii=False, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"Wrote {len(records)} chunk(s) to {args.output}")
|
|
print(f"Wrote RAG manifest to {manifest_path}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|