Spaces:
Sleeping
Sleeping
File size: 1,385 Bytes
3927a42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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
|