sqlseed-ai 0.1.0__py3-none-any.whl

sqlseed_ai/__init__.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from sqlseed.plugins.hookspecs import hookimpl
+from sqlseed_ai.analyzer import SchemaAnalyzer
+
+_SIMPLE_COL_RE = re.compile(
+    r'(^|[_\s])('
+    r'name|email|phone|address|url|uuid|'
+    r'date|time|datetime|timestamp|boolean|'
+    r'int|float|double|real|text|string|'
+    r'char|varchar|blob|byte|id|code|title|'
+    r'description|status|type|category|count|'
+    r'amount|price|value|number|index|order|level'
+    r')($|[_\s])',
+    re.IGNORECASE,
+)
+
+
+class AISqlseedPlugin:
+    def __init__(self) -> None:
+        self._analyzer: SchemaAnalyzer | None = None
+
+    def _get_analyzer(self) -> SchemaAnalyzer:
+        if self._analyzer is None:
+            self._analyzer = SchemaAnalyzer()
+        return self._analyzer
+
+    def _is_simple_column(self, column_name: str, column_type: str) -> bool:
+        return bool(
+            _SIMPLE_COL_RE.search(column_name) or _SIMPLE_COL_RE.search(column_type)
+        )
+
+    @hookimpl
+    def sqlseed_ai_analyze_table(
+        self,
+        table_name: str,
+        columns: list[Any],
+        indexes: list[dict[str, Any]],
+        sample_data: list[dict[str, Any]],
+        foreign_keys: list[Any],
+        all_table_names: list[str],
+    ) -> dict[str, Any] | None:
+        analyzer = self._get_analyzer()
+        return analyzer.analyze_table(
+            table_name=table_name,
+            columns=columns,
+            indexes=indexes,
+            sample_data=sample_data,
+            foreign_keys=foreign_keys,
+            all_table_names=all_table_names,
+        )
+
+    @hookimpl
+    def sqlseed_pre_generate_templates(
+        self,
+        table_name: str,
+        column_name: str,
+        column_type: str,
+        count: int,
+        sample_data: list[Any],
+    ) -> list[Any] | None:
+        if self._is_simple_column(column_name, column_type):
+            return None
+
+        analyzer = self._get_analyzer()
+        try:
+            return analyzer.generate_template_values(
+                column_name=column_name,
+                column_type=column_type,
+                count=min(count, 50),
+                sample_data=sample_data,
+            )
+        except Exception:
+            return None
+
+    @hookimpl
+    def sqlseed_register_providers(self, registry: Any) -> None:
+        pass
+
+    @hookimpl
+    def sqlseed_register_column_mappers(self, mapper: Any) -> None:
+        pass
+
+
+plugin = AISqlseedPlugin()

sqlseed_ai/_client.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlseed._utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_openai_client(config: Any | None = None) -> Any:
+    try:
+        from openai import OpenAI
+
+        from sqlseed_ai.config import AIConfig
+
+        if config is None:
+            config = AIConfig.from_env()
+
+        api_key = config.api_key if hasattr(config, "api_key") else None
+        base_url = config.base_url if hasattr(config, "base_url") else None
+
+        if not api_key:
+            raise ValueError(
+                "AI API key not configured. Set SQLSEED_AI_API_KEY or OPENAI_API_KEY environment variable."
+            )
+
+        return OpenAI(api_key=api_key, base_url=base_url)
+    except ImportError:
+        raise ImportError(
+            "openai is not installed. Install it with: pip install sqlseed-ai"
+        ) from None

sqlseed_ai/_json_utils.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+
+def parse_json_response(content: str) -> dict[str, Any]:
+    cleaned = content.strip()
+    if cleaned.startswith("```json"):
+        cleaned = cleaned[7:]
+    if cleaned.startswith("```"):
+        cleaned = cleaned[3:]
+    if cleaned.endswith("```"):
+        cleaned = cleaned[:-3]
+    cleaned = cleaned.strip()
+
+    try:
+        result = json.loads(cleaned)
+        if isinstance(result, dict):
+            return result
+    except json.JSONDecodeError:
+        pass
+
+    return {}

