test_query.py

#!/usr/bin/env python3
"""
Test the query_codebase method to get structured results.
"""
import os
import json
from app.main import CodeEmbedder

def main():
    """Run the test query."""
    # Create a data directory for the database
    data_dir = os.path.join(os.path.dirname(__file__), "data")
    os.makedirs(data_dir, exist_ok=True)
    
    # Initialize the code embedder
    print("Initializing CodeEmbedder...")
    embedder = CodeEmbedder(data_dir=data_dir)
    
    # Use the query_codebase method
    query = "how to process CSV data"
    print(f"Querying codebase with: '{query}'")
    
    results = embedder.query_codebase(query, top_k=3)
    
    # Pretty print the structured results
    print("\nStructured Results:")
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()