| import os |
| import json |
| import pandas as pd |
| from dotenv import load_dotenv |
| from os.path import join, dirname |
| from llm import embedding_llm |
| from langchain_community.vectorstores import SupabaseVectorStore |
| from langchain_core.documents import Document |
| from supabase.client import Client, create_client |
| from langchain.tools.retriever import create_retriever_tool |
|
|
|
|
| dotenv_path = join(dirname(__file__), '.env') |
| load_dotenv(dotenv_path) |
|
|
| def get_supabase_client(): |
| supabase: Client = create_client( |
| os.environ.get("SUPABASE_URL"), |
| os.environ.get("SUPABASE_SERVICE_KEY") |
| ) |
| return supabase |
|
|
|
|
| def __upload_gaia_data(supabase_client): |
| supabase_client = get_supabase_client() |
| df = pd.read_csv('supabase_docs.csv')[['content', 'metadata']] |
| df['metadata'] = df['metadata'].apply(lambda x: x.replace("'", '"')).apply(json.loads) |
| docs_dict = df.to_dict(orient='records') |
| docs = [Document(page_content=d['content'], metadata=d['metadata']) for d in docs_dict] |
| vector_store = SupabaseVectorStore.from_documents( |
| docs, |
| embedding_llm, |
| client=supabase_client, |
| table_name="documents", |
| query_name="match_documents_langchain", |
| ) |
| return vector_store |
|
|
| def get_gaia_data_vector_store(supabase_client): |
| return SupabaseVectorStore( |
| client=supabase_client, |
| embedding=embedding_llm, |
| table_name="documents", |
| query_name="match_documents_langchain" |
| ) |
|
|
| def get_gaia_data_retriever_tool(supabase_client): |
| vector_store = get_gaia_data_vector_store(supabase_client) |
| return create_retriever_tool( |
| retriever=vector_store.as_retriever(), |
| name="Question Search", |
| description="A tool to retrieve similar questions from a vector store.", |
| ) |
|
|
|
|
|
|
| if __name__ == '__main__': |
| |
| |
| supabase_client = get_supabase_client() |
| |
| |
| |
| vector_store = get_gaia_data_vector_store(supabase_client) |
| print(vector_store.similarity_search("Box Office Mojo's 2020 Worldwide Box Office list, how")) |
| |