sqlseed_ai/analyzer.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlseed._utils.logger import get_logger
+from sqlseed_ai._client import get_openai_client
+from sqlseed_ai.config import AIConfig
+
+logger = get_logger(__name__)
+
+SYSTEM_PROMPT = """You are an expert database test data engineer.
+You analyze SQLite table schemas and recommend data generation configurations for the sqlseed toolkit.
+
+## Available Generators
+- string (params: min_length, max_length, charset)
+- integer (params: min_value, max_value)
+- float (params: min_value, max_value, precision)
+- boolean
+- bytes (params: length)
+- name, first_name, last_name
+- email, phone, address, company
+- url, ipv4, uuid
+- date (params: start_year, end_year)
+- datetime (params: start_year, end_year)
+- timestamp
+- text (params: min_length, max_length)
+- sentence, password
+- choice (params: choices)
+- json (params: schema)
+- pattern (params: regex) — generates strings matching a regex pattern
+
+## Key Rules
+1. INTEGER PRIMARY KEY AUTOINCREMENT columns → do NOT include (auto-skip)
+2. Columns with DEFAULT values → do NOT include (auto-skip)
+3. Nullable columns → do NOT include unless they have semantic meaning
+4. Use `pattern` generator with regex for card numbers, codes, IDs with specific formats
+5. Use `derive_from` + `expression` when one column is computed from another
+6. Use `constraints.unique: true` for columns that must be unique
+7. Detect cross-column dependencies: if CutCard4byte = last 8 chars of sCardNo, use derive_from
+8. Detect implicit business associations: if sUserNo appears in multiple tables, note it
+
+## Output Format
+You MUST respond with a valid JSON object (NOT YAML, NOT markdown fences).
+The JSON object must have this exact structure:
+{
+  "name": "table_name",
+  "count": 1000,
+  "columns": [
+    {
+      "name": "column_name",
+      "generator": "generator_name",
+      "params": {"key": "value"}
+    },
+    {
+      "name": "derived_column",
+      "derive_from": "source_column",
+      "expression": "value[-8:]",
+      "constraints": {"unique": true}
+    }
+  ]
+}
+
+IMPORTANT: Do NOT include columns that are PRIMARY KEY AUTOINCREMENT or have DEFAULT values."""
+
+
+class SchemaAnalyzer:
+    def __init__(self, config: AIConfig | None = None) -> None:
+        self._config = config
+
+    def analyze_table(
+        self,
+        table_name: str,
+        columns: list[Any],
+        indexes: list[dict[str, Any]],
+        sample_data: list[dict[str, Any]],
+        foreign_keys: list[Any],
+        all_table_names: list[str],
+    ) -> dict[str, Any] | None:
+        if self._config is None:
+            self._config = AIConfig.from_env()
+
+        if not self._config.api_key:
+            logger.warning("AI API key not configured, skipping analysis")
+            return None
+
+        messages = self.build_initial_messages(
+            table_name=table_name,
+            columns=columns,
+            indexes=indexes,
+            sample_data=sample_data,
+            foreign_keys=foreign_keys,
+            all_table_names=all_table_names,
+        )
+
+        try:
+            return self.call_llm(messages)
+        except Exception as e:
+            logger.warning("AI analysis failed", table_name=table_name, error=str(e))
+            return None
+
+    def build_initial_messages(
+        self,
+        table_name: str,
+        columns: list[Any],
+        indexes: list[dict[str, Any]],
+        sample_data: list[dict[str, Any]],
+        foreign_keys: list[Any],
+        all_table_names: list[str],
+        distribution_profiles: list[dict[str, Any]] | None = None,
+    ) -> list[dict[str, str]]:
+        context = self._build_context(
+            table_name=table_name,
+            columns=columns,
+            indexes=indexes,
+            sample_data=sample_data,
+            foreign_keys=foreign_keys,
+            all_table_names=all_table_names,
+            distribution_profiles=distribution_profiles,
+        )
+        messages: list[dict[str, str]] = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+        ]
+
+        from sqlseed_ai.examples import FEW_SHOT_EXAMPLES
+        for example in FEW_SHOT_EXAMPLES:
+            messages.append({"role": "user", "content": example["input"]})
+            messages.append({"role": "assistant", "content": example["output"]})
+
+        messages.append({"role": "user", "content": context})
+
+        return messages
+
+    def call_llm(self, messages: list[dict[str, str]]) -> dict[str, Any]:
+        if self._config is None:
+            self._config = AIConfig.from_env()
+        if not self._config.api_key:
+            raise ValueError("AI API key not configured")
+
+        client = get_openai_client(self._config)
+        try:
+            response = client.chat.completions.create(
+                model=self._config.model,
+                messages=messages,
+                max_tokens=self._config.max_tokens,
+                temperature=self._config.temperature,
+                response_format={"type": "json_object"},
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"LLM API call failed (model={self._config.model}): {e}"
+            ) from e
+
+        if not response.choices:
+            raise RuntimeError(
+                f"LLM returned no choices (model={self._config.model}). "
+                "The API key or model may be invalid."
+            )
+        content = response.choices[0].message.content
+        if content is None:
+            return {}
+        return self._parse_json_response(content)
+
+    TEMPLATE_SYSTEM_PROMPT = (
+        "You are a data generation assistant. Generate realistic sample values "
+        "for the given database column. Return a JSON object with a 'values' "
+        "array containing the requested number of unique, realistic values. "
+        "Each value must be valid for the column type. Do NOT include explanations."
+    )
+
+    def generate_template_values(
+        self,
+        column_name: str,
+        column_type: str,
+        count: int,
+        sample_data: list[Any],
+    ) -> list[Any]:
+        prompt = (
+            f"Generate {count} realistic sample values for a database column "
+            f"named '{column_name}' with type '{column_type}'."
+        )
+        if sample_data:
+            prompt += f"\nExisting sample values: {sample_data[:5]}"
+        prompt += (
+            f"\nRespond with a JSON object: {{\"values\": [...]}}."
+            f"\nEach value should be a valid {column_type} value."
+        )
+
+        messages = [
+            {"role": "system", "content": self.TEMPLATE_SYSTEM_PROMPT},
+            {"role": "user", "content": prompt},
+        ]
+        result = self.call_llm(messages)
+        return result.get("values", [])
+
+    def _build_context(
+        self,
+        table_name: str,
+        columns: list[Any],
+        indexes: list[dict[str, Any]],
+        sample_data: list[dict[str, Any]],
+        foreign_keys: list[Any],
+        all_table_names: list[str],
+        distribution_profiles: list[dict[str, Any]] | None = None,
+    ) -> str:
+        lines: list[str] = []
+        lines.append(f"# Table: {table_name}")
+        lines.append("")
+
+        self._append_columns_info(lines, columns)
+
+        if indexes:
+            self._append_indexes_info(lines, indexes)
+
+        if foreign_keys:
+            lines.append("")
+            lines.append("## Foreign Keys")
+            for fk in foreign_keys:
+                lines.append(f"- {fk.column} → {fk.ref_table}.{fk.ref_column}")
+
+        if all_table_names:
+            lines.append("")
+            lines.append("## All Tables in Database")
+            lines.append(", ".join(all_table_names))
+
+        if sample_data:
+            lines.append("")
+            lines.append("## Sample Data (existing rows)")
+            for i, row in enumerate(sample_data[:5]):
+                row_str = ", ".join(f"{k}={v}" for k, v in row.items())
+                lines.append(f"  Row {i + 1}: {row_str}")
+
+        if distribution_profiles:
+            self._append_distribution_info(lines, distribution_profiles)
+
+        lines.append("")
+        lines.append(
+            "Please analyze this table schema and recommend "
+            "a complete sqlseed JSON configuration for generating test data."
+        )
+
+        return "\n".join(lines)
+
+    def _append_columns_info(
+        self,
+        lines: list[str],
+        columns: list[Any],
+    ) -> None:
+        lines.append("## Columns")
+        for col in columns:
+            parts = [f"- {col.name}: {col.type}"]
+            if col.is_primary_key:
+                parts.append("PRIMARY KEY")
+            if col.is_autoincrement:
+                parts.append("AUTOINCREMENT")
+            if col.nullable:
+                parts.append("NULLABLE")
+            if col.default is not None:
+                parts.append(f"DEFAULT={col.default}")
+            if not col.nullable and col.default is None and not col.is_primary_key:
+                parts.append("NOT NULL")
+            lines.append(" ".join(parts))
+
+    def _append_indexes_info(
+        self,
+        lines: list[str],
+        indexes: list[dict[str, Any]],
+    ) -> None:
+        lines.append("")
+        lines.append("## Indexes")
+        for idx in indexes:
+            unique_str = "UNIQUE " if idx.get("unique") else ""
+            cols_str = ", ".join(idx.get("columns", []))
+            lines.append(f"- {unique_str}INDEX ({cols_str})")
+
+    def _append_distribution_info(
+        self,
+        lines: list[str],
+        distribution_profiles: list[dict[str, Any]],
+    ) -> None:
+        lines.append("")
+        lines.append("## Column Distribution (from existing data)")
+        for profile in distribution_profiles:
+            col = profile["column"]
+            distinct = profile.get("distinct_count", "?")
+            null_ratio = profile.get("null_ratio", 0)
+            lines.append(
+                f"- {col}: {distinct} distinct values, {null_ratio:.1%} null"
+            )
+            top_values = profile.get("top_values", [])
+            if top_values:
+                top_str = ", ".join(
+                    f"{tv['value']}({tv['frequency']:.0%})"
+                    for tv in top_values[:3]
+                )
+                lines.append(f"  Top values: {top_str}")
+            vr = profile.get("value_range")
+            if vr:
+                lines.append(f"  Range: [{vr['min']}, {vr['max']}]")
+
+    def _parse_json_response(self, content: str) -> dict[str, Any]:
+        from sqlseed_ai._json_utils import parse_json_response
+
+        return parse_json_response(content)
+

