in skills/text_to_sql/evaluation/prompts.py [0:0]
def generate_prompt_with_rag(context):
from vectordb import VectorDB
# Load the vector database
vectordb = VectorDB()
vectordb.load_db()
user_query = context['vars']['user_query']
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)
relevant_schema = vectordb.search(user_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])
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.