File size: 3,000 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 78 79 80 81 82 83 | import os
import uuid
import datetime
from pymongo import MongoClient
from supabase import create_client, Client
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
load_dotenv()
# MongoDB Configuration
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017/")
MONGO_DB = "diet_db"
# Supabase Configuration
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY") or os.getenv("SUPABASE_KEY")
def migrate_kb():
print("Starting Knowledge Base migration from MongoDB to Supabase...")
if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
print("Error: Supabase configuration missing.")
return
mongo_client = MongoClient(MONGO_URI)
supabase: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
# Load embedding model
print("Loading embedding model (all-MiniLM-L6-v2)...")
model = SentenceTransformer('all-MiniLM-L6-v2')
# We need to migrate multiple collections possibly
# In kb_manager.py, default was 'documents' in 'medical_kb'
# In dietary_tools.py, it was 'diabetes' in 'test_index'
collections_to_migrate = [
{"db": "medical_kb", "coll": "documents"},
{"db": "test_index", "coll": "diabetes"},
{"db": "diet_db", "coll": "diabetes"}
]
for item in collections_to_migrate:
db_name = item["db"]
coll_name = item["coll"]
print(f"Migrating {db_name}.{coll_name}...")
try:
mongo_docs = list(mongo_client[db_name][coll_name].find())
if not mongo_docs:
print(f"No documents found in {db_name}.{coll_name}. Skipping.")
continue
# Batch processing for efficiency
batch_size = 50
for i in range(0, len(mongo_docs), batch_size):
batch = mongo_docs[i:i+batch_size]
texts = [doc.get("text", "") for doc in batch]
# Generate embeddings
embeddings = model.encode(texts).tolist()
supabase_data = []
for doc, emb in zip(batch, embeddings):
supabase_data.append({
"id": str(uuid.uuid4()),
"content": doc.get("text", ""),
"metadata": doc.get("metadata", {}),
"embedding": emb,
"collection_name": f"{db_name}_{coll_name}"
})
supabase.table("knowledge_base").insert(supabase_data).execute()
print(f"Uploaded {len(supabase_data)} documents...")
print(f"Successfully migrated {len(mongo_docs)} documents from {db_name}.{coll_name}.")
except Exception as e:
print(f"Error migrating {db_name}.{coll_name}: {e}")
print("Knowledge Base migration complete!")
if __name__ == "__main__":
migrate_kb()
|