# python chromadb_add.py "e:/insert.json" "e:/chromadb/dbname" "collection_name"
import sys
import chromadb
import json
from chromadb.config import Settings

argvlen = len(sys.argv)
print(argvlen)
if argvlen == 4:
    jsonfilename = sys.argv[1]
    localdbpath = sys.argv[2]
    collectionname = sys.argv[3]
	
	# 设置本地数据库名称（Persistent本地持久化）和表名（Collection集合）
    CHROMA_CLIENT = chromadb.PersistentClient(
        path = localdbpath,
        settings = chromadb.Settings(anonymized_telemetry=False)
    )
    COLLECTION = CHROMA_CLIENT.get_or_create_collection(
        name = collectionname,
        metadata={"hnsw:space": "cosine"}
    )
	# 加载json文件
    with open(jsonfilename, 'r', encoding='utf-8') as f:
        jsoninsert = json.load(f)
        ids = jsoninsert['ids']
        embeddings = jsoninsert['embeddings']
        documents = jsoninsert['chunks']
        metadatas = jsoninsert['metadatas']
        # 插入到数据库中
        COLLECTION.upsert(
            ids=ids,
            embeddings=embeddings,
            documents=documents,
			metadatas=metadatas)
else:
    print("python chromadb_add.py \"e:/insert.json\" \"e:/chromadb/dbname\" \"collection_name\"") 
# python chromadb_add.py "e:/insert.json" "e:/chromadb/dbname" "collection_name"
"""
# The sentences to encode
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]

# 2. Calculate embeddings by calling model.encode()
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 384]

# 3. Calculate the embedding similarities
similarities = model.similarity(embeddings, embeddings)
# tensor([[1.0000, 0.6660, 0.1046],
#         [0.6660, 1.0000, 0.1411],
#         [0.1046, 0.1411, 1.0000]])
"""