Merge pull request #1 from rahulapjs/patch

Patch

Author: zccott

backend/app/api/query.py

@@ -29,7 +29,8 @@ def query_sql(payload: SQLQueryRequest):
 
     except ValueError as exc:
         raise HTTPException(status_code=400, detail=str(exc))
-    except Exception:
+    except Exception as exc:
+        print(exc)
         raise HTTPException(
             status_code=500,
             detail="Failed to process SQL query"

backend/app/api/schema.py

@@ -0,0 +1,10 @@
+from fastapi import APIRouter
+from app.services.schema_loader import load_database_schema
+
+router = APIRouter(tags=["Schema"])
+
+
+@router.get("/schema")
+def get_schema():
+    schema = load_database_schema()
+    return {"tables": schema}

backend/app/api/schema_create.py

@@ -0,0 +1,13 @@
+from fastapi import APIRouter, HTTPException
+from app.services.schema_creator import create_table
+
+router = APIRouter(tags=["Schema"])
+
+
+@router.post("/schema/create")
+def create_schema(payload: dict):
+    try:
+        create_table(payload["table_name"], payload["columns"])
+        return {"status": "table created"}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))

backend/app/services/schema_creator.py

@@ -0,0 +1,43 @@
+import psycopg2
+from app.core.config import settings
+
+ALLOWED_TYPES = {
+    "text": "TEXT",
+    "int": "INTEGER",
+    "float": "REAL",
+    "timestamp": "TIMESTAMP",
+}
+
+
+def create_table(table_name: str, columns: list[dict]):
+    cols_sql = []
+
+    for col in columns:
+        if col["type"] not in ALLOWED_TYPES:
+            raise ValueError(f"Unsupported column type: {col['type']}")
+
+        cols_sql.append(
+            f"{col['name']} {ALLOWED_TYPES[col['type']]}"
+        )
+
+    sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id SERIAL PRIMARY KEY,
+        {", ".join(cols_sql)}
+    );
+    """
+
+    conn = psycopg2.connect(
+        host=settings.DB_HOST,
+        port=settings.DB_PORT,
+        dbname=settings.DB_NAME,
+        user=settings.DB_USER,
+        password=settings.DB_PASSWORD,
+    )
+
+    try:
+        with conn.cursor() as cursor:
+            cursor.execute(sql)
+            conn.commit()
+    finally:
+        conn.close()

backend/app/services/sql_generator.py

@@ -21,9 +21,9 @@ def generate_sql(question: str) -> Dict[str, Any]:
     schema = load_database_schema()
     formatted_schema = format_schema_for_prompt(schema)
 
-    system_prompt = f"""
+    # We move the "Expert" rules to the system_instruction parameter
+    system_instruction = f"""
 You are an expert SQL assistant.
-
 You are given a PostgreSQL database schema.
 Your task is to generate a SINGLE, READ-ONLY SQL query
 that answers the user's question.
@@ -46,18 +46,12 @@ def generate_sql(question: str) -> Dict[str, Any]:
 }}
 """
 
+    # Corrected method call for the google-genai SDK
     response = client.models.generate_content(
         model=settings.GEMINI_MODEL,
-        contents=[
-            types.Content(
-                role="user",
-                parts=[
-                    types.Part(text=system_prompt),
-                    types.Part(text=f"\nUSER QUESTION:\n{question}")
-                ]
-            )
-        ],
-        generation_config=types.GenerationConfig(
+        contents=f"USER QUESTION: {question}",
+        config=types.GenerateContentConfig(
+            system_instruction=system_instruction,
             temperature=0.1,
             response_mime_type="application/json"
         )
@@ -68,9 +62,9 @@ def generate_sql(question: str) -> Dict[str, Any]:
     try:
         data = json.loads(raw_text)
     except json.JSONDecodeError as exc:
-        raise ValueError("Gemini returned invalid JSON") from exc
+        raise ValueError(f"Gemini returned invalid JSON: {raw_text}") from exc
 
     if "sql" not in data or "explanation" not in data:
         raise ValueError("Gemini response missing required fields")
 
-    return data
+    return data

backend/main.py

@@ -2,6 +2,8 @@
 from fastapi.middleware.cors import CORSMiddleware
 
 from app.api.query import router as query_router
+from app.api.schema import router as schema_router
+from app.api.schema_create import router as schema_create_router
 from app.core.config import settings
 
 app = FastAPI(
@@ -20,7 +22,10 @@
 )
 
 # Routes
+
+app.include_router(schema_router, prefix="/api")
 app.include_router(query_router, prefix="/api")
+app.include_router(schema_create_router, prefix="/api")
 
 
 @app.get("/health")