seed_memory.py

import os
from vanna_setup import get_agent

def seed_and_test():
    vn = get_agent()
    print("--- Starting Seeding ---")

    print("Memory Seeding Complete. Running Tests...")

    test_questions = [
        "How many patients do we have?", "List all doctors and their specializations",
        "Show me appointments for last month", "Which doctor has the most appointments?",
        "What is the total revenue?", "Show revenue by doctor",
        "How many cancelled appointments last quarter?", "Top 5 patients by spending",
        "Average treatment cost by specialization", "Show monthly appointment count",
        "Which city has the most patients?", "List patients who visited more than 3 times",
        "Show unpaid invoices", "What percentage of appointments are no-shows?",
        "Show the busiest day of the week for appointments", "Revenue trend by month",
        "Average appointment duration by doctor", "List patients with overdue invoices",
        "Compare revenue between departments", "Show patient registration trend"
    ]

    with open("RESULTS.md", "w") as f:
        f.write("# NL2SQL Test Results\n\n| # | Question | Generated SQL | Status |\n|---|---|---|---|\n")
        
        for i, q in enumerate(test_questions, 1):
            print(f"Testing {i}/20: {q}")
            try:
                prompt = f"System: You are a SQLite expert. Convert the user question to a valid SQL query. Return ONLY the SQL code.\nUser: {q}"
                
            
                payload = vn.llm_service._build_payload(prompt)
                response = vn.llm_service.send_request(payload)
                
                sql_out = vn.llm_service._parse_response(response)
                
                sql_clean = str(sql_out).replace('```sql', '').replace('```', '').replace('\n', ' ').strip()
                f.write(f"| {i} | {q} | `{sql_clean}` | Success |\n")
            except Exception as e:
                f.write(f"| {i} | {q} | N/A | Error: {str(e)} |\n")

    print("\nSUCCESS: RESULTS.md generated.")

if __name__ == "__main__":
    seed_and_test()