from pymilvus import FieldSchema, DataType

# milvus 向量数据库
fields = [
    FieldSchema(name="chunk_id", dtype=DataType.INT64, is_primary=True, auto_id=False),
    FieldSchema(name="doc_id", dtype=DataType.VARCHAR, max_length=64),
    # 三种向量字段
    FieldSchema(name="vector_text", dtype=DataType.FLOAT_VECTOR, dim=2560),
    FieldSchema(name="vector_summary", dtype=DataType.FLOAT_VECTOR, dim=2560),
    FieldSchema(name="vector_questions", dtype=DataType.FLOAT_VECTOR, dim=2560),
    # metadata
    FieldSchema(name="topic", dtype=DataType.VARCHAR, max_length=255),
    FieldSchema(name="domain", dtype=DataType.VARCHAR, max_length=100),
    FieldSchema(name="task_type", dtype=DataType.VARCHAR, max_length=100),
    FieldSchema(name="summary", dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(
        name="keywords",
        dtype=DataType.ARRAY,
        element_type=DataType.VARCHAR,
        max_length=100,
    ),
    FieldSchema(
        name="concepts",
        dtype=DataType.ARRAY,
        element_type=DataType.VARCHAR,
        max_length=100,
    ),
    FieldSchema(
        name="questions",
        dtype=DataType.ARRAY,
        element_type=DataType.VARCHAR,
        max_length=200,
    ),
    FieldSchema(name="topic_purity", dtype=DataType.FLOAT),
    FieldSchema(name="tokens", dtype=DataType.INT64),
]


__all__ = ["fields"]