retriever.py

import os
from dotenv import load_dotenv

from huggingface_hub import login
from supabase import Client, create_client
from supabase.client import ClientOptions

from langchain_community.vectorstores import SupabaseVectorStore
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain.tools.retriever import create_retriever_tool

load_dotenv()

MODEL_NAME = "BAAI/bge-base-en-v1.5"
TBL_NAME = "documents_tbl"
QUERY_NAME = "match_documents"

def get_retriever_tool():
    embedding_model = HuggingFaceEmbeddings(model_name = MODEL_NAME)
    DIMS_EMBEDDING = embedding_model._client.get_sentence_embedding_dimension()

    # Supabase client
    supabase: Client = create_client(
        os.environ.get("SUPABASE_URL"), 
        os.environ.get("SUPABASE_ANON_KEY"), 
        options = ClientOptions(schema = "public")
    )

    # Vector Store
    vector_store = SupabaseVectorStore(
        client = supabase,
        embedding = embedding_model,
        table_name = TBL_NAME,
        query_name = QUERY_NAME
    )

    vector_retriever = vector_store.as_retriever()

    retriever_tool = create_retriever_tool(
        retriever = vector_retriever,
        name = "question_search_tool",
        description = "A tool to retrieve similar questions based on embedding from Supabase vector store."
    )
    return vector_store, vector_retriever, retriever_tool