main.py

import argparse
from pathlib import Path
from src import Pipeline


def main():
    parser = argparse.ArgumentParser(description="REFRAG CLI")
    sub = parser.add_subparsers(dest="command")

    ix = sub.add_parser("index")
    ix.add_argument("--file", "-f", type=str)
    ix.add_argument("--dir", "-d", type=str)
    ix.add_argument("--namespace", "-n", default="")

    q = sub.add_parser("query")
    q.add_argument("question", type=str)
    q.add_argument("--namespace", "-n", default="")
    q.add_argument("--stream", "-s", action="store_true")
    q.add_argument("--no-rl", action="store_true")

    sub.add_parser("stats")

    args = parser.parse_args()

    if args.command == "index":
        pipe = Pipeline()
        if args.file:
            p = Path(args.file)
            if not p.exists():
                print(f"Not found: {p}")
                return
            pipe.index_document(p.read_text(encoding="utf-8"), p.stem, namespace=args.namespace)
        elif args.dir:
            d = Path(args.dir)
            if not d.exists():
                print(f"Not found: {d}")
                return
            docs = [{"id": f.stem, "text": f.read_text(encoding="utf-8")} for f in d.glob("*.txt")]
            if docs:
                pipe.index_documents(docs, namespace=args.namespace)
            else:
                print("No .txt files found")

    elif args.command == "query":
        pipe = Pipeline()
        if args.stream:
            print("\nAnswer: ", end="", flush=True)
            for tok in pipe.query_stream(args.question, args.namespace, use_rl=not args.no_rl):
                print(tok, end="", flush=True)
            print()
        else:
            r = pipe.query(args.question, args.namespace, use_rl=not args.no_rl)
            print(f"\nQ: {r['query']}")
            print(f"A: {r['answer']}")
            if "decompression_stats" in r:
                s = r["decompression_stats"]
                print(f"Decompressed: {s['decompressed_count']}/{s['total_chunks']}")
            if "token_savings" in r:
                print(f"Savings: {r['token_savings']['percentage_saved']}%")

    elif args.command == "stats":
        pipe = Pipeline()
        s = pipe.get_stats()
        print(f"Encoder: {s['encoder']['model']} (dim={s['encoder']['dim']})")
        print(f"Decoder: {s['decoder']['model']}")
        print(f"Vectors: {s['vector_store']['total_vectors']}")

    else:
        parser.print_help()


if __name__ == "__main__":
    main()