db-chat-widget 0.1.0__py3-none-any.whl

db_chat_widget/__init__.py

@@ -0,0 +1,6 @@
+from db_chat_widget.chat.engine import ChatEngine, ChatResult
+from db_chat_widget.config import Settings
+
+__version__ = "0.1.0"
+
+__all__ = ["ChatEngine", "ChatResult", "Settings", "__version__"]

db_chat_widget/chat/engine.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from sqlalchemy import Engine
+
+from db_chat_widget.config import Settings
+from db_chat_widget.db.connection import build_engine
+from db_chat_widget.db.executor import run_safe_query
+from db_chat_widget.db.introspection import get_schema, schema_to_prompt
+from db_chat_widget.db.safety import UnsafeQueryError
+from db_chat_widget.llm.base import LLMProvider
+from db_chat_widget.llm.factory import create_llm_provider
+
+
+@dataclass
+class ChatResult:
+    question: str
+    sql: str | None
+    columns: list[str]
+    rows: list[dict[str, Any]]
+    truncated: bool
+    answer: str
+    error: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "question": self.question,
+            "sql": self.sql,
+            "columns": self.columns,
+            "rows": self.rows,
+            "truncated": self.truncated,
+            "answer": self.answer,
+            "error": self.error,
+        }
+
+
+class ChatEngine:
+    """Orchestrates: question -> SQL generation -> safety check -> execution -> answer."""
+
+    def __init__(
+        self,
+        settings: Settings,
+        llm_provider: LLMProvider | None = None,
+        engine: Engine | None = None,
+    ):
+        self.settings = settings
+        self.engine = engine or build_engine(settings.db_url)
+        self.llm = llm_provider or create_llm_provider(settings)
+        self._dialect = self.engine.dialect.name
+
+    def ask(self, question: str) -> ChatResult:
+        schema_tables = get_schema(self.engine, allowed_tables=self.settings.allowed_tables)
+        schema_prompt = schema_to_prompt(schema_tables)
+
+        try:
+            sql = self.llm.generate_sql(question, schema_prompt, self._dialect)
+        except Exception as exc:  # noqa: BLE001
+            return ChatResult(
+                question=question, sql=None, columns=[], rows=[], truncated=False,
+                answer="Sorry, I couldn't process that question.", error=f"LLM error: {exc}",
+            )
+
+        if sql.strip().upper() == "CANNOT_ANSWER" or not sql.strip():
+            return ChatResult(
+                question=question, sql=None, columns=[], rows=[], truncated=False,
+                answer="I can't answer that from the available database schema.",
+            )
+
+        try:
+            result = run_safe_query(
+                self.engine,
+                sql,
+                read_only=self.settings.read_only,
+                allowed_tables=self.settings.allowed_tables,
+                max_rows=self.settings.max_rows,
+                query_timeout_seconds=self.settings.query_timeout_seconds,
+            )
+        except UnsafeQueryError as exc:
+            return ChatResult(
+                question=question, sql=sql, columns=[], rows=[], truncated=False,
+                answer="The generated query was blocked for safety reasons.", error=str(exc),
+            )
+        except Exception as exc:  # noqa: BLE001
+            return ChatResult(
+                question=question, sql=sql, columns=[], rows=[], truncated=False,
+                answer="The query failed to execute.", error=str(exc),
+            )
+
+        answer = self.llm.summarize(question, result.sql, result.columns, result.rows)
+        return ChatResult(
+            question=question,
+            sql=result.sql,
+            columns=result.columns,
+            rows=result.rows,
+            truncated=result.truncated,
+            answer=answer,
+        )

