import streamlit as st
from langchain_community.vectorstores import FAISS
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_openai import OpenAIEmbeddings

# Example queries remain the same
examples = [
    {
        "input": "Retrieve the user  who have placed the highest total value of orders.",
        "query": "SELECT u.username, SUM(p.price) AS total_order_value FROM users u JOIN orders o ON u.user_id = o.user_id JOIN products p ON o.product_id = p.product_id GROUP BY u.username ORDER BY total_order_value DESC LIMIT 10;"
    },
    # ... (rest of the examples)
]

@st.cache_resource
def get_example_selector(api_key):
    embeddings = OpenAIEmbeddings(api_key=api_key)
    example_selector = SemanticSimilarityExampleSelector.from_examples(
        examples,
        embeddings,
        FAISS,
        k=2,
        input_variables=["input"],
    )
    return example_selector