qazws345/vector-migration
0
1from fastapi import FastAPI, HTTPException2from pinecone import Pinecone3from qdrant_client import QdrantClient4from qdrant_client.http.models import Distance, VectorParams, PointStruct5from tqdm import tqdm6from pinecone import Pinecone, ServerlessSpec 7from pinecone import ServerlessSpec8from langchain_pinecone import PineconeVectorStore9from fastapi import FastAPI, HTTPException10from pydantic import BaseModel11import logging12from qdrant_client import QdrantClient13from qdrant_client.http.models import Distance, VectorParams14from chromadb import Client as ChromaClient15from chromadb.config import Settings16 17import chromadb18 19import numpy as np20from pydantic import BaseModel21import logging22 23app = FastAPI()24 25# Configure logging26logging.basicConfig(level=logging.INFO)27logger = logging.getLogger(__name__)28 29class PineconToQdratMigrateRequest(BaseModel):30 pinecone_api_key: str31 qdrant_api_key: str32 qdrant_url: str33 pinecone_db: str34 collection_name: str35 namespace: str = None # Optional Pinecone namespace36 batch_size: int = 100 # Default batch size37 38class PineconetoChromaMigrateRequest(BaseModel):39 pinecone_api_key:str40 chroma_cloud_api:str41 pinecone_db:str42 chroma_db:str43 44@app.post("/")45async def main():46 return {"message": "Hello, this is a simple POST route!"}47@app.post('/chroma_to_pinecone')48async def migrate(req:PineconetoChromaMigrateRequest):49 try:50 chroma_client = chromadb.HttpClient(51 host="https://api.trychroma.com", # This is the default for Chroma Cloud52 api_key=req.chroma_cloud_api53 )54 collection = chroma_client.get_or_create_collection(name=req.chroma_db)55 pc = Pinecone(api_key=req.pinecone_api_key)56 57 results = collection.get(include=["embeddings", "documents", "metadatas"])58 59 pinecone_vector = []60 for id, vectors, metadata in zip(results["ids"], results["embeddings"], results["metadatas"]):61 vector = {62 "id":id,63 "values":vectors,64 "metadata":metadata65 }66 pinecone_vector.append(vector)67 68 pc.create_index(69 name=req.pinecone_Db,70 dimension=384,71 metric='cosine',72 spec=ServerlessSpec(73 cloud='aws',74 region='us-east-1'75 )76 )77 index = pc.Index(req.pinecone_db)78 79 all_ids = []80 for page in index.list(limit=100, namespace=None): # adjust namespace as needed81 all_ids.extend(page)82 print(f"Found {len(all_ids)} vector IDs.") 83 84 85 except Exception as e:86 raise HTTPException(status_code=500, detail=str(e))87 88@app.post('/pinecone_to_qdrant')89async def pinecone_to_qdrant(req: PineconToQdratMigrateRequest):90 try:91 # Initialize Pinecone92 pc = Pinecone(api_key=req.pinecone_api_key)93 index = pc.Index(req.pinecone_db)94 95 # Initialize Qdrant with longer timeout96 qdrant_client = QdrantClient(97 url=req.qdrant_url,98 api_key=req.qdrant_api_key,99 prefer_grpc=True,100 timeout=60.0 # Increased timeout to 60 seconds101 )102 103 # Get index stats104 stats = index.describe_index_stats()105 vector_dimension = stats['dimension']106 distance = stats['metric']107 108 # Create collection if it doesn't exist109 if not qdrant_client.collection_exists(req.collection_name):110 qdrant_client.create_collection(111 collection_name=req.collection_name,112 vectors_config=VectorParams(113 size=vector_dimension,114 distance=map_pinecone_to_qdrant_distance(distance)115 )116 )117 118 # Fetch all vector IDs119 all_ids = []120 for page in index.list(limit=100, namespace=req.namespace):121 all_ids.extend(page)122 logger.info(f"Found {len(all_ids)} vector IDs to migrate")123 124 # Process in batches125 success_count = 0126 batch_size = min(req.batch_size, 500) # Cap at 500 for safety127 128 for i in tqdm(range(0, len(all_ids), batch_size), desc="Migrating vectors"):129 batch_ids = all_ids[i:i + batch_size]130 131 # Fetch vectors from Pinecone132 pinecone_vectors = index.fetch(ids=batch_ids, namespace=req.namespace).vectors133 134 # Prepare Qdrant points135 points = [136 PointStruct(137 id=vid,138 vector=vector.values,139 payload=vector.metadata or {}140 )141 for vid, vector in pinecone_vectors.items()142 if len(vector.values) == vector_dimension143 ]144 145 # Upsert to Qdrant146 if points:147 qdrant_client.upsert(148 collection_name=req.collection_name,149 points=points,150 wait=True # Wait for confirmation151 )152 success_count += len(points)153 154 # Verify migration155 collection_info = qdrant_client.get_collection(req.collection_name)156 157 return {158 "status": "success",159 "migrated_vectors": success_count,160 "collection_count": collection_info.points_count,161 "vector_dimension": vector_dimension,162 "distance_metric": distance163 }164 165 except Exception as e:166 logger.error(f"Migration failed: {str(e)}")167 return {168 "status": "error",169 "message": str(e)170 }171 172def map_pinecone_to_qdrant_distance(metric: str) -> Distance:173 metric = metric.lower()174 if metric == "cosine":175 return Distance.COSINE176 elif metric == "euclidean":177 return Distance.EUCLID178 elif metric == "dotproduct":179 return Distance.DOT180 else:181 raise ValueError(f"Unsupported Pinecone metric: {metric}")182 183 184@app.post('/pinecone_to_chroma')185async def pinecone_to_chroma(req:PineconetoChromaMigrateRequest):186 try:187 pc = Pinecone(api_key=req.pinecone_api_key)188 index = pc.Index(req.pinecone_db) # your existing index189 190 all_ids = []191 for page in index.list(limit=100, namespace=None): # adjust namespace as needed192 all_ids.extend(page)193 print(f"Found {len(all_ids)} vector IDs.")194 195 def fetch_vectors(index, ids, batch_size=100):196 all_vectors = []197 for i in range(0, len(ids), batch_size):198 batch_ids = ids[i:i + batch_size]199 response = index.fetch(ids=batch_ids)200 201 for vid, vector in response.vectors.items():202 all_vectors.append({203 "id": vid,204 "vector": vector.values,205 "metadata": vector.metadata or {}206 })207 return all_vectors208 209 vectors = fetch_vectors(index, all_ids)210 211 chroma_client = chromadb.HttpClient(212 host="https://api.trychroma.com", # This is the default for Chroma Cloud213 api_key=req.chroma_cloud_api214 )215 collection = chroma_client.get_or_create_collection(name=req.chroma_db)216 217 ids = [item["id"] for item in vectors]218 embeddings = [item["vector"] for item in vectors]219 metadatas = [item["metadata"] for item in vectors]220 documents = [item["metadata"].get("text", "") for item in vectors]221 collection.add(222 ids=ids,223 embeddings=embeddings,224 documents=documents, 225 metadatas=metadatas226 )227 return {"status": "success", "message": f"Migrated {len(ids)} vectors."}228 229 except Exception as e:230 raise HTTPException(status_code=500, detail=str(e))231 232 233 234class MigrationRequest(BaseModel):235 chroma_api_key: str236 chroma_tenant: str237 chroma_database: str238 chroma_collection: str239 240 qdrant_url: str241 qdrant_api_key: str242 qdrant_collection: str243 244@app.post("/migrate/chroma-to-qdrant")245def migrate_chroma_to_qdrant(req: MigrationRequest):246 try:247 # Initialize Chroma client248 chroma_client = ChromaClient(249 Settings(250 chroma_api_impl="rest",251 chroma_server_host="api.chroma.cloud",252 chroma_server_http_port="443",253 chroma_server_ssl_enabled=True,254 tenant=req.chroma_tenant,255 database=req.chroma_database,256 anonymized_telemetry=False,257 chroma_api_key=req.chroma_api_key258 )259 )260 261 # Fetch collection data262 chroma_collection = chroma_client.get_collection(name=req.chroma_collection)263 chroma_data = chroma_collection.get(include=["embeddings", "metadatas", "documents"])264 265 if not chroma_data["ids"]:266 raise HTTPException(status_code=404, detail="No data found in Chroma collection.")267 268 logger.info(f"Fetched {len(chroma_data['ids'])} records from Chroma")269 270 # Initialize Qdrant client271 qdrant = QdrantClient(url=req.qdrant_url, api_key=req.qdrant_api_key)272 273 # Create Qdrant collection274 qdrant.recreate_collection(275 collection_name=req.qdrant_collection,276 vectors_config=VectorParams(size=len(chroma_data["embeddings"][0]), distance=Distance.COSINE),277 )278 279 # Upload to Qdrant280 qdrant.upload_collection(281 collection_name=req.qdrant_collection,282 ids=chroma_data["ids"],283 vectors=chroma_data["embeddings"],284 payload=chroma_data["metadatas"]285 )286 287 return {"status": "success", "message": f"{len(chroma_data['ids'])} vectors migrated to Qdrant."}288 289 except Exception as e:290 logger.exception("Migration failed.")291 raise HTTPException(status_code=500, detail=str(e))292 293 294 295@app.post("/migrate/qdrant-to-chroma")296def migrate_qdrant_to_chroma(req: MigrationRequest):297 try:298 # Connect to Qdrant299 qdrant_client = QdrantClient(300 url=req.qdrant_url,301 api_key=req.qdrant_api_key302 )303 304 # Fetch all vectors, payloads, and ids from Qdrant305 scroll_results = qdrant_client.scroll(306 collection_name=req.qdrant_collection,307 scroll_filter=None,308 limit=10000,309 with_payload=True,310 with_vectors=True311 )312 313 if not scroll_results[0]:314 raise Exception("No data found in Qdrant collection.")315 316 vectors = []317 metadatas = []318 documents = []319 ids = []320 321 for point in scroll_results[0]:322 vectors.append(point.vector)323 metadatas.append(point.payload if point.payload else {})324 documents.append(point.payload.get("document", "")) # Assumes doc is under 'document'325 ids.append(str(point.id))326 327 # Connect to Chroma Cloud328 chroma_client = ChromaClient(Settings(329 chroma_api_impl="rest",330 chroma_server_host="api.chroma.cloud",331 chroma_server_http_port="443",332 tenant=req.chroma_tenant,333 database="default",334 api_key=req.chroma_api_key,335 anonymized_telemetry=False336 ))337 338 # Get or create Chroma collection339 collection = chroma_client.get_or_create_collection(name=req.chroma_collection)340 341 # Push data to Chroma342 collection.add(343 ids=ids,344 documents=documents,345 metadatas=metadatas,346 embeddings=vectors347 )348 349 return {"status": "success", "message": f"Migrated {len(ids)} vectors from Qdrant to Chroma."}350 351 except Exception as e:352 raise HTTPException(status_code=500, detail=str(e))353 354 355class QdrantToPineconeMigrateRequest(BaseModel):356 qdrant_api_key: str357 qdrant_url: str358 qdrant_collection: str359 pinecone_api_key: str360 pinecone_db: str361 namespace: str = None362 batch_size: int = 100 # default batch size363 364@app.post("/qdrant_to_pinecone")365async def qdrant_to_pinecone(req: QdrantToPineconeMigrateRequest):366 try:367 # Init Qdrant368 qdrant_client = QdrantClient(369 url=req.qdrant_url,370 api_key=req.qdrant_api_key,371 prefer_grpc=True,372 timeout=60.0373 )374 375 # Get vectors from Qdrant376 scroll_results, _ = qdrant_client.scroll(377 collection_name=req.qdrant_collection,378 scroll_filter=None,379 limit=10000,380 with_payload=True,381 with_vectors=True382 )383 384 if not scroll_results:385 raise Exception("No data found in Qdrant collection.")386 387 # Organize data388 ids, vectors, metadatas = [], [], []389 for point in scroll_results:390 ids.append(str(point.id))391 vectors.append(point.vector)392 metadatas.append(point.payload or {})393 394 vector_dimension = len(vectors[0]) # Assumes all have same dim395 logger.info(f"Fetched {len(ids)} vectors of dim {vector_dimension} from Qdrant")396 397 # Init Pinecone398 pc = Pinecone(api_key=req.pinecone_api_key)399 pinecone_index = pc.Index(req.pinecone_db)400 401 # Batch upsert402 batch_size = min(req.batch_size, 500)403 success_count = 0404 for i in tqdm(range(0, len(ids), batch_size), desc="Uploading to Pinecone"):405 batch_ids = ids[i:i + batch_size]406 batch_vectors = vectors[i:i + batch_size]407 batch_metadatas = metadatas[i:i + batch_size]408 409 pinecone_index.upsert(vectors=[410 {411 "id": id_,412 "values": vec,413 "metadata": meta414 } for id_, vec, meta in zip(batch_ids, batch_vectors, batch_metadatas)415 ], namespace=req.namespace)416 417 success_count += len(batch_ids)418 419 return {420 "status": "success",421 "migrated_vectors": success_count,422 "vector_dimension": vector_dimension423 }424 425 except Exception as e:426 logger.error(f"Migration failed: {str(e)}")427 raise HTTPException(status_code=500, detail=str(e))