Byted Bytehouse Ai Query

# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse 知识库内容添加脚本
用于向已创建的ByteHouse知识库中添加内容

配置方式：设置以下环境变量
- BYTEHOUSE_HOST: ByteHouse 主机地址 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
- BYTEHOUSE_PASSWORD: 密码 (用作 Bearer token)
- KB_ID: 知识库ID (可选，如果未设置会从配置文件读取)
"""

import os
import requests
import json
import sys
import argparse
from create_knowledge_base import load_kb_config

# 从环境变量读取默认配置
BYTEHOUSE_HOST = os.environ.get('BYTEHOUSE_HOST', '')
BYTEHOUSE_PASSWORD = os.environ.get('BYTEHOUSE_PASSWORD', '')
KB_ID = os.environ.get('KB_ID', '')


def build_kb_api_url(host: str, endpoint: str) -> str:
    """构建知识库API URL"""
    if not host:
        print("Error: BYTEHOUSE_HOST environment variable is required", file=sys.stderr)
        sys.exit(1)
    
    # 处理 host - 可能是 host:port 格式
    if host.startswith('http'):
        base_url = host.rstrip('/')
    else:
        base_url = f"https://{host}"
    
    return f"{base_url}/matrix/v1/{endpoint.lstrip('/')}"


def get_kb_id() -> int:
    """获取知识库ID，优先从环境变量，其次从配置文件"""
    if KB_ID:
        try:
            return int(KB_ID)
        except ValueError:
            print(f"Error: KB_ID environment variable is not a valid integer: {KB_ID}", file=sys.stderr)
            sys.exit(1)
    
    # 从配置文件读取
    config = load_kb_config()
    kb_id = config.get('kb_id')
    
    if not kb_id:
        print("Error: No KB_ID found. Please set KB_ID environment variable or create a knowledge base first.", file=sys.stderr)
        print("Hint: Use create_knowledge_base.py to create a new knowledge base.", file=sys.stderr)
        sys.exit(1)
    
    return kb_id


def add_content_to_kb(content: str, kb_id: int = None, config: dict = None) -> dict:
    """
    向知识库添加内容
    
    Args:
        content: 要添加的内容文本
        kb_id: 知识库ID，如果不提供会自动获取
        config: 可选的配置 dict，支持:
            - url: 自定义 API URL
            - api_key: 自定义 API Key
    
    Returns:
        API返回结果
    """
    if kb_id is None:
        kb_id = get_kb_id()
    
    # 决定使用哪个配置：用户提供的 config > 环境变量
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base/add"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base/add')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    # 构建请求 payload
    payload = {
        "knowledgeBaseID": kb_id,
        "content": content
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=15)
    response.raise_for_status()
    
    return response.json()


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='ByteHouse 知识库内容添加工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  # 添加单条内容
  python3 add_content_to_kb.py "store_sales表包含销售数据，字段有ss_sold_date_sk, ss_item_sk, ss_quantity等"
  
  # 从文件添加内容
  python3 add_content_to_kb.py --file ./schema.md
  
  # 指定知识库ID
  python3 add_content_to_kb.py --kb-id 123 "表结构说明"
  
  # 使用自定义配置
  python3 add_content_to_kb.py "内容" --config '{"api_key": "xxx", "url": "https://xxx"}'

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
  BYTEHOUSE_PASSWORD  - 密码 (作为 Bearer token)
  KB_ID               - 知识库ID (可选，优先使用)
'''
    )
    
    parser.add_argument('content', nargs='?', help='要添加的内容文本')
    parser.add_argument('--file', type=str, help='从文件读取内容，支持markdown、txt等文本文件')
    parser.add_argument('--kb-id', type=int, help='指定知识库ID，优先级高于环境变量和配置文件')
    parser.add_argument('--config', type=str, help='JSON 格式的配置，包含 api_key, url')
    
    args = parser.parse_args()
    
    # 读取内容：优先从文件，其次从参数
    content = ""
    if args.file:
        try:
            with open(args.file, 'r', encoding='utf-8') as f:
                content = f.read()
        except Exception as e:
            print(f"Error: Failed to read file {args.file}: {e}", file=sys.stderr)
            sys.exit(1)
    elif args.content:
        content = args.content
    else:
        print("Error: Either content or --file is required", file=sys.stderr)
        sys.exit(1)
    
    # 解析 config 参数
    config = None
    if args.config:
        try:
            config = json.loads(args.config)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in --config: {e}", file=sys.stderr)
            sys.exit(1)
    
    try:
        kb_id = args.kb_id if args.kb_id else get_kb_id()
        result = add_content_to_kb(content, kb_id, config)
        print(f"内容添加成功！知识库ID: {kb_id}")
        print(f"返回结果: {json.dumps(result, indent=2, ensure_ascii=False)}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse 连接客户端 - 通用模块
根据不同实例自动选择合适的连接方式
"""

