import os
from pymilvus import connections


def connect():
    # Initialize a MilvusClient instance
    # Replace uri and API key with your own
    connections.connect(
        alias=os.getenv('MILVUS_DEFAULT'),
        uri=os.getenv("MILVUS_CLUSTER_ENDPOINT"), # Cluster endpoint obtained from the console
        token=os.getenv("MILVUS_TOKEN")  # API key or a colon-separated cluster username and password
    )

def disconnect():
    connections.disconnect("MILVUS_DEFAULT")