import shutil
import tempfile

import lancedb, pyarrow as pa
from lancedb.rerankers import LinearCombinationReranker

DB_PATH = tempfile.mkdtemp(prefix="testdb")

try:
    db = lancedb.connect(DB_PATH)

    data = pa.table({
        "text": ["apple orange", "banana grape", "cherry melon", "date fig"],
        "vector": [[0.1, 0.2], [0.9, 0.8], [0.5, 0.5], [0.3, 0.4]],
    })
    table = db.create_table("text", data=data, mode="overwrite")
    table.create_fts_index("text", replace=True)

    reranker = LinearCombinationReranker(weight=0.7, return_score="all")
    results = (
        table.search(query_type="hybrid")
        .vector([0.1, 0.2])
        .text("apple orange")
        .rerank(reranker)
        .to_list()
    )
    for r in results:
        print(r["text"], r["_relevance_score"])

finally:
    shutil.rmtree(DB_PATH, ignore_errors=True)

banana grape None 0.9999999403953552 0.699999988079071
apple orange 2.4079456329345703 0.0 0.30000001192092896
cherry melon None 0.25 0.17500002682209015
date fig None 0.08000001311302185 0.056000012904405594