54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
from langchain_community.vectorstores import FAISS
|
|
from langchain_openai import OpenAIEmbeddings
|
|
from langchain.schema import Document
|
|
|
|
|
|
class FAISSService:
|
|
"""
|
|
A service for creating and loading FAISS indexes for document embeddings.
|
|
"""
|
|
|
|
def __init__(self, openai_api_key, index_path="local_faiss_index"):
|
|
"""
|
|
Initialize the FAISS service.
|
|
|
|
:param openai_api_key: OpenAI API key for embeddings.
|
|
:param index_path: Path to save or load the FAISS index.
|
|
"""
|
|
self.openai_api_key = openai_api_key
|
|
self.index_path = index_path
|
|
|
|
def create_faiss_index(self, documents):
|
|
"""
|
|
Create a FAISS index from a list of documents.
|
|
|
|
:param documents: List of langchain Document objects.
|
|
:return: FAISS vectorstore instance.
|
|
"""
|
|
print("[INFO] Creating FAISS index...")
|
|
vectorstore = FAISS.from_documents(
|
|
documents,
|
|
OpenAIEmbeddings(
|
|
model="text-embedding-ada-002",
|
|
openai_api_key=self.openai_api_key
|
|
)
|
|
)
|
|
vectorstore.save_local(self.index_path)
|
|
print(f"[INFO] FAISS index saved to {self.index_path}.")
|
|
return vectorstore
|
|
|
|
def load_faiss_index(self):
|
|
"""
|
|
Load an existing FAISS index.
|
|
|
|
:return: Loaded FAISS vectorstore instance.
|
|
"""
|
|
print("[INFO] Loading FAISS index...")
|
|
vectorstore = FAISS.load_local(
|
|
self.index_path,
|
|
OpenAIEmbeddings(openai_api_key=self.openai_api_key),
|
|
allow_dangerous_deserialization=True
|
|
)
|
|
print(f"[INFO] FAISS index loaded from {self.index_path}.")
|
|
return vectorstore
|