sqlseed_ai/config.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import os
+
+from pydantic import BaseModel, Field
+
+
+class AIConfig(BaseModel):
+    api_key: str | None = None
+    model: str = "qwen3-coder-plus"
+    base_url: str | None = None
+    temperature: float = Field(default=0.3, ge=0.0, le=2.0)
+    max_tokens: int = Field(default=4096, gt=0)
+
+    @classmethod
+    def from_env(cls) -> AIConfig:
+        api_key = os.environ.get("SQLSEED_AI_API_KEY") or os.environ.get("OPENAI_API_KEY")
+        base_url = os.environ.get("SQLSEED_AI_BASE_URL") or os.environ.get("OPENAI_BASE_URL")
+        model = os.environ.get("SQLSEED_AI_MODEL", cls.model_fields["model"].default)
+        return cls(api_key=api_key, base_url=base_url, model=model)

sqlseed_ai/errors.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class ErrorSummary:
+    error_type: str
+    message: str
+    column: str | None
+    retryable: bool
+
+    def to_prompt_str(self) -> str:
+        parts = [f"Error Type: {self.error_type}", f"Message: {self.message}"]
+        if self.column:
+            parts.append(f"Affected Column: {self.column}")
+        return "\n".join(parts)
+
+
+def summarize_error(exc: Exception) -> ErrorSummary:
+    import json as _json
+
+    try:
+        from pydantic import ValidationError
+
+        if isinstance(exc, ValidationError):
+            first = exc.errors()[0]
+            loc = " → ".join(str(part) for part in first["loc"])
+            col_name = _extract_column_from_pydantic_loc(first["loc"])
+            return ErrorSummary(
+                error_type="pydantic_validation",
+                message=f"Field '{loc}': {first['msg']} (type={first['type']})",
+                column=col_name,
+                retryable=True,
+            )
+    except ImportError:
+        pass
+
+    if isinstance(exc, _json.JSONDecodeError):
+        return ErrorSummary(
+            error_type="json_syntax",
+            message=f"JSON parsing failed at position {exc.pos}: {exc.msg}",
+            column=None,
+            retryable=True,
+        )
+
+    if isinstance(exc, AttributeError) and "generate_" in str(exc):
+        gen_name = _extract_generator_name(str(exc))
+        return ErrorSummary(
+            error_type="unknown_generator",
+            message=(
+                f"Generator '{gen_name}' does not exist. "
+                "Use one of the available generators listed in the system prompt."
+            ),
+            column=None,
+            retryable=True,
+        )
+
+    try:
+        from sqlseed.generators.stream import UnknownGeneratorError
+        if isinstance(exc, UnknownGeneratorError):
+            return ErrorSummary(
+                error_type="unknown_generator",
+                message=(
+                    f"Generator '{exc.generator_name}' does not exist. "
+                    "Use one of the available generators listed in the system prompt."
+                ),
+                column=exc.column_name,
+                retryable=True,
+            )
+    except ImportError:
+        pass
+
+    exc_type_name = type(exc).__name__
+    exc_module = str(getattr(type(exc), "__module__", ""))
+    if "ExpressionTimeout" in exc_type_name or "simpleeval" in exc_module:
+        return ErrorSummary(
+            error_type="expression_error",
+            message=f"Expression evaluation failed: {str(exc)[:150]}",
+            column=_extract_column_from_message(str(exc)),
+            retryable=True,
+        )
+
+    if isinstance(exc, (FileNotFoundError, PermissionError)):
+        return ErrorSummary(
+            error_type="fatal",
+            message=str(exc)[:200],
+            column=None,
+            retryable=False,
+        )
+
+    return ErrorSummary(
+        error_type="runtime_error",
+        message=f"{exc_type_name}: {str(exc)[:200]}",
+        column=_extract_column_from_message(str(exc)),
+        retryable=True,
+    )
+
+
+def _extract_column_from_pydantic_loc(loc: tuple[Any, ...]) -> str | None:
+    if len(loc) >= 3 and loc[0] == "columns":
+        if hasattr(loc[2], "value"):
+            return loc[2].value.get("name") if isinstance(loc[2].value, dict) else str(loc[2].value)
+        return str(loc[2])
+    if len(loc) >= 2 and loc[0] == "columns":
+        return str(loc[1])
+    return None
+
+
+def _extract_column_from_message(msg: str) -> str | None:
+    match = re.search(r"column[:\s]+'?(\w+)'?", msg, re.IGNORECASE)
+    return match.group(1) if match else None
+
+
+def _extract_generator_name(msg: str) -> str:
+    match = re.search(r"generate_(\w+)", msg)
+    return match.group(1) if match else "unknown"

