from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_vector_store(vector_store_path):
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cpu'}
    )
    
    # Allow deserialization since we're loading files we just created
    return FAISS.load_local(
        vector_store_path, 
        embeddings, 
        allow_dangerous_deserialization=True
    )

def create_qa_chain(vector_store_path, model="gpt-4o", api_key=None):
    vector_store = load_vector_store(vector_store_path)
    llm = ChatOpenAI(
        model=model,
        api_key=api_key,
        temperature=0
    )
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 4}),
        return_source_documents=True,
        verbose=True
    )
    
    return qa_chain