Schema lookup implementation

.gitignore

@@ -37,4 +37,5 @@ ENV/
 # Misc
 .DS_Store
 
-__sql__
+__sql__
+fixtures

foundation_sql/db.py

@@ -10,7 +10,8 @@
 from sqlalchemy import create_engine, text
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import SQLAlchemyError
-
+from sqlalchemy import MetaData
+from sqlalchemy.schema import CreateTable
 from jinja2sql import Jinja2SQL
 from datetime import datetime
 
@@ -232,7 +233,28 @@ def is_empty(self) -> bool:
             True if no rows, False otherwise
         """
         return len(self.rows) == 0
+
+
+# Function to load the schema from the database
+def extract_schema_from_db(db_url: str) -> str:
+    """Extract the schema from the database.
+
+    Args:
+        db_url: Database URL to use
+
+    Returns:
+        Schema as a string
+    """
+    engine = create_engine(db_url)
+    metadata = MetaData()
+    metadata.reflect(bind=engine)
+
+    schema_lines = []
+    for table in metadata.sorted_tables:
+        ddl = str(CreateTable(table).compile(engine))
+        schema_lines.append(ddl + ";")
 
+    return "\n\n".join(schema_lines)
 
 
 def get_db(db_url: str) -> Database:
@@ -280,7 +302,8 @@ def parse_query_to_pydantic(data: Dict[str, Any], model_class: Type[BaseModel])
 
     # Check the response type and transform accordingly
     if model_class == int:
-        return int(unflattened_data["result"])
+        # FIX : STILL ONLY GETS FIRST LINE OF RESPONSE
+        return int(next(iter(unflattened_data.values())))
     elif model_class == NoneType:
         return None
 

foundation_sql/prompts.md

@@ -1,12 +1,12 @@
 You are an expert SQL developer. Write one or more SQL queries that can perform the actions as explained by the user. Ensure, the SQL query is usable across sqlite and postgresql. The SQL template generated is a jinja2 template - so jinja2 syntax can be used.
 
-1. Start with a -- comment to document the function name, parameters and docstring, explaining what the SQL query does.
+1. Start with a comment to document the function name, parameters and docstring, explaining what the SQL query does. Make sure to start comments with `--` (Only 2 dashes, no more , no less)
 2. Use jinja2 template to generate SQL
 3. When accessing nested fields handle cases if they aren't defined. Use default filter with None value for such cases e.g.
-{{user.zip_code|default(None)}}
+   {{user.zip_code|default(None)}}
 4. Ensure response rows can be parsed into Pydantic model. As long as the model fields are named the same as the columns in the SQL query. It also supports nested models by using double underscores to separate nested fields.
-5. For complex tasks, more than one queries can be run, separated by ';'
-6. Only respond with a single ```sql``` block which contains all queries.
+5. For complex tasks, more than one queries can be run, separated by ';', Make sure queries end with ';'.
+6. Only respond with a single `sql` block which contains all queries.
 7. No other explanation is necessary
 8. For insert queries, avoid any RETURNING clause. Let it return the default.
 9. We use jinja2 syntax to generate SQL - so parameters don't need to be quoted e.g. use {{user.zip_code|default(None)}} and not '{{user.zip_code|default(None)}}'
@@ -19,18 +19,18 @@ You are an expert SQL developer. Write one or more SQL queries that can perform
 Here is an example
 
 def get_task(workspace: schema.Workspace, task_no: int) -> schema.Task:
-    """
-        Creates and returns a Task object, for the provided workspace and task_no
-    """
-    pass
-
+"""
+Creates and returns a Task object, for the provided workspace and task_no
+"""
+pass
 
 The SQL generated would look like the following
+
 ```sql
-    --- def get_task(workspace: schema.Workspace, task_no: int) -> schema.Task
-    --- Creates and returns a Task object, for the provided workspace and task_no
-    --- Expects task_no and workspace.id are defined. If no tasks are found, returns None
-    SELECT 
+    -- def get_task(workspace: schema.Workspace, task_no: int) -> schema.Task;
+    -- Creates and returns a Task object, for the provided workspace and task_no;
+    -- Expects task_no and workspace.id are defined. If no tasks are found, returns None;
+    SELECT
         t.id as `id`,
         t.task_no as `task_no`,
         t.title as `title`,
@@ -55,7 +55,7 @@ The SQL generated would look like the following
     LEFT JOIN agents a ON t.agent_id = a.id
     LEFT JOIN models m ON a.model_id = m.id
     LEFT JOIN workspace_tasks wt ON t.id = wt.task_id
-    WHERE t.task_no = {{task_no}} AND wt.workspace_id = {{workspace.id}}
+    WHERE t.task_no = {{task_no}} AND wt.workspace_id = {{workspace.id}};
 ```
 
 Below are the real specifications for which query needs to be generated.

foundation_sql/query.py

@@ -57,15 +57,21 @@ def __init__(
         self.name = name
         self.regen = regen
         self.cache_dir = cache_dir
-        self.schema = schema or self.load_file(schema_path)
+        self.db_url = db_url or os.environ.get("DATABASE_URL")
+        if ( not self.db_url):
+            raise ValueError(f"Database URL not provided either through constructor or DATABASE_URL environment variable")
+
+        if (not schema and not schema_path):
+            # Load the schema from the database
+            self.schema = db.extract_schema_from_db(self.db_url)
+
+        else:
+            self.schema = schema or self.load_file(schema_path)
+
         if system_prompt or system_prompt_path:
             self.system_prompt = system_prompt or self.load_file(system_prompt_path)
         else:
             self.system_prompt = DEFAULT_SYSTEM_PROMPT
-
-        self.db_url = db_url
-        if not self.db_url:
-            raise ValueError(f"Database URL not provided either through constructor or {db_url_env} environment variable")
 
         # Initialize cache and SQL generator
         self.cache = SQLTemplateCache(cache_dir=cache_dir)

tests/common.py

@@ -2,33 +2,39 @@
 import os
 from foundation_sql import db
 from foundation_sql.query import SQLQueryDecorator
+from typing import Optional
 
 from dotenv import load_dotenv
 load_dotenv()
 
-DB_URL = os.environ.get("DATABSE_URL", "sqlite:///:memory:")
+# DB_URL = os.environ.get("DATABASE_URL", "sqlite:///:memory:")
 
-def create_query(schema):
+def create_query (schema: Optional[str] = None, db_url : Optional[str] = None):
+    final_db_url = (
+        db_url or 
+        os.environ.get("DATABASE_URL") or 
+        "sqlite:///:memory:"
+    )
     return SQLQueryDecorator(schema=schema, 
-                   db_url=DB_URL, 
+                   db_url=final_db_url, 
                    api_key=os.getenv("OPENAI_API_KEY"),
                    base_url=os.getenv("OPENAI_API_BASE_URL"),
                    model=os.getenv("OPENAI_MODEL"))
 
 class DatabaseTests(unittest.TestCase):
     """Base test class for database-driven tests with common setup and helper methods."""
 
-    db_url = DB_URL
+    db_url = "sqlite:///:memory:"
     schema_sql = None
     schema_path = None
 
     def setUp(self):
         """Create a fresh database connection for each test."""
-        # Re-initialize the schema for each test to ensure clean state
+        #Re-initialize the schema for each test to ensure clean state
         if (self.schema_sql or self.schema_path) and self.db_url:
             db.get_db(self.db_url).init_schema(schema_sql=self.schema_sql, schema_path=self.schema_path)
         else:
-            raise ValueError("At least one of schema_sql, schema_path must be provided along with db_url")
+            pass
 
 
     def tearDown(self):

tests/test_schema_discovery.py

@@ -0,0 +1,54 @@
+from typing import List
+from tests import common
+from pydantic import BaseModel
+from tests.utils import BIKES_DB_PATH, create_bike_db
+
+class Bike(BaseModel):
+    make: str
+    model: str
+    price: int
+
+
+create_bike_db()
+
+query = common.create_query(db_url=f"sqlite:///{BIKES_DB_PATH}")
+
+@query
+def get_bikes() -> List[Bike]:
+    """
+    Gets all bikes.
+    """
+    pass
+
+@query
+def create_bike(bike: Bike) -> Bike:
+    """
+    Creates a new bike.
+    """
+    pass
+
+@query
+def get_total_price() -> int:
+    """
+    Get the total price of all the bikes
+    """
+    pass
+
+
+class TestSchemaDiscovery(common.DatabaseTests):
+    db_url = f"sqlite:///{BIKES_DB_PATH}"
+    schema_sql = None
+
+    def test_schema_discovery(self):
+
+        re_bike = Bike(make="RE", model="Classic", price=600)
+        create_bike(bike=re_bike)
+
+        harley_bike = Bike(make="Harley", model="A very good one", price=500)
+        create_bike(bike = harley_bike)
+
+        bikes = get_bikes()
+        self.assertEqual(len(bikes), 2)
+
+        price = get_total_price()
+        self.assertEqual(price,1100)

tests/utils.py

@@ -0,0 +1,26 @@
+# utils.py
+import os
+import sqlite3
+
+BIKES_DB_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), "fixtures", "bikes.db")
+)
+
+def create_bike_db():
+    os.makedirs(os.path.dirname(BIKES_DB_PATH), exist_ok=True) 
+
+
+    if os.path.exists(BIKES_DB_PATH):
+        os.remove(BIKES_DB_PATH)
+
+    conn = sqlite3.connect(BIKES_DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("""
+    CREATE TABLE bikes (
+        make TEXT NOT NULL,
+        model TEXT NOT NULL,
+        price INTEGER NOT NULL
+    );
+    """)
+    conn.commit()
+    conn.close()