feat: refactor schema handling and enhance column/table management

- Removed deprecated Column and Table classes from the schema module, introducing ColumnRef for better reference management.
- Updated the SchemaRetrieverNode to build tables from schema snapshots, incorporating relationships and metadata.
- Enhanced the LogicalValidatorNode to enforce join relationships and validate against column statistics.
- Improved error handling and logging in various pipeline nodes for better traceability.
- Refactored imports across the codebase to utilize the new schema structure, ensuring consistency and clarity.

Author: nadeem4

packages/adapter-sdk/src/nl2sql_adapter_sdk/__init__.py

@@ -4,8 +4,6 @@
 from .contracts import AdapterRequest, ResultError, ResultFrame
 from .protocols import DatasourceAdapterProtocol
 from .schema import (
-    Column,
-    Table,
     TableRef,
     ColumnStatistics,
     ColumnMetadata,
@@ -16,6 +14,7 @@
     SchemaContract,
     SchemaMetadata,
     SchemaSnapshot,
+    ColumnRef   
 )
 
 __all__ = [
@@ -24,8 +23,6 @@
     "ResultError",
     "ResultFrame",
     "DatasourceAdapterProtocol",
-    "Column",
-    "Table",
     "TableRef",
     "ColumnStatistics",
     "ColumnMetadata",
@@ -36,4 +33,5 @@
     "SchemaContract",
     "SchemaMetadata",
     "SchemaSnapshot",
+    "ColumnRef",
 ]

packages/adapter-sdk/src/nl2sql_adapter_sdk/schema.py

@@ -8,22 +8,10 @@
 JsonValue = Union[Scalar, List[Scalar], Dict[str, Scalar]]
 
 
-class Column(BaseModel):
-    """Lightweight column schema for routing/planning."""
-
-    name: str
-    type: Optional[str] = None
-
-    model_config = ConfigDict(extra="allow")
-
-
-class Table(BaseModel):
-    """Lightweight table schema for routing/planning."""
-
-    name: str
-    columns: List[Column] = Field(default_factory=list)
+class ColumnRef(BaseModel):
+    table: TableRef
+    column_name: str
 
-    model_config = ConfigDict(extra="allow")
 
 
 class TableRef(BaseModel):

packages/core/src/nl2sql/indexing/chunk_builder.py

@@ -7,10 +7,9 @@
     ColumnChunk,
     RelationshipChunk,
     MetricChunk,
-    TableRef,
-    ColumnRef,
 )
 from nl2sql.schema import SchemaSnapshot
+from nl2sql_adapter_sdk.schema import TableRef, ColumnRef
 
 
 class SchemaChunkBuilder:
@@ -150,14 +149,13 @@ def _build_column_chunks(self) -> List[ColumnChunk]:
                 column_md = table_md.columns.get(column_name) if table_md else None
 
                 column_ref = ColumnRef(
-                    schema_name=table_ref.schema_name,
-                    table_name=table_ref.table_name,
+                    table=table_ref,
                     column_name=column_name,
                 )
 
                 chunks.append(
                     ColumnChunk(
-                        id=f"schema.column:{column_ref.full_name}:{self.schema_version}",
+                        id=f"schema.column:{column_ref.table.full_name}:{column_ref.column_name}:{self.schema_version}",
                         datasource_id=self.ds_id,
                         column=column_ref,
                         dtype=column_contract.data_type,

packages/core/src/nl2sql/indexing/models.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 from typing import List, Dict, Optional, Literal, Any
 from pydantic import BaseModel, Field
+from nl2sql_adapter_sdk.schema import TableRef, ColumnRef
 
 
 class BaseChunk(BaseModel):
@@ -20,28 +21,6 @@ def get_metadata(self) -> Dict[str, Any]:
             "type": self.type,
         }
 
-class TableRef(BaseModel):
-    schema_name: str
-    table_name: str
-
-    @property
-    def full_name(self) -> str:
-        return f"{self.schema_name}.{self.table_name}"
-
-
-class ColumnRef(BaseModel):
-    schema_name: str
-    table_name: str
-    column_name: str
-
-    @property
-    def full_name(self) -> str:
-        return f"{self.schema_name}.{self.table_name}.{self.column_name}"
-
-    @property
-    def table_full_name(self) -> str:
-        return f"{self.schema_name}.{self.table_name}"
-
 
 class DatasourceChunk(BaseChunk):
     type: Literal["schema.datasource"] = Field(
@@ -106,6 +85,9 @@ def get_metadata(self) -> Dict[str, Any]:
             "table": self.table.full_name,
             "row_count": self.row_count,
             "schema_version": self.schema_version,
+            "description": self.description,
+            "primary_key": ','.join(self.primary_key),
+            "foreign_keys": ','.join(self.foreign_keys),
         }
 
 
@@ -131,7 +113,8 @@ def get_page_content(self) -> str:
             else ""
         )
         return (
-            f"Column: {self.column.full_name}\n"
+            f"Table: {self.column.table.full_name}\n"
+            f"Column: {self.column.column_name}\n"
             f"Type: {self.dtype}\n"
             f"{self.description or ''}\n"
             f"{stats}\n"
@@ -142,11 +125,12 @@ def get_metadata(self) -> Dict[str, Any]:
         return {
             **super().get_metadata(),
             "datasource_id": self.datasource_id,
-            "column": self.column.full_name,
-            "table": self.column.table_full_name,
+            "column": self.column.column_name,
+            "table": self.column.table.full_name,
             "dtype": self.dtype,
             "pii": self.pii,
             "schema_version": self.schema_version,
+            "description": self.description,
         }
 
 class RelationshipChunk(BaseChunk):

packages/core/src/nl2sql/indexing/vector_store.py

@@ -50,6 +50,16 @@ def _initialize_vector_store(self) -> None:
             persist_directory=self.persist_directory,
         )
 
