| """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) |
