Spaces:
Running
Running
| import os | |
| import chromadb | |
| from db import SanatanDatabase | |
| from config import SanatanConfig | |
| import logging | |
| import argparse | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO, force=True) | |
| logger = logging.getLogger(__name__) | |
| logger.setLevel(logging.INFO) | |
| def get_target_field(scripture_config): | |
| """Identifies the hierarchical index field from scripture config.""" | |
| allowed_indices = ["relative_verse_number", "sloka_number", "sloka_index", "verse_number", "position_in_chapter"] | |
| for field in scripture_config.get("metadata_fields", []): | |
| if field["name"] in allowed_indices: | |
| return field["name"] | |
| return None | |
| def migrate_position_in_chapter(full_run=False): | |
| config = SanatanConfig() | |
| db = SanatanDatabase() | |
| # Use HttpClient with CHROMADB_SPACE_URL | |
| chroma_url = os.getenv("CHROMADB_SPACE_URL") | |
| if not chroma_url: | |
| logger.error("CHROMADB_SPACE_URL not set in environment.") | |
| return | |
| chroma_client = chromadb.HttpClient(host=chroma_url) | |
| BATCH_SIZE = 500 | |
| for scripture in config.scriptures: | |
| scripture_name = scripture["name"] | |
| collection_name = scripture["collection_name"] | |
| target_field = get_target_field(scripture) | |
| logger.info(f"Processing {scripture_name} (Chroma: {collection_name}, Field: {target_field or '_global_index'})") | |
| # If not a full run, check if any verses already have the property | |
| if not full_run: | |
| count_query = f""" | |
| MATCH (v:Verse)-[:PART_OF]->(s:Scripture {{name: '{scripture_name}'}}) | |
| WHERE v.position_in_chapter IS NOT NULL | |
| RETURN count(v) as count | |
| """ | |
| result = db.run_cypher(count_query) | |
| existing_count = result[0]["count"] if result and "count" in result[0] else 0 | |
| logger.info(f"existing_count = {existing_count}") | |
| if existing_count > 0: | |
| logger.info(f"Skipping {scripture_name}: already has {existing_count} verses populated. Use --full to override.") | |
| continue | |
| # Check if collection exists exactly as named | |
| coll_name = scripture["collection_name"] | |
| try: | |
| collection = chroma_client.get_collection(coll_name) | |
| except: | |
| # Fallback for _openai suffix | |
| try: | |
| coll_name = f"{coll_name}_openai" | |
| collection = chroma_client.get_collection(coll_name) | |
| except: | |
| logger.warning(f"Could not find collection for {scripture_name}") | |
| continue | |
| logger.info(f"Processing {scripture_name} (Chroma: {coll_name}, Field: {target_field or '_global_index'})") | |
| # Use a loop to fetch all records in case of pagination limits | |
| data = {"metadatas": [], "ids": []} | |
| offset = 0 | |
| limit = 1000 | |
| while True: | |
| chunk = collection.get(offset=offset, limit=limit, include=["metadatas"]) | |
| if not chunk["ids"]: | |
| break | |
| data["metadatas"].extend(chunk["metadatas"]) | |
| data["ids"].extend(chunk["ids"]) | |
| offset += limit | |
| metadatas = data["metadatas"] | |
| logger.info(f"Retrieved {len(metadatas)} records for {scripture_name}") | |
| updates = [] | |
| for i, meta in enumerate(metadatas): | |
| # The global_id in ArcadeDB is consistently scripture_name + "_" + _global_index | |
| global_index = meta.get("_global_index") | |
| if global_index is None: | |
| continue | |
| global_id = f"{scripture_name}_{global_index}" | |
| # Use dictionary get safely for relative_verse_number etc | |
| val = meta.get(target_field) if target_field and target_field in meta else global_index | |
| if i == 0: | |
| logger.info(f"DEBUG: scripture={scripture_name}, global_id={global_id}, target={target_field}, val={val}") | |
| if val is not None: | |
| updates.append({"global_id": global_id, "position": val}) | |
| if updates: | |
| # DEBUG: Sample the first update to ensure it's not None | |
| logger.info(f"DEBUG: Sample update: {updates[0]}") | |
| logger.info(f"Updating {len(updates)} records for {scripture_name} in batches of {BATCH_SIZE}...") | |
| for i in range(0, len(updates), BATCH_SIZE): | |
| batch = updates[i:i + BATCH_SIZE] | |
| query = """ | |
| UNWIND $batch AS row | |
| MATCH (v:Verse {global_id: row.global_id}) | |
| SET v.position_in_chapter = row.position | |
| """ | |
| try: | |
| db.run_cypher(query, {"batch": batch}) | |
| logger.info(f"Batch {i//BATCH_SIZE + 1} for {scripture_name} successful.") | |
| except Exception as e: | |
| logger.error(f"Failed to update batch starting at {i}: {e}") | |
| else: | |
| logger.info(f"No updates found for {scripture_name}") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Migrate position_in_chapter to ArcadeDB") | |
| parser.add_argument("--full", action="store_true", help="Perform a full run, overwriting existing data") | |
| args = parser.parse_args() | |
| migrate_position_in_chapter(full_run=args.full) | |