from sentence_transformers import (
    SentenceTransformer
)

from qdrant_client import (
    QdrantClient
)


class BNSRetriever:

    def __init__(
        self,
        collection_name="bns",
        model_name="BAAI/bge-large-en-v1.5"
    ):

        self.collection_name = (
            collection_name
        )

        self.client = (
            QdrantClient(
                host="localhost",
                port=6333
            )
        )

        self.model = (
            SentenceTransformer(
                model_name
            )
        )

    def search(
        self,
        query: str,
        limit: int = 10
    ):

        query_vector = (
            self.model.encode(
                query,
                normalize_embeddings=True
            ).tolist()
        )

        hits = (
            self.client.query_points(
                collection_name=
                    self.collection_name,

                query=
                    query_vector,

                limit=
                    limit
            )
        )

        return hits
    

if __name__=="__main__":
    query = "punishment for trafficking of person"
    
    bns_search=BNSRetriever()
    print(bns_search.search(query))