| 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.") |
| |
| |
| 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() |
| |
| |
| 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 |
| }) |
| |
| |
| 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__": |
| |
| 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") |
