stock_analysis_rag_project / vector_store.py
genaibeauty's picture
Create vector_store.py
3b0bc72 verified
raw
history blame
648 Bytes
import faiss
import numpy as np
from langchain.vectorstores import FAISS
from langchain.schema import Document
from utils.embeddings import embedding_model
# Initialize FAISS index
dim = 384 # MiniLM embedding size
index = faiss.IndexFlatL2(dim)
vector_db = FAISS(embedding_function=embedding_model, index=index)
def add_document(text, doc_id):
"""Add document embeddings to FAISS."""
doc = Document(page_content=text, metadata={"id": doc_id})
vector_db.add_documents([doc])
def search(query, top_k=3):
"""Search FAISS index for similar documents."""
results = vector_db.similarity_search(query, k=top_k)
return results