from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA
import os

def build_vector_store(directory: str):
    """Load text files from directory, split into chunks, and build FAISS store."""
    loader = TextLoader(directory, glob='**/*.txt')
    documents = loader.load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = splitter.split_documents(documents)
    embeddings = OpenAIEmbeddings()
    vector_store = FAISS.from_documents(docs, embeddings)
    return vector_store

def main():
    # Load documents and build vector store
    vector_store = build_vector_store('data')
    # Create RetrievalQA chain
    llm = OpenAI(temperature=0)
    qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
    # Sample query
    query = "What is the capital of France?"
    result = qa_chain.run(query)
    print("Query:", query)
    print("Answer:", result)

if __name__ == "__main__":
    main()