import json
import chromadb

def test_chromadb():
    # 初始化
    client = chromadb.HttpClient()
    # 创建之前先删除 collection，并且二次运行代码报错
    client.delete_collection("sample_collection")
    # 创建 collection
    collection = client.create_collection("sample_collection")
    # Add docs to the collection. Can also update and delete. Row-based API coming soon!
    collection.add(
        documents=["This is document1", "This is document2"],  # we embed for you, or bring your own
        metadatas=[{"source": "notion"}, {"source": "google-docs"}],  # filter on arbitrary metadata!
        ids=["doc1", "doc2"],  # must be unique for each doc
    )
    results = collection.query(
        query_texts=["This is a query document"],
        n_results=2,
        # where={"metadata_field": "is_equal_to_this"}, # optional filter
        # where_document={"$contains":"search_string"}  # optional filter
    )
    print(json.dumps(results))


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    test_chromadb()

def test_chromadb():
    client = chromadb.HttpClient()
    client.delete_collection("sample_collection")
    collection = client.create_collection("sample_collection")
    collection.add(
    # 增加一条新的文档 video
        documents=["This is document1", "This is document2", "This is a video"],
        metadatas=[{"source": "notion"}, {"source": "google-docs"}, {"source": "test-dir"}],
        ids=["doc1", "doc2", 'videoDir'],
    )
    results = collection.query(
        query_texts=["This is a query document"],
        # 查询结果，最多返回 10 条
        n_results=10,
    )
    print(json.dumps(results))

def test_chromadb():
    # Use a breakpoint in the code line below to debug your script.
    client = chromadb.HttpClient()
    client.delete_collection("sample_collection")
    collection = client.create_collection("sample_collection", metadata={
        "hnsw:space": "cosine", # 改成使用余弦相似度
    })
    # Add docs to the collection. Can also update and delete. Row-based API coming soon!
    collection.add(
        documents=["This is document1", "This is document2", "This is a video"],  # we embed for you, or bring your own
        metadatas=[{"source": "notion"}, {"source": "google-docs"}, {"source": "test-dir"}],
        # filter on arbitrary metadata!
        ids=["doc1", "doc2", 'videoDir'],  # must be unique for each doc
    )
    results = collection.query(
        query_texts=["This is a query document"],
        n_results=10,
        # where={"metadata_field": "is_equal_to_this"}, # optional filter
        # where_document={"$contains":"search_string"}  # optional filter
    )
    print(json.dumps(results))