db_chat_widget/cli.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+from db_chat_widget.config import Settings
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="db-chat-widget", description="Embeddable database chatbot widget."
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    serve = subparsers.add_parser("serve", help="Start the chatbot API + demo server.")
+    serve.add_argument(
+        "--db-url", default=os.environ.get("DB_CHAT_DB_URL"), help="SQLAlchemy database URL."
+    )
+    serve.add_argument(
+        "--llm-provider",
+        default=os.environ.get("DB_CHAT_LLM_PROVIDER", "anthropic"),
+        choices=["anthropic", "openai", "groq", "ollama"],
+    )
+    serve.add_argument("--llm-model", default=os.environ.get("DB_CHAT_LLM_MODEL"))
+    serve.add_argument("--llm-api-key", default=os.environ.get("DB_CHAT_LLM_API_KEY"))
+    serve.add_argument("--llm-base-url", default=os.environ.get("DB_CHAT_LLM_BASE_URL"))
+    serve.add_argument("--host", default=os.environ.get("DB_CHAT_HOST", "127.0.0.1"))
+    serve.add_argument("--port", type=int, default=int(os.environ.get("DB_CHAT_PORT", "8000")))
+    serve.add_argument(
+        "--max-rows", type=int, default=int(os.environ.get("DB_CHAT_MAX_ROWS", "200"))
+    )
+    serve.add_argument(
+        "--allow-write",
+        action="store_true",
+        help="Allow non-SELECT statements (disabled by default for safety).",
+    )
+    serve.add_argument(
+        "--allowed-table",
+        action="append",
+        dest="allowed_tables",
+        default=None,
+        help="Restrict the chatbot to this table; repeat the flag for multiple tables.",
+    )
+    serve.add_argument("--cors-origin", action="append", dest="cors_origins", default=None)
+    serve.add_argument("--reload", action="store_true")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "serve":
+        if not args.db_url:
+            parser.error("--db-url is required (or set DB_CHAT_DB_URL).")
+
+        settings = Settings(
+            db_url=args.db_url,
+            llm_provider=args.llm_provider,
+            llm_model=args.llm_model,
+            llm_api_key=args.llm_api_key,
+            llm_base_url=args.llm_base_url,
+            read_only=not args.allow_write,
+            max_rows=args.max_rows,
+            allowed_tables=args.allowed_tables,
+        )
+
+        try:
+            import uvicorn
+
+            from db_chat_widget.server.app import create_app
+        except ImportError as exc:
+            parser.error(
+                "The 'serve' command requires the optional server dependencies. "
+                f"Install them with: pip install db-chat-widget[server] ({exc})"
+            )
+            return 1
+
+        app = create_app(settings, cors_origins=args.cors_origins)
+        uvicorn.run(app, host=args.host, port=args.port, reload=args.reload)
+        return 0
+
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

db_chat_widget/config.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Settings:
+    """Runtime configuration for the chatbot engine and server.
+
+    db_url: SQLAlchemy connection URL, e.g. "postgresql://user:pass@host/db"
+    llm_provider: one of "anthropic", "openai", "ollama"
+    llm_model: provider-specific model name; a sane default is picked per provider if omitted
+    llm_api_key: API key for the provider (not required for "ollama")
+    read_only: when True (default), only SELECT statements are allowed
+    max_rows: maximum number of rows returned from any query
+    allowed_tables: optional allowlist restricting which tables the LLM may query
+    """
+
+    db_url: str
+    llm_provider: str = "anthropic"
+    llm_model: str | None = None
+    llm_api_key: str | None = None
+    llm_base_url: str | None = None
+    read_only: bool = True
+    max_rows: int = 200
+    allowed_tables: list[str] | None = field(default=None)
+    query_timeout_seconds: int = 15

