use gemini + read from file and write to file

Author: cded

generateInserts.py

@@ -1,9 +1,17 @@
 import re
-import requests
 import logging
+from dotenv import load_dotenv
+import os
+import google.generativeai as genai
+
+load_dotenv()
+
 
 logging.basicConfig(level=logging.NOTSET)
 
+genai.configure(api_key=os.environ['GEMINI_API_KEY'])
+model = genai.GenerativeModel('gemini-pro')
+
 
 def extract_table_columns(sql_content):
     create_table_pattern = re.compile(
@@ -50,7 +58,15 @@ def generate_insert_statements(table_columns):
             f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({values_placeholder});")
 
 
-def generate_insert_statement_from_ai(sql_content):
+def extract_inserts(response):
+    sql_pattern = re.compile(r'```sql(.*?)```', re.DOTALL)
+    sql_matches = sql_pattern.findall(response)
+    return sql_matches
+
+
+def generate_insert_statement_from_ai(sql_content: str):
+    if not sql_content.lower().startswith('create table'):
+        return
     logging.info(f"Generating INSERT Statements for table {sql_content}")
     prompt = (
         f"Generate me random fake data for the following tables:"
@@ -61,10 +77,25 @@ def generate_insert_statement_from_ai(sql_content):
         f"Values (...)"
         f"(...)"
         f"Make 10 statements"
+        f"Note that bigserial is a number. Each row table has an id."
     )
-    res = requests.get(f"https://api.freegpt4.ddns.net/?text={prompt}")
-    print(res)
-    print(res.text)
+
+    response = model.generate_content(prompt)
+    sql_inserts = ""
+    try:
+        sql_inserts = extract_inserts(response.text)
+    except:
+        print("EXCEPTTTT.....")
+        print(response)
+        sql_inserts = extract_inserts(
+            response.candidates[0].content.parts)
+
+    # Write the extracted SQL statements to a file
+    output_file_path = 'extracted_sql_statements.sql'
+    with open(output_file_path, 'a') as output_file:
+        for sql_statement in sql_inserts:
+            output_file.write(sql_statement.strip())
+
     return ''
 
 
@@ -73,23 +104,4 @@ def generate_insert_statement_from_ai(sql_content):
     sql_content = []
     with open(file_path, 'r') as sql_file:
         for line in sql_file:
-            sql_content.append(line)
-    print(sql_content)
-
-    insert_statement = generate_insert_statement_from_ai(sql_content)
-
-    # sql_content = """
-    #     create table "address" (
-    #         "customer_id" bigint,
-    #         "id" bigint not null,
-    #         "address_line" varchar(255),
-    #         "city" varchar(255),
-    #         "coordinates" varchar(255),
-    #         "country" varchar(255),
-    #         "postal_code" varchar(255),
-    #         "street_number" varchar(255),
-    #         "unit_number" varchar(255),
-    #         primary key ("id"),
-    #         foreign key ("customer_id") references "customer"
-    #     );
-    #     """
+            generate_insert_statement_from_ai(line)

my-schema.sql

@@ -1,3 +1,4 @@
+create table "product" ("price" float(53) not null, "category_id" bigint, "id" bigserial not null, "menu_id" bigint, "description" varchar(255), "name" varchar(255), primary key ("id"));
 create table "address" ("id" bigint not null, "address_line" varchar(255), "city" varchar(255), "coordinates" varchar(255), "country" varchar(255), "postal_code" varchar(255), "street_number" varchar(255), "unit_number" varchar(255), primary key ("id"));
 create table "billing_info" ("id" bigserial not null, "order_id" bigint unique, "credit_card_number" varchar(255), primary key ("id"));
 create table "customer" ("user_id" bigint not null, primary key ("user_id"));
@@ -8,7 +9,6 @@ create table "menu" ("category_id" bigint, "id" bigserial not null, "restaurant_
 create table "order" ("customer_id" bigint, "delivery_id" bigint unique, "id" bigserial not null, "order_date" timestamp(6), "restaurant_id" bigint, "order_status" varchar(255), primary key ("id"));
 create table "order_item" ("quantity" integer not null, "id" bigserial not null, "order_id" bigint, "product_id" bigint, primary key ("id"));
 create table "payment_transaction" ("amount" numeric(38,2), "id" bigserial not null, "order_id" bigint unique, "payment_gateway" varchar(255), primary key ("id"));
-create table "product" ("price" float(53) not null, "category_id" bigint, "id" bigserial not null, "menu_id" bigint, "description" varchar(255), "name" varchar(255), primary key ("id"));
 create table "product_category" ("id" bigserial not null, "name" varchar(255), primary key ("id"));
 create table "product_image" ("id" bigserial not null, "product_id" bigint, "image_url" varchar(255), primary key ("id"));
 create table "restaurant" ("address_id" bigint unique, "id" bigint not null, "name" varchar(255), "operating_hours" varchar(255), "phone_number" varchar(255), primary key ("id"));

requirements.txt

@@ -1,2 +1,21 @@
+cachetools==5.3.2
+certifi==2023.11.17
+charset-normalizer==3.3.2
+google-ai-generativelanguage==0.4.0
+google-api-core==2.15.0
+google-auth==2.26.2
+google-generativeai==0.3.2
+googleapis-common-protos==1.62.0
+grpcio==1.60.0
+grpcio-status==1.60.0
+idna==3.6
+proto-plus==1.23.0
+protobuf==4.25.2
+pyasn1==0.5.1
+pyasn1-modules==0.3.0
+python-dotenv==1.0.0
 requests==2.31.0
-urllib3==1.26.18
+rsa==4.9
+tqdm==4.66.1
+typing_extensions==4.9.0
+urllib3==2.1.0

roadmap

@@ -1,7 +1,18 @@
+Dec 28
 Read from .sql file
 Iterate through rows
 Validate constraints
 Useful logging with extract columns
 Try other models - llama
 Make it parallel with threads
-How to deal better with FKs
+How to deal better with FKs
+
+Jan 12
+Some rows don't have ids
+
+- Make it parallel with threads
+  How to deal better with FKs
+
+* less important
+  Validate constraints
+  Useful logging with extract columns