local-coze 0.0.1__py3-none-any.whl

local_coze/core/exceptions.py

@@ -0,0 +1,67 @@
+from typing import Optional, Any
+
+
+class CozeSDKError(Exception):
+    def __init__(
+        self,
+        message: str,
+        code: Optional[str] = None,
+        details: Optional[dict] = None
+    ):
+        self.message = message
+        self.code = code
+        self.details = details or {}
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        if self.code:
+            return f"[{self.code}] {self.message}"
+        return self.message
+
+
+class ConfigurationError(CozeSDKError):
+    def __init__(self, message: str, missing_key: Optional[str] = None):
+        self.missing_key = missing_key
+        super().__init__(
+            message=message,
+            code="CONFIGURATION_ERROR",
+            details={"missing_key": missing_key}
+        )
+
+
+class APIError(CozeSDKError):
+    def __init__(
+        self,
+        message: str,
+        code: Optional[str] = None,
+        status_code: Optional[int] = None,
+        response_data: Optional[dict] = None
+    ):
+        self.status_code = status_code
+        self.response_data = response_data or {}
+        super().__init__(
+            message=message,
+            code=code,
+            details={"status_code": status_code, "response_data": response_data}
+        )
+
+
+class NetworkError(CozeSDKError):
+    def __init__(self, message: str, original_error: Optional[Exception] = None):
+        self.original_error = original_error
+        super().__init__(
+            message=f"网络请求失败: {message}",
+            code="NETWORK_ERROR",
+            details={"original_error": str(original_error) if original_error else None}
+        )
+
+
+class ValidationError(CozeSDKError):
+    def __init__(self, message: str, field: Optional[str] = None, value: Optional[Any] = None):
+        self.field = field
+        self.value = value
+        super().__init__(
+            message=message,
+            code="VALIDATION_ERROR",
+            details={"field": field, "value": str(value) if value else None}
+        )

local_coze/database/__init__.py

@@ -0,0 +1,29 @@
+"""
+数据库模块
+提供 PostgreSQL 数据库连接、会话管理、ORM 基类和迁移功能
+"""
+
+from .client import (
+    Base,
+    get_db_url,
+    get_engine,
+    get_session,
+    get_sessionmaker,
+)
+from .migration import (
+    generate_models,
+    upgrade,
+)
+
+__all__ = [
+    # ORM 基类
+    "Base",
+    # 数据库连接
+    "get_db_url",
+    "get_engine",
+    "get_session",
+    "get_sessionmaker",
+    # 迁移功能
+    "generate_models",
+    "upgrade",
+]

local_coze/database/client.py