db_chat_widget/db/connection.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from sqlalchemy import Engine, create_engine
+from sqlalchemy.engine import make_url
+
+# Default client encoding applied when the database URL doesn't specify one.
+# MySQL/MariaDB servers commonly default to latin1, which mangles accented
+# characters (e.g. French text) unless the client explicitly asks for UTF-8.
+_DEFAULT_MYSQL_CHARSET = "utf8mb4"
+
+
+def build_engine(db_url: str, **engine_kwargs: object) -> Engine:
+    """Create a SQLAlchemy engine, applying a sane default client encoding.
+
+    For MySQL/MariaDB URLs without an explicit `charset` query parameter,
+    `charset=utf8mb4` is added so text data round-trips correctly by default.
+    Any encoding explicitly set in `db_url` is left untouched.
+    """
+    url = make_url(db_url)
+
+    if url.get_backend_name() == "mysql" and "charset" not in url.query:
+        url = url.update_query_dict({"charset": _DEFAULT_MYSQL_CHARSET})
+
+    return create_engine(url, **engine_kwargs)

db_chat_widget/db/executor.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import sqlglot
+from sqlalchemy import Engine, text
+from sqlglot import exp
+
+from db_chat_widget.db.safety import UnsafeQueryError, validate_select_only
+
+
+@dataclass
+class QueryResult:
+    columns: list[str]
+    rows: list[dict[str, Any]]
+    truncated: bool
+    sql: str
+
+
+def _apply_row_limit(sql: str, max_rows: int, dialect: str | None) -> tuple[str, str]:
+    """Cap the query at `max_rows` results.
+
+    Returns (exec_sql, display_sql). exec_sql requests one extra row (max_rows + 1) so
+    truncation can be detected; display_sql reports the effective LIMIT (max_rows) so the
+    off-by-one fetch trick isn't leaked to end users. If the query already has a smaller
+    LIMIT, that LIMIT is left untouched in both and is never reported as truncated.
+    """
+    parsed = sqlglot.parse_one(sql, read=dialect)
+    existing_limit = parsed.args.get("limit")
+    if existing_limit is not None:
+        try:
+            current = int(existing_limit.expression.this)
+            if current <= max_rows:
+                unchanged = parsed.sql(dialect=dialect)
+                return unchanged, unchanged
+        except (TypeError, ValueError, AttributeError):
+            pass
+    exec_sql = parsed.copy().limit(max_rows + 1).sql(dialect=dialect)
+    display_sql = parsed.copy().limit(max_rows).sql(dialect=dialect)
+    return exec_sql, display_sql
+
+
+def run_safe_query(
+    engine: Engine,
+    sql: str,
+    *,
+    read_only: bool = True,
+    allowed_tables: list[str] | None = None,
+    max_rows: int = 200,
+    query_timeout_seconds: int = 15,
+) -> QueryResult:
+    """Validate and execute a single SELECT statement, capping returned rows."""
+    dialect = engine.dialect.name if engine.dialect.name != "sqlite" else None
+
+    if read_only:
+        safe_sql = validate_select_only(sql, allowed_tables=allowed_tables, dialect=dialect)
+    else:
+        parsed = sqlglot.parse_one(sql, read=dialect)
+        if not isinstance(parsed, (exp.Select, exp.Union, exp.With)):
+            raise UnsafeQueryError("Only a single statement is allowed.")
+        safe_sql = parsed.sql(dialect=dialect)
+
+    exec_sql, display_sql = _apply_row_limit(safe_sql, max_rows, dialect)
+
+    with engine.connect() as conn:
+        if engine.dialect.name == "postgresql" and read_only:
+            conn = conn.execution_options(postgresql_readonly=True)
+        cursor_result = conn.execute(text(exec_sql))
+        columns = list(cursor_result.keys())
+        fetched = cursor_result.fetchmany(max_rows + 1)
+
+    truncated = len(fetched) > max_rows
+    rows = [dict(zip(columns, row, strict=True)) for row in fetched[:max_rows]]
+
+    return QueryResult(columns=columns, rows=rows, truncated=truncated, sql=display_sql)