import os
from typing import Optional

try:
    import clickhouse_connect
except ImportError:
    raise ImportError("clickhouse-connect not installed. Run: pip install clickhouse-connect")


def create_client(
    host: str = None,
    port: str = None,
    user: str = None,
    password: str = None,
    database: str = None
) -> 'clickhouse_connect.Client':
    """
    创建 ByteHouse 连接客户端
    
    自动处理不同实例的连接差异：
    - 8123 端口：需要 secure=True
    - 8443 端口：默认 secure=True
    - password 可能是 "user:password" 格式或纯密码
    
    Args:
        host: ByteHouse 主机
        port: 端口 (默认自动判断)
        user: 用户名
        password: 密码
        database: 数据库名 (可选)
    
    Returns:
        clickhouse_connect 客户端
    """
    # 从环境变量读取
    host = host or os.environ.get('BYTEHOUSE_HOST', '')
    port = port or os.environ.get('BYTEHOUSE_PORT', '8123')
    user = user or os.environ.get('BYTEHOUSE_USER', 'bytehouse')
    password = password or os.environ.get('BYTEHOUSE_PASSWORD', '')
    database = database or os.environ.get('BYTEHOUSE_DATABASE', '')
    
    if not host:
        raise ValueError("BYTEHOUSE_HOST is required")
    if not user or not password:
        raise ValueError("BYTEHOUSE_USER and BYTEHOUSE_PASSWORD are required")
    
    # 解析端口
    port = int(port) if port else None
    
    # 自动判断端口和加密设置
    if port:
        # 明确指定端口
        secure = port in (8123, 8443, 443)
    else:
        # 自动判断：根据 host 判断
        if 'bytehouse-ce' in host:
            # CE 版本用 8443
            port = 8443
        else:
            # 公有云版本用 8123
            port = 8123
        secure = True
    
    return clickhouse_connect.get_client(
        host=host,
        port=port,
        username=user,
        password=password,
        database=database if database else None,
        secure=secure
    )


def query(client: 'clickhouse_connect.Client', sql: str) -> list:
    """执行查询并返回结果"""
    result = client.query(sql)
    return result.result_rows


def execute(client: 'clickhouse_connect.Client', sql: str):
    """执行查询并打印结果"""
    result = client.query(sql)
    print(result.result_set)


if __name__ == "__main__":
    # 测试连接
    client = create_client()
    result = client.query("SELECT 1 as test")
    print("Connected! Result:", result.result_rows)
    client.close()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse 知识库创建脚本
用于创建ByteHouse知识库并返回知识库ID

