101-0731/sherry/semantic_search.py

### Python = 3.9
import os
from dotenv import load_dotenv
load_dotenv()

import openai 
openai_api_key = os.getenv("OPENAI_API_KEY")
openai.api_key = openai_api_key

from langchain_openai import OpenAIEmbeddings
embeddings_model = OpenAIEmbeddings()

from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_chroma import Chroma

from supabase import create_client, Client 
supabase_url = os.getenv("SUPABASE_URL")
supabase_key = os.getenv("SUPABASE_KEY")
supabase: Client = create_client(supabase_url, supabase_key)

############# Load data #############
# def extract_field(doc, field_name):
#     for line in doc.page_content.split('\n'):
#         if line.startswith(f"{field_name}:"):
#             return line.split(':', 1)[1].strip()
#     return None

# loader = CSVLoader(file_path="video_cache_rows.csv")
# data = loader.load()
# field_name = "question"
# question = [extract_field(doc, field_name) for doc in data]

# ####### load data from supabase #######
# embeddings_model = OpenAIEmbeddings()

def generated(language:str ="ch"):
    global response,count,question,ids,question_id_map,vectorstore
    response,count = supabase.table("video_cache").select("question","id").eq("language",language).order("id").execute()
    data = response[1]
    question = [item['question'] for item in data if 'question' in item]
    ids = [item['id'] for item in data if 'id' in item]
    question_id_map = {item['question']: item['id'] for item in data if 'id' in item and 'question' in item}

    ########## generate embedding ###########
    embedding = embeddings_model.embed_documents(question)

    ########## Write embedding to the supabase table  #######
    # for id, new_embedding in zip(ids, embedding):
    #     supabase.table("video_cache_rows_duplicate").insert({"embedding": embedding.tolist()}).eq("id", id).execute()

    ######### Vector Store ##########
    # Put pre-compute embeddings to vector store. ## save to disk
    vectorstore = Chroma.from_texts(
        texts=question,
        embedding=embeddings_model,
        persist_directory="./chroma_db"
        )

    vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings_model)

    print("gernerate")

def get_id_by_question(question):
    return question_id_map.get(question)

# print(question)
# created_at = []
# question = []
# ids = []
# answer = []
# video_url = []

# for item in data:
#     ids.append(item['id'])
#     created_at.append(item['created_at'])
#     question.append(item['question'])
#     answer.append(item['answer'])
#     video_url.append(item['video_url'])


def ask_question(question:str, SIMILARITY_THRESHOLD:int = 0.83,language:str ="ch"):
    generated(language=language)
    docs_and_scores = vectorstore.similarity_search_with_relevance_scores(question, k=1)
    doc, score = docs_and_scores[0]
    print(doc,score)
    
    if score >= SIMILARITY_THRESHOLD:
        id = get_id_by_question(doc.page_content)
        data,count = supabase.table("video_cache").select("*").eq("id",id).execute()

        if data[1][0]["answer"] == None :
            return None

        return data[1]
    else:
        return None


if __name__ == "__main__" :
    ####### load from disk  #######
    query = "美食街在哪裡"
    docs = vectorstore.similarity_search(query)
    print(f"Query: {query}  | 最接近文檔：{docs[0].page_content}")

    ####### Query it #########
    query = "101可以帶狗嗎"
    docs = vectorstore.similarity_search(query)
    print(f"Query: {query}  | 最接近文檔：{docs[0].page_content}")