File size: 2,826 Bytes
76962bf | 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 | import os
import sys
import uuid
import datetime
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
from supabase import create_client, Client
from src.utils.logger import setup_logger
logger = setup_logger("KBManager")
def _get_supabase_client() -> Client:
url = os.getenv("SUPABASE_URL")
key = os.getenv("SUPABASE_SERVICE_ROLE_KEY") or os.getenv("SUPABASE_KEY")
if not url or not key:
raise ValueError("Supabase configuration missing")
return create_client(url, key)
def create_knowledge_base_supabase(pdf_path: str, collection_name: str = "general"):
"""
Creates a new knowledge base in Supabase (pgvector) from a PDF document.
"""
if not os.path.exists(pdf_path):
logger.error(f"Error: File not found at {pdf_path}")
return
logger.info(f"Loading document: {pdf_path}")
loader = PyPDFLoader(pdf_path)
docs = loader.load()
for doc in docs:
if 'page' in doc.metadata:
doc.metadata['page_index'] = doc.metadata['page']
logger.info(f"Successfully loaded {len(docs)} pages.")
logger.info("Splitting text into chunks...")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
logger.info(f"Created {len(splits)} chunks.")
# Generate embeddings
logger.info("Generating embeddings...")
model = SentenceTransformer('all-MiniLM-L6-v2')
texts = [s.page_content for s in splits]
embeddings = model.encode(texts).tolist()
supabase = _get_supabase_client()
# Store in Supabase
logger.info(f"Storing {len(splits)} chunks in Supabase knowledge_base...")
supabase_data = []
for chunk, emb in zip(splits, embeddings):
supabase_data.append({
"content": chunk.page_content,
"metadata": chunk.metadata,
"embedding": emb,
"collection_name": collection_name
})
# Bulk insert (limited by payload size, so batching)
batch_size = 100
for i in range(0, len(supabase_data), batch_size):
supabase.table("knowledge_base").insert(supabase_data[i:i+batch_size]).execute()
logger.info("Knowledge base integration to Supabase complete.")
return True
if __name__ == "__main__":
# Example usage for Supabase migration
sources_dir = "data/sources"
if os.path.exists(sources_dir):
for filename in os.listdir(sources_dir):
if filename.endswith(".pdf"):
pdf_path = os.path.join(sources_dir, filename)
create_knowledge_base_supabase(pdf_path, collection_name="diabetes_guidelines") |