+    def initialize_if_not_exists(self) -> None:
+        """
+        Initializes the vector store if it does not exist.
+        """
+        try:
+            _ = self.vectorstore._collection.count()
+        except Exception:
+            logger.info("Vector store not found, initializing new store.")
+            self._initialize_vector_store()
+
     def is_empty(self) -> bool:
         """
         Checks whether the vector store is empty.
@@ -185,6 +195,7 @@ def retrieve_datasource_candidates(
         Returns:
             Retrieved datasource documents.
         """
+        self.initialize_if_not_exists()
         from nl2sql.common.resilience import VECTOR_BREAKER
 
         @VECTOR_BREAKER
@@ -216,8 +227,10 @@ def retrieve_schema_context(
         Returns:
             Retrieved schema documents.
         """
+        self.initialize_if_not_exists()
         from nl2sql.common.resilience import VECTOR_BREAKER
 
+
         @VECTOR_BREAKER
         def _execute():
             return self.vectorstore.max_marginal_relevance_search(
@@ -254,6 +267,8 @@ def retrieve_column_candidates(
         """
         from nl2sql.common.resilience import VECTOR_BREAKER
 
+        self.initialize_if_not_exists()
+
         @VECTOR_BREAKER
         def _execute():
             return self.vectorstore.max_marginal_relevance_search(
@@ -292,6 +307,8 @@ def retrieve_planning_context(
         """
         from nl2sql.common.resilience import VECTOR_BREAKER
 
+        self.initialize_if_not_exists()
+    
         @VECTOR_BREAKER
         def _execute():
             return self.vectorstore.max_marginal_relevance_search(

packages/core/src/nl2sql/pipeline/nodes/ast_planner/node.py

@@ -74,6 +74,8 @@ def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
                 }
             )
 
+            logger.info(f"Generated Plan: {plan.model_dump_json(indent=2)}")
+
             return {
                 "ast_planner_response": ASTPlannerResponse(plan=plan),
                 "reasoning": [

packages/core/src/nl2sql/pipeline/nodes/ast_planner/prompts.py

@@ -104,18 +104,22 @@
 
     "[INSTRUCTIONS]\n"
     "1. Analyze [USER_QUERY] and [SEMANTIC_CONTEXT].\n"
-    "2. Select tables from [RELEVANT_TABLES]. Assign strict 'ordinal' positions 0..N.\n"
-    "3. Define joins using ONLY table aliases (left_alias/right_alias).\n"
-    "4. Build Expr trees using:\n"
+    "2. Select ONLY tables from [RELEVANT_TABLES]. Assign strict 'ordinal' positions 0..N.\n"
+    "3. When joining, use relationships listed in [RELEVANT_TABLES]. If no relationship exists, do not join.\n"
+    "4. Define joins using ONLY table aliases (left_alias/right_alias).\n"
+    "5. Build Expr trees using:\n"
     "   literal | column | func | binary | unary | case\n"
-    "5. Every list MUST contain `ordinal` fields in ascending order starting at 0.\n"
-    "6. Order lists to match ordinals (0..N) exactly.\n\n"
+    "6. Every list MUST contain `ordinal` fields in ascending order starting at 0.\n"
+    "7. Order lists to match ordinals (0..N) exactly.\n"
+    "8. For literal values on '=' or 'IN', choose values from column stats if available.\n"
+    "9. If no exact match is available, fall back to LIKE but keep the pattern derived from stats/synonyms.\n\n"
 
     "[OUTPUT CONTRACT]\n"
     "- If [EXPECTED_SCHEMA] is provided and non-empty:\n"
     "  - select_items length MUST equal expected_schema length.\n"
     "  - select_items aliases MUST match expected_schema names in the same order.\n"
     "- All table/column references MUST come from [RELEVANT_TABLES].\n"
+    "- Joins MUST use relationships provided in [RELEVANT_TABLES].\n"
     "- The [EXAMPLES] are illustrative; always follow [EXPECTED_SCHEMA] when provided.\n\n"
 
     "[CONSTRAINTS]\n"