sqlseed_ai/examples.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import json
+
+FEW_SHOT_EXAMPLES: list[dict[str, str]] = [
+    {
+        "input": """# Table: users
+## Columns
+- id: INTEGER PRIMARY KEY AUTOINCREMENT
+- name: VARCHAR(50) NOT NULL
+- email: VARCHAR(100) NOT NULL
+- status: INTEGER DEFAULT 1
+- created_at: DATETIME
+## Indexes
+- UNIQUE INDEX (email)
+## All Tables in Database
+users, orders""",
+
+        "output": json.dumps({
+            "name": "users",
+            "count": 1000,
+            "columns": [
+                {"name": "name", "generator": "name"},
+                {"name": "email", "generator": "email", "constraints": {"unique": True}},
+                {"name": "created_at", "generator": "datetime", "params": {"start_year": 2020, "end_year": 2025}},
+            ]
+        }, indent=2),
+    },
+    {
+        "input": """# Table: card_info
+## Columns
+- cardId: INTEGER PRIMARY KEY AUTOINCREMENT
+- sCardNo: VARCHAR(20) NOT NULL
+- sUserNo: VARCHAR(32) NOT NULL
+- CutCard4byte: VARCHAR(8)
+- nStatus: INTEGER DEFAULT 0
+- dCreateTime: DATETIME
+## Indexes
+- UNIQUE INDEX (sCardNo)
+- UNIQUE INDEX (sUserNo)
+## All Tables in Database
+card_info, user_info""",
+
+        "output": json.dumps({
+            "name": "card_info",
+            "count": 1000,
+            "columns": [
+                {
+                    "name": "sCardNo",
+                    "generator": "pattern",
+                    "params": {"regex": "[0-9]{16}"},
+                    "constraints": {"unique": True},
+                },
+                {
+                    "name": "sUserNo",
+                    "generator": "pattern",
+                    "params": {"regex": "U[0-9]{10}"},
+                    "constraints": {"unique": True},
+                },
+                {"name": "CutCard4byte", "derive_from": "sCardNo", "expression": "value[-8:]"},
+                {
+                    "name": "dCreateTime",
+                    "generator": "datetime",
+                    "params": {"start_year": 2023, "end_year": 2025},
+                },
+            ]
+        }, indent=2),
+    },
+    {
+        "input": """# Table: orders
+## Columns
+- id: INTEGER PRIMARY KEY AUTOINCREMENT
+- user_id: INTEGER NOT NULL
+- product_name: VARCHAR(100) NOT NULL
+- quantity: INTEGER NOT NULL
+- unit_price: FLOAT NOT NULL
+- order_status: VARCHAR(20) NOT NULL
+- order_date: DATE
+- notes: TEXT
+## Foreign Keys
+- user_id → users.id
+## Indexes
+- INDEX (user_id)
+## All Tables in Database
+users, orders""",
+
+        "output": json.dumps({
+            "name": "orders",
+            "count": 5000,
+            "columns": [
+                {
+                    "name": "user_id",
+                    "generator": "foreign_key",
+                    "params": {"ref_table": "users", "ref_column": "id"},
+                },
+                {
+                    "name": "product_name",
+                    "generator": "string",
+                    "params": {"min_length": 5, "max_length": 50},
+                },
+                {
+                    "name": "quantity",
+                    "generator": "integer",
+                    "params": {"min_value": 1, "max_value": 100},
+                },
+                {
+                    "name": "unit_price",
+                    "generator": "float",
+                    "params": {"min_value": 0.99, "max_value": 999.99, "precision": 2},
+                },
+                {
+                    "name": "order_status",
+                    "generator": "choice",
+                    "params": {
+                        "choices": [
+                            "pending", "confirmed",
+                            "shipped", "delivered", "cancelled",
+                        ],
+                    },
+                },
+                {
+                    "name": "order_date",
+                    "generator": "date",
+                    "params": {"start_year": 2023, "end_year": 2025},
+                },
+            ]
+        }, indent=2),
+    },
+    {
+        "input": """# Table: employees
+## Columns
+- emp_id: INTEGER PRIMARY KEY AUTOINCREMENT
+- dept_id: INTEGER NOT NULL
+- first_name: VARCHAR(50) NOT NULL
+- last_name: VARCHAR(50) NOT NULL
+- hire_date: DATE NOT NULL
+- salary: INTEGER NOT NULL
+- is_active: BOOLEAN
+- metadata: TEXT
+## Foreign Keys
+- dept_id → departments.id
+## Indexes
+- UNIQUE INDEX (first_name, last_name)
+## All Tables in Database
+departments, employees""",
+
+        "output": json.dumps({
+            "name": "employees",
+            "count": 2000,
+            "columns": [
+                {
+                    "name": "dept_id",
+                    "generator": "foreign_key",
+                    "params": {"ref_table": "departments", "ref_column": "id"},
+                },
+                {"name": "first_name", "generator": "first_name"},
+                {"name": "last_name", "generator": "last_name"},
+                {
+                    "name": "hire_date",
+                    "generator": "date",
+                    "params": {"start_year": 2015, "end_year": 2025},
+                },
+                {
+                    "name": "salary",
+                    "generator": "integer",
+                    "params": {"min_value": 30000, "max_value": 200000},
+                },
+                {"name": "is_active", "generator": "boolean"},
+            ]
+        }, indent=2),
+    },
+]

