mvp coded to bare minimum of the app working

Author: Arsalanjdev

.idea/dataSources.xml

@@ -23,3 +23,18 @@ def send_message_to_api(self, message: str) -> str:
             ],
         )
         return completion.choices[0].message["content"]
+
+    def prompt_fake_data(self, db_schema: str, num_rows: int = 10) -> str:
+        """
+        Generates a prompt based on the DB information it recieves
+        :return: str. the response of ai api.
+        """
+        prompt = f"""
+        Generate {num_rows} rows of fake sql data for the following table:column schema. Make the data look believable and realistic (not john doe or anything like that). 
+        Your response should be only in plain json with no words before or after. table:column schema: {db_schema}
+        """
+        return self.send_message_to_api(prompt)
+
+    def prompt_insert_statements(self, jsondata: str, table_name: str):
+        prompt = f"""Given the following jsondata that is to be inserted into a SQL table (table_name:{table_name}, generate insert statements for it. Your output should only be insert statements without any sentences before or after. data: {jsondata}"""
+        return self.send_message_to_api(prompt)

src/rowgen/hf_api.py

@@ -3,6 +3,9 @@
 from hf_api import HFapi
 import argparse
 
+from rowgen.extract_from_db import DBconnect
+from rowgen.parser import parse_sql_from_code_block
+
 
 def main():
     parser = argparse.ArgumentParser(description="RowGen CLI")
@@ -25,10 +28,20 @@ def main():
     else:
         db_url = "{args.db_type}://{args.user}:{args.password}@{args.host}:{args.port}/{args.database}"
 
-    engine = create_engine(db_url)
-    with engine.connect() as conn:
-        result = conn.execute(text("SELECT 1"))
-        print("Connected to database:", result.scalar())
+    dbc = DBconnect(db_url)
+    hf = HFapi()
+    data = hf.prompt_fake_data(dbc.table_columns, 20)
+    inserts = hf.prompt_insert_statements(data, table_name="users")
+    # print(inserts)
+    sql_parser = parse_sql_from_code_block(inserts)
+    sql_parser = sql_parser.split("\n")
+    print(sql_parser)
+    # execute
+    engine = create_engine("sqlite:///testrowgendb.sqlite")  # replace with your DB URL
+    with engine.connect() as connection:
+        for sql in sql_parser:
+            connection.execute(text(sql))
+        connection.commit()
 
 
 if __name__ == "__main__":

src/rowgen/json_parser.py

@@ -0,0 +1,72 @@
+import json
+
+
+def trim_code_block(text: str, language: str = "") -> str:
+    """
+    Removes markdown code block fences like ```json ... ```
+    """
+    prefix = f"```{language}" if language else "```"
+
+    if text.startswith(prefix):
+        text = text.removeprefix(prefix).strip()
+    elif text.startswith("```"):
+        text = text.removeprefix("```").strip()
+
+    if text.endswith("```"):
+        text = text.removesuffix("```").strip()
+
+    return text
+
+
+def parse_json_from_code_block(text: str) -> dict | None:
+    """
+    Trim and parse JSON text from code block
+    """
+    cleaned = trim_code_block(text, "json")
+    try:
+        return json.loads(cleaned)
+    except json.JSONDecodeError as e:
+        print("Failed to parse JSON:", e)
+        print(cleaned)
+        return None
+
+
+def parse_sql_from_code_block(text: str) -> str:
+    """
+    Return cleaned SQL query string
+    """
+    return trim_code_block(text, "sql")
+
+
+def save_json(data: dict, file_path: str = "output.json"):
+    with open(file_path, "w") as f:
+        json.dump(data, f)
+
+
+# print(
+#     parse_sql_from_code_block(
+#         """```sql
+# INSERT INTO users (id, name, username, rank) VALUES
+# (1, 'Alice Johnson', 'alicej', 3),
+# (2, 'Bob Smith', 'bobsmith', 1),
+# (3, 'Charlie Brown', 'charlieb', 2),
+# (4, 'Diana Prince', 'dianap', 4),
+# (5, 'Evan Wright', 'evanw', 5),
+# (6, 'Fiona Gallagher', 'fionag', 2),
+# (7, 'George King', 'georgek', 3),
+# (8, 'Hannah Montana', 'hannahm', 1),
+# (9, 'Ian McKellen', 'ianm', 4),
+# (10, 'Jessica Jones', 'jessicaj', 5),
+# (11, 'Kevin Hart', 'kevinh', 2),
+# (12, 'Laura Palmer', 'laurap', 3),
+# (13, 'Michael Scott', 'michaels', 1),
+# (14, 'Nancy Drew', 'nancyd', 4),
+# (15, 'Oscar Wilde', 'oscarw', 5),
+# (16, 'Pam Beesly', 'pamb', 2),
+# (17, 'Quincy Jones', 'quincyj', 3),
+# (18, 'Rachel Green', 'rachelg', 1),
+# (19, 'Steve Rogers', 'stever', 4),
+# (20, 'Tina Fey', 'tinaf', 5);
+# ```"""
+#     )
+# )

src/rowgen/main.py

@@ -1,5 +1,5 @@
 import json
-from rowgen.json_parser import JsonParse
+from rowgen.parser import parse_sql_from_code_block, parse_json_from_code_block
 from rowgen.hf_api import HFapi
 from pytest import fixture
 
@@ -13,9 +13,9 @@ def test_rows_json(hfapi):
     raw_query = hfapi.send_message_to_api(
         "This is my database: id, name, username, email, access_level. generate 5 fake data rows for it. don't say anything just generate in json. i want your whole message to be just a json"
     )
-    jsp = JsonParse(raw_query)
+    jsp = parse_json_from_code_block(raw_query)
     try:
-        json.loads(jsp.processed_query)
+        json.loads(jspy)
         return True
     except json.JSONDecodeError:
         return False