Files
llm/scripts/check_1c_rag_quality.py
T

88 lines
3.0 KiB
Python

from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from common import read_json, search_lexical_index
from rag_profiles import resolve_rag_profile
ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INDEX = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_index.json"
DEFAULT_CASES = ROOT / "plugins" / "1c" / "rag" / "quality-smoke.json"
def load_cases(path: Path) -> list[dict]:
data = read_json(path)
cases = data.get("cases")
if not isinstance(cases, list):
raise ValueError(f"{path} must contain a cases list")
return cases
def case_text(results: list[dict]) -> str:
return "\n".join((result["document"].get("content") or "") for result in results).lower()
def run_case(index: dict, case: dict, limit: int) -> dict:
profile = resolve_rag_profile(case.get("profile") or "auto", str(case["query"]))
results = search_lexical_index(
index,
str(case["query"]),
limit=limit or int(profile["limit"]),
candidate_limit=int(profile["candidate_limit"]),
dedupe_by_document=bool(profile["dedupe_by_document"]),
min_score=float(profile["min_score"]),
source_types=profile["source_types"],
)
text = case_text(results)
missing = [word for word in case.get("must_contain") or [] if str(word).lower() not in text]
return {
"id": case.get("id"),
"profile": profile["id"],
"query": case.get("query"),
"status": "passed" if results and not missing else "failed",
"missing": missing,
"top_sources": [
{
"source_path": result["document"].get("source_path"),
"title": result["document"].get("title"),
"chunk_index": result["document"].get("chunk_index"),
"score": round(float(result.get("score") or 0), 4),
}
for result in results
],
}
def main() -> int:
parser = argparse.ArgumentParser(description="Run smoke quality checks for the 1C RAG index.")
parser.add_argument("--index", type=Path, default=DEFAULT_INDEX)
parser.add_argument("--cases", type=Path, default=DEFAULT_CASES)
parser.add_argument("--limit", type=int, default=5)
parser.add_argument("--print", action="store_true")
args = parser.parse_args()
index = read_json(args.index)
results = [run_case(index, case, args.limit) for case in load_cases(args.cases)]
report = {
"status": "ok" if all(result["status"] == "passed" for result in results) else "failed",
"case_count": len(results),
"results": results,
}
if args.print:
print(json.dumps(report, ensure_ascii=False, indent=2))
else:
print(f"1C RAG quality status: {report['status']}")
if report["status"] != "ok":
for result in results:
if result["status"] != "passed":
print(f"- {result['id']}: missing {result['missing']}", file=sys.stderr)
return 0 if report["status"] == "ok" else 1
if __name__ == "__main__":
raise SystemExit(main())