tests/test_db.py

import json
import logging

from db import SanatanDatabase
from metadata import MetadataFilter, MetadataWhereClause


if __name__ == "__main__":
    logging.basicConfig()
    collection_name = "divya_prabandham"
    database = SanatanDatabase()
    print("count = ", database.count(collection_name))
    while True:
        query = input("Search for: ")
        if query.strip() == "":
            break
        # response = database.search_for_literal(
        #     collection_name=collection_name, literal_to_search_for=query, n_results=1
        # )
        metadata_where_clause = MetadataWhereClause(
            filters=[
                MetadataFilter(
                    metadata_field="prabandham_code",
                    metadata_search_operator="$eq",
                    metadata_value="NAT",
                ),
                # MetadataFilter(
                #     metadata_field="decade",
                #     metadata_search_operator="$eq",
                #     metadata_value=4,
                # ),
                MetadataFilter(
                    metadata_field="position_in_chapter",
                    metadata_search_operator="$eq",
                    metadata_value=5,
                ),
            ]
        )
        response = database.search(
            collection_name=collection_name,
            query=query,
            metadata_where_clause=metadata_where_clause,
            n_results=5,
        )
        print("Matches", json.dumps(response["metadatas"][0], ensure_ascii=False, indent=1))
        # print("Document: ")
        # print(response["documents"][0][0])
        # print("Metadata: ")
        # print(response["metadatas"][0][0])