sqlseed_ai/nl_config.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlseed._utils.logger import get_logger
+from sqlseed_ai._client import get_openai_client
+from sqlseed_ai.config import AIConfig
+
+logger = get_logger(__name__)
+
+
+class NLConfigGenerator:
+
+    def __init__(self, config: Any | None = None) -> None:
+        self._config = config
+
+    def generate(self, description: str, db_path: str | None = None) -> dict[str, Any]:
+        try:
+            client = get_openai_client(self._config)
+            model = AIConfig.model_fields["model"].default
+            if self._config is not None and hasattr(self._config, "model"):
+                model = self._config.model
+
+            schema_info = ""
+            if db_path:
+                schema_info = self._read_schema(db_path)
+
+            prompt = (
+                f"Generate a sqlseed configuration based on this description:\n"
+                f'"{description}"\n\n'
+            )
+            if schema_info:
+                prompt += f"Database schema:\n{schema_info}\n\n"
+
+            prompt += (
+                "The JSON should follow this structure:\n"
+                '{"db_path": "path", "provider": "mimesis", "locale": "en_US", '
+                '"tables": [{"name": "table_name", "count": 1000, '
+                '"columns": [{"name": "column_name", "generator": "generator_name", "params": {}}]}]}\n\n'
+                "Respond with valid JSON only."
+            )
+
+            response = client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=2048,
+                temperature=0.5,
+                response_format={"type": "json_object"},
+            )
+
+            content = response.choices[0].message.content if response.choices else None
+            if content is None:
+                return {}
+
+            from sqlseed_ai._json_utils import parse_json_response
+
+            return parse_json_response(content)
+
+        except Exception as e:
+            logger.warning("NL config generation failed", error=e)
+            return {}
+
+    def _read_schema(self, db_path: str) -> str:
+        try:
+            from sqlseed.database.raw_sqlite_adapter import RawSQLiteAdapter
+
+            adapter = RawSQLiteAdapter()
+            adapter.connect(db_path)
+            tables = adapter.get_table_names()
+            lines: list[str] = []
+            for table in tables:
+                columns = adapter.get_column_info(table)
+                col_desc = ", ".join(
+                    f"{c.name}({c.type})" for c in columns
+                )
+                lines.append(f"  {table}: {col_desc}")
+            adapter.close()
+            return "\n".join(lines)
+        except Exception:
+            return ""

