from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel

# Model loading with the new model name
llm = AutoModelForCausalLM.from_pretrained("sqlcoder-7b.Q4_K_S.gguf")

class Validation(BaseModel):
    prompt: str  # Assuming this includes both user_question and table_metadata_string

app = FastAPI()

@app.post("/generate_sql")
async def generate_sql(item: Validation):
    # Updated system prompt
    system_prompt = """### Task
Generate a SQL query to answer the following question:
`{question}`

### Database Schema
The query will run on a database with the following schema:
{schema}

### Answer
Given the database schema, here is the SQL query that answers `{question}`:
```sql
"""
    # Format the actual prompt using item.prompt
    prompt = system_prompt.format(user_question="Your question here", table_metadata_string="Your schema here")
    completion = llm(prompt)
    return completion