from pymilvus import CollectionSchema, FieldSchema, DataType
from pymilvus import Collection, db, connections,utility

def create_collections(collections_name,database_name : str ="JXTest"):
  db.using_database(database_name)

  id = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True)
  text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096)
  header_1 = FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  header_2 = FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  header_3 = FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024,)

  schema = CollectionSchema(
      fields=[id, header_1, header_2, header_3, text, vector],
      auto_id=True,
      enable_dynamic_field=True
  )
  index_params = {
      "metric_type": "L2",
      "index_type": "IVF_FLAT", #FLAT
      "params": {"nlist": 128}
  }
  collection = Collection(name=collections_name, schema=schema, using='default', shards_num=2)
  collection.create_index(field_name="vector", index_params=index_params)

  utility.index_building_progress("word_vector")
  return "collections created"