外部 AI 函数

通过将 TiDB Cloud Lake 与您自己的基础设施连接，构建强大的 AI/ML 能力。外部函数让您能够部署自定义模型、利用 GPU 加速，并与任何 ML 框架集成，同时确保数据安全。

关键能力

工作原理

1. Create AI Server：使用 Python 和 tidbcloudlake-udf 构建您的 AI/ML 服务器
2. Register Function：使用 CREATE FUNCTION 将您的服务器连接到 TiDB Cloud Lake
3. Use in SQL：直接在 SQL 查询中调用您的自定义 AI 函数

示例：文本嵌入函数

# Simple embedding UDF server demo
from tidbcloudlake_udf import udf, UDFServer
from sentence_transformers import SentenceTransformer

# Load pre-trained model
model = SentenceTransformer('all-mpnet-base-v2')  # 768-dimensional vectors

@udf(
    input_types=["STRING"],
    result_type="ARRAY(FLOAT)",
)
def ai_embed_768(inputs: list[str], headers) -> list[list[float]]:
    """Generate 768-dimensional embeddings for input texts"""
    try:
        # Process inputs in a single batch
        embeddings = model.encode(inputs)
        # Convert to list format
        return [embedding.tolist() for embedding in embeddings]
    except Exception as e:
        print(f"Error generating embeddings: {e}")
        # Return empty lists in case of error
        return [[] for _ in inputs]

if __name__ == '__main__':
    print("Starting embedding UDF server on port 8815...")
    server = UDFServer("0.0.0.0:8815")
    server.add_function(ai_embed_768)
    server.serve()

-- Register the external function in TiDB Cloud Lake
CREATE OR REPLACE FUNCTION ai_embed_768 (STRING)
    RETURNS ARRAY(FLOAT)
    LANGUAGE PYTHON
    HANDLER = 'ai_embed_768'
    ADDRESS = 'https://your-ml-server.example.com';

-- Use the custom embedding in queries
SELECT
    id,
    title,
    cosine_distance(
        ai_embed_768(content),
        ai_embed_768('machine learning techniques')
    ) AS similarity
FROM articles
ORDER BY similarity ASC
LIMIT 5;