demeter

  1import sys
  2import os
  3import uuid
  4import pypdf
  5from qdrant_client import models
  6from fastembed import TextEmbedding
  7
  8current_dir = os.path.dirname(os.path.abspath(__file__))
  9project_root = os.path.abspath(os.path.join(current_dir, '../../'))
 10sys.path.append(project_root)
 11
 12from Qdrant.Client import client  
 13
 14# --- CONFIGURATION ---
 15COLLECTION_NAME = "Knowledge_Base"
 16VECTOR_SIZE = 384  
 17DOCS_FOLDER = os.path.join(project_root, "Knowledge_Base") 
 18
 19def init_collection():
 20    """
 21    Creates the collection if it doesn't exist.
 22    """
 23    if client.collection_exists(COLLECTION_NAME):
 24        print(f"ℹ️  Collection '{COLLECTION_NAME}' already exists. Appending data...")
 25    else:
 26        print(f"🔨 Creating new collection: {COLLECTION_NAME}")
 27        client.create_collection(
 28            collection_name=COLLECTION_NAME,
 29            vectors_config=models.VectorParams(
 30                size=VECTOR_SIZE, 
 31                distance=models.Distance.COSINE
 32            )
 33        )
 34        print("✅ Collection created.")
 35
 36def extract_text_from_pdf(pdf_path):
 37    """
 38    Reads a PDF file page by page and returns the full text.
 39    """
 40    text = ""
 41    try:
 42        reader = pypdf.PdfReader(pdf_path)
 43        for page in reader.pages:
 44            page_text = page.extract_text()
 45            if page_text:
 46                text += page_text + "\n"
 47    except Exception as e:
 48        print(f"❌ Error reading PDF {pdf_path}: {e}")
 49    return text
 50
 51def chunk_text(text, chunk_size=500, overlap=50):
 52    """
 53    Splits long text into smaller overlapping pieces.
 54    Overlap helps preserve context between chunks.
 55    """
 56    if not text:
 57        return []
 58    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size - overlap)]
 59
 60def ingest_docs():
 61    # 1. Setup Collection & Model
 62    init_collection()
 63    
 64    print("🧠 Loading Embedding Model (bge-small-en)...")
 65    model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
 66    
 67    # 2. Check if folder exists
 68    if not os.path.exists(DOCS_FOLDER):
 69        os.makedirs(DOCS_FOLDER)
 70        print(f"⚠️ Created folder '{DOCS_FOLDER}'. Please put your PDFs there and run this script again!")
 71        return
 72
 73    # 3. Scan for files
 74    files = [f for f in os.listdir(DOCS_FOLDER) if f.endswith(('.pdf', '.txt'))]
 75    if not files:
 76        print(f"📭 No files found in '{DOCS_FOLDER}'. Add some PDFs!")
 77        return
 78
 79    print(f"📚 Found {len(files)} documents. Starting ingestion...")
 80    total_chunks = 0
 81
 82    for file_name in files:
 83        file_path = os.path.join(DOCS_FOLDER, file_name)
 84        print(f"   📄 Processing: {file_name}")
 85
 86        # A. Extract Text
 87        content = ""
 88        if file_name.endswith('.pdf'):
 89            content = extract_text_from_pdf(file_path)
 90        else:
 91            with open(file_path, 'r', encoding='utf-8') as f:
 92                content = f.read()
 93
 94        if not content.strip():
 95            print(f"      ⚠️ Skipping empty file.")
 96            continue
 97
 98        # B. Chunk Text
 99        chunks = chunk_text(content)
100        if not chunks:
101            continue
102            
103        # C. Convert to Vectors (Embed)
104        # FastEmbed handles the list of strings automatically
105        embeddings = list(model.embed(chunks))
106
107        # D. Prepare Points for Qdrant
108        points = []
109        for i, (text_chunk, vector) in enumerate(zip(chunks, embeddings)):
110            points.append(models.PointStruct(
111                id=str(uuid.uuid4()),  
112                vector=vector.tolist(),
113                payload={
114                    "text": text_chunk,
115                    "source": file_name,
116                    "chunk_id": i
117                }
118            ))
119
120        # E. Upload Batch
121        client.upsert(collection_name=COLLECTION_NAME, points=points)
122        total_chunks += len(points)
123        print(f"      ✅ Uploaded {len(points)} chunks.")
124
125    print(f"\n🎉 Success! Knowledge Base now contains {total_chunks} searchable segments.")
126
127if __name__ == "__main__":
128    ingest_docs()