diff --git a/rag_agent.py b/rag_agent.py new file mode 100644 index 0000000..5199a0b --- /dev/null +++ b/rag_agent.py @@ -0,0 +1,32 @@ +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()