code: import sys
  import os
  
  # this is needed to import classes from the API. it will be removed when the worker is refactored
  sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))
  
  import time
  import pika
  import json
  import pinecone
  import logging
  import weaviate
  import redis
  import lancedb
  import pymongo
  import pyarrow as pa
  import numpy as np
  import worker.config as config
  import services.database.batch_service as batch_service
  import services.database.job_service as job_service
  from services.database.database import get_db, safe_db_operation
  from shared.job_status import JobStatus
  from shared.batch_status import BatchStatus
  from qdrant_client import QdrantClient
  from qdrant_client.models import PointStruct
  from pymilvus import Collection, connections
  from shared.embeddings_type import EmbeddingsType
  from shared.vector_db_type import VectorDBType
  from shared.utils import generate_uuid_from_tuple
  from urllib.parse import quote_plus
  from services.rabbitmq.rabbit_service import create_connection_params
  from pika.exceptions import AMQPConnectionError
  
  logging.basicConfig(filename='./vdb-log.txt', level=logging.INFO)
  logging.basicConfig(filename='./vdb-errors.txt', level=logging.ERROR)
  
  def upload_batch(batch_id, chunks_with_embeddings):
      batch = safe_db_operation(batch_service.get_batch, batch_id)
      if batch.batch_status == BatchStatus.FAILED:
          safe_db_operation(batch_service.update_batch_retry_count, batch.id, batch.retries+1)
          logging.info(f"Retrying vector db upload of batch {batch.id}")
  
      batch = safe_db_operation(batch_service.get_batch, batch_id)
      vectors_uploaded = write_embeddings_to_vector_db(chunks_with_embeddings, batch.vector_db_metadata, batch.id, batch.job_id)
  
      if vectors_uploaded:
          status = safe_db_operation(batch_service.update_batch_status_with_successful_minibatch, batch.id)
          update_batch_and_job_status(batch.job_id, status, batch.id)
      else:
          update_batch_and_job_status(batch.job_id, BatchStatus.FAILED, batch.id)
  
  def write_embeddings_to_vector_db(chunks, vector_db_metadata, batch_id, job_id):
      # NOTE: the legacy code expects a list of tuples, (text_chunk, embedding) of form (str, list[float])
      text_embeddings_list = [(chunk['text'], chunk['vector']) for chunk in chunks]
      
      job = safe_db_operation(job_service.get_job, job_id)
      source_filename = job.source_filename
      
      if vector_db_metadata.vector_db_type == VectorDBType.PINECONE:
          upsert_list = create_pinecone_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_pinecone(upsert_list, vector_db_metadata)
      elif vector_db_metadata.vector_db_type == VectorDBType.QDRANT:
          upsert_list = create_qdrant_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_qdrant(upsert_list, vector_db_metadata)
      elif vector_db_metadata.vector_db_type == VectorDBType.WEAVIATE:
          return write_embeddings_to_weaviate(text_embeddings_list, vector_db_metadata, batch_id, job_id, source_filename)
      elif vector_db_metadata.vector_db_type == VectorDBType.MILVUS:
          upsert_list = create_milvus_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_milvus(upsert_list, vector_db_metadata)
      elif vector_db_metadata.vector_db_type == VectorDBType.REDIS:
          upsert_list = create_redis_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_redis(upsert_list, vector_db_metadata)
      elif vector_db_metadata.vector_db_type == VectorDBType.LANCEDB:
          upsert_list = create_lancedb_source_chunks(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_lancedb(upsert_list, batch_id)
      elif vector_db_metadata.vector_db_type == VectorDBType.MONGODB:
          upsert_list = create_mongodb_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename)
          return write_embeddings_to_mongodb(upsert_list, vector_db_metadata)
      else:
          logging.error('Unsupported vector DB type: %s', vector_db_metadata.vector_db_type.value)
  
  def create_mongodb_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename):
      upsert_list = []
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          upsert_list.append(
              {"_id": generate_uuid_from_tuple((job_id, batch_id, i)), 
              "values": embedding, 
              "source_text": source_text,
              "source_document": source_filename
              })
      return upsert_list
  
  def write_embeddings_to_mongodb(upsert_list, vector_db_metadata):
      mongo_conn_uri = vector_db_metadata.environment
      mongo_password = quote_plus(os.getenv('VECTOR_DB_KEY'))
      mongo_conn_uri = mongo_conn_uri.replace("<password>", mongo_password)
  
      mongo_client = pymongo.MongoClient(mongo_conn_uri)
      db_name, collection = vector_db_metadata.index_name.split(".")
      db = mongo_client[db_name]
  
      try:
          db.command("ping")
      except Exception as e:
          logging.error(f"Error connecting to MongoDB via python client: {e}")
          return None
  
      if collection not in db.list_collection_names():
          logging.error(f"Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}")
          return None
      
      index = db.get_collection(collection)
      
      logging.info(f"Starting MongoDB upsert for {len(upsert_list)} vectors")
  
      batch_size = config.PINECONE_BATCH_SIZE
      vectors_uploaded = 0
  
      for i in range(0,len(upsert_list), batch_size):
          try:
              upsert_batch = upsert_list[i:i+batch_size]
              upsert_response = index.insert_many(upsert_batch)
              vectors_uploaded += len(upsert_batch)
          except Exception as e:
              logging.error('Error writing embeddings to Mongo:', e)
              return None
      
      logging.info(f"Successfully uploaded {vectors_uploaded} vectors to MongoDB")
      return vectors_uploaded
  
  def create_pinecone_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename):
      upsert_list = []
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          upsert_list.append(
              {"id": generate_uuid_from_tuple((job_id, batch_id, i)), 
              "values": embedding, 
              "metadata": {"source_text": source_text, "source_document": source_filename}})
      return upsert_list
  
  def write_embeddings_to_pinecone(upsert_list, vector_db_metadata):
      pinecone_api_key = os.getenv('VECTOR_DB_KEY')
      pinecone.init(api_key=pinecone_api_key, environment=vector_db_metadata.environment)
      index = pinecone.GRPCIndex(vector_db_metadata.index_name)
      if not index:
          logging.error(f"Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}")
          return None
      
      logging.info(f"Starting pinecone upsert for {len(upsert_list)} vectors")
  
      batch_size = config.PINECONE_BATCH_SIZE
      vectors_uploaded = 0
  
      for i in range(0,len(upsert_list), batch_size):
          try:
              upsert_response = index.upsert(vectors=upsert_list[i:i+batch_size])
              vectors_uploaded += upsert_response.upserted_count
          except Exception as e:
              logging.error('Error writing embeddings to pinecone:', e)
              return None
      
      logging.info(f"Successfully uploaded {vectors_uploaded} vectors to pinecone")
      return vectors_uploaded
  
  def create_redis_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename):
      ids = []
      source_texts = []
      source_documents = []
      embeddings = []
  
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          ids.append(generate_uuid_from_tuple((job_id, batch_id, i)))
          source_texts.append(source_text)
          embeddings.append(embedding)
          source_documents.append(source_filename)
  
      return [ids, source_texts, embeddings, source_documents]
  
  def write_embeddings_to_redis(upsert_list, vector_db_metadata):
      redis_client = redis.from_url(url=vector_db_metadata.environment, password=os.getenv('VECTOR_DB_KEY'), decode_responses=True)
      
      try:
          redis_client.ft(vector_db_metadata.index_name).info()
      except redis.exceptions.ResponseError as e:
          if "Unknown Index name" in str(e):
              logging.error(f"Index {vector_db_metadata.index_name} does not exist at redis URL {vector_db_metadata.environment}")
              return None
      
      logging.info(f"Starting redis upsert for {len(upsert_list)} vectors")
  
      redis_pipeline = redis_client.pipeline()
  
      for i in range(0,len(upsert_list[0])):
          key = f'{vector_db_metadata.collection}:{upsert_list[0][i]}'
          obj = {"source_data": upsert_list[1][i], "embeddings": np.array(upsert_list[2][i]).tobytes(), "source_document": upsert_list[3][i]}
  
          redis_pipeline.hset(key, mapping=obj)
  
      res = redis_pipeline.execute()
  
      logging.info(f"Successfully uploaded {len(res)} vectors to redis")
      return len(res)
  
  def create_qdrant_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename):
      upsert_list = []
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          upsert_list.append(
              PointStruct(
                  id=generate_uuid_from_tuple((job_id, batch_id, i)),
                  vector=embedding,
                  payload={"source_text": source_text, "source_document": source_filename}
              )
          )
      return upsert_list
  
  def write_embeddings_to_qdrant(upsert_list, vector_db_metadata):
      qdrant_client = QdrantClient(
          url=vector_db_metadata.environment, 
          api_key=os.getenv('VECTOR_DB_KEY'),
          grpc_port=6334, 
          prefer_grpc=True,
          timeout=5
      ) if vector_db_metadata.environment != os.getenv('LOCAL_VECTOR_DB') else QdrantClient(os.getenv('LOCAL_VECTOR_DB'), port=6333)
  
      index = qdrant_client.get_collection(collection_name=vector_db_metadata.index_name)
      if not index:
          logging.error(f"Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}")
          return None
      
      logging.info(f"Starting qdrant upsert for {len(upsert_list)} vectors")
  
      batch_size = config.PINECONE_BATCH_SIZE
  
      for i in range(0, len(upsert_list), batch_size):
          try:
              qdrant_client.upsert(
                  collection_name=vector_db_metadata.index_name,
                  points=upsert_list[i:i+batch_size]
              )
          except Exception as e:
              logging.error('Error writing embeddings to qdrant:', e)
              return None
      
      logging.info(f"Successfully uploaded {len(upsert_list)} vectors to qdrant")
      return len(upsert_list)
  
  def write_embeddings_to_weaviate(text_embeddings_list, vector_db_metadata,  batch_id, job_id, source_filename):
      client = weaviate.Client(
          url=vector_db_metadata.environment,
          auth_client_secret=weaviate.AuthApiKey(api_key=os.getenv('VECTOR_DB_KEY')),
      ) if vector_db_metadata.environment != os.getenv('LOCAL_VECTOR_DB') else weaviate.Client(url=vector_db_metadata.environment)
  
      index = client.schema.get()
      class_list = [class_dict["class"] for class_dict in index["classes"]]
      if not index or not vector_db_metadata.index_name in class_list:
          logging.error(f"Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}")
          return None
      
      logging.info(f"Starting Weaviate upsert for {len(text_embeddings_list)} vectors")
      try:
          with client.batch(batch_size=config.PINECONE_BATCH_SIZE, dynamic=True, num_workers=2) as batch:
              for i, (text, vector) in enumerate(text_embeddings_list):
                  properties = {
                      "source_data": text,
                      "vectoflow_id": generate_uuid_from_tuple((job_id, batch_id, i)),
                      "source_document": source_filename
                  }
  
                  client.batch.add_data_object(
                      properties,
                      vector_db_metadata.index_name,
                      vector=vector
                  )
      except Exception as e:
          logging.error('Error writing embeddings to weaviate: %s', e)
          return None
      
      logging.info(f"Successfully uploaded {len(text_embeddings_list)} vectors to Weaviate")
      return len(text_embeddings_list)
  
  def create_milvus_source_chunk_dict(text_embeddings_list, batch_id, job_id, source_filename):
      ids = []
      source_texts = []
      embeddings = []
      source_filenames = []
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          ids.append(generate_uuid_from_tuple((job_id, batch_id, i)))
          source_texts.append(source_text)
          embeddings.append(embedding)
          source_filenames.append(source_filename)
      return [ids, source_texts, embeddings, source_filenames]
  
  def write_embeddings_to_milvus(upsert_list, vector_db_metadata):
      if vector_db_metadata.environment != os.getenv('LOCAL_VECTOR_DB'):
          connections.connect("default", 
              uri = vector_db_metadata.environment,
              token = os.getenv('VECTOR_DB_KEY')
          )
      else:
          connections.connect("default", 
              host = vector_db_metadata.environment
          )
  
      collection = Collection(vector_db_metadata.index_name)
      if not collection:
          logging.error(f"Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}")
          return None
      
      logging.info(f"Starting Milvus insert for {len(upsert_list)} vectors")
      batch_size = config.PINECONE_BATCH_SIZE
      vectors_uploaded = 0
  
      for i in range(0,len(upsert_list), batch_size):
          try:
              insert_response = collection.insert(upsert_list[i:i+batch_size])
              vectors_uploaded += insert_response.insert_count
          except Exception as e:
              logging.error('Error writing embeddings to milvus: %s', e)
              return None
      
      logging.info(f"Successfully uploaded {vectors_uploaded} vectors to milvus")
      return vectors_uploaded
  
  def create_lancedb_source_chunks(text_embeddings_list, batch_id, job_id, source_filename):
      upsert_list = []
      for i, (source_text, embedding) in enumerate(text_embeddings_list):
          upsert_list.append(
              {
                  "id": generate_uuid_from_tuple((job_id, batch_id, i)),
                  "vector": embedding,
                  "source_text": source_text, 
                  "source_document": source_filename
              }
          )
      return upsert_list
  
  def write_embeddings_to_lancedb(upsert_list, batch_id):
      # right now only local connection, since its serverless and their cloud is in beta
      batch = safe_db_operation(batch_service.get_batch, batch_id)
      db = lancedb.connect(batch.vector_db_metadata.environment)
      try:
          table = db.open_table(batch.vector_db_metadata.index_name)
      except FileNotFoundError as e:
          logging.info(f"Table {batch.vector_db_metadata.index_name} does not exist in environment {batch.vector_db_metadata.environment}.")
  
          if batch.embeddings_metadata.embeddings_type == EmbeddingsType.OPEN_AI:
              schema = pa.schema(
                  [
                      pa.field("id", pa.string()),
                      pa.field("vector", pa.list_(pa.float32(), 1536)),
                      pa.field("source_text", pa.string()),
                      pa.field("source_document", pa.string()),
                  ])
              table = db.create_table(batch.vector_db_metadata.index_name, schema=schema)
              logging.info(f"Created table {batch.vector_db_metadata.index_name} in environment {batch.vector_db_metadata.environment}.")
          else:
              logging.error(f"Embeddings type {batch.embeddings_metadata.embeddings_type} not supported for LanceDB. Only Open AI")
              return None
  
      logging.info(f"Starting LanceDB upsert for {len(upsert_list)} vectors")
  
      batch_size = config.PINECONE_BATCH_SIZE
      vectors_uploaded = 0
  
      for i in range(0,len(upsert_list), batch_size):
          try:
              table.add(data=upsert_list[i:i+batch_size])
              vectors_uploaded += batch_size
          except Exception as e:
              logging.error('Error writing embeddings to lance db:', e)
              return None
      
      logging.info(f"Successfully uploaded {vectors_uploaded} vectors to lance db")
      return vectors_uploaded
  
  # TODO: refactor into utils
  def update_batch_and_job_status(job_id, batch_status, batch_id):
      try:
          if not job_id and batch_id:
              job = safe_db_operation(batch_service.get_batch, batch_id)
              job_id = job.job_id
          updated_batch_status = safe_db_operation(batch_service.update_batch_status, batch_id, batch_status)
          job = safe_db_operation(job_service.update_job_with_batch, job_id, updated_batch_status)
          if job.job_status == JobStatus.COMPLETED:
              logging.info(f"Job {job_id} completed successfully")
          elif job.job_status == JobStatus.PARTIALLY_COMPLETED:
              logging.info(f"Job {job_id} partially completed. {job.batches_succeeded} out of {job.total_batches} batches succeeded")
                  
      except Exception as e:
          logging.error('Error updating job and batch status: %s', e)
          safe_db_operation(job_service.update_job_status, job_id, JobStatus.FAILED)
  
  def callback(ch, method, properties, body):
      # do these outside the try-catch so it can update the batch status if there's an error
      # if this parsing logic fails, the batch shouldn't be marked as failed
      data = json.loads(body)
      batch_id, chunks_with_embeddings, vector_db_key = data
  
      if vector_db_key:
          os.environ["VECTOR_DB_KEY"] = vector_db_key
      else:
          logging.info("No vector DB key provided")
      
      try:
          logging.info("Batch retrieved successfully")
          upload_batch(batch_id, chunks_with_embeddings)
          logging.info("Batch processed successfully")
      except Exception as e:
          logging.error('Error processing batch: %s', e)
          update_batch_and_job_status(None, batch_id, BatchStatus.FAILED)
  
      ch.basic_ack(delivery_tag=method.delivery_tag)
  
  def start_connection(max_retries=5, retry_delay=5):
      for attempt in range(max_retries):
          try:
              connection_params = create_connection_params()
              connection = pika.BlockingConnection(connection_params)
              channel = connection.channel()
  
              queue_name = os.getenv('VDB_UPLOAD_QUEUE')
              channel.queue_declare(queue=queue_name)
  
              channel.basic_consume(queue=queue_name, on_message_callback=callback)
  
              logging.info('Waiting for messages.')
              channel.start_consuming()
              return  # If successful, exit the function
              
          except AMQPConnectionError as e:
              logging.error('AMQP Connection Error: %s', e)
          except Exception as e:
              logging.error('Unexpected error: %s', e)
          finally:
              if connection and not connection.is_closed:
                  connection.close()
  
          logging.info('Retrying to connect in %s seconds (Attempt %s/%s)', retry_delay, attempt + 1, max_retries)
          time.sleep(retry_delay)
  
  if __name__ == "__main__":
      while True:
          try:
              start_connection()
          except Exception as e:
              logging.error('Error in start_connection: %s', e)
              logging.info('Restarting start_connection after encountering an error.')
              time.sleep(config.PIKA_RETRY_INTERVAL) 
  
