|
|
"""Vector store module for document embedding and retrieval""" |
|
|
|
|
|
from typing import List |
|
|
from langchain_community.vectorstores import FAISS |
|
|
from langchain_openai import OpenAIEmbeddings |
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings |
|
|
|
|
|
from langchain_core.documents import Document |
|
|
|
|
|
|
|
|
class VectorStore: |
|
|
"""Manages vector store operations""" |
|
|
|
|
|
def __init__(self): |
|
|
"""Initialize vector store with OpenAI embeddings""" |
|
|
|
|
|
|
|
|
self.embeddings = HuggingFaceEmbeddings( |
|
|
model_name="sentence-transformers/all-MiniLM-L6-v2", |
|
|
model_kwargs={"device": "cpu"}, |
|
|
encode_kwargs={"normalize_embeddings": True} |
|
|
) |
|
|
self.vectorstore = None |
|
|
self.retriever = None |
|
|
|
|
|
def create_vectorstore(self, documents: List[Document]): |
|
|
""" |
|
|
Create vector store from documents |
|
|
|
|
|
Args: |
|
|
documents: List of documents to embed |
|
|
""" |
|
|
self.vectorstore = FAISS.from_documents(documents, self.embedding) |
|
|
self.retriever = self.vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 5, "lambda_mult":0.25}) |
|
|
|
|
|
def get_retriever(self): |
|
|
""" |
|
|
Get the retriever instance |
|
|
|
|
|
Returns: |
|
|
Retriever instance |
|
|
""" |
|
|
if self.retriever is None: |
|
|
raise ValueError("Vector store not initialized. Call create_vectorstore first.") |
|
|
return self.retriever |
|
|
|
|
|
def retrieve(self, query: str, k: int = 4) -> List[Document]: |
|
|
""" |
|
|
Retrieve relevant documents for a query |
|
|
|
|
|
Args: |
|
|
query: Search query |
|
|
k: Number of documents to retrieve |
|
|
|
|
|
Returns: |
|
|
List of relevant documents |
|
|
""" |
|
|
if self.retriever is None: |
|
|
raise ValueError("Vector store not initialized. Call create_vectorstore first.") |
|
|
return self.retriever.invoke(query) |
