| | 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}") |
