project/lab.py

"""Run the local quote-only lab. No API credentials or dependencies required."""

import argparse
import json
from pathlib import Path

from retrieval import Index, evaluate, evidence_packet, load_documents


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    search = commands.add_parser("search", help="retrieve verbatim synthetic passages")
    search.add_argument("question")
    search.add_argument("--k", type=int, default=2)
    search.add_argument("--max-quote-chars", type=int, default=1200)
    evaluation = commands.add_parser("evaluate", help="evaluate ten synthetic teaching cases")
    evaluation.add_argument("--k", type=int, default=2)
    arguments = parser.parse_args()
    data = Path(__file__).resolve().parent / "data"
    try:
        index = Index(load_documents(data / "corpus.json"))
        if arguments.command == "search":
            result = evidence_packet(index.search(arguments.question, arguments.k), arguments.max_quote_chars)
        else:
            cases = json.loads((data / "questions.json").read_text(encoding="utf-8"))
            result = evaluate(index, cases, arguments.k)
    except (ValueError, OSError) as error:
        parser.error(str(error))
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()