@@ -0,0 +1,170 @@
+"""
+数据库连接模块
+提供 PostgreSQL 数据库连接、会话管理和 ORM 基类
+"""
+
+import os
+import time
+import logging
+from typing import Optional
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import sessionmaker, DeclarativeBase, Session
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.engine import Engine
+
+logger = logging.getLogger(__name__)
+
+MAX_RETRY_TIME = 20  # 连接最大重试时间（秒）
+
+_engine: Optional[Engine] = None
+_SessionLocal = None
+
+
+class Base(DeclarativeBase):
+    """SQLAlchemy ORM 模型基类"""
+
+    pass
+
+
+def _load_env() -> None:
+    """加载环境变量（内部使用）"""
+    try:
+        from dotenv import load_dotenv
+
+        load_dotenv()
+    except ImportError:
+        pass
+
+
+def get_db_url() -> str:
+    """
+    获取数据库连接 URL
+
+    优先从环境变量 PGDATABASE_URL 获取，
+    如果不存在则尝试从 coze_workload_identity 获取。
+
+    Returns:
+        str: 数据库连接 URL
+
+    Raises:
+        ValueError: 如果无法获取数据库 URL
+    """
+    _load_env()
+
+    url = os.getenv("PGDATABASE_URL") or ""
+    if url:
+        return url
+
+    try:
+        from coze_workload_identity import Client
+
+        client = Client()
+        env_vars = client.get_project_env_vars()
+        client.close()
+        for env_var in env_vars:
+            if env_var.key == "PGDATABASE_URL":
+                url = env_var.value.replace("'", "'\\''")
+                return url
+    except Exception as e:
+        logger.error(f"Error loading PGDATABASE_URL: {e}")
+        raise e
+
+    if not url:
+        logger.error("PGDATABASE_URL is not set")
+        raise ValueError("PGDATABASE_URL is not set")
+
+    return url
+
+
+def _create_engine_with_retry() -> Engine:
+    """创建数据库引擎（带重试）"""
+    url = get_db_url()
+    if not url:
+        logger.error("PGDATABASE_URL is not set")
+        raise ValueError("PGDATABASE_URL is not set")
+
+    size = 100
+    overflow = 100
+    recycle = 1800
+    timeout = 30
+
+    engine = create_engine(
+        url,
+        pool_size=size,
+        max_overflow=overflow,
+        pool_pre_ping=True,
+        pool_recycle=recycle,
+        pool_timeout=timeout,
+    )
+
+    # 验证连接，带重试
+    start_time = time.time()
+    last_error = None
+    while time.time() - start_time < MAX_RETRY_TIME:
+        try:
+            with engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+            return engine
+        except OperationalError as e:
+            last_error = e
+            elapsed = time.time() - start_time
+            logger.warning(f"Database connection failed, retrying... (elapsed: {elapsed:.1f}s)")
+            time.sleep(min(1, MAX_RETRY_TIME - elapsed))
+
+    logger.error(f"Database connection failed after {MAX_RETRY_TIME}s: {last_error}")
+    raise last_error
+
+
+def get_engine() -> Engine:
+    """
+    获取数据库引擎（单例）
+
+    Returns:
+        Engine: SQLAlchemy 数据库引擎
+    """
+    global _engine
+    if _engine is None:
+        _engine = _create_engine_with_retry()
+    return _engine
+
+
+def get_sessionmaker():
+    """
+    获取 sessionmaker（单例）
+
+    Returns:
+        sessionmaker: SQLAlchemy sessionmaker
+    """
+    global _SessionLocal
+    if _SessionLocal is None:
+        _SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=get_engine())
+    return _SessionLocal
+
+
+def get_session() -> Session:
+    """
+    创建新的数据库会话
+
+    Returns:
+        Session: SQLAlchemy 会话实例
+
+    Example:
+        >>> from local_coze.database import get_session
+        >>> db = get_session()
+        >>> try:
+        ...     # 执行数据库操作
+        ...     pass
+        ... finally:
+        ...     db.close()
+    """
+    return get_sessionmaker()()
+
+
+__all__ = [
+    "Base",
+    "get_db_url",
+    "get_engine",
+    "get_sessionmaker",
+    "get_session",
+]

local_coze/database/migration.py

