src/embedder.py

from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from src.config import CHUNK_SIZE, CHUNK_OVERLAP, EMBEDDING_MODEL
from src.persistence import get_vectorstore_path, save_paper_metadata, load_paper_metadata
import uuid

def chunk_text(md_text: str) -> list[str]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", " ", ""]
    )
    return splitter.split_text(md_text)

def get_embeddings():
    """Get embeddings with GPU fallback to CPU."""
    try:
        return HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            model_kwargs={"device": "cuda"},
        )
    except Exception:
        print("⚠️ GPU not available, falling back to CPU")
        return HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            model_kwargs={"device": "cpu"},
        )

def build_vectorstore(chunks: list[str], paper_id: str = None) -> tuple[FAISS, str]:
    """Build and save vectorstore. Returns (vectorstore, paper_id)."""
    if paper_id is None:
        paper_id = str(uuid.uuid4())[:8]
    
    embeddings = get_embeddings()
    vectorstore = FAISS.from_texts(chunks, embeddings)
    
    # Save to disk
    vectorstore_path = get_vectorstore_path(paper_id)
    vectorstore.save_local(str(vectorstore_path))
    
    return vectorstore, paper_id

def load_vectorstore(paper_id: str) -> FAISS:
    """Load a previously saved vectorstore."""
    vectorstore_path = get_vectorstore_path(paper_id)
    embeddings = get_embeddings()
    
    if vectorstore_path.exists():
        vectorstore = FAISS.load_local(
            str(vectorstore_path),
            embeddings,
            allow_dangerous_deserialization=True
        )
        return vectorstore
    
    raise FileNotFoundError(f"Vectorstore for paper {paper_id} not found")