%pip install -q anthropic pandas voyageai matplotlib seaborn

import os
import sqlite3
 
import pandas as pd
from anthropic import Anthropic
from IPython.display import display
 
# Set your Claude API key
os.environ["ANTHROPIC_API_KEY"] = "YOUR_ANTHROPIC_API_KEY"
os.environ["VOYAGE_API_KEY"] = "YOUR_VOYAGE_API_KEY"
 
# Initialize the Anthropic client
client = Anthropic()
MODEL_NAME = "claude-sonnet-4-5"
 
# Filepath to the SQLite database
DATABASE_PATH = "data/data.db"

import random
from datetime import datetime, timedelta
 
if not os.path.exists(DATABASE_PATH):
    print("Database does not exist. Creating and populating...")
    # Create a new SQLite database and tables
    with sqlite3.connect(DATABASE_PATH) as conn:
        cursor = conn.cursor()
 
        cursor.executescript("""
        CREATE TABLE IF NOT EXISTS departments (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            location TEXT
        );
        CREATE TABLE IF NOT EXISTS employees (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            age INTEGER,
            department_id INTEGER,
            salary REAL,
            hire_date DATE,
            FOREIGN KEY (department_id) REFERENCES departments (id)
        );
        """)
 
        # Insert sample data
        cursor.executemany(
            "INSERT INTO departments VALUES (?,?,?)",
            [
                (1, "HR", "New York"),
                (2, "Engineering", "San Francisco"),
                (3, "Marketing", "Chicago"),
                (4, "Sales", "Los Angeles"),
                (5, "Finance", "Boston"),
                (6, "Customer Support", "Dallas"),
                (7, "Research", "Seattle"),
                (8, "Legal", "Washington D.C."),
                (9, "Product", "Austin"),
                (10, "Operations", "Denver"),
            ],
        )
 
        first_names = [
            "John",
            "Jane",
            "Bob",
            "Alice",
            "Charlie",
            "Diana",
            "Edward",
            "Fiona",
            "George",
            "Hannah",
            "Ian",
            "Julia",
            "Kevin",
            "Laura",
            "Michael",
            "Nora",
            "Oliver",
            "Patricia",
            "Quentin",
            "Rachel",
            "Steve",
            "Tina",
            "Ulysses",
            "Victoria",
            "William",
            "Xena",
            "Yannick",
            "Zoe",
        ]
        last_names = [
            "Smith",
            "Johnson",
            "Williams",
            "Jones",
            "Brown",
            "Davis",
            "Miller",
            "Wilson",
            "Moore",
            "Taylor",
            "Anderson",
            "Thomas",
            "Jackson",
            "White",
            "Harris",
            "Martin",
            "Thompson",
            "Garcia",
            "Martinez",
            "Robinson",
            "Clark",
            "Rodriguez",
            "Lewis",
            "Lee",
            "Walker",
            "Hall",
            "Allen",
            "Young",
            "King",
        ]
 
        employees_data = []
        for i in range(1, 201):  # Generate 200 employees
            name = f"{random.choice(first_names)} {random.choice(last_names)}"
            age = random.randint(22, 65)
            department_id = random.randint(1, 10)
            salary = round(random.uniform(40000, 200000), 2)
            hire_date = (datetime.now() - timedelta(days=random.randint(0, 3650))).strftime(
                "%Y-%m-%d"
            )
            employees_data.append((i, name, age, department_id, salary, hire_date))
 
        cursor.executemany("INSERT INTO employees VALUES (?,?,?,?,?,?)", employees_data)
 
    print("Database created and populated successfully.")
else:
    print("Database already exists. Skipping creation and population.")
 
# Display table contents
with sqlite3.connect(DATABASE_PATH) as conn:
    for table in ["departments", "employees"]:
        df = pd.read_sql_query(f"SELECT * FROM {table}", conn)
        print(f"\n{table.capitalize()} table:")
        display(df)

