refactored hf_api and added tests

Author: Arsalanjdev

src/rowgen/hf_api.py

@@ -0,0 +1,25 @@
+from huggingface_hub import InferenceClient
+
+from .utils import API_KEY
+
+
+class HFapi:
+    def __init__(self, provider: str = "novita", api_key: str = API_KEY):
+        self.client = InferenceClient(provider=provider, api_key=api_key)
+
+    def send_message_to_api(self, message: str) -> str:
+        """
+        Sends a message to the AI and returns the response
+        :param message: The prompt to be sent.
+        :return: the response of AI model.
+        """
+        completion = self.client.chat.completions.create(
+            model="deepseek-ai/DeepSeek-V3",
+            messages=[
+                {
+                    "role": "user",
+                    "content": message,
+                }
+            ],
+        )
+        return completion.choices[0].message["content"]

src/rowgen/insert_todb.py

@@ -6,8 +6,6 @@
 from dotenv import load_dotenv
 from utils import API_KEY
 
-print()
-
 
 class JsonDB:
     def __init__(self):
@@ -23,19 +21,6 @@ def __getitem__(self, item):
 
 print(JsonDB()["columns"])
 
-client = InferenceClient(provider="novita", api_key=API_KEY)
-completion = client.chat.completions.create(
-    model="deepseek-ai/DeepSeek-V3",
-    messages=[
-        {
-            "role": "user",
-            "content": f"generate 10 rows for the following schema: {JsonDB().get_column_names()}. Use real looking data (no john doe or @example.com. also make emails beliavble such as using numbers underscores. sometimes even use unrelated to emails that are unrelate to names.). do not say anything. just generate data in a json format.",
-        }
-    ],
-)
-
-print(completion.choices[0])
-
 
 class JsonParse:
     def __init__(self, query: str):
@@ -64,6 +49,3 @@ def save_to_json(self, file_path: str = "output.json"):
 
         with open(file_path, "w") as f:
             json.dump(data, f)
-
-
-JsonParse(completion.choices[0].message["content"]).save_to_json()

src/rowgen/main.py

@@ -6,6 +6,7 @@
 
 from rowgen.utils import API_KEY
 from huggingface_hub import InferenceClient, HfApi
+from rowgen.hf_api import HFapi
 
 
 DEEP_SEEK_MODEL = "deepseek-ai/DeepSeek-V3"
@@ -23,6 +24,12 @@ def hf_api_client():
     return HfApi()
 
 
+@pytest.fixture
+def hf_api():
+    client = HFapi()
+    return client
+
+
 def test_hf_hub_ping(hf_api_client):
     try:
         info = hf_api_client.model_info(DEEP_SEEK_MODEL)
@@ -48,3 +55,17 @@ def test_hf_inference_connection(inference_client):
 
     except Exception as d:
         pytest.fail(f"Failed inference call: {d}")
+
+
+def test_hf_api_send_message_to_api(hf_api):
+    prompt_message = "Hi! This is a test. say banana if you can hear me."
+    try:
+        response_message = hf_api.send_message_to_api(prompt_message)
+        assert len(response_message) > 0
+        assert "banana" in response_message.casefold()
+
+    except BadRequestError as e:
+        print(e)
+
+    except Exception as d:
+        pytest.fail(f"Failed inference call: {d}")