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 连接客户端 - 通用模块
根据不同实例自动选择合适的连接方式
"""

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: 用户名（默认为bytehouse）
        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,
        verify=False
    )


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()

# 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.

"""
多模态向量化模块
基于豆包多模态向量化模型实现文本、图片、视频的向量化
"""

import os
import json
import numpy as np
from volcenginesdkarkruntime import Ark
from typing import List, Union, Dict


class MultimodalEmbedding:
    """多模态向量化客户端"""
    
    def __init__(self):
        # 先尝试从环境变量读取配置
        api_key = os.environ.get("BH_ARK_API_KEY")
        base_url = os.environ.get("BH_ARK_BASE_URL")
        
        # 如果环境变量没有配置，尝试从OpenClaw配置文件读取
        if not api_key or not base_url:
            config_path = os.path.expanduser("~/.openclaw/openclaw.json")
            if os.path.exists(config_path):
                try:
                    with open(config_path, "r", encoding="utf-8") as f:
                        config = json.load(f)
                    ark_config = config.get("models", {}).get("providers", {}).get("ark", {})
                    if not api_key:
                        api_key = ark_config.get("apiKey")
                    if not base_url:
                        base_url = ark_config.get("baseUrl")
                except Exception as e:
                    print(f"读取OpenClaw配置文件失败: {str(e)}")
        
        # 检查必要配置是否存在
        if not api_key:
            raise ValueError("未找到ARK API Key，请配置BH_ARK_API_KEY环境变量或在openclaw.json中配置models.providers.ark.apiKey")
        if not base_url:
            raise ValueError("未找到ARK Base URL，请配置BH_ARK_BASE_URL环境变量或在openclaw.json中配置models.providers.ark.baseUrl")
        
        self.client = Ark(
            api_key=api_key,
            base_url=base_url
        )
        self.model = os.environ.get("BH_EMBEDDING_MODEL", "doubao-embedding-vision-251215")
    
    def encode(self, 
               input_data: Union[str, List[Dict]], 
               modality: str = "text",
               instruction: str = None) -> List[float]:
        """
        多模态向量化接口
        
        Args:
            input_data: 输入数据
                - 文本：直接传入字符串
                - 图片/视频：传入 {"type": "image_url"/"video_url", "url": "xxx"} 格式
            modality: 数据类型，可选 text/image/video
            instruction: 自定义指令，用于提升特定场景检索精度
        
        Returns:
            向量列表
        """
        try:
            if isinstance(input_data, str):
                input_item = {"type": "text", "text": input_data}
            else:
                input_item = input_data
            
            # 输入格式校验
            if modality in ["image", "video"] and not isinstance(input_item, dict):
                raise ValueError(f"{modality}类型输入必须为包含url的字典格式")
            
            # 构造请求参数
            request_params = {
                "model": self.model,
                "encoding_format": "float",
                "input": [input_item]
            }
            
            # 添加自定义指令（251215及以上版本支持）
            if instruction and "251215" in self.model:
                request_params["instructions"] = instruction
            
            # 调用 API
            resp = self.client.multimodal_embeddings.create(**request_params)
            
            if hasattr(resp, 'data'):
                embedding = resp.data.embedding
                vec = np.array(embedding).flatten().tolist()
                
                
                return vec
            else:
                raise ValueError("API响应格式错误，未找到embedding字段")
        
        except Exception as e:
            error_msg = str(e).lower()
            if "api key" in error_msg or "unauthorized" in error_msg or "permission" in error_msg:
                raise PermissionError(f"向量化失败：API密钥无效或权限不足。错误详情：{e}")
            elif "connection" in error_msg or "timeout" in error_msg or "network" in error_msg:
                raise ConnectionError(f"向量化失败：网络连接异常。错误详情：{e}")
            elif "invalid" in error_msg or "parameter" in error_msg or "format" in error_msg:
                raise ValueError(f"向量化失败：输入参数或格式错误。错误详情：{e}")
            else:
                raise Exception(f"向量化失败：{e}")
    
    def encode_text(self, text: str, instruction: str = None) -> List[float]:
        """文本向量化"""
        return self.encode(text, "text", instruction)
    
    def encode_image(self, image_url: str, instruction: str = None) -> List[float]:
        """图片URL向量化"""
        input_item = {"type": "image_url", "image_url": {"url": image_url}}
        return self.encode(input_item, "image", instruction)
    
    def encode_video(self, video_url: str, instruction: str = None) -> List[float]:
        """视频URL向量化"""
        input_item = {"type": "video_url", "video_url": {"url": video_url}}
        return self.encode(input_item, "video", instruction)

def main():
    """命令行入口"""
    print("Hello, World!")
    multimodal_embedding = MultimodalEmbedding()
    print(multimodal_embedding.encode_text("你好"))

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 查询')
    parser.add_argument('--force', action='store_true', help='强制执行非DQL查询（跳过确认）')
    
    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)
        
    # 程序化控制：检查是否为非DQL查询（create无需用户确认，也算作DQL）
    sql_upper = sql.strip().upper()
    is_dql = any(sql_upper.startswith(prefix) for prefix in ("SELECT", "SHOW", "DESC", "DESCRIBE", "EXPLAIN", "WITH", "CREATE"))
    
    if not is_dql and not args.force:
        if sys.stdin.isatty():
            confirm = input(f"警告: 检测到非DQL查询:\n{sql}\n是否确认执行？(y/N): ")
            if confirm.lower() not in ['y', 'yes']:
                print("已取消执行。", file=sys.stderr)
                sys.exit(0)
        else:
            print("Error: 这是一个非DQL查询（如INSERT/UPDATE/DROP等）。作为AI助手，请立即中断执行，并向用户明确询问是否确认执行此操作（展示出具体的SQL）。如果用户明确同意，请在命令中添加 --force 参数再次执行。", 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()

#!/bin/bash

load_config() {
  # 检查 jq 是否安装
  if ! command -v jq &> /dev/null; then
    echo "Error: jq is not installed. Please install jq first (e.g. brew install jq or sudo apt install jq)."
    return 1
  fi

  # 解析 json，将每个 key-value 转成 export KEY="VALUE" 的格式
  local exports
  exports=$(jq -r 'to_entries | .[] | "export \(.key)=\(.value | @sh)"' ~/.bytehouse_config.json)

  # 执行生成的 export 命令
  eval "$exports"
  echo "Configuration loaded from ~/.bytehouse_config.json"
}

load_config

#!/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）
  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()

# 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
import json
import asyncio
from typing import List, Dict, Any

from embedding import MultimodalEmbedding


class ByteHouseMultimodalSearch:
    """ByteHouse多模态检索客户端"""
    
    def __init__(self, 
                 connection_type: str = "http",
                 secure: bool = True,
                 compress: str = "zstd",
                 connect_timeout: int = 300,
                 send_receive_timeout: int = 1000,
                 prefer_mcp: bool = True):
        """
        初始化ByteHouse多模态检索客户端
        
        Args:
            connection_type: 连接方式，可选 http/tcp
            secure: 是否启用加密连接
            compress: 压缩方式，可选 zstd/lz4/False
            connect_timeout: 连接超时时间，单位秒
            send_receive_timeout: 请求超时时间，单位秒
            prefer_mcp: 是否优先使用ByteHouse MCP Skill
        """
        self.connection_type = connection_type
        self.dimensions = int(os.environ.get("EMBEDDING_DIMENSIONS", 2048))
        self.embedding = MultimodalEmbedding()
        self.use_mcp = False
        self.mcp_client = None
        
        # 优先尝试使用MCP连接
        if prefer_mcp:
            try:
                from mcp_client import ByteHouseMCPClient
                
                async def test_mcp_connection():
                    async with ByteHouseMCPClient() as client:
                        await client.connect()
                        return client
                
                self.mcp_client = asyncio.run(test_mcp_connection())
                self.use_mcp = True
            except Exception:
                self.use_mcp = False
        
        # MCP不可用时使用原生驱动连接
        if not self.use_mcp:
            if connection_type == "http":
                import clickhouse_connect
                self.client = clickhouse_connect.get_client(
                    host=os.environ.get("BYTEHOUSE_HOST"),
                    port=int(os.environ.get("BYTEHOUSE_PORT", 8123)),
                    username=os.environ.get("BYTEHOUSE_USER"),
                    password=os.environ.get("BYTEHOUSE_PASSWORD"),
                    database=os.environ.get("BYTEHOUSE_DATABASE", "default"),
                    secure=secure,
                    compress=compress,
                    send_receive_timeout=send_receive_timeout
                )
            elif connection_type == "tcp":
                from clickhouse_driver import Client
                self.client = Client(
                    host=os.environ.get("BYTEHOUSE_HOST"),
                    port=int(os.environ.get("BYTEHOUSE_PORT", 9000)),
                    user=os.environ.get("BYTEHOUSE_USER"),
                    password=os.environ.get("BYTEHOUSE_PASSWORD"),
                    database=os.environ.get("BYTEHOUSE_DATABASE", "default"),
                    connect_timeout=connect_timeout,
                    send_receive_timeout=send_receive_timeout,
                    compression=compress if compress else False,
                    secure=secure,
                    client_revision=54430
                )
            else:
                raise ValueError(f"不支持的连接类型: {connection_type}")
    
    
    def _check_dql(self, sql: str, force: bool = False):
        """检查SQL是否为DQL，如果不是则要求确认"""
        import sys
        sql_upper = sql.strip().upper()
        is_dql = any(sql_upper.startswith(prefix) for prefix in ("SELECT", "SHOW", "DESC", "DESCRIBE", "EXPLAIN", "WITH", "CREATE"))
        
        if not is_dql and not force:
            if sys.stdin.isatty():
                confirm = input(f"警告: 检测到非DQL查询:\n{sql}\n是否确认执行？(y/N): ")
                if confirm.lower() not in ['y', 'yes']:
                    print("已取消执行。", file=sys.stderr)
                    sys.exit(0)
            else:
                print("Error: 这是一个非DQL查询（如INSERT/UPDATE/DROP等）。作为AI助手，请立即中断执行，并向用户明确询问是否确认执行此操作（展示出具体的SQL）。如果用户明确同意，请在调用时传递 force=True 参数再次执行。", file=sys.stderr)
                sys.exit(1)

    def _execute_sql(self, sql: str, query_type: str = "select", force: bool = False):
        """内部通用SQL执行方法，自动适配MCP和原生驱动"""
        self._check_dql(sql, force)
        try:
            if self.use_mcp:
                tool_name = "run_select_query" if query_type == "select" else "run_dml_ddl_query"
                
                async def run_mcp_query():
                    return await self.mcp_client.call_tool(tool_name, {"query": sql})
                
                result = asyncio.run(run_mcp_query())
                if result and len(result) > 0:
                    try:
                        return [list(item.values()) for item in json.loads(result[0])]
                    except:
                        return [line.split('\t') for line in result[0].strip().split('\n')]
                return []
            else:
                if query_type == "select":
                    result = self.client.query(sql)
                    return result.result_rows if hasattr(result, 'result_rows') else result
                else:
                    return self.client.command(sql)
        
        except Exception as e:
            error_msg = str(e).lower()
            if "connection" in error_msg or "timeout" in error_msg:
                raise ConnectionError(f"数据库连接异常：{e}")
            elif "syntax" in error_msg or "parse" in error_msg:
                raise ValueError(f"SQL语法错误：{e}")
            elif "permission" in error_msg or "auth" in error_msg:
                raise PermissionError(f"权限不足：{e}")
            else:
                raise Exception(f"数据库操作失败：{e}")
    
    def create_multimodal_table(self, 
                               table_name: str,
                               enable_text_search: bool = True,
                               index_type: str = "HNSW",
                               metric: str = "COSINE",
                               hnsw_m: int = 32,
                               hnsw_ef_construction: int = 512,
                               force: bool = False):
        """
        创建多模态检索表
        
        Args:
            table_name: 表名
            enable_text_search: 是否开启全文检索
            index_type: 索引类型，可选 HNSW/HNSW_SQ/IVF_FLAT/IVF_PQ/IVF_PQ_FS
            metric: 距离度量，可选 COSINE/L2
            hnsw_m: HNSW 每个节点最大连接数
            hnsw_ef_construction: HNSW 构建时探索因子
        """
        if index_type in ["HNSW", "HNSW_SQ"]:
            index_config = f"TYPE {index_type}('DIM={self.dimensions}, METRIC={metric}, M={hnsw_m}, EF_CONSTRUCTION={hnsw_ef_construction}')"
        else:
            index_config = f"TYPE {index_type}('dim={self.dimensions}', 'metric={metric}')"
        
        text_index = f"INDEX text_idx (title, content) TYPE inverted('standard', '{{\"version\":\"v4\"}}')" if enable_text_search else ""
        
        create_sql = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            id UInt64 COMMENT '唯一ID',
            content_type Enum('text' = 1, 'image' = 2, 'video' = 3) COMMENT '内容类型',
            content String COMMENT '原始内容或URL',
            title String COMMENT '标题/描述',
            embedding Array(Float32) COMMENT '向量',
            CONSTRAINT cons_vec_len CHECK length(embedding) = {self.dimensions},
            metadata Map(String, String) COMMENT '元数据',
            create_time DateTime DEFAULT now() COMMENT '创建时间',
            INDEX vec_idx embedding {index_config},
            {text_index}
        ) ENGINE = CnchMergeTree
        ORDER BY id
        SETTINGS 
            index_granularity = 1024,
            index_granularity_bytes = 0,
            enable_vector_index_preload = 1
        """
        
        self._execute_sql(create_sql, query_type="ddl", force=True)
    
    def insert_document(self,
                       table_name: str,
                       doc_id: int,
                       content_type: str,
                       content: str,
                       title: str = "",
                       metadata: Dict = None,
                       embedding: List[float] = None,
                       instruction: str = None) -> bool:
        """插入单条文档"""
        if embedding is None:
            if content_type == "text":
                embedding = self.embedding.encode_text(content, instruction)
            elif content_type == "image":
                embedding = self.embedding.encode_image(content, instruction)
            elif content_type == "video":
                embedding = self.embedding.encode_video(content, instruction)
            else:
                raise ValueError(f"不支持的内容类型: {content_type}")
        
        if len(embedding) != self.dimensions:
            raise ValueError(f"向量维度错误：期望{self.dimensions}维，实际{len(embedding)}维")
        
        metadata_str = json.dumps(metadata).replace("'", "''") if metadata else "{}"
        
        insert_sql = f"""
        INSERT INTO {table_name}
        (id, content_type, content, title, embedding, metadata)
        VALUES
        ({doc_id}, '{content_type}', '{content.replace("'", "''")}', 
         '{title.replace("'", "''")}', {embedding}, '{metadata_str}')
        """
        
        self._execute_sql(insert_sql, query_type="dml", force=True)
        return True
    
    def insert_batch_documents(self,
                              table_name: str,
                              documents: List[Dict],
                              instruction: str = None,
                              skip_error: bool = True) -> Dict:
        """批量插入文档"""
        rows = []
        failed = []
        
        for idx, doc in enumerate(documents):
            try:
                required_fields = ["doc_id", "content_type", "content"]
                for field in required_fields:
                    if field not in doc:
                        raise ValueError(f"缺少必填字段: {field}")
                
                if doc.get('embedding'):
                    embedding = doc['embedding']
                    if len(embedding) != self.dimensions:
                        raise ValueError(f"向量维度错误")
                else:
                    if doc['content_type'] == "text":
                        embedding = self.embedding.encode_text(doc['content'], instruction)
                    elif doc['content_type'] == "image":
                        embedding = self.embedding.encode_image(doc['content'], instruction)
                    elif doc['content_type'] == "video":
                        embedding = self.embedding.encode_video(doc['content'], instruction)
                    else:
                        raise ValueError(f"不支持的内容类型: {doc['content_type']}")
                
                metadata = doc.get('metadata', {})
                metadata_str = json.dumps(metadata).replace("'", "''")
                
                rows.append([
                    doc['doc_id'],
                    doc['content_type'],
                    doc['content'].replace("'", "''"),
                    doc.get('title', '').replace("'", "''"),
                    embedding,
                    metadata_str
                ])
            
            except Exception as e:
                failed.append({"doc_id": doc.get("doc_id", idx), "error": str(e)})
                if not skip_error:
                    raise
        
        if rows:
            try:
                if self.use_mcp:
                    values_str = [f"({row[0]}, '{row[1]}', '{row[2]}', '{row[3]}', {row[4]}, '{row[5]}')" for row in rows]
                    insert_sql = f"INSERT INTO {table_name} VALUES {','.join(values_str)}"
                    self._execute_sql(insert_sql, query_type="dml", force=True)
                else:
                    self._check_dql(f"INSERT INTO {table_name} (批量插入 {len(rows)} 条数据)", True)
                    if self.connection_type == "http":
                        self.client.insert(
                            table_name,
                            rows,
                            column_names=['id', 'content_type', 'content', 'title', 'embedding', 'metadata'],
                            column_type_names=['UInt64', 'Enum', 'String', 'String', 'Array(Float32)', 'Map(String, String)']
                        )
                    else:
                        self.client.execute(
                            f'INSERT INTO {table_name} VALUES',
                            rows
                        )
                success_count = len(rows)
            except Exception as e:
                for row in rows:
                    failed.append({"doc_id": row[0], "error": f"批量插入失败: {str(e)}"})
                success_count = 0
        else:
            success_count = 0
        
        return {
            "success_count": success_count,
            "failed_count": len(failed),
            "failed_details": failed
        }
    
    def vector_search(self,
                     table_name: str,
                     query_embedding: List[float],
                     top_k: int = 10,
                     filter_condition: str = None,
                     metric: str = "COSINE",
                     hnsw_ef_s: int = 200) -> List[Dict]:
        """纯向量检索"""
        distance_func = "cosineDistance" if metric == "COSINE" else "L2Distance"
        
        sql = f"""
        SELECT 
            id, content_type, content, title, metadata, create_time,
            {distance_func}(embedding, {query_embedding}) AS score
        FROM {table_name}
        {f"WHERE {filter_condition}" if filter_condition else ""}
        ORDER BY score ASC
        LIMIT {top_k}
        SETTINGS enable_new_ann = 1, hnsw_ef_s = {hnsw_ef_s}
        """
        
        rows = self._execute_sql(sql, query_type="select")
        columns = ['id', 'content_type', 'content', 'title', 'metadata', 'create_time', 'score']
        return [dict(zip(columns, row)) for row in rows]
    
    def hybrid_search(self,
                     table_name: str,
                     query_text: str,
                     query_embedding: List[float] = None,
                     top_k: int = 10,
                     filter_condition: str = None,
                     vector_weight: float = 0.7,
                     text_weight: float = 0.3,
                     metric: str = "COSINE") -> List[Dict]:
        """混合检索：向量检索 + 全文检索"""
        if query_embedding is None:
            query_embedding = self.embedding.encode_text(query_text)
        
        vector_results = self.vector_search(table_name, query_embedding, top_k * 2, filter_condition, metric)
        
        query_escaped = self.transform_string(query_text)
        text_search_sql = f"""
        SELECT id, content_type, content, title, metadata, create_time, _text_search_score AS text_score
        FROM {table_name}
        WHERE textSearch(content, '{query_escaped}')
        {f'AND {filter_condition}' if filter_condition else ''}
        ORDER BY text_score DESC
        LIMIT {top_k * 2}
        """
        text_rows = self._execute_sql(text_search_sql, query_type="select")
        text_results = [dict(zip(['id', 'content_type', 'content', 'title', 'metadata', 'create_time', 'text_score'], row)) 
                       for row in text_rows]
        
        # RRF 融合算法
        all_results = {}
        k = 60
        
        for rank, item in enumerate(vector_results):
            doc_id = item['id']
            if doc_id not in all_results:
                all_results[doc_id] = item
            all_results[doc_id]['vector_rank'] = rank
        
        for rank, item in enumerate(text_results):
            doc_id = item['id']
            if doc_id not in all_results:
                all_results[doc_id] = item
            all_results[doc_id]['text_rank'] = rank
        
        for doc_id, item in all_results.items():
            vector_score = 1.0 / (k + item.get('vector_rank', 10000))
            text_score = 1.0 / (k + item.get('text_rank', 10000))
            item['final_score'] = vector_weight * vector_score + text_weight * text_score
        
        sorted_results = sorted(all_results.values(), key=lambda x: x['final_score'], reverse=True)
        return sorted_results[:top_k]
    
    def text_search_image(self, table_name: str, query_text: str, top_k: int = 10, **kwargs) -> List[Dict]:
        """以文搜图"""
        filter_cond = "content_type = 'image'"
        if 'filter_condition' in kwargs:
            filter_cond += f" AND {kwargs.pop('filter_condition')}"
        
        instruction = kwargs.pop('instruction', None)
        query_embedding = self.embedding.encode_text(query_text, instruction)
        return self.vector_search(table_name, query_embedding, top_k, filter_condition=filter_cond, **kwargs)
    
    def image_search_image(self, table_name: str, image_url: str, top_k: int = 10, **kwargs) -> List[Dict]:
        """以图搜图"""
        filter_cond = "content_type = 'image'"
        if 'filter_condition' in kwargs:
            filter_cond += f" AND {kwargs.pop('filter_condition')}"
        
        instruction = kwargs.pop('instruction', None)
        query_embedding = self.embedding.encode_image(image_url, instruction)
        return self.vector_search(table_name, query_embedding, top_k, filter_condition=filter_cond, **kwargs)
    
    def text_search_video(self, table_name: str, query_text: str, top_k: int = 10, **kwargs) -> List[Dict]:
        """以文搜视频"""
        filter_cond = "content_type = 'video'"
        if 'filter_condition' in kwargs:
            filter_cond += f" AND {kwargs.pop('filter_condition')}"
        
        instruction = kwargs.pop('instruction', None)
        query_embedding = self.embedding.encode_text(query_text, instruction)
        return self.vector_search(table_name, query_embedding, top_k, filter_condition=filter_cond, **kwargs)
    
    def transform_string(self, s: str) -> str:
        # Remove all single quotes
        s = s.replace("'", "")
        # Split each character with |
        return "|".join(s)

#!/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 查询，并判断 SQL 是否为 DQL

配置方式：设置以下环境变量
- 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

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


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}"
    
    # 构建请求 payload
    payload = {
        "systemHints": system_hints,
        "input": natural_language,
        "knowledgeBaseIDsString": ["*"],
        "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()