配置方式：设置以下环境变量
- BYTEHOUSE_HOST: ByteHouse 主机地址 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
- BYTEHOUSE_PASSWORD: 密码 (用作 Bearer token)
"""

import os
import requests
import json
import sys
import argparse

# 从环境变量读取默认配置
BYTEHOUSE_HOST = os.environ.get('BYTEHOUSE_HOST', '')
BYTEHOUSE_PASSWORD = os.environ.get('BYTEHOUSE_PASSWORD', '')
KB_CONFIG_PATH = os.path.expanduser('~/.bytehouse_kb_config.json')
IDENTITY_PATH = '/root/.openclaw/workspace/IDENTITY.md'


def build_kb_api_url(host: str, endpoint: str) -> str:
    """构建知识库API URL"""
    if not host:
        print("Error: BYTEHOUSE_HOST environment variable is required", file=sys.stderr)
        sys.exit(1)
    
    # 处理 host - 可能是 host:port 格式
    if host.startswith('http'):
        base_url = host.rstrip('/')
    else:
        base_url = f"https://{host}"
    
    return f"{base_url}/matrix/v1/{endpoint.lstrip('/')}"


def create_knowledge_base(name: str, config: dict = None) -> int:
    """
    创建ByteHouse知识库
    
    Args:
        name: 知识库名称
        config: 可选的配置 dict，支持:
            - url: 自定义 API URL
            - api_key: 自定义 API Key
    
    Returns:
        创建成功的知识库ID
    """
    # 决定使用哪个配置：用户提供的 config > 环境变量
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"

    if not name:
        name = get_claw_name()
        # kb_name加上时间戳
        import time
        name += f"_{int(time.time())}"
    
    # 构建请求 payload
    payload = {
        "name": name
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=15)
    print(response.text)
    response.raise_for_status()
    
    result = response.json()
    kb_id = result.get('id') or (result.get('data', {}).get('id') if isinstance(result.get('data'), dict) else None)

    if not kb_id:
        print(f"Error: Failed to get knowledge base ID from response: {result}", file=sys.stderr)
        sys.exit(1)
    
    # 保存知识库ID到配置文件
    save_kb_config(kb_id)
    
    return kb_id


def get_claw_name() -> str:
    """从IDENTITY.md获取当前Claw的名字"""
    default_name = "ByteHouse Text2SQL 知识库"
    
    if not os.path.exists(IDENTITY_PATH):
        return default_name
    
    try:
        with open(IDENTITY_PATH, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 查找Name字段
        import re
        match = re.search(r'- \*\*Name:\*\*\s*(\w+)', content)
        if match:
            claw_name = match.group(1)
            return f"{claw_name} Text2SQL 知识库"
        else:
            return default_name
    except Exception:
        return default_name


def save_kb_config(kb_id: int):
    """保存知识库ID到配置文件"""
    config = {}
    if os.path.exists(KB_CONFIG_PATH):
        try:
            with open(KB_CONFIG_PATH, 'r') as f:
                config = json.load(f)
        except:
            pass
    
    config['kb_id'] = kb_id
    
    with open(KB_CONFIG_PATH, 'w') as f:
        json.dump(config, f, indent=2)


def load_kb_config() -> dict:
    """加载知识库配置"""
    if os.path.exists(KB_CONFIG_PATH):
        try:
            with open(KB_CONFIG_PATH, 'r') as f:
                return json.load(f)
        except:
            pass
    return {}


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='ByteHouse 知识库创建工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  python3 create_knowledge_base.py "我的SQL知识库"
  python3 create_knowledge_base.py "自定义知识库" --config '{"api_key": "xxx", "url": "https://xxx"}'

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
  BYTEHOUSE_PASSWORD  - 密码 (作为 Bearer token)

自定义 Config (通过 --config 参数):
  api_key - 自定义 API Key
  url     - 自定义 API URL
'''
    )
    
    parser.add_argument('name', nargs='?', help='知识库名称（可选，默认从IDENTITY.md获取）')
    parser.add_argument('--config', type=str, help='JSON 格式的配置，包含 api_key, url')
    
    args = parser.parse_args()
    
    # 解析 config 参数
    config = None
    if args.config:
        try:
            config = json.loads(args.config)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in --config: {e}", file=sys.stderr)
            sys.exit(1)
    
    import time
    # 获取知识库名称
    kb_name = args.name if args.name else get_claw_name()  + f"_{int(time.time())}"
    
    try:
        kb_id = create_knowledge_base(kb_name, config)
        print(f"知识库创建成功！知识库名称: {kb_name}")
        print(f"知识库ID: {kb_id}")
        print(f"配置已保存到: {KB_CONFIG_PATH}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse Execute SQL 脚本
执行 SQL 查询并返回结果

依赖: clickhouse-connect
安装: pip install clickhouse-connect

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机地址
  BYTEHOUSE_PORT     - 端口 (默认自动判断)
  BYTEHOUSE_USER     - 用户名
  BYTEHOUSE_PASSWORD - 密码
  BYTEHOUSE_DATABASE - 默认数据库 (可选)
"""

import sys
import argparse
from client import create_client, query


