import logging import time import csv from django.core.management.base import BaseCommand from core.text2sql.handler import QueryDataHandler from core.text2sql.prompt import get_prompt from core.text2sql.eval_queries import queries logger = logging.getLogger(__name__) class Command(BaseCommand): help = "Morningstar API to save the JSON response to a file which contains secIds with other details" def handle(self, *args, **options) -> None: t1 = time.perf_counter() q=[] count =1 for query in queries[26:]: print("count: ", query["Query Number"]) prompt = get_prompt(query["Query Description"]) logger.info(f"Prompt: {prompt}") generated_query, data = QueryDataHandler().get_data_from_query(prompt) print(f"Description: {query['Query Description']}, Query: {query.get('SQL Statement')}, Generated: {generated_query} ") q.append({ "Query Number": query["Query Number"], "Complexity Level": query["Complexity Level"], "Description": query["Query Description"], "Query": query.get("SQL Statement", "-"), "Generated": generated_query, }) count+=1 time.sleep(1) csv_file_path = 'queries_data.csv' # Writing data to CSV with open(csv_file_path, 'w', newline='', encoding='utf-8') as csv_file: fieldnames = q[0].keys() print(fieldnames) writer = csv.DictWriter(csv_file, fieldnames=fieldnames) # Write the header writer.writeheader() # Write the data writer.writerows(q) print(f'Data has been written to {csv_file_path}.') self.stdout.write(f"Time taken for evaluation: {time.perf_counter() - t1}")