import json
from langchain.vectorstores import FAISS
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

# Load the FAISS vector store from the directory 'faiss_index'
db = FAISS.load_local(
    "faiss_index",
    HuggingFaceEmbeddings(model_name='sentence-transformers/paraphrase-MiniLM-L6-v2'),
    allow_dangerous_deserialization=True
)

#Load attrributes along with their supported operators
attribute_data = json.load(open("attribute_data.json")) 

# Connect query to FAISS index using a retriever
retriever = db.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5}
)

# Modify fetch function to include operators
def fetch(query):
    res = retriever.get_relevant_documents(query)
    docs = []
    for i in res:
        attribute_line = i.page_content.split(",")[0]  # "attribute: X"
        attribute = attribute_line.split(": ")[1]  # Extract "X"
        
        # Find the matching operators from attribute_data
        operators = next((item["operators"] for item in attribute_data if item["attribute"] == attribute), [])
        operators_list = f"operators: {operators}"
        docs.append(f"{i.page_content}, {operators_list}")
    return docs