Spaces:
Sleeping
Sleeping
File size: 3,104 Bytes
454f118 923d3a8 454f118 | 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 84 85 86 87 88 89 90 91 92 93 94 | import os
from qdrant_client import QdrantClient
from dotenv import load_dotenv
from qdrant_client.models import PointStruct
def migrate():
# Load .env manually in case this is run directly
load_dotenv()
# 1. Local Client
local_path = "./data/qdrant_store"
if not os.path.exists(local_path):
print(f"Error: Local Qdrant store not found at {local_path}")
return
local_client = QdrantClient(path=local_path)
collection_name = "financial_rag"
# 2. Cloud Client
url = os.environ.get("QDRANT_URL")
api_key = os.environ.get("QDRANT_API_KEY")
if not url or not api_key:
print("Error: QDRANT_URL or QDRANT_API_KEY not found in environment.")
print("Please add them to your .env file.")
return
cloud_client = QdrantClient(url=url, port=443, api_key=api_key, timeout=60)
print(f"Connected to Local ({local_path}) and Cloud ({url})")
# 3. Create Cloud Collection if missing
try:
local_info = local_client.get_collection(collection_name)
except Exception as e:
print(f"Error reading local collection: {e}")
return
existing_cloud = [c.name for c in cloud_client.get_collections().collections]
if collection_name not in existing_cloud:
print(f"Creating collection '{collection_name}' in Cloud...")
cloud_client.create_collection(
collection_name=collection_name,
vectors_config=local_info.config.params.vectors,
)
from qdrant_client.models import PayloadSchemaType
for f in ['chunk_id', 'parent_chunk_id', 'company', 'ticker', 'source_file', 'fiscal_year', 'section', 'subsection', 'content_type', 'chunk_level']:
cloud_client.create_payload_index(collection_name=collection_name, field_name=f, field_schema=PayloadSchemaType.KEYWORD)
cloud_client.create_payload_index(collection_name=collection_name, field_name='page_number', field_schema=PayloadSchemaType.INTEGER)
# 4. Migrate points in batches
print("Starting migration...")
batch_size = 100
offset = None
total_migrated = 0
while True:
# Scroll through local points
records, next_offset = local_client.scroll(
collection_name=collection_name,
limit=batch_size,
offset=offset,
with_payload=True,
with_vectors=True
)
if not records:
break
# Convert Record objects to PointStruct objects
point_structs = [
PointStruct(id=record.id, vector=record.vector, payload=record.payload)
for record in records
]
# Upsert to cloud
cloud_client.upsert(
collection_name=collection_name,
points=point_structs
)
total_migrated += len(records)
print(f"Migrated {total_migrated} points...")
offset = next_offset
if offset is None:
break
print(f"\nMigration complete! Successfully moved {total_migrated} vectors to Qdrant Cloud.")
if __name__ == "__main__":
migrate()
|