LanceDB

.gitignore

@@ -0,0 +1 @@
+app2.py

app.py

@@ -2,12 +2,15 @@ import os
 import torch
 import duckdb
 import spaces
+import lancedb
 import gradio as gr
 import pandas as pd
+import pyarrow as pa
+from langchain import hub
+from langsmith import traceable
+from sentence_transformers import SentenceTransformer
 from langchain_huggingface.llms import HuggingFacePipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
-from langsmith import traceable
-from langchain import hub
 
 # Height of the Tabs Text Area
 TAB_LINES = 8
@@ -16,8 +19,8 @@ TAB_LINES = 8
 #----------CONNECT TO DATABASE----------
 md_token = os.getenv('MD_TOKEN')
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
-
 #---------------------------------------
+
 if torch.cuda.is_available():
     device = torch.device("cuda")
     print(f"Using GPU: {torch.cuda.get_device_name(device)}")
@@ -26,6 +29,25 @@ else:
     print("Using CPU")
 #---------------------------------------
 
+#--------------LanceDB-------------
+
+lance_db = lancedb.connect(
+        uri=os.getenv('lancedb_uri'),
+        api_key=os.getenv('lancedb_api_key'),
+        region=os.getenv('lancedb_region')
+        )
+
+lance_schema = pa.schema([
+    pa.field("vector", pa.list_(pa.float32())),
+    pa.field("sql-query", pa.utf8())
+])
+
+try:
+  table = lance_db.create_table(name="SQL-Queries", schema=lance_schema)
+except:
+  table = lance_db.open_table(name="SQL-Queries")
+#---------------------------------------
+
 #-------LOAD HUGGINGFACE PIPELINE-------
 tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
 
@@ -46,7 +68,9 @@ hf = HuggingFacePipeline(pipeline=pipe)
 prompt = hub.pull("sql-agent-prompt")
 #---------------------------------------
 
-
+#-----LOAD EMBEDDING MODEL-----
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2", device=device)
+#---------------------------------------
 
 #--------------ALL UTILS----------------
 # Get Databases
@@ -91,6 +115,20 @@ def get_prompt(schema, query_input):
 def generate_sql(prompt):
     result = hf.invoke(prompt)
     return result.strip()
+@spaces.GPU(duration=10)
+def embed_query(sql_query):
+    print(f'Creating Emebeddings {sql_query}')
+    if sql_query is not None:
+        embeddings = embedding_model.encode(sql_query, normalize_embeddings=True).tolist()
+    return embeddings
+
+def log2lancedb(embeddings, sql_query):
+    data = [{
+        "sql-query": sql_query,
+        "vector": embeddings
+    }]
+    table.add(data)
+    print(f'Added to Lance DB.')
 #---------------------------------------
 
 
@@ -108,6 +146,7 @@ def text2sql(table, query_input):
     print(f'Schema Generated...')
     prompt = get_prompt(schema, query_input)
     print(f'Prompt Generated...')
+    
     try:
         print(f'Generating SQL... {model.device}')
         result = generate_sql(prompt)
@@ -119,6 +158,14 @@ def text2sql(table, query_input):
             generated_query: "",
             result_output:pd.DataFrame([{"error": f"❌ Unable to get the SQL query based on the text. {e}"}])
         }
+        
+    try:
+        embeddings = embed_query(result)
+        log2lancedb(embeddings, result)
+    except Exception as e:
+        print("Error Generating and Logging Embeddings...")
+        print(e)
+    
     try:
         query_result = conn.sql(result).df()
         

requirements.txt

@@ -4,4 +4,7 @@ transformers==4.44.2
 duckdb==1.1.1
 langsmith==0.1.135
 langchain==0.3.4
+lancedb==0.15.0
+sentence-transformers==3.2.1
+pyarrow==17.0.0
 langchain-huggingface