def main():
    parser = argparse.ArgumentParser(
        description='执行 ByteHouse SQL 查询',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  python3 execute_sql.py "SELECT * FROM tpcds.call_center LIMIT 5"
  python3 execute_sql.py "SELECT count(*) FROM tpcds.store_sales"
  python3 execute_sql.py "SHOW TABLES FROM tpcds"

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机
  BYTEHOUSE_PORT     - 端口 (默认自动判断)
  BYTEHOUSE_USER     - 用户名
  BYTEHOUSE_PASSWORD - 密码
  BYTEHOUSE_DATABASE - 默认数据库 (可选)
'''
    )
    
    parser.add_argument('query', nargs='?', help='SQL 查询语句')
    parser.add_argument('--file', '-i', help='从文件读取 SQL 查询')
    
    args = parser.parse_args()
    
    # 获取查询语句
    sql = args.query
    if args.file:
        with open(args.file, 'r') as f:
            sql = f.read().strip()
    
    if not sql:
        print("Error: Please provide a query or use --file", file=sys.stderr)
        sys.exit(1)
    
    try:
        client = create_client()
        result = query(client, sql)
        client.close()
        
        # 打印结果
        for row in result:
            print('\t'.join(str(v) for v in row))
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse List Tables 脚本
列出指定数据库中的所有表

依赖: clickhouse-connect
安装: pip install clickhouse-connect

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机地址
  BYTEHOUSE_PORT     - 端口 (默认自动判断)
  BYTEHOUSE_USER     - 用户名
  BYTEHOUSE_PASSWORD - 密码
  BYTEHOUSE_DATABASE - 默认数据库 (可选)
"""

import sys
import argparse
from client import create_client, query


def list_databases():
    """列出所有数据库"""
    client = create_client()
    result = query(client, "SHOW DATABASES")
    client.close()
    return [row[0] for row in result]


def list_tables(database: str):
    """列出指定数据库中的所有表"""
    client = create_client()
    # 用系统表查询更稳定
    result = query(client, f"SELECT name FROM system.tables WHERE database = '{database}' LIMIT 100")
    client.close()
    return [row[0] for row in result]


def main():
    parser = argparse.ArgumentParser(description='列出 ByteHouse 数据库中的表')
    parser.add_argument('--database', '-d', help='数据库名')
    parser.add_argument('--databases', '-D', action='store_true', help='列出所有数据库')
    
    args = parser.parse_args()
    
    try:
        if args.databases:
            dbs = list_databases()
            print("Databases:")
            for db in dbs:
                print(f"  - {db}")
        else:
            database = args.database
            if not database:
                print("Error: Please specify --database", file=sys.stderr)
                sys.exit(1)
            tables = list_tables(database)
            print(f"Tables in '{database}':")
            for table in tables:
                print(f"  - {table}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse 知识库查询脚本
用于在ByteHouse知识库中搜索相关内容

配置方式：设置以下环境变量
- BYTEHOUSE_HOST: ByteHouse 主机地址 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
- BYTEHOUSE_PASSWORD: 密码 (用作 Bearer token)
- KB_ID: 知识库ID (可选，如果未设置会从配置文件读取)
"""

import os
import requests
import json
import sys
import argparse
from create_knowledge_base import load_kb_config

# 从环境变量读取默认配置
BYTEHOUSE_HOST = os.environ.get('BYTEHOUSE_HOST', '')
BYTEHOUSE_PASSWORD = os.environ.get('BYTEHOUSE_PASSWORD', '')
KB_ID = os.environ.get('KB_ID', '')


def build_kb_api_url(host: str, endpoint: str) -> str:
    """构建知识库API URL"""
    if not host:
        print("Error: BYTEHOUSE_HOST environment variable is required", file=sys.stderr)
        sys.exit(1)
    
    # 处理 host - 可能是 host:port 格式
    if host.startswith('http'):
        base_url = host.rstrip('/')
    else:
        base_url = f"https://{host}"
    
    return f"{base_url}/matrix/v1/{endpoint.lstrip('/')}"


def get_kb_id() -> int:
    """获取知识库ID，优先从环境变量，其次从配置文件"""
    if KB_ID:
        try:
            return int(KB_ID)
        except ValueError:
            print(f"Error: KB_ID environment variable is not a valid integer: {KB_ID}", file=sys.stderr)
            sys.exit(1)
    
    # 从配置文件读取
    config = load_kb_config()
    kb_id = config.get('kb_id')
    
    if not kb_id:
        print("Error: No KB_ID found. Please set KB_ID environment variable or create a knowledge base first.", file=sys.stderr)
        print("Hint: Use create_knowledge_base.py to create a new knowledge base.", file=sys.stderr)
        sys.exit(1)
    
    return kb_id


def search_knowledge_base(
    search_query: str, 
    kb_id: int = None, 
    top_k: int = 5,
    min_distance: int = 0,
    config: dict = None
) -> list:
    """
    在知识库中搜索相关内容
    
    Args:
        search_query: 搜索查询文本
        kb_id: 知识库ID，如果不提供会自动获取
        top_k: 返回最相关的前K条结果，默认5
        min_distance: 最小距离阈值，默认0
        config: 可选的配置 dict，支持:
            - url: 自定义 API URL
            - api_key: 自定义 API Key
    
    Returns:
        搜索结果列表
    """
    if kb_id is None:
        kb_id = get_kb_id()
    
    # 决定使用哪个配置：用户提供的 config > 环境变量
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base/evaluate"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base/evaluate')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    # 构建请求 payload
    payload = {
        "knowledgeBaseID": kb_id,
        "search": search_query,
        "topK": top_k,
        "minimumDistance": min_distance
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=15)
    if not response.ok:
        response.raise_for_status()
    result = response.json()
    # 兼容不同返回格式：results可能在根节点或data字段里
    return result.get('results', []) or (result.get('data', {}).get('results', []) if isinstance(result.get('data'), dict) else [])


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='ByteHouse 知识库查询工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  # 简单搜索
  python3 search_knowledge_base.py "销售表结构"
  
  # 返回前10条结果
  python3 search_knowledge_base.py "store_sales字段" --top-k 10
  
  # 指定知识库ID
  python3 search_knowledge_base.py "查询语法" --kb-id 123
  
  # 自定义阈值
  python3 search_knowledge_base.py "日期字段" --min-distance 50
  
  # 使用自定义配置
  python3 search_knowledge_base.py "查询" --config '{"api_key": "xxx", "url": "https://xxx"}'

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
  BYTEHOUSE_PASSWORD  - 密码 (作为 Bearer token)
  KB_ID               - 知识库ID (可选，优先使用)
'''
    )
    
    parser.add_argument('query', help='搜索查询文本')
    parser.add_argument('--top-k', type=int, default=5, help='返回最相关的前K条结果，默认5')
    parser.add_argument('--min-distance', type=int, default=0, help='最小距离阈值，默认0')
    parser.add_argument('--kb-id', type=int, help='指定知识库ID，优先级高于环境变量和配置文件')
    parser.add_argument('--raw', action='store_true', help='输出原始JSON结果')
    parser.add_argument('--config', type=str, help='JSON 格式的配置，包含 api_key, url')
    
    args = parser.parse_args()
    
    # 解析 config 参数
    config = None
    if args.config:
        try:
            config = json.loads(args.config)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in --config: {e}", file=sys.stderr)
            sys.exit(1)
    
    try:
        kb_id = args.kb_id if args.kb_id else get_kb_id()
        results = search_knowledge_base(args.query, kb_id, args.top_k, args.min_distance, config)
        
        if args.raw:
            print(json.dumps(results, indent=2, ensure_ascii=False))
        else:
            print(f"搜索结果 (知识库ID: {kb_id}, 查询: \"{args.query}\"):")
            print("=" * 80)
            for i, result in enumerate(results, 1):
                # 兼容相似度字段名：可能是distance或similarityScore
                similarity = result.get('distance') or result.get('similarityScore')
                print(f"\n{i}. 相似度: {similarity:.4f}" if similarity is not None else f"\n{i}. 相似度: N/A")
                print(f"内容: {result.get('content', '').strip()[:200]}..." if len(result.get('content', '')) > 200 else f"内容: {result.get('content', '').strip()}")
                print("-" * 60)
            
            print(f"\n共找到 {len(results)} 条结果")
    
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse Text2SQL API 客户端
用于调用 ByteHouse 的 text2sql 接口，将自然语言转换为 SQL 查询

