jsonparser added along tests

Author: Arsalanjdev

src/rowgen/db.json

@@ -0,0 +1,3 @@
+"""
+Extracts information (schema) from a SQL database.
+"""

src/rowgen/extract_from_db.py

@@ -1,6 +1,6 @@
 from huggingface_hub import InferenceClient
 
-from .utils import API_KEY
+from rowgen.utils import API_KEY
 
 
 class HFapi:

src/rowgen/hf_api.py

@@ -0,0 +1 @@
+""" """

src/rowgen/insert_todb.py

@@ -0,0 +1,46 @@
+import json
+
+
+# class JsonDB:
+#     def __init__(self):
+#         with open("db.json", "r") as f:
+#             self._json_raw = json.load(f)
+#
+#     def get_column_names(self):
+#         return self._json_raw["columns"]
+#
+#     def __getitem__(self, item):
+#         return self._json_raw[item]
+
+
+class JsonParse:
+    def __init__(self, query: str):
+        self._raw_query = query
+        self.processed_query = JsonParse.trun_query(self._raw_query)
+
+    @staticmethod
+    def trun_query(query: str) -> str:
+        if query.startswith("```json"):
+            query = query.removeprefix("```json").strip()
+        elif query.startswith("```"):
+            query = query.removeprefix("```").strip()
+
+        if query.endswith("```"):
+            query = query.removesuffix("```").strip()
+
+        return query
+
+    @property
+    def get_processed_query(self):
+        return self.processed_query
+
+    def save_to_json(self, file_path: str = "output.json"):
+        try:
+            data = json.loads(self.processed_query)
+        except json.JSONDecodeError as e:
+            print("failed to parse json", e)
+            print(self.processed_query)
+            return
+
+        with open(file_path, "w") as f:
+            json.dump(data, f)

src/rowgen/json_parser.py

@@ -1,51 +0,0 @@
-import json
-import os
-
-from huggingface_hub import InferenceClient
-import sqlalchemy
-from dotenv import load_dotenv
-from utils import API_KEY
-
-
-class JsonDB:
-    def __init__(self):
-        with open("db.json", "r") as f:
-            self._json_raw = json.load(f)
-
-    def get_column_names(self):
-        return self._json_raw["columns"]
-
-    def __getitem__(self, item):
-        return self._json_raw[item]
-
-
-print(JsonDB()["columns"])
-
-
-class JsonParse:
-    def __init__(self, query: str):
-        self._raw_query = query
-        self.processed_query = JsonParse.trun_query(self._raw_query)
-
-    @staticmethod
-    def trun_query(query: str) -> str:
-        if query.startswith("```json"):
-            query = query.removeprefix("```json").strip()
-        elif query.startswith("```"):
-            query = query.removeprefix("```").strip()
-
-        if query.endswith("```"):
-            query = query.removesuffix("```").strip()
-
-        return query
-
-    def save_to_json(self, file_path: str = "output.json"):
-        try:
-            data = json.loads(self.processed_query)
-        except json.JSONDecodeError as e:
-            print("failed to parse json", e)
-            print(self.processed_query)
-            return
-
-        with open(file_path, "w") as f:
-            json.dump(data, f)

src/rowgen/main.py

@@ -0,0 +1,22 @@
+import pytest
+import json
+from rowgen.json_parser import JsonParse
+from rowgen.hf_api import HFapi
+from pytest import fixture
+
+
+@fixture
+def hfapi():
+    return HFapi()
+
+
+def test_single_row_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)
+    try:
+        json.loads(jsp.processed_query)
+        return True
+    except json.JSONDecodeError:
+        return False