llama-index-vector-stores-faiss

import os import faiss from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader, load_index_from_storage from llama_index.vector_stores.faiss import FAISSVectorStore from llama_index.core.embeddings import resolve_embed_model # --- Setup: Install necessary packages and configure API key --- # pip install llama-index-vector-stores-faiss faiss-cpu llama-index-llms-openai llama-index-embeddings-openai # Set up OpenAI API key for default embedding model, or configure another model os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "sk-YOUR_OPENAI_KEY") # Create a dummy directory with a text file for demonstration if not os.path.exists("./data"): os.makedirs("./data") with open("./data/example.txt", "w") as f: f.write("The quick brown fox jumps over the lazy dog. LlamaIndex is powerful.\n") f.write("FAISS is a library for efficient similarity search and clustering of dense vectors.\n") f.write("RAG combines retrieval with large language models.") # 1. Load documents from the dummy directory documents = SimpleDirectoryReader("./data").load_data() # 2. Determine embedding dimension using the default embedding model # This is crucial for initializing the FAISS index correctly. embed_model = resolve_embed_model("default") # Uses OpenAI by default dummy_embedding = embed_model.get_text_embedding("hello world") d = len(dummy_embedding) # e.g., 1536 for OpenAI's text-embedding-ada-002 # 3. Initialize FAISS index (e.g., a simple L2 distance index) faiss_index = faiss.IndexFlatL2(d) # 4. Initialize FAISSVectorStore with the created FAISS index vector_store = FAISSVectorStore(faiss_index=faiss_index) # 5. Create StorageContext, linking it to the FAISSVectorStore storage_context = StorageContext.from_defaults(vector_store=vector_store) # 6. Create VectorStoreIndex from documents, using the defined storage context index = VectorStoreIndex.from_documents( documents, storage_context=storage_context, ) # 7. Query the index query_engine = index.as_query_engine() response = query_engine.query("What is RAG?") print(f"\nQuery Response: {response}\n") # 8. Example of persistence and loading persist_dir = "./faiss_storage" if not os.path.exists(persist_dir): os.makedirs(persist_dir) # Persist the index, including the FAISS index file index.storage_context.persist(persist_dir=persist_dir) print(f"Index persisted to {persist_dir}\n") # Load the index back from storage # First, load the FAISS index itself loaded_faiss_index = faiss.read_index(os.path.join(persist_dir, "vector_store.faiss")) loaded_vector_store = FAISSVectorStore(faiss_index=loaded_faiss_index) # Then, load the LlamaIndex StorageContext and the index loaded_storage_context = StorageContext.from_defaults( vector_store=loaded_vector_store, persist_dir=persist_dir ) loaded_index = load_index_from_storage(loaded_storage_context) loaded_query_engine = loaded_index.as_query_engine() loaded_response = loaded_query_engine.query("What did the fox do?") print(f"Loaded Index Query Response: {loaded_response}\n") # --- Clean up dummy data and directories --- import shutil shutil.rmtree("./data") shutil.rmtree(persist_dir) print("Cleaned up dummy data and storage directories.")