@@ -0,0 +1,342 @@
+"""
+数据库迁移模块
+提供 Alembic 迁移功能的 Python API 封装
+"""
+
+import os
+import sys
+import tempfile
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Alembic script.py.mako 模板
+SCRIPT_MAKO_TEMPLATE = '''"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}
+'''
+
+# 全局缓存的 alembic 目录
+_alembic_dir: Optional[str] = None
+
+
+def _load_env() -> None:
+    """加载环境变量"""
+    try:
+        from dotenv import load_dotenv
+        load_dotenv()
+    except ImportError:
+        pass
+
+    try:
+        from coze_workload_identity import Client
+        client = Client()
+        env_vars = client.get_project_env_vars()
+        client.close()
+        for env_var in env_vars:
+            if env_var.key not in os.environ:
+                os.environ[env_var.key] = env_var.value
+    except Exception as e:
+        logger.debug(f"coze_workload_identity not available: {e}")
+
+
+def _get_db_url() -> str:
+    """获取数据库 URL"""
+    _load_env()
+
+    url = os.getenv("PGDATABASE_URL")
+    if url:
+        # 兼容 postgres:// 前缀
+        if url.startswith("postgres://"):
+            url = "postgresql://" + url[len("postgres://"):]
+        return url
+
+    raise ValueError(
+        "Database URL not configured. Set PGDATABASE_URL environment variable.\n"
+        "Did you create a database? You can create one via the Coze Coding platform."
+    )
+
+
+def _get_alembic_dir() -> str:
+    """获取或创建 alembic 临时目录"""
+    global _alembic_dir
+    if _alembic_dir and os.path.exists(_alembic_dir):
+        return _alembic_dir
+
+    # 使用固定的临时目录，避免每次创建新目录
+    _alembic_dir = os.path.join(tempfile.gettempdir(), "coze_sdk_alembic")
+    os.makedirs(_alembic_dir, exist_ok=True)
+    os.makedirs(os.path.join(_alembic_dir, "versions"), exist_ok=True)
+    return _alembic_dir
+
+
+def _ensure_alembic_env(model_import_path: str) -> str:
+    """
+    确保 alembic 环境已初始化
+
+    Args:
+        model_import_path: 模型导入路径，如 "storage.database.shared.model"
+
+    Returns:
+        str: alembic 脚本目录路径
+    """
+    script_location = _get_alembic_dir()
+
+    # 创建 env.py（支持 -x 参数传递 version_table 和 version_table_schema）
+    env_py_content = f'''"""Alembic Environment Configuration (Auto-generated)"""
+
+import os
+import sys
+from logging.config import fileConfig
+from sqlalchemy import engine_from_config, pool
+from alembic import context
+
+config = context.config
+
+if config.config_file_name:
+    fileConfig(config.config_file_name)
+
+# 获取 -x 参数
+x_args = context.get_x_argument(as_dictionary=True)
+version_table = x_args.get("version_table")
+version_table_schema = x_args.get("version_table_schema")
+
+# 导入 Base
+from local_coze.database import Base
+
+# 导入用户模型以注册到 metadata
+from {model_import_path} import *
+
+target_metadata = Base.metadata
+
+
+def run_migrations_offline():
+    url = config.get_main_option("sqlalchemy.url")
+    kwargs = {{
+        "url": url,
+        "target_metadata": target_metadata,
+        "literal_binds": True,
+        "compare_type": True,
+        "compare_server_default": True,
+        "compare_nullable": True,
+    }}
+    if version_table:
+        kwargs["version_table"] = version_table
+    if version_table_schema:
+        kwargs["version_table_schema"] = version_table_schema
+    context.configure(**kwargs)
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online():
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+    with connectable.connect() as connection:
+        kwargs = {{
+            "connection": connection,
+            "target_metadata": target_metadata,
+            "compare_type": True,
+            "compare_server_default": True,
+            "compare_nullable": True,
+        }}
+        if version_table:
+            kwargs["version_table"] = version_table
+        if version_table_schema:
+            kwargs["version_table_schema"] = version_table_schema
+        context.configure(**kwargs)
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
+'''
+
+    env_py_path = os.path.join(script_location, "env.py")
+    with open(env_py_path, "w") as f:
+        f.write(env_py_content)
+
+    # 创建 script.py.mako
+    mako_path = os.path.join(script_location, "script.py.mako")
+    with open(mako_path, "w") as f:
+        f.write(SCRIPT_MAKO_TEMPLATE)
+
+    return script_location
+
+
+def _get_alembic_config(model_import_path: str, model_path: Optional[str] = None):
+    """
+    创建 Alembic 配置对象
+
+    Args:
+        model_import_path: 模型导入路径
+        model_path: 模型文件所在目录（会添加到 sys.path）
+    """
+    from alembic.config import Config
+
+    # 添加模型目录到 sys.path
+    if model_path and model_path not in sys.path:
+        sys.path.insert(0, model_path)
+
+    script_location = _ensure_alembic_env(model_import_path)
+
+    config = Config()
+    config.set_main_option("script_location", script_location)
+    config.set_main_option("sqlalchemy.url", _get_db_url())
+
+    return config
+
+
+def generate_models(output_path: str, verbose: bool = False) -> None:
+    """
+    从数据库生成 ORM 模型
+
+    Args:
+        output_path: 输出文件路径，如 "src/storage/database/shared/model.py"
+        verbose: 是否输出详细信息
+
+    Example:
+        >>> from local_coze.database import generate_models
+        >>> generate_models("src/storage/database/shared/model.py")
+    """
+    db_url = _get_db_url()
+
+    # 确保输出目录存在
+    output_dir = os.path.dirname(output_path)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    # 使用 subprocess 调用 sqlacodegen
+    import subprocess
+
+    cmd = ["sqlacodegen", db_url, "--outfile", output_path]
+    if verbose:
+        print(f"Running: {' '.join(cmd)}")
+
+    result = subprocess.run(cmd, capture_output=True, text=True)
+
+    if result.returncode != 0:
+        raise RuntimeError(f"sqlacodegen failed: {result.stderr}")
+
+    # 在生成的文件头部添加 Base 导入
+    with open(output_path, "r") as f:
+        content = f.read()
+
+    # 替换默认的 Base 定义为从 SDK 导入
+    import re
+
+    # 删除 "class Base(DeclarativeBase):\n    pass\n"
+    content = re.sub(r"class Base\(DeclarativeBase\):\s*\n\s*pass\s*\n*", "", content)
+
+    # 删除 DeclarativeBase 相关导入
+    content = re.sub(r",\s*DeclarativeBase", "", content)
+    content = re.sub(r"DeclarativeBase,\s*", "", content)
+    content = re.sub(r"from sqlalchemy\.orm import DeclarativeBase\n", "", content)
+
+    # 在文件开头添加 SDK 的 Base 导入
+    content = "from local_coze.database import Base\n\n" + content
+
+    with open(output_path, "w") as f:
+        f.write(content)
+
+    if verbose:
+        print(f"Models generated at {output_path}")
+
+
+def upgrade(
+    model_import_path: str = "storage.database.shared.model",
+    model_path: Optional[str] = None,
+    verbose: bool = False,
+    version_table: str = "schema_version",
+    version_table_schema: str = "internal",
+) -> None:
+    """
+    执行数据库迁移（包含自动生成迁移版本）
+
+    相当于执行：
+    1. alembic revision --autogenerate
+    2. alembic upgrade head
+
+    Args:
+        model_import_path: 模型导入路径，如 "storage.database.shared.model"
+        model_path: 模型文件所在目录（会添加到 sys.path）
+        verbose: 是否输出详细信息
+        version_table: 版本表名，默认 "schema_version"
+        version_table_schema: 版本表所在 schema，默认 "internal"
+
+    Example:
+        >>> from local_coze.database import upgrade
+        >>> upgrade()  # 使用默认配置
+        >>> upgrade(model_import_path="myapp.models", model_path="/path/to/src")
+    """
+    from alembic import command
+    from sqlalchemy import create_engine, text
+
+    # 确保 version_table_schema 存在
+    db_url = _get_db_url()
+    engine = create_engine(db_url, isolation_level="AUTOCOMMIT")
+    with engine.connect() as conn:
+        conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{version_table_schema}"'))
+    engine.dispose()
+
+    config = _get_alembic_config(model_import_path, model_path)
+
+    # 设置 -x 参数
+    config.cmd_opts = type("obj", (object,), {"x": [
+        f"version_table={version_table}",
+        f"version_table_schema={version_table_schema}",
+    ]})()
+
+    # 去除 alembic 输出的前导空格
+    config.print_stdout = lambda msg, *args: print(msg.lstrip() % args if args else msg.lstrip())
+
+    # 自动生成迁移版本
+    try:
+        command.revision(config, message="auto migration", autogenerate=True)
+        if verbose:
+            print("Migration revision generated")
+    except Exception as e:
+        # 如果没有变更，revision 可能会失败，忽略
+        logger.debug(f"Revision generation skipped: {e}")
+
+    # 执行升级
+    command.upgrade(config, "head")
+    if verbose:
+        print("Database upgraded to head")
+    logger.info("Database upgraded to head")
+
+
+__all__ = [
+    "generate_models",
+    "upgrade",
+]

local_coze/embedding/__init__.py

@@ -0,0 +1,31 @@
+from .client import EmbeddingClient
+from .models import (
+    EmbeddingConfig,
+    EmbeddingInputItem,
+    EmbeddingInputImageURL,
+    EmbeddingInputVideoURL,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    EmbeddingData,
+    EmbeddingUsage,
+    MultiEmbeddingConfig,
+    SparseEmbeddingConfig,
+    SparseEmbeddingItem,
+    PromptTokensDetails,
+)
+
+__all__ = [
+    "EmbeddingClient",
+    "EmbeddingConfig",
+    "EmbeddingInputItem",
+    "EmbeddingInputImageURL",
+    "EmbeddingInputVideoURL",
+    "EmbeddingRequest",
+    "EmbeddingResponse",
+    "EmbeddingData",
+    "EmbeddingUsage",
+    "MultiEmbeddingConfig",
+    "SparseEmbeddingConfig",
+    "SparseEmbeddingItem",
+    "PromptTokensDetails",
+]