src/infer.py

import os
import argparse
from pathlib import Path

from .context_retriever import query_index
from .llm import generate_answer
from .embeddings import ingest_files


def ensure_index(data_dir: str, index_path: str, meta_path: str, model: str = "nomic-ai/nomic-embed-text-v1"):
    # If index or meta missing, instruct user to run ingest; try to call ingest if available.
    if Path(index_path).exists() and Path(meta_path).exists():
        return True

    print("Index or metadata not found. Attempting to create index from PDFs in", data_dir)
    try:

        ingest_files(data_dir, os.path.dirname(index_path), model_name=model)
        return Path(index_path).exists() and Path(meta_path).exists()
    except Exception as e:
        print("Automatic ingest failed:", e)
        print("Please run: python src/ingest.py --input-dir", data_dir, "--output-dir", os.path.dirname(index_path))
        return False


def run_inference(index_path: str, meta_path: str, question: str, k: int = 5, use_llm: bool = True, backend: str = "ollama", model_name: str = "gpt-oss:20b"):
    ok = ensure_index("data", index_path, meta_path)
    if not ok:
        return

    # Retrieve top-k chunks
    results = query_index(index_path, meta_path, question, top_k=k)
    print("\nRetrieved contexts:")
    contexts = []
    for i, r in enumerate(results, 1):
        print(f"[{i}] source={r['source']} score={r['score']:.4f}")
        txt = r.get("text") or ""
        snippet = txt.replace('\n', ' ')[:800]
        print(snippet)
        print("---")
        contexts.append(txt)

    if use_llm:
        print("\nGenerating answer using LLM (backend=", backend, ", model=", model_name, ")")
        ans = generate_answer(question, contexts, backend=backend, model_name=model_name)
        print("\n--- ANSWER ---\n")
        print(ans)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--index", default="data/faiss.index")
    parser.add_argument("--meta", default="data/meta.pkl")
    parser.add_argument("--query", required=True)
    parser.add_argument("--k", type=int, default=5)
    parser.add_argument("--no-llm", action="store_true", help="Only retrieve contexts, don't call LLM")
    parser.add_argument("--backend", default="ollama", choices=["hf", "ollama"], help="LLM backend")
    parser.add_argument("--model", default="gpt-oss:20b", help="Model name for the chosen backend")
    args = parser.parse_args()

    run_inference(args.index, args.meta, args.query, k=args.k, use_llm=not args.no_llm, backend=args.backend, model_name=args.model)


if __name__ == "__main__":
    main()