66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
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"
|
|
|
|
|
|
def print_result(result: dict, index: int) -> None:
|
|
document = result["document"]
|
|
content = (document.get("content") or "").strip().replace("\n", " ")
|
|
if len(content) > 500:
|
|
content = content[:497].rstrip() + "..."
|
|
|
|
print(f"{index}. score={result['score']:.4f}")
|
|
print(f" source={document.get('source_path')} chunk={document.get('chunk_index')}")
|
|
print(f" title={document.get('title')}")
|
|
if (document.get("metadata") or {}).get("heading"):
|
|
print(f" heading={(document.get('metadata') or {}).get('heading')}")
|
|
print(f" content={content}")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Search the local 1C RAG lexical index.")
|
|
parser.add_argument("query")
|
|
parser.add_argument("--index", type=Path, default=DEFAULT_INDEX)
|
|
parser.add_argument("--profile", default="auto")
|
|
parser.add_argument("--limit", type=int)
|
|
parser.add_argument("--candidate-limit", type=int)
|
|
parser.add_argument("--dedupe-by-document", action="store_true")
|
|
parser.add_argument("--min-score", type=float)
|
|
parser.add_argument("--source-type", action="append", dest="source_types")
|
|
parser.add_argument("--json", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
index = read_json(args.index)
|
|
profile = resolve_rag_profile(args.profile, args.query)
|
|
source_types = args.source_types if args.source_types is not None else profile["source_types"]
|
|
results = search_lexical_index(
|
|
index,
|
|
query=args.query,
|
|
limit=args.limit or int(profile["limit"]),
|
|
candidate_limit=args.candidate_limit or int(profile["candidate_limit"]),
|
|
dedupe_by_document=args.dedupe_by_document or bool(profile["dedupe_by_document"]),
|
|
min_score=float(args.min_score if args.min_score is not None else profile["min_score"]),
|
|
source_types=source_types,
|
|
)
|
|
|
|
if args.json:
|
|
print(json.dumps(results, ensure_ascii=False, indent=2))
|
|
elif results:
|
|
for position, result in enumerate(results, start=1):
|
|
print_result(result, position)
|
|
else:
|
|
print("No matches.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|