Update app.py

app.py

@@ -1,17 +1,19 @@
 import json
-import openai
 import gradio as gr
 import duckdb
 import tempfile
 from functools import lru_cache
 import os
+import replicate
 
 # =========================
 # Configuration and Setup
 # =========================
 
-openai.api_key = os.getenv("OPENAI_API_KEY")
-dataset_path = 'sample_contract_df.parquet'  # Update with your Parquet file path
+replicate_api_token = os.getenv("REPLICATE_API_TOKEN")
+replicate.Client(api_token=replicate_api_token)
+
+dataset_path = 'sample_contract_df.parquet'
 
 schema = [
     {"column_name": "department_ind_agency", "column_type": "VARCHAR"},
@@ -59,26 +61,46 @@ COLUMN_TYPES = {col['column_name']: col['column_type'] for col in get_schema()}
 # OpenAI API Integration
 # =========================
 
+import replicate
+import json
+
 def parse_query(nl_query):
-    messages = [
-        {"role": "system", "content": "You are an assistant that converts natural language queries into SQL queries for the 'contract_data' table."},
-        {"role": "user", "content": f"Schema:\n{json.dumps(schema, indent=2)}\n\nQuery:\n\"{nl_query}\"\n\nSQL:"}
-    ]
+    system_prompt = (
+        "You are an assistant that converts natural language into SQL queries for a DuckDB database. "
+        "The table is named `contract_data`. Only return the SQL query, no explanations."
+    )
+
+    user_prompt = f"""\
+Schema:
+{json.dumps(schema, indent=2)}
+
+Natural Language Query:
+"{nl_query}"
+
+SQL:"""
 
+    full_prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{system_prompt}<|eot|>" \
+                  f"<|start_header_id|>user<|end_header_id|>\n{user_prompt}<|eot|>" \
+                  f"<|start_header_id|>assistant<|end_header_id|>\n"
+
+    output = ""
     try:
-        response = openai.chat.completions.create(
-            model="gpt-4o-mini",
-            messages=messages,
-            temperature=0,
-            max_tokens=150,
-        )
-        sql_query = response.choices[0].message.content.strip()
-        # Remove surrounding backticks and formatting artifacts
-        if sql_query.startswith("```") and sql_query.endswith("```"):
-            sql_query = sql_query[sql_query.find('\n')+1:sql_query.rfind('\n')].strip()
-        return sql_query, ""
+        for event in replicate.stream(
+            "meta/meta-llama-3-8b-instruct",
+            input={
+                "top_p": 0.9,
+                "temperature": 0.2,
+                "max_tokens": 300,
+                "prompt": full_prompt
+            }
+        ):
+            output += str(event)
+
+        return output.strip()
+
     except Exception as e:
-        return "", f"Error generating SQL query: {e}"
+        return f"-- Error generating SQL: {e}"
+
 
 # =========================
 # Database Interaction