pratham0011/QueryMate_Text-to-SQL-CSV

1from fastapi import FastAPI, HTTPException
2from pydantic import BaseModel
3import sqlite3
4import pandas as pd
5import os
6from dotenv import load_dotenv
7import google.generativeai as genai
8
9app = FastAPI()
10
11# Load environment variables and configure Genai
12load_dotenv()
13genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
14
15class Query(BaseModel):
16    question: str
17    data_source: str
18
19def get_gemini_response(question, prompt):
20    model = genai.GenerativeModel('gemini-pro')
21    response = model.generate_content([prompt, question])
22    return response.text
23
24def get_csv_columns():
25    df = pd.read_csv('employee.csv')
26    return df.columns.tolist()
27
28csv_columns = get_csv_columns()
29
30sql_prompt = """
31You are an expert in converting English questions to SQL code!
32The SQL database has the name STUDENT and has the following Columns - NAME, CLASS, SECTION
33
34For example:
35- How many entries of records are present? SQL command: SELECT COUNT(*) FROM STUDENT;
36- Tell me all the students studying in Data Science class? SQL command: SELECT * FROM STUDENT where CLASS="Data Science";
37
38Also, the SQL code should not have ''' in the beginning or at the end, and SQL word in output.
39Ensure that you only generate valid SQL queries, not pandas or Python code.
40"""
41
42csv_prompt = f"""
43You are an expert in analyzing CSV data and converting English questions to pandas query syntax.
44The CSV file is named 'employee.csv' and contains employee information.
45The available columns in the CSV file are: {', '.join(csv_columns)}
46
47For example:
48- How many employees are there? Pandas query: len(df)
49- List all employees in the Sales department. Pandas query: df[df['Department'] == 'Sales']
50- Show employees with a specific ID. Pandas query: df[df['ID'] == specific_id]
51
52Provide only the pandas query syntax without any additional explanation or markdown formatting.
53Do not include 'df = ' or any variable assignment in your response.
54Make sure to use only the columns that are available in the CSV file.
55Ensure that you only generate valid pandas queries, not SQL or other types of code.
56"""
57
58def execute_sql_query(query):
59    conn = sqlite3.connect('student.db')
60    try:
61        cursor = conn.cursor()
62        cursor.execute(query)
63        result = cursor.fetchall()
64        return result
65    except sqlite3.Error as e:
66        raise HTTPException(status_code=400, detail=f"SQL Error: {str(e)}")
67    finally:
68        conn.close()
69
70def execute_pandas_query(query):
71    df = pd.read_csv('employee.csv')
72    try:
73        result = eval(query, {'df': df, 'pd': pd})
74        if isinstance(result, pd.DataFrame):
75            return result.to_dict(orient='records')
76        elif isinstance(result, pd.Series):
77            return result.to_dict()
78        else:
79            return result
80    except Exception as e:
81        raise HTTPException(status_code=400, detail=f"Pandas Error: {str(e)}")
82
83@app.post("/query")
84async def process_query(query: Query):
85    if query.data_source == "SQL Database":
86        ai_response = get_gemini_response(query.question, sql_prompt)
87        try:
88            result = execute_sql_query(ai_response)
89            return {"query": ai_response, "result": result}
90        except HTTPException as e:
91            raise HTTPException(status_code=400, detail=f"Error in SQL query: {e.detail}")
92    else:  # CSV Data
93        ai_response = get_gemini_response(query.question, csv_prompt)
94        try:
95            result = execute_pandas_query(ai_response)
96            return {"query": ai_response, "result": result, "columns": csv_columns}
97        except HTTPException as e:
98            raise HTTPException(status_code=400, detail=f"Error in pandas query: {e.detail}")