seekdb-cli 0.1.0__py3-none-any.whl

seekdb_cli/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

seekdb_cli/commands/ai.py

@@ -0,0 +1,306 @@
+"""seekdb ai command — use database AI (DBMS_AI_SERVICE, AI_COMPLETE)."""
+
+from __future__ import annotations
+
+import json
+
+import click
+import pymysql
+
+from seekdb_cli import output
+from seekdb_cli.connection import get_connection
+from seekdb_cli.logger import log_operation
+
+# OceanBase AI model types (for CREATE_AI_MODEL)
+AI_MODEL_TYPES = ("dense_embedding", "completion", "rerank")
+
+
+@click.group("ai")
+@click.pass_context
+def ai_cmd(ctx: click.Context) -> None:
+    """AI model management and completion (via database DBMS_AI_SERVICE)."""
+
+
+# ---------------------------------------------------------------------------
+# seekdb ai model ...
+# ---------------------------------------------------------------------------
+
+@ai_cmd.group("model")
+@click.pass_context
+def model_group(ctx: click.Context) -> None:
+    """Manage AI models registered in the database (DBMS_AI_SERVICE)."""
+
+
+@model_group.command("list")
+@click.pass_context
+def model_list(ctx: click.Context) -> None:
+    """List all registered AI models (from oceanbase.DBA_OB_AI_MODELS)."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai model list", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute(
+                "SELECT MODEL_ID, NAME, TYPE, MODEL_NAME "
+                "FROM oceanbase.DBA_OB_AI_MODELS ORDER BY NAME"
+            )
+            rows = cur.fetchall()
+            # DictCursor: keys may be uppercase from view
+            result = []
+            for r in rows:
+                row = dict(r)
+                result.append({
+                    "name": row.get("NAME") or row.get("name"),
+                    "type": row.get("TYPE") or row.get("type"),
+                    "model_name": row.get("MODEL_NAME") or row.get("model_name"),
+                    "model_id": row.get("MODEL_ID") or row.get("model_id"),
+                })
+
+        log_operation("ai model list", ok=True, time_ms=timer.elapsed_ms)
+        output.success(result, time_ms=timer.elapsed_ms, fmt=fmt)
+    except pymysql.err.ProgrammingError as exc:
+        log_operation("ai model list", ok=False, error_code="SQL_ERROR")
+        output.error(
+            "SQL_ERROR",
+            str(exc) + ". Ensure the database supports DBA_OB_AI_MODELS (OceanBase AI).",
+            fmt=fmt,
+        )
+    except pymysql.err.Error as exc:
+        log_operation("ai model list", ok=False, error_code="SQL_ERROR")
+        output.error("SQL_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()
+
+
+@model_group.command("create")
+@click.argument("name")
+@click.option(
+    "--type",
+    "model_type",
+    type=click.Choice(AI_MODEL_TYPES),
+    required=True,
+    help="Model type: dense_embedding, completion, or rerank.",
+)
+@click.option(
+    "--model",
+    "provider_model_name",
+    required=True,
+    help="Provider model name (e.g. BAAI/bge-m3, THUDM/GLM-4-9B-0414).",
+)
+@click.pass_context
+def model_create(
+    ctx: click.Context,
+    name: str,
+    model_type: str,
+    provider_model_name: str,
+) -> None:
+    """Register an AI model via DBMS_AI_SERVICE.CREATE_AI_MODEL. Create an endpoint separately to use it."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai model create", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    config = json.dumps({"type": model_type, "model_name": provider_model_name})
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute(
+                "CALL DBMS_AI_SERVICE.CREATE_AI_MODEL(%s, %s)",
+                (name, config),
+            )
+            conn.commit()
+
+        log_operation("ai model create", ok=True, time_ms=timer.elapsed_ms)
+        output.success(
+            {"name": name, "type": model_type, "model_name": provider_model_name},
+            time_ms=timer.elapsed_ms,
+            fmt=fmt,
+        )
+    except pymysql.err.Error as exc:
+        log_operation("ai model create", ok=False, error_code="SQL_ERROR")
+        output.error("SQL_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()
+
+
+@model_group.command("delete")
+@click.argument("name")
+@click.pass_context
+def model_delete(ctx: click.Context, name: str) -> None:
+    """Drop an AI model via DBMS_AI_SERVICE.DROP_AI_MODEL. Drop endpoints first if any."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai model delete", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute("CALL DBMS_AI_SERVICE.DROP_AI_MODEL(%s)", (name,))
+            conn.commit()
+
+        log_operation("ai model delete", ok=True, time_ms=timer.elapsed_ms)
+        output.success({"deleted": name}, time_ms=timer.elapsed_ms, fmt=fmt)
+    except pymysql.err.Error as exc:
+        log_operation("ai model delete", ok=False, error_code="SQL_ERROR")
+        output.error("SQL_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()
+
+
+# ---------------------------------------------------------------------------
+# seekdb ai model endpoint ...
+# ---------------------------------------------------------------------------
+
+@model_group.group("endpoint")
+@click.pass_context
+def endpoint_group(ctx: click.Context) -> None:
+    """Create or delete AI model endpoints (DBMS_AI_SERVICE)."""
+
+
+@endpoint_group.command("create")
+@click.argument("endpoint_name")
+@click.argument("ai_model_name")
+@click.option("--url", required=True, help="API URL (e.g. https://api.siliconflow.cn/v1/chat/completions).")
+@click.option("--access-key", required=True, help="API key for the service.")
+@click.option("--provider", default="siliconflow", help="Provider name (siliconflow, openai, dashscope, etc.).")
+@click.pass_context
+def endpoint_create(
+    ctx: click.Context,
+    endpoint_name: str,
+    ai_model_name: str,
+    url: str,
+    access_key: str,
+    provider: str,
+) -> None:
+    """Create an AI model endpoint via DBMS_AI_SERVICE.CREATE_AI_MODEL_ENDPOINT."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai model endpoint create", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    config = json.dumps({
+        "ai_model_name": ai_model_name,
+        "url": url,
+        "access_key": access_key,
+        "provider": provider,
+    })
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute(
+                "CALL DBMS_AI_SERVICE.CREATE_AI_MODEL_ENDPOINT(%s, %s)",
+                (endpoint_name, config),
+            )
+            conn.commit()
+
+        log_operation("ai model endpoint create", ok=True, time_ms=timer.elapsed_ms)
+        output.success(
+            {"endpoint": endpoint_name, "ai_model": ai_model_name, "provider": provider},
+            time_ms=timer.elapsed_ms,
+            fmt=fmt,
+        )
+    except pymysql.err.Error as exc:
+        log_operation("ai model endpoint create", ok=False, error_code="SQL_ERROR")
+        output.error("SQL_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()
+
+
+@endpoint_group.command("delete")
+@click.argument("endpoint_name")
+@click.pass_context
+def endpoint_delete(ctx: click.Context, endpoint_name: str) -> None:
+    """Drop an AI model endpoint via DBMS_AI_SERVICE.DROP_AI_MODEL_ENDPOINT."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai model endpoint delete", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute(
+                "CALL DBMS_AI_SERVICE.DROP_AI_MODEL_ENDPOINT(%s)",
+                (endpoint_name,),
+            )
+            conn.commit()
+
+        log_operation("ai model endpoint delete", ok=True, time_ms=timer.elapsed_ms)
+        output.success({"deleted": endpoint_name}, time_ms=timer.elapsed_ms, fmt=fmt)
+    except pymysql.err.Error as exc:
+        log_operation("ai model endpoint delete", ok=False, error_code="SQL_ERROR")
+        output.error("SQL_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()
+
+
+# ---------------------------------------------------------------------------
+# seekdb ai complete ...
+# ---------------------------------------------------------------------------
+
+@ai_cmd.command("complete")
+@click.argument("prompt")
+@click.option("--model", "model_name", required=True, help="Registered completion model name (from ai model list).")
+@click.pass_context
+def ai_complete(ctx: click.Context, prompt: str, model_name: str) -> None:
+    """Run text completion using the database AI_COMPLETE function."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        conn = get_connection(dsn)
+    except Exception as exc:
+        log_operation("ai complete", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer, conn.cursor() as cur:
+            cur.execute(
+                "SELECT AI_COMPLETE(%s, %s) AS response",
+                (model_name, prompt),
+            )
+            row = cur.fetchone()
+        response_text = (row.get("response") or row.get("RESPONSE") or "") if row else ""
+
+        log_operation("ai complete", ok=True, time_ms=timer.elapsed_ms)
+        output.success(
+            {"model": model_name, "response": response_text},
+            time_ms=timer.elapsed_ms,
+            fmt=fmt,
+        )
+    except pymysql.err.Error as exc:
+        log_operation("ai complete", ok=False, error_code="AI_ERROR")
+        output.error("AI_ERROR", str(exc), fmt=fmt)
+    finally:
+        conn.close()

seekdb_cli/commands/collection.py

@@ -0,0 +1,188 @@
+"""seekdb collection command — manage vector collections."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import click
+
+from seekdb_cli import output
+from seekdb_cli.logger import log_operation
+from seekdb_cli.vecconnection import get_vec_client
+
+
+@click.group("collection")
+@click.pass_context
+def collection_cmd(ctx: click.Context) -> None:
+    """Manage vector collections."""
+
+
+@collection_cmd.command("list")
+@click.pass_context
+def list_collections(ctx: click.Context) -> None:
+    """List all collections in the current database."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        client = get_vec_client(dsn)
+    except (ImportError, ValueError) as exc:
+        log_operation("collection list", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer:
+            collections = client.list_collections()
+            result = []
+            for col in collections:
+                info: dict[str, Any] = {"name": col.name}
+                try:
+                    info["count"] = col.count()
+                except Exception:
+                    info["count"] = None
+                result.append(info)
+
+        log_operation("collection list", ok=True, time_ms=timer.elapsed_ms)
+        output.success(result, time_ms=timer.elapsed_ms, fmt=fmt)
+    except Exception as exc:
+        log_operation("collection list", ok=False, error_code="COLLECTION_ERROR")
+        output.error("COLLECTION_ERROR", str(exc), fmt=fmt)
+
+
+@collection_cmd.command("create")
+@click.argument("name")
+@click.option("--dimension", "-d", type=int, default=384, help="Vector dimension (default: 384).")
+@click.option("--distance", type=click.Choice(["cosine", "l2", "ip"]), default="cosine", help="Distance metric.")
+@click.pass_context
+def create_collection(ctx: click.Context, name: str, dimension: int, distance: str) -> None:
+    """Create a new collection."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        client = get_vec_client(dsn)
+    except (ImportError, ValueError) as exc:
+        log_operation("collection create", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        from pyseekdb import HNSWConfiguration
+
+        with timer:
+            config = HNSWConfiguration(dimension=dimension, distance=distance)
+            client.create_collection(
+                name=name,
+                configuration=config,
+                embedding_function=None,
+            )
+
+        log_operation("collection create", ok=True, time_ms=timer.elapsed_ms)
+        output.success(
+            {"name": name, "dimension": dimension, "distance": distance},
+            time_ms=timer.elapsed_ms,
+            fmt=fmt,
+        )
+    except Exception as exc:
+        log_operation("collection create", ok=False, error_code="COLLECTION_ERROR")
+        output.error("COLLECTION_ERROR", str(exc), fmt=fmt)
+
+
+@collection_cmd.command("delete")
+@click.argument("name")
+@click.pass_context
+def delete_collection(ctx: click.Context, name: str) -> None:
+    """Delete a collection."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        client = get_vec_client(dsn)
+    except (ImportError, ValueError) as exc:
+        log_operation("collection delete", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer:
+            if not client.has_collection(name):
+                log_operation("collection delete", ok=False, error_code="COLLECTION_NOT_FOUND")
+                output.error(
+                    "COLLECTION_NOT_FOUND",
+                    f"Collection '{name}' does not exist",
+                    fmt=fmt,
+                    extra={"collections": [c.name for c in client.list_collections()]},
+                )
+                return
+            client.delete_collection(name)
+
+        log_operation("collection delete", ok=True, time_ms=timer.elapsed_ms)
+        output.success({"deleted": name}, time_ms=timer.elapsed_ms, fmt=fmt)
+    except SystemExit:
+        raise
+    except Exception as exc:
+        log_operation("collection delete", ok=False, error_code="COLLECTION_ERROR")
+        output.error("COLLECTION_ERROR", str(exc), fmt=fmt)
+
+
+@collection_cmd.command("info")
+@click.argument("name")
+@click.pass_context
+def info_collection(ctx: click.Context, name: str) -> None:
+    """Show collection details (count, peek)."""
+    fmt: str = ctx.obj["format"]
+    dsn: str | None = ctx.obj["dsn"]
+
+    try:
+        client = get_vec_client(dsn)
+    except (ImportError, ValueError) as exc:
+        log_operation("collection info", ok=False, error_code="CONNECTION_ERROR")
+        output.error("CONNECTION_ERROR", str(exc), fmt=fmt)
+        return
+
+    timer = output.Timer()
+    try:
+        with timer:
+            if not client.has_collection(name):
+                log_operation("collection info", ok=False, error_code="COLLECTION_NOT_FOUND")
+                output.error(
+                    "COLLECTION_NOT_FOUND",
+                    f"Collection '{name}' does not exist",
+                    fmt=fmt,
+                    extra={"collections": [c.name for c in client.list_collections()]},
+                )
+                return
+
+            col = client.get_collection(name, embedding_function=None)
+            count = col.count()
+            preview = col.peek(limit=3)
+
+            data: dict[str, Any] = {
+                "name": name,
+                "count": count,
+                "preview": {
+                    "ids": preview.get("ids", []),
+                    "documents": preview.get("documents", []),
+                    "metadatas": preview.get("metadatas", []),
+                },
+            }
+            # Dimension and distance come from list_collections() metadata (pyseekdb API)
+            for coll in client.list_collections():
+                if getattr(coll, "name", None) == name:
+                    if getattr(coll, "dimension", None) is not None:
+                        data["dimension"] = coll.dimension
+                    if getattr(coll, "distance", None) is not None:
+                        data["distance"] = coll.distance
+                    break
+
+        log_operation("collection info", ok=True, time_ms=timer.elapsed_ms)
+        output.success(data, time_ms=timer.elapsed_ms, fmt=fmt)
+    except SystemExit:
+        raise
+    except Exception as exc:
+        log_operation("collection info", ok=False, error_code="COLLECTION_ERROR")
+        output.error("COLLECTION_ERROR", str(exc), fmt=fmt)