配置方式：设置以下环境变量
- BYTEHOUSE_HOST: ByteHouse 主机地址，同时也作为 API URL 的主机部分
- BYTEHOUSE_PASSWORD: 密码 (同时也用作API的 Bearer token)

或通过 --config 参数传入自定义配置：
- reasoningModel: 自定义模型ID
- reasoningAPIKey: 自定义 API Key
- url: 自定义 API URL
"""

import os
import requests
import json
import re
import sys
import argparse

# 导入知识库相关功能
from create_knowledge_base import create_knowledge_base, load_kb_config, save_kb_config

# 从环境变量读取默认配置
BYTEHOUSE_HOST = os.environ.get('BYTEHOUSE_HOST', '')
BYTEHOUSE_PASSWORD = os.environ.get('BYTEHOUSE_PASSWORD', '')
KB_ID = os.environ.get('KB_ID', '')


def get_kb_id(config: dict = None) -> str:
    """
    获取知识库ID
    优先级：环境变量KB_ID > 配置文件 > 自动创建新知识库
    """
    # 先检查环境变量
    if KB_ID:
        return KB_ID
    
    # 检查配置文件
    kb_config = load_kb_config()
    if 'kb_id' in kb_config:
        return str(kb_config['kb_id'])
    
    # 没有找到，自动创建知识库
    print("未检测到知识库ID，正在自动创建新的知识库...", file=sys.stderr)
    try:
        kb_id = create_knowledge_base("", config)
        print(f"自动创建知识库成功，知识库ID: {kb_id}", file=sys.stderr)
        return str(kb_id)
    except Exception as e:
        print(f"自动创建知识库失败: {e}, 将使用默认知识库(*)", file=sys.stderr)
        return "*"


def build_text2sql_url(host: str) -> str:
    """构建 Text2SQL API URL"""
    if not host:
        print("Error: BYTEHOUSE_HOST environment variable is required", file=sys.stderr)
        sys.exit(1)
    
    # 处理 host - 可能是 host:port 格式
    if host.startswith('http'):
        base_url = host.rstrip('/')
    else:
        base_url = f"https://{host}"
    
    return f"{base_url}/matrix/v1/conversation"


def call_text2sql(
    natural_language: str, 
    tables: list, 
    system_hints: str = "TEXT2SQL",
    config: dict = None
) -> str:
    """
    调用 ByteHouse Text2SQL 接口，将自然语言转换为 SQL
    
    Args:
        natural_language: 自然语言描述的查询需求
        tables: 要查询的表名列表，如 ["bytehouse.query_history"]
        system_hints: 系统提示词，默认为 "TEXT2SQL"
        config: 可选的配置 dict，支持:
            - reasoningModel: 自定义模型ID
            - reasoningAPIKey: 自定义 API Key
            - url: 自定义 API URL
    
    Returns:
        转换后的 SQL 语句
    """
    # 决定使用哪个配置：用户提供的 config > 环境变量
    if config:
        # 使用用户提供的自定义配置
        base_url = config.get('url', '')
        auth_token = config.get('reasoningAPIKey', '')
        reasoning_model = config.get('reasoningModel', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        # 用户提供的 URL 可能已经包含完整路径，直接使用
        # 如果 URL 已经包含任何 path（如 /api/v3），则不再追加
        if '/' in base_url.split('://')[-1]:
            url = base_url.rstrip('/')
        else:
            url = f"{base_url.rstrip('/')}/matrix/v1/conversation"
    else:
        # 使用环境变量中的默认配置
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_text2sql_url(BYTEHOUSE_HOST)
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
        reasoning_model = ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    # 获取知识库ID
    kb_id = get_kb_id(config)
    
    # 构建请求 payload
    payload = {
        "systemHints": system_hints,
        "input": natural_language,
        "knowledgeBaseIDsString": [kb_id] if kb_id != "*" else ["*"],
        "tables": tables
    }
    
    # 如果用户指定了 reasoningModel，添加到 payload
    if reasoning_model:
        payload["config"] = {
            "reasoningModel": reasoning_model
        }
    
    # 使用流式请求
    response = requests.post(url, headers=headers, json=payload, stream=True, timeout=30)
    if response.status_code != 200:
        print(response.text)
    response.raise_for_status()
    
    # 收集所有内容片段
    full_content = ""
    
    for line in response.iter_lines():
        if line:
            line = line.decode('utf-8')
            if line.startswith('data: '):
                data_str = line[6:]  # 去掉 'data: ' 前缀
                try:
                    data = json.loads(data_str)
                    event_type = data.get('event_type')
                    
                    if event_type == 'DELTA':
                        event_data = json.loads(data.get('event_data', '{}'))
                        message = event_data.get('message', {})
                        content = message.get('content', '')
                        full_content += content
                    elif event_type == 'DONE':
                        break
                except json.JSONDecodeError:
                    continue
    
    # 清理和提取 SQL
    sql = extract_sql(full_content)
    return sql


def extract_sql(content: str) -> str:
    """
    从内容中提取 SQL 语句
    移除 markdown 代码块标记，清理空白字符
    """
    if not content:
        return ""
    
    # 移除 markdown 代码块标记
    sql = content.strip()
    sql = re.sub(r'^```\w*\n', '', sql)
    sql = re.sub(r'\n```$', '', sql)
    
    # 规范化空白字符
    sql = re.sub(r'\s+', ' ', sql)
    sql = sql.strip()
    
    return sql


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='ByteHouse Text2SQL - 将自然语言转换为 SQL 查询',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  python3 text2sql.py "get count of queries" "bytehouse.query_history"
  python3 text2sql.py "查看最近10条记录" "bytehouse.query_history" --config '{"reasoningModel": "ep-xxx", "reasoningAPIKey": "xxx", "url": "https://xxx"}'

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
  BYTEHOUSE_PASSWORD  - 密码 (作为 Bearer token)

自定义 Config (通过 --config 参数):
  reasoningModel  - 自定义模型ID
  reasoningAPIKey - 自定义 API Key
  url            - 自定义 API URL
'''
    )
    
    parser.add_argument('query', help='自然语言查询')
    parser.add_argument('tables', nargs='+', help='要查询的表名列表')
    parser.add_argument('--config', type=str, help='JSON 格式的配置，包含 reasoningModel, reasoningAPIKey, url')
    parser.add_argument('--system-hints', type=str, default='TEXT2SQL', help='系统提示词 (默认: TEXT2SQL)')
    
    args = parser.parse_args()
    
    # 解析 config 参数
    config = None
    if args.config:
        try:
            config = json.loads(args.config)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in --config: {e}", file=sys.stderr)
            sys.exit(1)
    
    # 检查必要配置
    if not config and not BYTEHOUSE_HOST:
        print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
        sys.exit(1)
    
    try:
        sql = call_text2sql(args.query, args.tables, args.system_hints, config)
        print(sql)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
ByteHouse 知识库文件上传脚本
用于上传文件到ByteHouse知识库并自动进行切片处理
支持格式：md, txt, pdf, docx, xlsx, csv等

配置方式：设置以下环境变量
- BYTEHOUSE_HOST: ByteHouse 主机地址 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
- BYTEHOUSE_PASSWORD: 密码 (用作 Bearer token)
- KB_ID: 知识库ID (可选，如果未设置会从配置文件读取)
"""

