Upload app.py

app.py

@@ -1,11 +1,22 @@
-from fastapi import FastAPI, HTTPException
+import logging
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from transformers import pipeline
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 app = FastAPI()
 
 # Initialize the text generation pipeline
-pipe = pipeline("text-generation", model="defog/sqlcoder-7b-2", pad_token_id=2)
+try:
+    pipe = pipeline("text-generation", model="defog/sqlcoder-7b-2", pad_token_id=2)
+    logger.info("Model loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load the model: {str(e)}")
+    raise
 
 class QueryRequest(BaseModel):
     text: str
@@ -15,9 +26,11 @@ def home():
     return {"message": "SQL Generation Server is running"}
 
 @app.post("/generate")
-def generate(request: QueryRequest):
+async def generate(request: QueryRequest):
     try:
         text = request.text
+        logger.info(f"Received request: {text}")
+        
         prompt = f"Generate a valid SQL query for the following request. Only return the SQL query, nothing else:\n\n{text}\n\nSQL query:"
         output = pipe(prompt, max_new_tokens=100)
         
@@ -35,16 +48,29 @@ def generate(request: QueryRequest):
         allowed_keywords = {
             'select', 'insert', 'update', 'delete', 'show', 'describe', 'from', 'where', 'and', 'or', 'like', 'limit', 'order by', 'group by', 'join', 'inner join', 'left join', 'right join', 'full join', 'on', 'using', 'union', 'union all', 'distinct', 'having', 'into', 'values', 'set', 'create', 'alter', 'drop', 'table', 'database', 'index', 'view', 'trigger', 'procedure', 'function', 'if', 'exists', 'primary key', 'foreign key', 'references', 'check', 'constraint', 'default', 'auto_increment', 'null', 'not null', 'in', 'is', 'is not', 'between', 'case', 'when', 'then', 'else', 'end', 'asc', 'desc', 'count', 'sum', 'avg', 'min', 'max', 'timestamp', 'date', 'time', 'varchar', 'char', 'int', 'integer', 'smallint', 'bigint', 'decimal', 'numeric', 'float', 'real', 'double', 'boolean', 'enum', 'text', 'blob', 'clob'
         }
+        
         # Ensure the query only contains allowed keywords
         tokens = sql_query.lower().split()
         for token in tokens:
             if not any(token.startswith(keyword) for keyword in allowed_keywords):
-                raise ValueError("Generated text contains invalid SQL syntax")
+                raise ValueError(f"Generated text contains invalid SQL syntax: {token}")
 
+        logger.info(f"Generated SQL query: {sql_query}")
         return {"output": sql_query}
+    except ValueError as ve:
+        logger.warning(f"Validation error: {str(ve)}")
+        raise HTTPException(status_code=400, detail=str(ve))
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Error in generate endpoint: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail="An error occurred while generating the SQL query")
+
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"message": exc.detail},
+    )
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)