def get_schema_info(db_path):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
 
    schema_info = []
 
    # Get all tables
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = cursor.fetchall()
 
    for (table_name,) in tables:
        # Get columns for this table
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()
 
        table_info = f"Table: {table_name}\n"
        table_info += "\n".join(f"  - {col[1]} ({col[2]})" for col in columns)
        schema_info.append(table_info)
 
    conn.close()
    return "\n\n".join(schema_info)
 
 
# Get the schema info
schema = get_schema_info(DATABASE_PATH)
print(schema)

def generate_prompt(schema, query):
    return f"""
        You are an AI assistant that converts natural language queries into SQL.
        Given the following SQL database schema:
 
        <schema>
        {schema}
        </schema>
 
        Convert the following natural language query into SQL:
        <query>
        {query}
        </query>
 
        Provide only the SQL query in your response, enclosed within <sql> tags.
    """
 
 
# Test the prompt
user_query = "What are the names of all employees in the Engineering department?"
prompt = generate_prompt(schema, user_query)
print(prompt)

def generate_sql(prompt):
    response = client.messages.create(
        model=MODEL_NAME,
        max_tokens=1000,
        temperature=0,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[0].text.strip()
 
 
# Generate SQL
result = generate_sql(prompt)
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("Generated SQL:")
print(sql)

def run_sql(sql):
    conn = sqlite3.connect(DATABASE_PATH)
    result = pd.read_sql_query(sql, conn)
    conn.close()
    return result
 
 
result = run_sql(sql)
print("Query result:")
display(result)

def generate_prompt_with_examples(schema, query):
    examples = """
        Example 1:
        <query>List all employees in the HR department.</<query>
        <output>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</output>
 
        Example 2:
        User: What is the average salary of employees in the Engineering department?
        SQL: SELECT AVG(e.salary) FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'Engineering';
 
        Example 3:
        User: Who is the oldest employee?
        SQL: SELECT name, age FROM employees ORDER BY age DESC LIMIT 1;
    """
 
    return f"""
        You are an AI assistant that converts natural language queries into SQL.
        Given the following SQL database schema:
 
        <schema>
        {schema}
        </schema>
 
        Here are some examples of natural language queries and their corresponding SQL:
 
        <examples>
        {examples}
        </examples>
 
        Now, convert the following natural language query into SQL:
        <query>
        {query}
        </query>
 
        Provide only the SQL query in your response, enclosed within <sql> tags.
    """
 
 
# Test the new prompt
user_query = "What are the names and salaries of employees in the Marketing department?"
prompt = generate_prompt_with_examples(schema, user_query)
print(prompt)

# Generate SQL using the improved prompt
result = generate_sql(prompt)
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("Generated SQL:")
print(sql)
 
# Run the generated SQL
result = run_sql(sql)
print("\nQuery result:")
display(result)

def generate_prompt_with_cot(schema, query):
    examples = """
    <example>
    <query>List all employees in the HR department.</query>
    <thought_process>
    1. We need to join the employees and departments tables.
    2. We'll match employees.department_id with departments.id.
    3. We'll filter for the HR department.
    4. We only need to return the employee names.
    </thought_process>
    <sql>SELECT e.name FROM employees e JOIN departments d ON e.department_id = d.id WHERE d.name = 'HR';</sql>
    </example>
 
    <example>
    <query>What is the average salary of employees hired in 2022?</query>
    <thought_process>
    1. We need to work with the employees table.
    2. We need to filter for employees hired in 2022.
    3. We'll use the YEAR function to extract the year from the hire_date.
    4. We'll calculate the average of the salary column for the filtered rows.
    </thought_process>
    <sql>SELECT AVG(salary) FROM employees WHERE YEAR(hire_date) = 2022;</sql>
    </example>
    """
 
    return f"""You are an AI assistant that converts natural language queries into SQL.
    Given the following SQL database schema:
 
    <schema>
    {schema}
    </schema>
 
    Here are some examples of natural language queries, thought processes, and their corresponding SQL:
 
    <examples>
    {examples}
    </examples>
 
    Now, convert the following natural language query into SQL:
    <query>
    {query}
    </query>
 
    Within <thought_process> tags, explain your thought process for creating the SQL query.
    Then, within <sql> tags, provide your output SQL query.
    """
 
 
# Test the new prompt
user_query = "What are the names and hire dates of employees in the Engineering department, ordered by their salary?"
prompt = generate_prompt_with_cot(schema, user_query)
print(prompt)

# Generate SQL using the chain-of-thought prompt
result = generate_sql(prompt)
print("Raw response from Claude:")
print(result)
 
# Extract thought process and SQL query
thought_process = result.split("<thought_process>")[1].split("</thought_process>")[0].strip()
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
 
print("\nThought Process:")
print(thought_process)
 
print("\nGenerated SQL:")
print(sql)
 
# Run the generated SQL
query_result = run_sql(sql)
print("\nQuery result:")
display(query_result)

import json
import os
import pickle
import sqlite3
 
import numpy as np
import voyageai
 
 
class VectorDB:
    def __init__(self, db_path="./data/vector_db.pkl"):
        self.client = voyageai.Client(api_key=os.getenv("VOYAGE_API_KEY"))
        self.db_path = db_path
        self.load_db()
 
    def load_db(self):
        if os.path.exists(self.db_path):
            with open(self.db_path, "rb") as file:
                data = pickle.load(file)
            self.embeddings, self.metadata, self.query_cache = (
                data["embeddings"],
                data["metadata"],
                json.loads(data["query_cache"]),
            )
        else:
            self.embeddings, self.metadata, self.query_cache = [], [], {}
 
    def load_data(self, data):
        if not self.embeddings:
            texts = [item["text"] for item in data]
            self.embeddings = [
                emb
                for batch in range(0, len(texts), 128)
                for emb in self.client.embed(
                    texts[batch : batch + 128], model="voyage-2"
                ).embeddings
            ]
            self.metadata = [item["metadata"] for item in data]  # Store only the inner metadata
            self.save_db()
 
    def search(self, query, k=5, similarity_threshold=0.3):
        if query not in self.query_cache:
            self.query_cache[query] = self.client.embed([query], model="voyage-2").embeddings[0]
            self.save_db()
 
        similarities = np.dot(self.embeddings, self.query_cache[query])
        top_indices = np.argsort(similarities)[::-1]
 
        return [
            {"metadata": self.metadata[i], "similarity": similarities[i]}
            for i in top_indices
            if similarities[i] >= similarity_threshold
        ][:k]
 
    def save_db(self):
        with open(self.db_path, "wb") as file:
            pickle.dump(
                {
                    "embeddings": self.embeddings,
                    "metadata": self.metadata,
                    "query_cache": json.dumps(self.query_cache),
                },
                file,
            )
 
 
# Initialize and load schema data
vectordb = VectorDB()
if not vectordb.embeddings:
    with sqlite3.connect(DATABASE_PATH) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        schema_data = [
            {
                "text": f"Table: {table[0]}, Column: {col[1]}, Type: {col[2]}",
                "metadata": {"table": table[0], "column": col[1], "type": col[2]},
            }
            for table in cursor.fetchall()
            for col in cursor.execute(f"PRAGMA table_info({table[0]})").fetchall()
        ]
    vectordb.load_data(schema_data)
 
# Test the search functionality
test_query = "What is the average salary of employees in each department?"
results = vectordb.search(test_query)
print("Search results:")
for result in results:
    print(f"Similarity: {result['similarity']}, Metadata: {result['metadata']}")

def generate_prompt_with_rag(query):
    relevant_schema = vectordb.search(query, k=10, similarity_threshold=0.3)
    schema_info = "\n".join(
        [
            f"Table: {item['metadata']['table']}, Column: {item['metadata']['column']}, Type: {item['metadata']['type']}"
            for item in relevant_schema
        ]
    )
 
    return generate_prompt_with_cot(schema_info, query)
 
 
# Test the RAG-based prompt
user_query = "What is the average salary of employees in each department?"
prompt = generate_prompt_with_rag(user_query)
print("Generated prompt:")
print(prompt)
 
# Generate and execute SQL
result = generate_sql(prompt)
print("\nGenerated result:")
print(result)
 
# Extract and run the SQL query
sql = result.split("<sql>")[1].split("</sql>")[0].strip()
print("\nExtracted SQL:")
print(sql)
 
query_result = run_sql(sql)
print("\nQuery result:")
display(query_result)

def execute_sql_with_feedback(sql):
    try:
        result = run_sql(sql)
        return True, result, "Query executed successfully."
    except Exception as e:
        return False, None, str(e)
 
 
def generate_prompt_with_self_improvement(query, max_attempts=3):
    feedback = None
    sql = None
 
    for attempt in range(max_attempts):
        if attempt == 0:
            prompt = generate_prompt_with_rag(query)
        else:
            prompt = f"""
            The previous SQL query resulted in this error: {feedback}
            Analyze the error and provide an improved SQL query.
            Original query: {sql}
            Explain your changes in <thought_process> tags and provide the corrected SQL in <sql> tags.
            """
 
        response = generate_sql(prompt)
        sql = response.split("<sql>")[1].split("</sql>")[0].strip()
 
        print(f"\nAttempt {attempt + 1}:")
 
        success, result, feedback = execute_sql_with_feedback(sql)
        if success:
            print("SQL executed successfully!")
            return sql, result
        else:
            print("SQL failed to execute")
 
    print("Maximum attempts reached. Could not generate a valid SQL query.")
    return None, None
 
 
# Test the self-improving SQL generation
user_query = "For each department, show the ratio of the highest paid employee's salary to the lowest paid employee's salary, but only for departments where this ratio is greater than 3"
final_sql, result = generate_prompt_with_self_improvement(user_query)
 
if final_sql:
    print("\nFinal SQL query:")
    print(final_sql)
    print("\nQuery result:")
    display(result)
else:
    print("Failed to generate a valid SQL query.")

import matplotlib.pyplot as plt
import numpy as np
 
# Data
models = [
    "3H Basic",
    "3H Few-Shot",
    "3H CoT+Few-Shot",
    "3H RAG+Few-Shot+CoT",
    "3.5S Basic",
    "3.5S Few-Shot",
    "3.5S CoT+Few-Shot",
    "3.5S RAG+Few-Shot+CoT",
]
pass_rates = [83.3, 77.8, 88.9, 88.9, 88.9, 94.4, 100, 100]
avg_latencies = [1561, 1758, 2187, 2564, 1887, 2297, 3900, 4614]
avg_tokens = [383, 572, 765, 1001, 309, 496, 765, 984]
costs = [0.0023, 0.0027, 0.0034, 0.0044, 0.020, 0.024, 0.041, 0.050]
 
# Normalize the data
max_latency = max(avg_latencies)
max_tokens = max(avg_tokens)
max_cost = max(costs)
norm_latencies = [latency / max_latency * 100 for latency in avg_latencies]
norm_tokens = [t / max_tokens * 100 for t in avg_tokens]
norm_costs = [c / max_cost * 100 for c in costs]
 
# Set up the plot
x = np.arange(len(models))
width = 0.2
 
fig, ax = plt.subplots(figsize=(15, 8))
 
# Create the bars
ax.bar(x - 1.5 * width, pass_rates, width, label="Pass Rate (%)", color="green")
ax.bar(x - 0.5 * width, norm_latencies, width, label="Avg. Latency (normalized)", color="blue")
ax.bar(x + 0.5 * width, norm_tokens, width, label="Avg. Tokens (normalized)", color="orange")
ax.bar(x + 1.5 * width, norm_costs, width, label="Cost (normalized)", color="red")
 
# Customize the plot
ax.set_ylabel("Percentage / Normalized Value")
ax.set_title("Comparison of Model Performance")
ax.set_xticks(x)
ax.set_xticklabels(models, rotation=45, ha="right")
ax.legend()
 
 
# Add value labels on the bars
def add_labels(rects):
    for rect in rects:
        height = rect.get_height()
        ax.annotate(
            f"{height:.1f}",
            xy=(rect.get_x() + rect.get_width() / 2, height),
            xytext=(0, 3),  # 3 points vertical offset
            textcoords="offset points",
            ha="center",
            va="bottom",
            rotation=90,
        )
 
 
add_labels(ax.containers[0])  # Only add labels for pass rates
 
plt.tight_layout()
plt.show()