|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
import sqlite3
|
|
import pandas as pd
|
|
import os
|
|
from dotenv import load_dotenv
|
|
import google.generativeai as genai
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
load_dotenv()
|
|
genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
|
|
|
|
class Query(BaseModel):
|
|
question: str
|
|
data_source: str
|
|
|
|
def get_gemini_response(question, prompt):
|
|
model = genai.GenerativeModel('gemini-pro')
|
|
response = model.generate_content([prompt, question])
|
|
return response.text
|
|
|
|
def get_csv_columns():
|
|
df = pd.read_csv('employee.csv')
|
|
return df.columns.tolist()
|
|
|
|
csv_columns = get_csv_columns()
|
|
|
|
sql_prompt = """
|
|
You are an expert in converting English questions to SQL code!
|
|
The SQL database has the name STUDENT and has the following Columns - NAME, CLASS, SECTION
|
|
|
|
For example:
|
|
- How many entries of records are present? SQL command: SELECT COUNT(*) FROM STUDENT;
|
|
- Tell me all the students studying in Data Science class? SQL command: SELECT * FROM STUDENT where CLASS="Data Science";
|
|
|
|
Also, the SQL code should not have ''' in the beginning or at the end, and SQL word in output.
|
|
Ensure that you only generate valid SQL queries, not pandas or Python code.
|
|
"""
|
|
|
|
csv_prompt = f"""
|
|
You are an expert in analyzing CSV data and converting English questions to pandas query syntax.
|
|
The CSV file is named 'employee.csv' and contains employee information.
|
|
The available columns in the CSV file are: {', '.join(csv_columns)}
|
|
|
|
For example:
|
|
- How many employees are there? Pandas query: len(df)
|
|
- List all employees in the Sales department. Pandas query: df[df['Department'] == 'Sales']
|
|
- Show employees with a specific ID. Pandas query: df[df['ID'] == specific_id]
|
|
|
|
Provide only the pandas query syntax without any additional explanation or markdown formatting.
|
|
Do not include 'df = ' or any variable assignment in your response.
|
|
Make sure to use only the columns that are available in the CSV file.
|
|
Ensure that you only generate valid pandas queries, not SQL or other types of code.
|
|
"""
|
|
|
|
def execute_sql_query(query):
|
|
conn = sqlite3.connect('student.db')
|
|
try:
|
|
cursor = conn.cursor()
|
|
cursor.execute(query)
|
|
result = cursor.fetchall()
|
|
return result
|
|
except sqlite3.Error as e:
|
|
raise HTTPException(status_code=400, detail=f"SQL Error: {str(e)}")
|
|
finally:
|
|
conn.close()
|
|
|
|
def execute_pandas_query(query):
|
|
df = pd.read_csv('employee.csv')
|
|
try:
|
|
result = eval(query, {'df': df, 'pd': pd})
|
|
if isinstance(result, pd.DataFrame):
|
|
return result.to_dict(orient='records')
|
|
elif isinstance(result, pd.Series):
|
|
return result.to_dict()
|
|
else:
|
|
return result
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=f"Pandas Error: {str(e)}")
|
|
|
|
@app.post("/query")
|
|
async def process_query(query: Query):
|
|
if query.data_source == "SQL Database":
|
|
ai_response = get_gemini_response(query.question, sql_prompt)
|
|
try:
|
|
result = execute_sql_query(ai_response)
|
|
return {"query": ai_response, "result": result}
|
|
except HTTPException as e:
|
|
raise HTTPException(status_code=400, detail=f"Error in SQL query: {e.detail}")
|
|
else:
|
|
ai_response = get_gemini_response(query.question, csv_prompt)
|
|
try:
|
|
result = execute_pandas_query(ai_response)
|
|
return {"query": ai_response, "result": result, "columns": csv_columns}
|
|
except HTTPException as e:
|
|
raise HTTPException(status_code=400, detail=f"Error in pandas query: {e.detail}") |