apis: lancedb.connect
extract_api: [((1090, 1155), 'logging.basicConfig', 'logging.basicConfig', ([], {'filename': '"""./vdb-log.txt"""', 'level': 'logging.INFO'}), "(filename='./vdb-log.txt', level=logging.INFO)\n", (1109, 1155), False, 'import logging\n'), ((1156, 1225), 'logging.basicConfig', 'logging.basicConfig', ([], {'filename': '"""./vdb-errors.txt"""', 'level': 'logging.ERROR'}), "(filename='./vdb-errors.txt', level=logging.ERROR)\n", (1175, 1225), False, 'import logging\n'), ((1291, 1343), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.get_batch', 'batch_id'], {}), '(batch_service.get_batch, batch_id)\n', (1308, 1343), False, 'from services.database.database import get_db, safe_db_operation\n'), ((1570, 1622), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.get_batch', 'batch_id'], {}), '(batch_service.get_batch, batch_id)\n', (1587, 1622), False, 'from services.database.database import get_db, safe_db_operation\n'), ((2324, 2370), 'services.database.database.safe_db_operation', 'safe_db_operation', (['job_service.get_job', 'job_id'], {}), '(job_service.get_job, job_id)\n', (2341, 2370), False, 'from services.database.database import get_db, safe_db_operation\n'), ((4945, 4980), 'pymongo.MongoClient', 'pymongo.MongoClient', (['mongo_conn_uri'], {}), '(mongo_conn_uri)\n', (4964, 4980), False, 'import pymongo\n'), ((6003, 6079), 'logging.info', 'logging.info', (['f"""Successfully uploaded {vectors_uploaded} vectors to MongoDB"""'], {}), "(f'Successfully uploaded {vectors_uploaded} vectors to MongoDB')\n", (6015, 6079), False, 'import logging\n'), ((6634, 6660), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (6643, 6660), False, 'import os\n'), ((6665, 6753), 'pinecone.init', 'pinecone.init', ([], {'api_key': 'pinecone_api_key', 'environment': 'vector_db_metadata.environment'}), '(api_key=pinecone_api_key, environment=vector_db_metadata.\n    environment)\n', (6678, 6753), False, 'import pinecone\n'), ((6761, 6810), 'pinecone.GRPCIndex', 'pinecone.GRPCIndex', (['vector_db_metadata.index_name'], {}), '(vector_db_metadata.index_name)\n', (6779, 6810), False, 'import pinecone\n'), ((7471, 7548), 'logging.info', 'logging.info', (['f"""Successfully uploaded {vectors_uploaded} vectors to pinecone"""'], {}), "(f'Successfully uploaded {vectors_uploaded} vectors to pinecone')\n", (7483, 7548), False, 'import logging\n'), ((13446, 13487), 'pymilvus.Collection', 'Collection', (['vector_db_metadata.index_name'], {}), '(vector_db_metadata.index_name)\n', (13456, 13487), False, 'from pymilvus import Collection, connections\n'), ((14146, 14221), 'logging.info', 'logging.info', (['f"""Successfully uploaded {vectors_uploaded} vectors to milvus"""'], {}), "(f'Successfully uploaded {vectors_uploaded} vectors to milvus')\n", (14158, 14221), False, 'import logging\n'), ((14884, 14936), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.get_batch', 'batch_id'], {}), '(batch_service.get_batch, batch_id)\n', (14901, 14936), False, 'from services.database.database import get_db, safe_db_operation\n'), ((14946, 14999), 'lancedb.connect', 'lancedb.connect', (['batch.vector_db_metadata.environment'], {}), '(batch.vector_db_metadata.environment)\n', (14961, 14999), False, 'import lancedb\n'), ((16477, 16554), 'logging.info', 'logging.info', (['f"""Successfully uploaded {vectors_uploaded} vectors to lance db"""'], {}), "(f'Successfully uploaded {vectors_uploaded} vectors to lance db')\n", (16489, 16554), False, 'import logging\n'), ((17756, 17772), 'json.loads', 'json.loads', (['body'], {}), '(body)\n', (17766, 17772), False, 'import json\n'), ((1401, 1492), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.update_batch_retry_count', 'batch.id', '(batch.retries + 1)'], {}), '(batch_service.update_batch_retry_count, batch.id, batch.\n    retries + 1)\n', (1418, 1492), False, 'from services.database.database import get_db, safe_db_operation\n'), ((1494, 1556), 'logging.info', 'logging.info', (['f"""Retrying vector db upload of batch {batch.id}"""'], {}), "(f'Retrying vector db upload of batch {batch.id}')\n", (1506, 1556), False, 'import logging\n'), ((1793, 1886), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.update_batch_status_with_successful_minibatch', 'batch.id'], {}), '(batch_service.\n    update_batch_status_with_successful_minibatch, batch.id)\n', (1810, 1886), False, 'from services.database.database import get_db, safe_db_operation\n'), ((4823, 4849), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (4832, 4849), False, 'import os\n'), ((5302, 5430), 'logging.error', 'logging.error', (['f"""Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}"""'], {}), "(\n    f'Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}'\n    )\n", (5315, 5430), False, 'import logging\n'), ((6837, 6965), 'logging.error', 'logging.error', (['f"""Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}"""'], {}), "(\n    f'Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}'\n    )\n", (6850, 6965), False, 'import logging\n'), ((10133, 10266), 'logging.error', 'logging.error', (['f"""Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}"""'], {}), "(\n    f'Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}'\n    )\n", (10146, 10266), False, 'import logging\n'), ((11215, 11266), 'weaviate.Client', 'weaviate.Client', ([], {'url': 'vector_db_metadata.environment'}), '(url=vector_db_metadata.environment)\n', (11230, 11266), False, 'import weaviate\n'), ((11451, 11584), 'logging.error', 'logging.error', (['f"""Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}"""'], {}), "(\n    f'Collection {vector_db_metadata.index_name} does not exist at cluster URL {vector_db_metadata.environment}'\n    )\n", (11464, 11584), False, 'import logging\n'), ((13141, 13169), 'os.getenv', 'os.getenv', (['"""LOCAL_VECTOR_DB"""'], {}), "('LOCAL_VECTOR_DB')\n", (13150, 13169), False, 'import os\n'), ((13336, 13403), 'pymilvus.connections.connect', 'connections.connect', (['"""default"""'], {'host': 'vector_db_metadata.environment'}), "('default', host=vector_db_metadata.environment)\n", (13355, 13403), False, 'from pymilvus import Collection, connections\n'), ((13519, 13647), 'logging.error', 'logging.error', (['f"""Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}"""'], {}), "(\n    f'Index {vector_db_metadata.index_name} does not exist in environment {vector_db_metadata.environment}'\n    )\n", (13532, 13647), False, 'import logging\n'), ((16856, 16932), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.update_batch_status', 'batch_id', 'batch_status'], {}), '(batch_service.update_batch_status, batch_id, batch_status)\n', (16873, 16932), False, 'from services.database.database import get_db, safe_db_operation\n'), ((16947, 17033), 'services.database.database.safe_db_operation', 'safe_db_operation', (['job_service.update_job_with_batch', 'job_id', 'updated_batch_status'], {}), '(job_service.update_job_with_batch, job_id,\n    updated_batch_status)\n', (16964, 17033), False, 'from services.database.database import get_db, safe_db_operation\n'), ((17925, 17966), 'logging.info', 'logging.info', (['"""No vector DB key provided"""'], {}), "('No vector DB key provided')\n", (17937, 17966), False, 'import logging\n'), ((17989, 18033), 'logging.info', 'logging.info', (['"""Batch retrieved successfully"""'], {}), "('Batch retrieved successfully')\n", (18001, 18033), False, 'import logging\n'), ((18097, 18141), 'logging.info', 'logging.info', (['"""Batch processed successfully"""'], {}), "('Batch processed successfully')\n", (18109, 18141), False, 'import logging\n'), ((19271, 19379), 'logging.info', 'logging.info', (['"""Retrying to connect in %s seconds (Attempt %s/%s)"""', 'retry_delay', '(attempt + 1)', 'max_retries'], {}), "('Retrying to connect in %s seconds (Attempt %s/%s)',\n    retry_delay, attempt + 1, max_retries)\n", (19283, 19379), False, 'import logging\n'), ((19384, 19407), 'time.sleep', 'time.sleep', (['retry_delay'], {}), '(retry_delay)\n', (19394, 19407), False, 'import time\n'), ((168, 193), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (183, 193), False, 'import os\n'), ((5151, 5219), 'logging.error', 'logging.error', (['f"""Error connecting to MongoDB via python client: {e}"""'], {}), "(f'Error connecting to MongoDB via python client: {e}')\n", (5164, 5219), False, 'import logging\n'), ((7844, 7891), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (7868, 7891), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((8226, 8252), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (8235, 8252), False, 'import os\n'), ((9930, 9958), 'os.getenv', 'os.getenv', (['"""LOCAL_VECTOR_DB"""'], {}), "('LOCAL_VECTOR_DB')\n", (9939, 9958), False, 'import os\n'), ((9977, 10005), 'os.getenv', 'os.getenv', (['"""LOCAL_VECTOR_DB"""'], {}), "('LOCAL_VECTOR_DB')\n", (9986, 10005), False, 'import os\n'), ((11181, 11209), 'os.getenv', 'os.getenv', (['"""LOCAL_VECTOR_DB"""'], {}), "('LOCAL_VECTOR_DB')\n", (11190, 11209), False, 'import os\n'), ((12316, 12376), 'logging.error', 'logging.error', (['"""Error writing embeddings to weaviate: %s"""', 'e'], {}), "('Error writing embeddings to weaviate: %s', e)\n", (12329, 12376), False, 'import logging\n'), ((12797, 12844), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (12821, 12844), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((15119, 15259), 'logging.info', 'logging.info', (['f"""Table {batch.vector_db_metadata.index_name} does not exist in environment {batch.vector_db_metadata.environment}."""'], {}), "(\n    f'Table {batch.vector_db_metadata.index_name} does not exist in environment {batch.vector_db_metadata.environment}.'\n    )\n", (15131, 15259), False, 'import logging\n'), ((16740, 16792), 'services.database.database.safe_db_operation', 'safe_db_operation', (['batch_service.get_batch', 'batch_id'], {}), '(batch_service.get_batch, batch_id)\n', (16757, 16792), False, 'from services.database.database import get_db, safe_db_operation\n'), ((17092, 17144), 'logging.info', 'logging.info', (['f"""Job {job_id} completed successfully"""'], {}), "(f'Job {job_id} completed successfully')\n", (17104, 17144), False, 'import logging\n'), ((17391, 17450), 'logging.error', 'logging.error', (['"""Error updating job and batch status: %s"""', 'e'], {}), "('Error updating job and batch status: %s', e)\n", (17404, 17450), False, 'import logging\n'), ((17459, 17533), 'services.database.database.safe_db_operation', 'safe_db_operation', (['job_service.update_job_status', 'job_id', 'JobStatus.FAILED'], {}), '(job_service.update_job_status, job_id, JobStatus.FAILED)\n', (17476, 17533), False, 'from services.database.database import get_db, safe_db_operation\n'), ((18177, 18223), 'logging.error', 'logging.error', (['"""Error processing batch: %s"""', 'e'], {}), "('Error processing batch: %s', e)\n", (18190, 18223), False, 'import logging\n'), ((18485, 18511), 'services.rabbitmq.rabbit_service.create_connection_params', 'create_connection_params', ([], {}), '()\n', (18509, 18511), False, 'from services.rabbitmq.rabbit_service import create_connection_params\n'), ((18537, 18579), 'pika.BlockingConnection', 'pika.BlockingConnection', (['connection_params'], {}), '(connection_params)\n', (18560, 18579), False, 'import pika\n'), ((18649, 18678), 'os.getenv', 'os.getenv', (['"""VDB_UPLOAD_QUEUE"""'], {}), "('VDB_UPLOAD_QUEUE')\n", (18658, 18678), False, 'import os\n'), ((18827, 18864), 'logging.info', 'logging.info', (['"""Waiting for messages."""'], {}), "('Waiting for messages.')\n", (18839, 18864), False, 'import logging\n'), ((4463, 4510), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (4487, 4510), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((5915, 5969), 'logging.error', 'logging.error', (['"""Error writing embeddings to Mongo:"""', 'e'], {}), "('Error writing embeddings to Mongo:', e)\n", (5928, 5969), False, 'import logging\n'), ((6345, 6392), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (6369, 6392), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((7380, 7437), 'logging.error', 'logging.error', (['"""Error writing embeddings to pinecone:"""', 'e'], {}), "('Error writing embeddings to pinecone:', e)\n", (7393, 7437), False, 'import logging\n'), ((8456, 8582), 'logging.error', 'logging.error', (['f"""Index {vector_db_metadata.index_name} does not exist at redis URL {vector_db_metadata.environment}"""'], {}), "(\n    f'Index {vector_db_metadata.index_name} does not exist at redis URL {vector_db_metadata.environment}'\n    )\n", (8469, 8582), False, 'import logging\n'), ((9790, 9816), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (9799, 9816), False, 'import os\n'), ((10674, 10729), 'logging.error', 'logging.error', (['"""Error writing embeddings to qdrant:"""', 'e'], {}), "('Error writing embeddings to qdrant:', e)\n", (10687, 10729), False, 'import logging\n'), ((13281, 13307), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (13290, 13307), False, 'import os\n'), ((14054, 14112), 'logging.error', 'logging.error', (['"""Error writing embeddings to milvus: %s"""', 'e'], {}), "('Error writing embeddings to milvus: %s', e)\n", (14067, 14112), False, 'import logging\n'), ((14499, 14546), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (14523, 14546), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((15739, 15872), 'logging.info', 'logging.info', (['f"""Created table {batch.vector_db_metadata.index_name} in environment {batch.vector_db_metadata.environment}."""'], {}), "(\n    f'Created table {batch.vector_db_metadata.index_name} in environment {batch.vector_db_metadata.environment}.'\n    )\n", (15751, 15872), False, 'import logging\n'), ((15889, 16016), 'logging.error', 'logging.error', (['f"""Embeddings type {batch.embeddings_metadata.embeddings_type} not supported for LanceDB. Only Open AI"""'], {}), "(\n    f'Embeddings type {batch.embeddings_metadata.embeddings_type} not supported for LanceDB. Only Open AI'\n    )\n", (15902, 16016), False, 'import logging\n'), ((16386, 16443), 'logging.error', 'logging.error', (['"""Error writing embeddings to lance db:"""', 'e'], {}), "('Error writing embeddings to lance db:', e)\n", (16399, 16443), False, 'import logging\n'), ((17219, 17348), 'logging.info', 'logging.info', (['f"""Job {job_id} partially completed. {job.batches_succeeded} out of {job.total_batches} batches succeeded"""'], {}), "(\n    f'Job {job_id} partially completed. {job.batches_succeeded} out of {job.total_batches} batches succeeded'\n    )\n", (17231, 17348), False, 'import logging\n'), ((19024, 19069), 'logging.error', 'logging.error', (['"""AMQP Connection Error: %s"""', 'e'], {}), "('AMQP Connection Error: %s', e)\n", (19037, 19069), False, 'import logging\n'), ((19113, 19153), 'logging.error', 'logging.error', (['"""Unexpected error: %s"""', 'e'], {}), "('Unexpected error: %s', e)\n", (19126, 19153), False, 'import logging\n'), ((19539, 19588), 'logging.error', 'logging.error', (['"""Error in start_connection: %s"""', 'e'], {}), "('Error in start_connection: %s', e)\n", (19552, 19588), False, 'import logging\n'), ((19601, 19673), 'logging.info', 'logging.info', (['"""Restarting start_connection after encountering an error."""'], {}), "('Restarting start_connection after encountering an error.')\n", (19613, 19673), False, 'import logging\n'), ((19686, 19724), 'time.sleep', 'time.sleep', (['config.PIKA_RETRY_INTERVAL'], {}), '(config.PIKA_RETRY_INTERVAL)\n', (19696, 19724), False, 'import time\n'), ((8898, 8925), 'numpy.array', 'np.array', (['upsert_list[2][i]'], {}), '(upsert_list[2][i])\n', (8906, 8925), True, 'import numpy as np\n'), ((9410, 9457), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (9434, 9457), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((11977, 12024), 'shared.utils.generate_uuid_from_tuple', 'generate_uuid_from_tuple', (['(job_id, batch_id, i)'], {}), '((job_id, batch_id, i))\n', (12001, 12024), False, 'from shared.utils import generate_uuid_from_tuple\n'), ((11109, 11135), 'os.getenv', 'os.getenv', (['"""VECTOR_DB_KEY"""'], {}), "('VECTOR_DB_KEY')\n", (11118, 11135), False, 'import os\n'), ((15416, 15427), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (15425, 15427), True, 'import pyarrow as pa\n'), ((15544, 15555), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (15553, 15555), True, 'import pyarrow as pa\n'), ((15606, 15617), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (15615, 15617), True, 'import pyarrow as pa\n'), ((4137, 4230), 'logging.error', 'logging.error', (['"""Unsupported vector DB type: %s"""', 'vector_db_metadata.vector_db_type.value'], {}), "('Unsupported vector DB type: %s', vector_db_metadata.\n    vector_db_type.value)\n", (4150, 4230), False, 'import logging\n'), ((15478, 15490), 'pyarrow.float32', 'pa.float32', ([], {}), '()\n', (15488, 15490), True, 'import pyarrow as pa\n')]