# python chromadb_query.py "question" "e:/search.json" "e:/chromadb/dbname1" "collection_name" 3
import sys
import chromadb
import json
from sentence_transformers import SentenceTransformer
from chromadb.config import Settings

argvlen = len(sys.argv)
retnum = 3
if argvlen == 7:
    questionfile = sys.argv[1]
    filenameout = sys.argv[2]
    modepath = sys.argv[3]
    localdbpaht = sys.argv[4]
    collectionname = sys.argv[5]
    retnum = int(sys.argv[6])
    
    # 设置本地数据库名称（Persistent本地持久化）和表名（Collection集合）
    model = SentenceTransformer(modepath)
    CHROMA_CLIENT = chromadb.PersistentClient(
        path = localdbpaht,
        settings=chromadb.Settings(anonymized_telemetry=False)
    )
    COLLECTION = CHROMA_CLIENT.get_or_create_collection(collectionname)
    with open(questionfile, "r", encoding="utf-8") as f:
        question = f.read()
    # 读取查询的问题
    if len(question) > 0:
        #print(question)
        query_embedding = model.encode([question]).tolist()
        # Chroma检索
        results = COLLECTION.query(
	        query_embeddings=query_embedding,
            n_results = retnum
        )
        
        json_string = json.dumps(results, ensure_ascii=False)
        #print(json_string)
		# 将查询的json结果写入输出文件
        with open(filenameout, 'wb') as filewrite:
           textuft8 = json_string.encode('utf-8')
           filewrite.write(textuft8)
else:
    print("python chromadb_query.py \"e:/question.txt\" \"e:/search.json\" \"e:/chromadb/dbname1\" \"collection_name\" retnum")          

# python chromadb_query.py "question" "e:/search.json" "e:/chromadb/dbname1" "collection_name" 3