from smart_chart.common.smartvector import SmartVectorDB, Text2VecEmbeddingFunction
text_vector = Text2VecEmbeddingFunction()

def dataset(*args, **kwargs):
    promote = args[0][0]
    db_config = args[1]
    table = db_config.get('table', 'vectors_key')
    return SmartVectorDB(db_config=db_config, text_vector=text_vector, table=table).get(promote)


def insert_dataset(*args, **kwargs):
    contents = args[0]
    table = args[1]
    connect_dict = args[3]
    docIndex = contents[0].index('document')
    contents[0].append('embedding')
    for item in contents[1:]:
        item.append(str(text_vector(item[docIndex])[0]))
    SmartVectorDB(db_config=connect_dict, text_vector=text_vector)._execute_load(contents, table)
    return len(contents) - 1