import tempfile
from pathlib import Path

import lancedb
from lancedb.embeddings import get_registry
from lancedb.pydantic import LanceModel, Vector
from lancedb.rerankers import OpenaiReranker

embedder = get_registry().get("sentence-transformers").create()
db = lancedb.connect(str(Path(tempfile.mkdtemp()) / "openai-rerank"))

class Schema(LanceModel):
    text: str = embedder.SourceField()
    vector: Vector(embedder.ndims()) = embedder.VectorField()

data = [
    {"text": "hello world"},
    {"text": "goodbye world"},
]
tbl = db.create_table("test", schema=Schema)
tbl.add(data)
reranker = OpenaiReranker(api_key="your-api-key")

# Run vector search with a reranker
vector_results = tbl.search("hello").rerank(reranker=reranker).to_list()
print(vector_results)

# Build the full-text index before either FTS or hybrid search.
tbl.create_fts_index("text", replace=True)

# Run FTS search with a reranker
fts_results = (
    tbl.search("hello", query_type="fts").rerank(reranker=reranker).to_list()
)
print(fts_results)

# Run hybrid search with a reranker
hybrid_results = (
    tbl.search("hello", query_type="hybrid").rerank(reranker=reranker).to_list()
)
print(hybrid_results)