from langchain_pinecone import PineconeVectorStore
from src.helper import download_embeddings
from dotenv import load_dotenv

load_dotenv()

embedding_model = download_embeddings()

INDEX_NAME = "assessment-recommender-agent"

vectorstore = PineconeVectorStore.from_existing_index(index_name=INDEX_NAME, embedding=embedding_model)

def get_retriever(k=10):
    return vectorstore.as_retriever(search_kwargs={"k": k})

def get_metadata_retriever(job_level=None, assessment_focus=None, languages=None, k=10):
    if job_level and assessment_focus :
        if languages:
            return vectorstore.as_retriever(search_kwargs={"k": k,
                                                        "filter" : {
                                                            "job_levels": {"$in": job_level} if job_level else None,
                                                            "languages": {"$in": languages} if languages else None,
                                                            "keys": {"$in": assessment_focus} if assessment_focus else None
                                            }})
        else:
            return vectorstore.as_retriever(search_kwargs={"k": k,
                                                        "filter" : {
                                                            "job_levels": {"$in": job_level} if job_level else None,
                                                            "keys": {"$in": assessment_focus} if assessment_focus else None
                                            }})
    else:
        return None