def testdb()

in devai-cli/src/devai/commands/rag/load.py [0:0]

• 27 lines of code
• 7 McCabe index (conditional complexity)

def testdb(db_path, qry):
    # 1. Load existing ChromaDB if it exists
    db = None
    persist_directory = db_path
    if Path(persist_directory).exists():
        # Assuming same embeddings were used to create the DB
        EMBEDDING_QPM = 100
        EMBEDDING_NUM_BATCH = 5
        embeddings = VertexAIEmbeddings(
            requests_per_minute=EMBEDDING_QPM,
            num_instances_per_batch=EMBEDDING_NUM_BATCH,
            model_name="textembedding-gecko@latest",
        )
        db = Chroma(
            persist_directory=persist_directory,
            embedding_function=embeddings,
            collection_name="source_code_embeddings"
        )

    # 2. Simple test if DB loaded
    if db:
        # Sample query
        sample_query = qry
        docs = db.similarity_search(sample_query)

        if docs:
            print("Database loaded successfully. Sample query results:")
            for doc in docs:
                # Print a snippet of each relevant document
                print(doc.page_content[:100])
        else:
            print("Database loaded, but no results for sample query.")
    else:
        print("Database could not be loaded.")