sqlseed_ai/provider.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class AIProvider:
+
+    @property
+    def name(self) -> str:
+        return "ai"
+
+    def set_locale(self, locale: str) -> None:
+        pass
+
+    def set_seed(self, seed: int) -> None:
+        pass
+
+    def generate_string(self, **kwargs: Any) -> str:
+        return ""
+
+    def generate_integer(self, **kwargs: Any) -> int:
+        return 0
+
+    def generate_float(self, **kwargs: Any) -> float:
+        return 0.0
+
+    def generate_boolean(self) -> bool:
+        return False
+
+    def generate_bytes(self, **kwargs: Any) -> bytes:
+        return b""
+
+    def generate_name(self) -> str:
+        return ""
+
+    def generate_first_name(self) -> str:
+        return ""
+
+    def generate_last_name(self) -> str:
+        return ""
+
+    def generate_email(self) -> str:
+        return ""
+
+    def generate_phone(self) -> str:
+        return ""
+
+    def generate_address(self) -> str:
+        return ""
+
+    def generate_company(self) -> str:
+        return ""
+
+    def generate_url(self) -> str:
+        return ""
+
+    def generate_ipv4(self) -> str:
+        return ""
+
+    def generate_uuid(self) -> str:
+        return ""
+
+    def generate_date(self, **kwargs: Any) -> str:
+        return ""
+
+    def generate_datetime(self, **kwargs: Any) -> str:
+        return ""
+
+    def generate_timestamp(self) -> int:
+        return 0
+
+    def generate_text(self, **kwargs: Any) -> str:
+        return ""
+
+    def generate_sentence(self) -> str:
+        return ""
+
+    def generate_password(self, **kwargs: Any) -> str:
+        return ""
+
+    def generate_choice(self, choices: list[Any]) -> Any:
+        return choices[0] if choices else None
+
+    def generate_json(self, **kwargs: Any) -> str:
+        return "{}"
+
+    def generate_pattern(self, *, regex: str, **kwargs: Any) -> str:
+        return ""

sqlseed_ai/refiner.py