import os
import requests
import json
import sys
import argparse
import os.path
from create_knowledge_base import load_kb_config

# 从环境变量读取默认配置
BYTEHOUSE_HOST = os.environ.get('BYTEHOUSE_HOST', '')
BYTEHOUSE_PASSWORD = os.environ.get('BYTEHOUSE_PASSWORD', '')
KB_ID = os.environ.get('KB_ID', '')


def build_kb_api_url(host: str, endpoint: str) -> str:
    """构建知识库API URL"""
    if not host:
        print("Error: BYTEHOUSE_HOST environment variable is required", file=sys.stderr)
        sys.exit(1)
    
    # 处理 host - 可能是 host:port 格式
    if host.startswith('http'):
        base_url = host.rstrip('/')
    else:
        base_url = f"https://{host}"
    
    return f"{base_url}/matrix/v1/{endpoint.lstrip('/')}"


def get_kb_id() -> int:
    """获取知识库ID，优先从环境变量，其次从配置文件"""
    if KB_ID:
        try:
            return int(KB_ID)
        except ValueError:
            print(f"Error: KB_ID environment variable is not a valid integer: {KB_ID}", file=sys.stderr)
            sys.exit(1)
    
    # 从配置文件读取
    config = load_kb_config()
    kb_id = config.get('kb_id')
    
    if not kb_id:
        print("Error: No KB_ID found. Please set KB_ID environment variable or create a knowledge base first.", file=sys.stderr)
        print("Hint: Use create_knowledge_base.py to create a new knowledge base.", file=sys.stderr)
        sys.exit(1)
    
    return kb_id