db_chat_widget/db/introspection.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from sqlalchemy import Engine, inspect
+
+
+@dataclass
+class ColumnInfo:
+    name: str
+    type: str
+    nullable: bool
+    primary_key: bool
+
+
+@dataclass
+class TableInfo:
+    name: str
+    columns: list[ColumnInfo]
+    foreign_keys: list[str]
+
+
+def get_schema(engine: Engine, allowed_tables: list[str] | None = None) -> list[TableInfo]:
+    """Introspect the database and return table/column metadata.
+
+    If allowed_tables is given, only those tables are included.
+    """
+    inspector = inspect(engine)
+    tables: list[TableInfo] = []
+
+    for table_name in inspector.get_table_names():
+        if allowed_tables is not None and table_name not in allowed_tables:
+            continue
+
+        pk_columns = set(inspector.get_pk_constraint(table_name).get("constrained_columns") or [])
+        columns = [
+            ColumnInfo(
+                name=col["name"],
+                type=str(col["type"]),
+                nullable=bool(col.get("nullable", True)),
+                primary_key=col["name"] in pk_columns,
+            )
+            for col in inspector.get_columns(table_name)
+        ]
+
+        fk_descriptions = []
+        for fk in inspector.get_foreign_keys(table_name):
+            if not fk.get("constrained_columns") or not fk.get("referred_table"):
+                continue
+            local_cols = ", ".join(fk["constrained_columns"])
+            ref_cols = ", ".join(fk.get("referred_columns") or [])
+            fk_descriptions.append(f"{local_cols} -> {fk['referred_table']}({ref_cols})")
+
+        tables.append(TableInfo(name=table_name, columns=columns, foreign_keys=fk_descriptions))
+
+    return tables
+
+
+def schema_to_prompt(tables: list[TableInfo]) -> str:
+    """Render schema metadata as compact text suitable for an LLM prompt."""
+    lines: list[str] = []
+    for table in tables:
+        col_parts = []
+        for col in table.columns:
+            marker = " PK" if col.primary_key else ""
+            null_marker = "" if col.nullable else " NOT NULL"
+            col_parts.append(f"{col.name} {col.type}{marker}{null_marker}")
+        lines.append(f"TABLE {table.name} ({', '.join(col_parts)})")
+        for fk in table.foreign_keys:
+            lines.append(f"  FK {fk}")
+    return "\n".join(lines)

db_chat_widget/db/safety.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import sqlglot
+from sqlglot import exp
+
+_WRITE_EXPRESSION_TYPES = (
+    exp.Insert,
+    exp.Update,
+    exp.Delete,
+    exp.Drop,
+    exp.Create,
+    exp.Alter,
+    exp.TruncateTable,
+    exp.Merge,
+    exp.Grant,
+)
+
+
+class UnsafeQueryError(ValueError):
+    """Raised when a generated SQL statement fails the safety checks."""
+
+
+def validate_select_only(
+    sql: str, *, allowed_tables: list[str] | None = None, dialect: str | None = None
+) -> str:
+    """Ensure `sql` contains exactly one read-only SELECT statement.
+
+    Raises UnsafeQueryError if the query writes data, runs DDL, contains
+    multiple statements, or references a table outside `allowed_tables`.
+    Returns the normalized (single) statement on success.
+    """
+    try:
+        statements = sqlglot.parse(sql, read=dialect)
+    except Exception as exc:  # noqa: BLE001 - surface parser errors as unsafe
+        raise UnsafeQueryError(f"Could not parse SQL: {exc}") from exc
+
+    statements = [s for s in statements if s is not None]
+    if len(statements) == 0:
+        raise UnsafeQueryError("No SQL statement found.")
+    if len(statements) > 1:
+        raise UnsafeQueryError("Only a single SQL statement is allowed.")
+
+    statement = statements[0]
+
+    if not isinstance(statement, (exp.Select, exp.Union, exp.With)):
+        kind = type(statement).__name__
+        raise UnsafeQueryError(f"Only SELECT statements are allowed, got: {kind}")
+
+    for write_type in _WRITE_EXPRESSION_TYPES:
+        if list(statement.find_all(write_type)):
+            raise UnsafeQueryError(f"Query contains a disallowed operation: {write_type.__name__}")
+
+    if allowed_tables is not None:
+        referenced = {table.name for table in statement.find_all(exp.Table)}
+        disallowed = referenced - set(allowed_tables)
+        if disallowed:
+            names = ", ".join(sorted(disallowed))
+            raise UnsafeQueryError(f"Query references disallowed table(s): {names}")
+
+    return statement.sql(dialect=dialect)

