Generative BI: Connecting LLMs to Your Enterprise Data Warehouse

import sqlglot
import sqlglot.expressions as exp
from openai import OpenAI
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue
import json
from typing import Optional
from dataclasses import dataclass


@dataclass
class UserContext:
    user_id: str
    tenant_id: str
    allowed_schemas: list[str]  # e.g. ["sales", "marketing"]
    row_level_filter: str       # e.g. "tenant_id = 'acme-corp'"


class GenerativeBIAgent:
    """
    Production Text-to-SQL agent with semantic layer retrieval,
    AST-based SQL validation, and access control enforcement.
    """
    FORBIDDEN_STATEMENT_TYPES = (exp.Insert, exp.Update, exp.Delete,
                                  exp.Drop, exp.Alter, exp.Create)

    def __init__(self, qdrant_host: str = "localhost"):
        self.openai = OpenAI()
        self.qdrant = QdrantClient(host=qdrant_host, port=6333)
        # Two collections: semantic_layer (metric defs), schema_docs (table/column docs)
        self.semantic_collection = "semantic_layer"
        self.schema_collection = "schema_docs"

    def _embed(self, text: str) -> list[float]:
        resp = self.openai.embeddings.create(
            input=text, model="text-embedding-3-small", dimensions=768
        )
        return resp.data[0].embedding

    def _retrieve_semantic_context(self, query: str, tenant_id: str) -> str:
        results = self.qdrant.search(
            collection_name=self.semantic_collection,
            query_vector=self._embed(query),
            query_filter=Filter(must=[
                FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id))
            ]),
            limit=5,
            with_payload=True,
        )
        if not results:
            return ""
        return "\n".join(
            f"METRIC: {r.payload['name']}\nDEFINITION: {r.payload['definition']}"
            for r in results
        )

    def _retrieve_schema_context(self, query: str, allowed_schemas: list[str]) -> str:
        results = self.qdrant.search(
            collection_name=self.schema_collection,
            query_vector=self._embed(query),
            query_filter=Filter(must=[
                FieldCondition(key="schema_name",
                               match=MatchValue(value=s)) for s in allowed_schemas
            ]),
            limit=10,
            with_payload=True,
        )
        return "\n".join(
            f"TABLE: {r.payload['table']} | COLUMN: {r.payload['column']} | "
            f"TYPE: {r.payload['type']} | DESC: {r.payload['description']}"
            for r in results
        )

    def _validate_sql(self, sql: str, user: UserContext) -> tuple[bool, str]:
        try:
            parsed = sqlglot.parse_one(sql, dialect="bigquery")
        except Exception as e:
            return False, f"SQL parse error: {e}"

        # Block DML/DDL statements
        for forbidden_type in self.FORBIDDEN_STATEMENT_TYPES:
            if isinstance(parsed, forbidden_type):
                return False, f"Forbidden statement type: {type(parsed).__name__}"

        # Ensure only allowed schemas are referenced
        for table in parsed.find_all(exp.Table):
            schema = table.db
            if schema and schema not in user.allowed_schemas:
                return False, f"Access denied to schema: {schema}"

        return True, "valid"

    def _enforce_access_control(self, sql: str, user: UserContext) -> str:
        """Rewrite SQL to add row-level security filter."""
        # Append tenant filter as a subquery wrapper
        return f"SELECT * FROM ({sql}) WHERE {user.row_level_filter}"

    def generate_sql(
        self, natural_language_query: str, user: UserContext, max_retries: int = 3
    ) -> Optional[str]:
        semantic_ctx = self._retrieve_semantic_context(
            natural_language_query, user.tenant_id
        )
        schema_ctx = self._retrieve_schema_context(
            natural_language_query, user.allowed_schemas
        )
        system_prompt = f"""You are a SQL generation assistant for BigQuery.
Generate a single SELECT query only. No DML or DDL.

Metric definitions (use these exact definitions for named metrics):
{semantic_ctx or 'No relevant metrics found.'}

Relevant schema:
{schema_ctx or 'No relevant schema found.'}

Respond with JSON: {{"sql": "your query here", "explanation": "one sentence"}}"""

        for attempt in range(max_retries):
            resp = self.openai.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": natural_language_query},
                ],
                temperature=0,
                response_format={"type": "json_object"},
            )
            result = json.loads(resp.choices[0].message.content)
            sql = result.get("sql", "")

            valid, reason = self._validate_sql(sql, user)
            if valid:
                return self._enforce_access_control(sql, user)
            else:
                # Feed validation error back for self-correction
                natural_language_query += f"\n\nPrevious attempt failed validation: {reason}. Correct and retry."

        return None  # Failed after max_retries


# Usage
if __name__ == "__main__":
    agent = GenerativeBIAgent()
    user = UserContext(
        user_id="user-123",
        tenant_id="acme-corp",
        allowed_schemas=["sales", "marketing"],
        row_level_filter="tenant_id = 'acme-corp'"
    )
    sql = agent.generate_sql(
        "What was our NRR trend over the last 6 quarters for enterprise accounts?",
        user=user
    )
    print(sql)