Update app.py

app.py

@@ -1,24 +1,61 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import pipeline
-
-## create a new FASTAPI app instance
-app=FastAPI()
-
-# Initialize the text generation pipeline
-pipe = pipeline("text-generation", model="defog/llama-3-sqlcoder-8b")
-
-
-@app.get("/")
-def home():
-    return {"message":"Hello World"}
-
-# Define a function to handle the GET request at `/generate`
-
-
-@app.get("/generate")
-def generate(text:str):
-    ## use the pipeline to generate text from given input text
-    output=pipe(text)
-
-    ## return the generate text in Json reposne
-    return {"output":output[0]['generated_text']}
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="SQL Coder API")
+
+# Initialize pipeline
+try:
+    pipe = pipeline("text-generation", 
+                   model="defog/llama-3-sqlcoder-8b",
+                   device_map="auto",
+                   torch_dtype="auto")
+    logger.info("Pipeline initialized successfully")
+except Exception as e:
+    logger.error(f"Error initializing pipeline: {str(e)}")
+    raise
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class QueryRequest(BaseModel):
+    messages: list[ChatMessage]
+    max_length: int = 1024
+    temperature: float = 0.7
+
+class QueryResponse(BaseModel):
+    generated_text: str
+
+@app.post("/generate", response_model=QueryResponse)
+async def generate(request: QueryRequest):
+    try:
+        # Format messages into a single string
+        formatted_prompt = "\n".join([f"{msg.role}: {msg.content}" for msg in request.messages])
+        
+        # Generate response using pipeline
+        response = pipe(
+            formatted_prompt,
+            max_length=request.max_length,
+            temperature=request.temperature,
+            do_sample=True,
+            num_return_sequences=1
+        )
+        
+        # Extract generated text
+        generated_text = response[0]['generated_text']
+        
+        return QueryResponse(generated_text=generated_text)
+    
+    except Exception as e:
+        logger.error(f"Error generating response: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}