db_chat_widget/llm/anthropic_provider.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from db_chat_widget.llm.base import SYSTEM_PROMPT_TEMPLATE, LLMProvider
+
+DEFAULT_MODEL = "claude-sonnet-5"
+
+
+class AnthropicProvider(LLMProvider):
+    def __init__(self, api_key: str, model: str | None = None):
+        try:
+            import anthropic
+        except ImportError as exc:  # pragma: no cover
+            raise ImportError(
+                "The 'anthropic' package is required for AnthropicProvider. "
+                "Install it with: pip install anthropic"
+            ) from exc
+
+        self._client = anthropic.Anthropic(api_key=api_key)
+        self._model = model or DEFAULT_MODEL
+
+    def generate_sql(self, question: str, schema_prompt: str, dialect: str) -> str:
+        system = SYSTEM_PROMPT_TEMPLATE.format(dialect=dialect, schema=schema_prompt)
+        response = self._client.messages.create(
+            model=self._model,
+            max_tokens=512,
+            system=system,
+            messages=[{"role": "user", "content": question}],
+        )
+        text = "".join(block.text for block in response.content if hasattr(block, "text"))
+        return self._strip_sql(text)

db_chat_widget/llm/base.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+SYSTEM_PROMPT_TEMPLATE = """You are a database assistant. Given a database schema and a \
+question in natural language, respond with a single {dialect} SQL SELECT statement that \
+answers the question.
+
+Rules:
+- Output ONLY the SQL statement, no explanation, no markdown code fences.
+- Only use SELECT statements. Never write, alter, or delete data.
+- Only reference tables and columns that appear in the schema below.
+- If the question cannot be answered with the given schema, output exactly: \
+CANNOT_ANSWER
+
+Schema:
+{schema}
+"""
+
+
+class LLMProvider(ABC):
+    """Abstract interface for a natural-language-to-SQL backend."""
+
+    @abstractmethod
+    def generate_sql(self, question: str, schema_prompt: str, dialect: str) -> str:
+        """Return a single SQL SELECT statement answering `question`, or "CANNOT_ANSWER"."""
+        raise NotImplementedError
+
+    def summarize(self, question: str, sql: str, columns: list[str], rows: list[dict]) -> str:
+        """Produce a short natural-language summary of the query result.
+
+        Default implementation is a plain, deterministic sentence (no extra LLM call).
+        Providers may override this to ask the model for a nicer summary.
+        """
+        count = len(rows)
+        if count == 0:
+            return "No rows matched your question."
+        if count == 1:
+            return f"Found 1 result: {rows[0]}"
+        return f"Found {count} result(s) across columns: {', '.join(columns)}."
+
+    @staticmethod
+    def _strip_sql(text: str) -> str:
+        cleaned = text.strip()
+        if cleaned.startswith("```"):
+            cleaned = cleaned.strip("`")
+            if cleaned.lower().startswith("sql"):
+                cleaned = cleaned[3:]
+        return cleaned.strip().rstrip(";").strip()