def generate_upload_url(
    kb_id: int, 
    file_name: str, 
    file_size: int,
    config: dict = None
) -> tuple[str, str, str]:
    """
    生成文件上传预签名URL
    
    Args:
        kb_id: 知识库ID
        file_name: 文件名
        file_size: 文件大小（字节）
        config: 可选的配置 dict，支持:
            - url: 自定义 API URL
            - api_key: 自定义 API Key
    
    Returns:
        (file_id, upload_url, upload_method)
    """
    # 决定使用哪个配置：用户提供的 config > 环境变量
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base/file/generate-upload-files-url"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base/file/generate-upload-files-url')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    # 构建请求 payload
    payload = {
        "knowledgeBaseID": kb_id,
        "files": [
            {
                "name": file_name,
                "sizeBytes": file_size
            }
        ]
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=30)
    response.raise_for_status()
    
    result = response.json()
    file_info = (result.get('data', []) or [])[0] if isinstance(result.get('data'), list) else {}
    
    file_id = file_info.get('fileID')
    upload_url = file_info.get('url')
    upload_method = file_info.get('method', 'PUT')
    
    if not file_id or not upload_url:
        print(f"Error: Failed to get upload URL from response: {result}", file=sys.stderr)
        sys.exit(1)
    
    return file_id, upload_url, upload_method


def upload_file(upload_url: str, upload_method: str, file_path: str) -> None:
    """上传文件到预签名URL"""
    with open(file_path, 'rb') as f:
        response = requests.request(
            upload_method.upper(),
            upload_url,
            data=f,
            timeout=120
        )
        response.raise_for_status()


def finish_upload(file_id: str, config: dict = None) -> None:
    """完成上传流程"""
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base/file/complete-upload"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base/file/complete-upload')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    payload = {
        "fileID": file_id
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=30)
    response.raise_for_status()


