BidiRag/examples/test_bid_field_extraction_rag.py

# -*- coding: utf-8 -*-
"""BidFieldExtractionRAG standalone test."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from bdirag.rag_methods.bid_field_extraction_rag import BidFieldExtractionRAG
from examples.rag_test_utils import build_vector_rag, make_documents, run_retrieval_test


if __name__ == "__main__":
    rag = build_vector_rag(BidFieldExtractionRAG)
    run_retrieval_test("BidFieldExtractionRAG", rag)

    print("\n" + "=" * 60)
    print("BidFieldExtractionRAG Field Extraction Test")
    print("=" * 60)
    rag = build_vector_rag(BidFieldExtractionRAG)
    rag.index_documents(make_documents())
    result = rag.extract_fields("extract project budget and evaluation method", k=3)
    print(result.answer)