Very interesting idea and implementation. Here are a few comments / observations.

1. Can we do the decorator on the class itself. Feels cleaner. e.g.

@schema_decorator
class User(BaseModel):

2. Lets follow the convention of keeping the prompts in a file (instead of .py). And then provide a way for user to override e.g.
   SQLTableSchemaDecorator(
   ... schema_prompt= OR schema_prompt=
   )

3. In-fact, following up on the previous comment, it might be a good idea to move prompts into jinja templates - so the function and object specifications can be passed to it, instead of constructing the prompt in prompt.py with string concat.

@anshum4n-git for the 1st point if we add the decorator above the class we might have to get the schema like

TABLE_SCHEMA = User.get_sql_schema() 

while if we define a function we could just use

TABLE_SCHEMA = user_table()

Updated implementation example -

class UserRole(str, Enum):
 ADMIN = "admin"
 USER = "user"
 GUEST = "guest"
decorator = SQLTableSchemaDecorator(
 api_key=self.api_key,
 base_url=self.base_url,
 model=self.model,
 cache_dir=perf_cache_dir
)
@decorator
class User(BaseModel):
 """User model with auto-generated SQL schema."""
 id: str
 name: str
 email: str
 role: UserRole
 created_at: Optional[datetime] = None
TABLE_SCHEMA = User.get_sql_schema()
# The TABLE_SCHEMA will contain something like:
# CREATE TABLE IF NOT EXISTS users (
# id VARCHAR(36) PRIMARY KEY NOT NULL,
# name VARCHAR(255) NOT NULL,
# email VARCHAR(255) NOT NULL,
# role VARCHAR(50) NOT NULL CHECK (role IN ('admin', 'user', 'guest')),
# created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
# );

• Implemented Jinja templating creating new templates folder and storing all the prompts as .j2 file.
• Moved all the system prompt into prompts folder
• Added overridding of the system prompts
• Update the query creating prompt, now instead of using backticks ( `` ), it will be generating the queries with single quotes ( ' ' ) as backticks were giving errors while using those queries in multiple places.
• Update the some test files as they were failing and giving error due to some prompt changes.

@anshum4n-git

I have added a yet another implementation. Below is how we can do it

from typing import Optional
from pydantic import BaseModel
from foundation_sql.query import SQLQueryDecorator
class User(BaseModel):
 id: str
 name: str
 email: str
 created_at: Optional[str] = None
create_user = SQLQueryDecorator(
 db_url="sqlite:///./app.db",
 cache_dir="__sql__",
 auto_schema=True, # enable automatic schema
 schema_validate=True, # apply/validate schema on the DB
 schema_regen=False, # do not force schema regeneration if cached
)
@create_user
def create_user(user: User) -> User:
 """Create a new user"""
 # body not used; SQL is LLM-generated
 ...

Now we can just add some new parameters in the decorator and then it will automatically track the schema and query based on that.

I have also updated the readme with the explanation of the edge cases