File size: 2,343 Bytes
e2640ca d8c20f3 7483ad4 e2640ca 1689c36 e2640ca 1689c36 e2640ca 1689c36 e2640ca 1689c36 7483ad4 1689c36 e2640ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import os
import time
import uuid
from pathlib import Path
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from core.settings import get_settings
load_dotenv()
UPLOAD_DIR = Path("./uploaded_docs")
UPLOAD_DIR.mkdir(exist_ok=True)
_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
# Lazy singleton — Pinecone is NOT connected at import time
_index = None
def get_index():
"""Connect to Pinecone on first call, then reuse the connection."""
global _index
if _index is not None:
return _index
settings = get_settings()
pc = Pinecone(api_key=settings.pinecone_api_key)
spec = ServerlessSpec(cloud="aws", region=settings.pinecone_region)
if settings.pinecone_index_name not in {i["name"] for i in pc.list_indexes()}:
pc.create_index(
name=settings.pinecone_index_name,
dimension=768,
metric="dotproduct",
spec=spec,
)
while not pc.describe_index(settings.pinecone_index_name).status["ready"]:
time.sleep(1)
_index = pc.Index(settings.pinecone_index_name)
return _index
def load_vectorstore(uploaded_files) -> int:
settings = get_settings()
embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
index = get_index()
total_chunks = 0
for file in uploaded_files:
save_path = UPLOAD_DIR / file.filename
save_path.write_bytes(file.file.read())
chunks = _splitter.split_documents(PyPDFLoader(str(save_path)).load())
if not chunks:
continue
texts = [c.page_content for c in chunks]
metadatas = [c.metadata for c in chunks]
embeddings = embed_model.embed_documents(texts)
vectors = [
{
"id": f"{save_path.stem}-{i}-{uuid.uuid4().hex[:8]}",
"values": embeddings[i],
"metadata": {**metadatas[i], "text": texts[i]},
}
for i in range(len(chunks))
]
index.upsert(vectors=vectors)
total_chunks += len(vectors)
return total_chunks
|