Update HF Model Engine > hf_engine.py

1. Implement custom wrapper for LangChain compatibility for later agent implementation.
2. Use InferenceClient to initialize the model with Inference API Provider
3. Tested with local code block
- Successfully loaded

src/nl2sql/hf_engine.py

@@ -1,98 +1,63 @@
-#"""Hugging Face inference helpers for SQL generation."""
-
+# Path: src/nl2sql/hf_engine.py
+# This module defines the HuggingFace-based engine for generating SQL queries from natural language questions.
 import os
-import re
-
-from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
-
-
-load_dotenv()
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    raise ValueError("Token Not Found!")
-
-client = InferenceClient(api_key=hf_token)
-MODEL_ID = "defog/llama-3-sqlcoder-8b:featherless-ai"
-
-
-def _build_messages(question: str, schema_context: str):
-    system_content = (
-        "You are an expert SQLite assistant that converts natural language into one "
-        "executable SQLite query.\n"
-        "Rules:\n"
-        "1. Use only tables, columns, and join paths present in the provided schema.\n"
-        "2. Generate valid SQLite syntax only.\n"
-        "3. Prefer exact column names from the schema, never invent columns.\n"
-        "4. Use explicit JOIN conditions when multiple tables are required.\n"
-        "5. Use GROUP BY for aggregates by entity, HAVING for aggregate filters, "
-        "ORDER BY for ranking, and LIMIT for top-N requests.\n"
-        "6. Return SQL only. No markdown, explanations, comments, or chain-of-thought.\n"
-        "7. If a join is needed, use short aliases that remain readable.\n"
-        "8. Produce a single SELECT statement."
-    )
-
-    user_content = f"""Database schema:
-{schema_context}
-
-Question:
-{question}
-
-Write the SQLite query that answers the question. Return only the SQL query."""
-
-    return [
-        {"role": "system", "content": system_content},
-        {"role": "user", "content": user_content},
-    ]
-
-
-def _extract_sql(raw_response: str) -> str:
-    text = raw_response.strip()
-    fenced_match = re.search(r"```(?:sql)?\s*(.*?)```", text, flags=re.IGNORECASE | re.DOTALL)
-    if fenced_match:
-        text = fenced_match.group(1).strip()
-
-    statement_match = re.search(
-        r"(?is)\b(WITH|SELECT)\b.*?(;|$)",
-        text,
-    )
-    if statement_match:
-        text = statement_match.group(0).strip()
-
-    lines = [
-        line.strip()
-        for line in text.splitlines()
-        if line.strip() and not line.strip().startswith(("--", "#"))
-    ]
-    sql = " ".join(lines).strip()
-    if sql and not sql.endswith(";"):
-        sql = f"{sql};"
-    return sql
-
-
-def generate_sql(question, ddl):
-    try:
-        completion = client.chat.completions.create(
-            model=MODEL_ID,
-            messages=_build_messages(question, ddl),
-            max_tokens=220,
-            temperature=0,
+from langchain_core.language_models.llms import LLM
+from typing import Any, List, Optional
+
+# Default Model
+DEFAULT_MODEL_ID = "defog/llama-3-sqlcoder-8b:featherless-ai"
+
+# Custom LangChain wrapper for HuggingFace Inference API
+class HFChatWrapper(LLM):
+    """
+    Custom LLM wrapper for HuggingFace Inference API to maintain compatibility with LangChain's LLM interface.
+    """
+    client: Any
+    model_id: str
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
+        completion = self.client.chat.completions.create(
+            model = self.model_id,
+            messages = [
+                {"role": "user", "content": prompt}
+            ],
+            temperature = 0.0,
+            max_tokens = 512
         )
-        raw_response = completion.choices[0].message.content or ""
-        sql = _extract_sql(raw_response)
-        return sql or raw_response.strip()
-    except Exception as error:
-        return f"Error: {error}"
-
-
+        return completion.choices[0].message.content
+    
+    @property
+    def _llm_type(self) -> str:
+        return "huggingface_inference_client"
+    
+# Initialize the HuggingFace endpoint using the InferenceClient
+def get_llm(model_id: str = DEFAULT_MODEL_ID):
+    """
+    Initializes the HuggingFace InferenceClient and returns an LLM instance for generating SQL queries.
+    """
+    # Load HuggingFace API token from environment variable
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HuggingFace API token not found!")
+    
+    print(f"Initializing HuggingFace InferenceClient with model: {model_id}")
+
+    # Initialize the HuggingFace InferenceClient
+    client = InferenceClient(api_key=hf_token)
+    llm = HFChatWrapper(client=client, model_id=model_id)
+
+    return llm
+
+# Local Test block
 if __name__ == "__main__":
-    my_ddl = "CREATE TABLE tracks (id INTEGER PRIMARY KEY, title TEXT, genre TEXT);"
-    my_question = "How many tracks are there in each genre?"
+    from dotenv import load_dotenv
+    load_dotenv()
 
-    print("Generating SQL query via Featherless AI...")
     try:
-        result = generate_sql(my_question, my_ddl)
-        print("-" * 20)
-        print(result)
-    except Exception as error:
-        print(f"An error occurred: {error}")
+        test_llm = get_llm()
+        print("Model loaded successfully! Running a quick ping...")
+        response = test_llm.invoke("Write a single SQL statement to count all rows in a table named 'Employee'.")
+        print(f"\nResponse:\n{response}")
+    except Exception as e:
+        print(f"Error during LLM initialization or invocation: {e}")