db_chat_widget/llm/factory.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from db_chat_widget.config import Settings
+from db_chat_widget.llm.base import LLMProvider
+
+_PROVIDERS = ("anthropic", "openai", "groq", "ollama")
+
+
+def create_llm_provider(settings: Settings) -> LLMProvider:
+    provider = settings.llm_provider.lower()
+
+    if provider == "anthropic":
+        from db_chat_widget.llm.anthropic_provider import AnthropicProvider
+
+        if not settings.llm_api_key:
+            raise ValueError("llm_api_key is required for the 'anthropic' provider.")
+        return AnthropicProvider(api_key=settings.llm_api_key, model=settings.llm_model)
+
+    if provider == "openai":
+        from db_chat_widget.llm.openai_provider import OpenAIProvider
+
+        if not settings.llm_api_key:
+            raise ValueError("llm_api_key is required for the 'openai' provider.")
+        return OpenAIProvider(api_key=settings.llm_api_key, model=settings.llm_model)
+
+    if provider == "groq":
+        from db_chat_widget.llm.groq_provider import GroqProvider
+
+        if not settings.llm_api_key:
+            raise ValueError("llm_api_key is required for the 'groq' provider.")
+        return GroqProvider(api_key=settings.llm_api_key, model=settings.llm_model)
+
+    if provider == "ollama":
+        from db_chat_widget.llm.ollama_provider import OllamaProvider
+
+        return OllamaProvider(base_url=settings.llm_base_url, model=settings.llm_model)
+
+    raise ValueError(
+        f"Unknown llm_provider '{settings.llm_provider}'. Expected one of {_PROVIDERS}."
+    )

db_chat_widget/llm/groq_provider.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from db_chat_widget.llm.base import SYSTEM_PROMPT_TEMPLATE, LLMProvider
+
+DEFAULT_MODEL = "llama-3.3-70b-versatile"
+
+
+class GroqProvider(LLMProvider):
+    def __init__(self, api_key: str, model: str | None = None):
+        try:
+            import groq
+        except ImportError as exc:  # pragma: no cover
+            raise ImportError(
+                "The 'groq' package is required for GroqProvider. "
+                "Install it with: pip install groq"
+            ) from exc
+
+        self._client = groq.Groq(api_key=api_key)
+        self._model = model or DEFAULT_MODEL
+
+    def generate_sql(self, question: str, schema_prompt: str, dialect: str) -> str:
+        system = SYSTEM_PROMPT_TEMPLATE.format(dialect=dialect, schema=schema_prompt)
+        response = self._client.chat.completions.create(
+            model=self._model,
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user", "content": question},
+            ],
+            temperature=0,
+        )
+        text = response.choices[0].message.content or ""
+        return self._strip_sql(text)

db_chat_widget/llm/ollama_provider.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from db_chat_widget.llm.base import SYSTEM_PROMPT_TEMPLATE, LLMProvider
+
+DEFAULT_MODEL = "llama3.1"
+DEFAULT_BASE_URL = "http://localhost:11434"
+
+
+class OllamaProvider(LLMProvider):
+    def __init__(self, base_url: str | None = None, model: str | None = None):
+        try:
+            import httpx
+        except ImportError as exc:  # pragma: no cover
+            raise ImportError(
+                "The 'httpx' package is required for OllamaProvider. "
+                "Install it with: pip install httpx"
+            ) from exc
+
+        self._httpx = httpx
+        self._base_url = (base_url or DEFAULT_BASE_URL).rstrip("/")
+        self._model = model or DEFAULT_MODEL
+
+    def generate_sql(self, question: str, schema_prompt: str, dialect: str) -> str:
+        system = SYSTEM_PROMPT_TEMPLATE.format(dialect=dialect, schema=schema_prompt)
+        response = self._httpx.post(
+            f"{self._base_url}/api/chat",
+            json={
+                "model": self._model,
+                "stream": False,
+                "options": {"temperature": 0},
+                "messages": [
+                    {"role": "system", "content": system},
+                    {"role": "user", "content": question},
+                ],
+            },
+            timeout=60,
+        )
+        response.raise_for_status()
+        text = response.json()["message"]["content"]
+        return self._strip_sql(text)