@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from sqlseed._utils.logger import get_logger
+from sqlseed_ai.errors import ErrorSummary, summarize_error
+
+if TYPE_CHECKING:
+    from sqlseed_ai.analyzer import SchemaAnalyzer
+
+logger = get_logger(__name__)
+
+
+class AISuggestionFailedError(Exception):
+    pass
+
+
+class AiConfigRefiner:
+
+    def __init__(
+        self,
+        analyzer: SchemaAnalyzer,
+        db_path: str,
+        *,
+        cache_dir: str | None = None,
+    ) -> None:
+        self._analyzer = analyzer
+        self._db_path = db_path
+        self._cache_dir = Path(cache_dir) if cache_dir else Path(".sqlseed_cache/ai_configs")
+
+    def generate_and_refine(
+        self,
+        table_name: str,
+        *,
+        max_retries: int = 3,
+        no_cache: bool = False,
+    ) -> dict[str, Any]:
+        from sqlseed.core.orchestrator import DataOrchestrator
+
+        with DataOrchestrator(self._db_path) as orch:
+            schema_hash = self._compute_schema_hash(orch, table_name)
+
+            if not no_cache:
+                cached = self.get_cached_config(table_name, schema_hash)
+                if cached is not None:
+                    logger.info("Using cached AI config", table_name=table_name)
+                    return cached
+
+            schema_ctx = orch.get_schema_context(table_name)
+
+            initial_messages = self._analyzer.build_initial_messages(
+                table_name=schema_ctx["table_name"],
+                columns=schema_ctx["columns"],
+                indexes=schema_ctx["indexes"],
+                sample_data=schema_ctx["sample_data"],
+                foreign_keys=schema_ctx["foreign_keys"],
+                all_table_names=schema_ctx["all_table_names"],
+                distribution_profiles=schema_ctx.get("distribution"),
+            )
+
+            messages_history = list(initial_messages)
+
+            for attempt in range(max_retries + 1):
+                messages = list(messages_history)
+                try:
+                    config_dict = self._analyzer.call_llm(messages)
+                except Exception as e:
+                    error = summarize_error(e)
+                    if not error.retryable:
+                        raise AISuggestionFailedError(
+                            f"Non-retryable error: {error.message}"
+                        ) from e
+                    if attempt == max_retries:
+                        raise AISuggestionFailedError(
+                            f"Failed after {max_retries} retries. Last error: {error.message}"
+                        ) from e
+                    logger.info(
+                        "LLM API call failed, retrying",
+                        attempt=attempt + 1,
+                        max_retries=max_retries,
+                        error=error.message,
+                    )
+                    continue
+
+                error = self._validate_config(orch, table_name, config_dict)
+
+                if error is None:
+                    logger.info(
+                        "AI config validated successfully",
+                        table_name=table_name,
+                        attempts=attempt + 1,
+                    )
+                    self._cache_successful_config(table_name, config_dict, schema_hash)
+                    return config_dict
+
+                if not error.retryable:
+                    raise AISuggestionFailedError(
+                        f"Non-retryable error: {error.message}"
+                    )
+
+                if attempt == max_retries:
+                    logger.warning(
+                        "AI config refinement exhausted all retries",
+                        table_name=table_name,
+                        last_error=error.error_type,
+                    )
+                    raise AISuggestionFailedError(
+                        f"Failed after {max_retries} retries. Last error: {error.message}"
+                    )
+
+                logger.info(
+                    "AI config refinement attempt",
+                    attempt=attempt + 1,
+                    max_retries=max_retries,
+                    error_type=error.error_type,
+                    column=error.column,
+                )
+
+                messages_history.append({
+                    "role": "assistant",
+                    "content": json.dumps(config_dict, ensure_ascii=False),
+                })
+                messages_history.append({
+                    "role": "user",
+                    "content": self._build_refinement_prompt(error, attempt, max_retries),
+                })
+
+        raise AISuggestionFailedError("Unexpected state")
+
+    def _compute_schema_hash(self, orch: Any, table_name: str) -> str:
+        column_names = orch.get_column_names(table_name)
+        raw = "|".join(sorted(column_names))
+        return hashlib.md5(raw.encode()).hexdigest()[:12]
+
+    def _validate_config(
+        self,
+        orch: Any,
+        table_name: str,
+        config_dict: dict[str, Any],
+    ) -> ErrorSummary | None:
+        from sqlseed.config.models import TableConfig
+
+        try:
+            table_config = TableConfig(**config_dict)
+        except Exception as e:
+            return summarize_error(e)
+
+        actual_columns = orch.get_column_names(table_name)
+        skippable_cols = orch.get_skippable_columns(table_name)
+        suggestable_cols = actual_columns - skippable_cols
+
+        for col_cfg in table_config.columns:
+            if col_cfg.name not in actual_columns:
+                return ErrorSummary(
+                    error_type="column_mismatch",
+                    message=(
+                        f"Column '{col_cfg.name}' does not exist in table "
+                        f"'{table_name}'. Available columns: "
+                        f"{sorted(actual_columns)}"
+                    ),
+                    column=col_cfg.name,
+                    retryable=True,
+                )
+
+        if suggestable_cols and len(table_config.columns) == 0:
+            return ErrorSummary(
+                error_type="empty_config",
+                message=(
+                    f"No column suggestions provided for table '{table_name}'. "
+                    f"There are {len(suggestable_cols)} suggestable columns: "
+                    f"{sorted(suggestable_cols)}. "
+                    "Please provide generator suggestions for at least the "
+                    "non-default, non-autoincrement columns."
+                ),
+                column=None,
+                retryable=True,
+            )
+
+        try:
+            orch.preview_table(
+                table_name=table_name,
+                count=5,
+                column_configs=table_config.columns,
+            )
+        except Exception as e:
+            return summarize_error(e)
+
+        return None
+
+    def _build_refinement_prompt(
+        self,
+        error: ErrorSummary,
+        attempt: int,
+        max_retries: int,
+    ) -> str:
+        parts = [
+            "Your previous configuration contained an error. Please fix it.",
+            "",
+            "## Error Details",
+            error.to_prompt_str(),
+            "",
+            "## Instructions",
+            "- Only fix the column(s) mentioned in the error.",
+            "- Do NOT modify other column configurations that were working correctly.",
+            "- Return the COMPLETE configuration JSON "
+            "with only the problematic parts corrected.",
+            "- If you are unsure how to fix the error, "
+            "use 'string' generator as a safe fallback.",
+            "",
+            f"This is refinement attempt {attempt + 1} of {max_retries}.",
+        ]
+
+        if attempt >= max_retries - 1:
+            parts.append(
+                "WARNING: This is the LAST attempt. "
+                "Use the simplest possible generators to ensure validity."
+            )
+
+        return "\n".join(parts)
+
+    def _cache_successful_config(
+        self,
+        table_name: str,
+        config_dict: dict[str, Any],
+        schema_hash: str,
+    ) -> None:
+        try:
+            self._cache_dir.mkdir(parents=True, exist_ok=True)
+            cache_file = self._cache_dir / f"{table_name}.json"
+            entry = {
+                "_meta": {
+                    "schema_hash": schema_hash,
+                    "created_at": time.time(),
+                },
+                "config": config_dict,
+            }
+            cache_file.write_text(
+                json.dumps(entry, ensure_ascii=False, indent=2),
+                encoding="utf-8",
+            )
+            logger.debug(
+                "Cached AI config",
+                table_name=table_name,
+                path=str(cache_file),
+                schema_hash=schema_hash,
+            )
+        except Exception as e:
+            logger.debug("Failed to cache AI config", error=str(e))
+
+    def get_cached_config(
+        self,
+        table_name: str,
+        schema_hash: str | None = None,
+    ) -> dict[str, Any] | None:
+        cache_file = self._cache_dir / f"{table_name}.json"
+        if cache_file.exists():
+            try:
+                entry = json.loads(cache_file.read_text(encoding="utf-8"))
+                if isinstance(entry, dict) and "_meta" in entry:
+                    cached_hash = entry["_meta"].get("schema_hash", "")
+                    if schema_hash and cached_hash != schema_hash:
+                        logger.debug(
+                            "Cache schema hash mismatch, invalidating",
+                            table_name=table_name,
+                            cached_hash=cached_hash,
+                            current_hash=schema_hash,
+                        )
+                        return None
+                    return entry.get("config")
+                return entry
+            except Exception:
+                pass
+        return None