def load_file(
    file_id: str,
    chunk_size: int = 512,
    delimiters: list = None,
    enable_image_ocr: bool = False,
    enable_chunk_auto_merge: bool = False,
    config: dict = None
) -> None:
    """加载文件到知识库，启动切片处理"""
    if delimiters is None:
        delimiters = ["#", "##"]
    
    if config:
        base_url = config.get('url', '')
        auth_token = config.get('api_key', '')
        
        if not base_url:
            print("Error: config.url is required when using custom config", file=sys.stderr)
            sys.exit(1)
        
        url = f"{base_url.rstrip('/')}/matrix/v1/knowledge-base/file/load"
    else:
        if not BYTEHOUSE_HOST:
            print("Error: Please set BYTEHOUSE_HOST environment variable or provide --config", file=sys.stderr)
            sys.exit(1)
        
        url = build_kb_api_url(BYTEHOUSE_HOST, 'knowledge-base/file/load')
        auth_token = BYTEHOUSE_PASSWORD if BYTEHOUSE_PASSWORD else ""
    
    headers = {
        "Content-Type": "application/json",
    }
    
    # 添加认证
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    # 构建切片配置
    chunk_settings = {
        "size": chunk_size,
        "delimiters": delimiters
    }
    
    if enable_image_ocr:
        chunk_settings["enableImageOcr"] = True
    
    if enable_chunk_auto_merge:
        chunk_settings["enableChunkAutoMerge"] = True
    
    payload = {
        "fileID": file_id,
        "chunkSettings": chunk_settings
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=30)
    response.raise_for_status()


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='ByteHouse 知识库文件上传工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
示例:
  # 上传文件，使用默认配置
  python3 upload_file_to_kb.py --file ./table_schema.md
  
  # 指定知识库ID
  python3 upload_file_to_kb.py --kb-id 123 --file ./business_rules.pdf
  
  # 自定义切片配置
  python3 upload_file_to_kb.py --file ./document.md --chunk-size 1024 --delimiters "#,##,###"
  
  # 启用图片OCR和自动合并
  python3 upload_file_to_kb.py --file ./report.pdf --enable-image-ocr --enable-chunk-auto-merge
  
  # 使用自定义配置
  python3 upload_file_to_kb.py --file ./data.csv --config '{"api_key": "xxx", "url": "https://xxx"}'

环境变量:
  BYTEHOUSE_HOST      - ByteHouse 主机 (如 tenant-xxx-cn-beijing-public.bytehouse.volces.com)
  BYTEHOUSE_PASSWORD  - 密码 (作为 Bearer token)
  KB_ID               - 知识库ID (可选，优先使用)
'''
    )
    
    parser.add_argument('--file', required=True, help='要上传的文件路径')
    parser.add_argument('--kb-id', type=int, help='指定知识库ID，优先级高于环境变量和配置文件')
    parser.add_argument('--chunk-size', type=int, default=512, help='切片大小（字节），默认512')
    parser.add_argument('--delimiters', type=str, default='#,##', help='切片分隔符，逗号分隔，默认"#,##"')
    parser.add_argument('--enable-image-ocr', action='store_true', help='启用图片OCR识别，默认关闭')
    parser.add_argument('--enable-chunk-auto-merge', action='store_true', help='启用切片自动合并，默认关闭')
    parser.add_argument('--config', type=str, help='JSON 格式的配置，包含 api_key, url')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.isfile(args.file):
        print(f"Error: File not found: {args.file}", file=sys.stderr)
        sys.exit(1)
    
    # 解析 config 参数
    config = None
    if args.config:
        try:
            config = json.loads(args.config)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in --config: {e}", file=sys.stderr)
            sys.exit(1)
    
    try:
        # 获取知识库ID
        kb_id = args.kb_id if args.kb_id else get_kb_id()
        
        # 获取文件信息
        file_name = os.path.basename(args.file)
        file_size = os.path.getsize(args.file)
        
        print("=========================================")
        print("ByteHouse 知识库文件上传")
        print("=========================================")
        print(f"知识库ID: {kb_id}")
        print(f"文件: {args.file}")
        print(f"文件名: {file_name}")
        print(f"文件大小: {file_size / 1024:.2f} KB")
        print()
        
        # Step 1: 生成上传URL
        print("Step 1/4: 生成预签名上传URL...")
        file_id, upload_url, upload_method = generate_upload_url(kb_id, file_name, file_size, config)
        print(f"  ✅ 成功，文件ID: {file_id}")
        print(f"  上传URL: {upload_url}")
        print(f"  上传方法: {upload_method}")
        
        # Step 2: 上传文件
        print("\nStep 2/4: 上传文件到对象存储...")
        upload_file(upload_url, upload_method, args.file)
        print("  ✅ 文件上传成功")
        
        # Step 3: 完成上传
        print("\nStep 3/4: 完成上传流程...")
        finish_upload(file_id, config)
        print("  ✅ 上传流程完成")
        
        # Step 4: 加载文件到知识库
        print("\nStep 4/4: 加载文件到知识库（启动切片处理）...")
        delimiters = [d.strip() for d in args.delimiters.split(',')]
        load_file(
            file_id, 
            args.chunk_size, 
            delimiters, 
            args.enable_image_ocr, 
            args.enable_chunk_auto_merge,
            config
        )
        print("  ✅ 文件加载成功，切片处理已启动")
        
        print("\n=========================================")
        print("✅ 所有步骤完成！文件已成功上传到知识库")
        print("=========================================")
        
    except Exception as e:
        print(f"\n❌ 上传失败: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()