sqlseed_ai/suggest.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlseed._utils.logger import get_logger
+from sqlseed_ai._client import get_openai_client
+from sqlseed_ai.config import AIConfig
+
+logger = get_logger(__name__)
+
+
+class ColumnSuggester:
+
+    def __init__(self, config: Any | None = None) -> None:
+        self._config = config
+
+    def suggest(
+        self,
+        column_name: str,
+        column_type: str,
+        table_name: str,
+        all_column_names: list[str],
+    ) -> dict[str, Any] | None:
+        try:
+            client = get_openai_client(self._config)
+            model = AIConfig.model_fields["model"].default
+            if self._config is not None and hasattr(self._config, "model"):
+                model = self._config.model
+
+            prompt = (
+                f"Given a SQLite table '{table_name}' with columns {all_column_names}, "
+                f"the column '{column_name}' has type '{column_type}'. "
+                f"Suggest the best data generator name and params for this column. "
+                f"Available generators: string, integer, float, boolean, bytes, "
+                f"name, first_name, last_name, email, phone, address, company, "
+                f"url, ipv4, uuid, date, datetime, timestamp, text, sentence, "
+                f"password, choice, json, foreign_key, pattern. "
+                f'Respond in JSON format: {{"generator": "...", "params": {{}}}}'
+            )
+
+            response = client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=256,
+                temperature=0.3,
+                response_format={"type": "json_object"},
+            )
+
+            content = response.choices[0].message.content if response.choices else None
+            if content is None:
+                return None
+
+            from sqlseed_ai._json_utils import parse_json_response
+
+            result = parse_json_response(content)
+            if "generator" in result:
+                return result
+
+        except Exception as e:
+            logger.warning("AI suggestion failed", column_name=column_name, error=e)
+
+        return None

sqlseed_ai-0.1.0.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: sqlseed-ai
+Version: 0.1.0
+Summary: AI-powered data generation plugin for sqlseed
+Project-URL: Homepage, https://github.com/sunbos/sqlseed
+Project-URL: Repository, https://github.com/sunbos/sqlseed/tree/main/plugins/sqlseed-ai
+Author-email: SunBo <1443584939@qq.com>
+License-Expression: AGPL-3.0-or-later
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+Requires-Dist: openai>=1.0
+Requires-Dist: pydantic>=2.0
+Requires-Dist: sqlseed

sqlseed_ai-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+sqlseed_ai/__init__.py,sha256=KD8UutkBS4td2IaeEcG1YlDPafiMEFYnX1A-2_X0qaw,2450
+sqlseed_ai/_client.py,sha256=0lLJD03sCvJPeCZXYMUYrdi_RVGq9Xhufa2oZt_TvL8,891
+sqlseed_ai/_json_utils.py,sha256=TByhP4wvxj17aiQe9-Lgx88k4tqIa2RoOBWR7i-xgB4,556
+sqlseed_ai/analyzer.py,sha256=aetDgwymMBDop14Qhu2kLIlTP6AWJ5m1xKPQNVX-qgU,10554
+sqlseed_ai/config.py,sha256=Zc3N1EOtbuAcfxo48nftKDJK1paR4Av8x0lYcYojTmg,713
+sqlseed_ai/errors.py,sha256=W43rjOWaP1zFUSltXRtpB9CXalEP_qgHzMfzdmD418c,3846
+sqlseed_ai/examples.py,sha256=Sfl9JVjGiMjwXm9NuuHZPmnKKGTvg_yl599w5z8O9Yk,5322
+sqlseed_ai/nl_config.py,sha256=Wcf8ujcoA98XAKjUA6HkReuQplodCC83afDL7UIFDP4,2781
+sqlseed_ai/provider.py,sha256=JNr4Gsr0-cXH3hRxVBogrXehIM4YHudWJhVJgwzEdiM,1853
+sqlseed_ai/refiner.py,sha256=br0wVQYGm-HuN3FW5x8wAyOgFqDGonFGsEx3INCQVnc,9896
+sqlseed_ai/suggest.py,sha256=EckLx5q5CgwK9RJRsslOJTjJG6DQNufdRqCOxmhr8VU,2196
+sqlseed_ai-0.1.0.dist-info/METADATA,sha256=RcpIcb4GFHQfGemI1UdxjgRcuSS7Ds5uDfNpnHs3U2c,826
+sqlseed_ai-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+sqlseed_ai-0.1.0.dist-info/entry_points.txt,sha256=zEHu6heoDb7qXrooVmJN10y-4LSWIhitMNdmcovm1UE,33
+sqlseed_ai-0.1.0.dist-info/RECORD,,

sqlseed_ai-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

sqlseed_ai-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[sqlseed]
+ai = sqlseed_ai:plugin