基于华为云FunctionGraph + OBS + GaussDB构建Serverless向量检索流水线：架构设计与代码实现

import psycopg2

def init_vector_table(conn):
    """初始化向量表和DiskANN索引"""
    with conn.cursor() as cur:
        # 启用向量数据库功能
        cur.execute("SET enable_vectordb = on")
        
        # 创建文档向量表
        cur.execute("""
            CREATE TABLE IF NOT EXISTS doc_vectors (
                id BIGSERIAL PRIMARY KEY,
                doc_name VARCHAR(512) NOT NULL,
                chunk_index INTEGER NOT NULL,
                content TEXT NOT NULL,
                embedding FloatVector(768),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # 创建DiskANN向量索引（余弦距离）
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_doc_embedding 
            ON doc_vectors 
            USING gsdiskann (embedding COSINE) 
            WITH (pq_nseg = 128, pq_nclus = 16, num_parallels = 32)
        """)
        
        # 创建doc_name普通索引用于过滤
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_doc_name 
            ON doc_vectors (doc_name)
        """)
    conn.commit()

# -*- coding: utf-8 -*-
import json
import io
import os
import logging
import re

import psycopg2
from psycopg2 import pool
from obs import ObsClient
from sentence_transformers import SentenceTransformer

logger = logging.getLogger()
logger.setLevel(logging.INFO)

# OBS客户端（全局复用）
obs_client = ObsClient(
    access_key_id=os.getenv('OBS_AK'),
    secret_access_key=os.getenv('OBS_SK'),
    server=os.getenv('OBS_ENDPOINT')
)

# 数据库连接池（全局复用）
db_pool = psycopg2.pool.ThreadedConnectionPool(
    minconn=2, maxconn=10,
    host=os.getenv('DB_HOST'),
    port=5432,
    database=os.getenv('DB_NAME'),
    user=os.getenv('DB_USER'),
    password=os.getenv('DB_PASSWORD'),
    sslmode='require'
)

# 嵌入模型（全局复用，避免每次冷启动重新加载）
_model = None

def get_model():
    global _model
    if _model is None:
        _model = SentenceTransformer('BAAI/bge-base-zh-v1.5')
    return _model


def handler(event, context):
    request_id = context.getRequestId()
    logger.info(f"Request {request_id} started, event: {json.dumps(event)}")

    # 1. 从OBS事件中提取桶名和对象名
    bucket_name, object_key = _parse_obs_event(event)
    if not bucket_name or not object_key:
        return _response(400, {'error': '无法解析OBS事件'})

    # 2. 从OBS读取文档内容
    resp = obs_client.getObject(bucketName=bucket_name, objectKey=object_key)
    if resp.status >= 300:
        return _response(500, {'error': f'OBS读取失败: {resp.errorMessage}'})

    content = resp.body.read().decode('utf-8', errors='ignore')
    content = re.sub(r'\s+', ' ', content).strip()
    if not content:
        return _response(400, {'error': '文档内容为空'})

    # 3. 文本分块（按固定长度切分，保留语义完整性）
    chunks = _split_chunks(content, max_chars=500, overlap=50)
    logger.info(f"Document split into {len(chunks)} chunks")

    # 4. 向量化并写入GaussDB
    model = get_model()
    conn = db_pool.getconn()
    try:
        # 启用向量功能
        with conn.cursor() as cur:
            cur.execute("SET enable_vectordb = on")
        
        # 批量插入
        with conn.cursor() as cur:
            for idx, chunk in enumerate(chunks):
                embedding = model.encode(chunk, normalize_embeddings=True)
                vector_str = '[' + ','.join(str(round(float(v), 6)) for v in embedding) + ']'
                cur.execute(
                    """INSERT INTO doc_vectors (doc_name, chunk_index, content, embedding)
                       VALUES (%s, %s, %s, %s::floatvector)""",
                    (object_key, idx, chunk, vector_str)
                )
        conn.commit()
    except Exception as e:
        conn.rollback()
        logger.error(f"DB insert failed: {str(e)}")
        return _response(500, {'error': f'入库失败: {str(e)}'})
    finally:
        db_pool.putconn(conn)

    return _response(200, {
        'message': '向量化完成',
        'doc_name': object_key,
        'chunks': len(chunks)
    })


def _parse_obs_event(event):
    """从OBS事件中解析桶名和对象名"""
    try:
        records = event.get('Records', [])
        if records:
            s3 = records[0].get('s3', {})
            bucket = s3.get('bucket', {}).get('name', '')
            obj = s3.get('object', {}).get('key', '')
            return bucket, obj
    except Exception as e:
        logger.error(f"Parse OBS event failed: {str(e)}")
    return None, None


def _split_chunks(text, max_chars=500, overlap=50):
    """按语义边界切分文本块"""
    sentences = re.split(r'(?<=[。！？.!?])\s*', text)
    chunks = []
    current = ''
    for sent in sentences:
        if not sent.strip():
            continue
        if len(current) + len(sent) <= max_chars:
            current += sent
        else:
            if current:
                chunks.append(current.strip())
            # 处理重叠
            if overlap > 0 and len(current) > overlap:
                current = current[-overlap:] + sent
            else:
                current = sent
    if current.strip():
        chunks.append(current.strip())
    return chunks


def _response(status_code, body):
    return {
        'statusCode': status_code,
        'headers': {'Content-Type': 'application/json'},
        'body': json.dumps(body, ensure_ascii=False)
    }

def handler(event, context):
    """向量检索函数入口"""
    query_text = event.get('queryStringParameters', {}).get('q', '')
    top_k = int(event.get('queryStringParameters', {}).get('top_k', '5'))
    
    if not query_text:
        return _response(400, {'error': '缺少查询参数q'})

    model = get_model()
    query_vec = model.encode(query_text, normalize_embeddings=True)
    vector_str = '[' + ','.join(str(round(float(v), 6)) for v in query_vec) + ']'

    conn = db_pool.getconn()
    try:
        with conn.cursor() as cur:
            cur.execute("SET enable_vectordb = on")
            # 余弦距离检索，按相似度降序
            cur.execute("""
                SELECT doc_name, chunk_index, content,
                       1 - (embedding <=> %s::floatvector) AS similarity
                FROM doc_vectors
                WHERE embedding <=> %s::floatvector IS NOT NULL
                ORDER BY embedding <=> %s::floatvector
                LIMIT %s
            """, (vector_str, vector_str, vector_str, top_k))
            
            results = []
            for row in cur.fetchall():
                results.append({
                    'doc_name': row[0],
                    'chunk_index': row[1],
                    'content': row[2][:200],
                    'similarity': round(float(row[3]), 4)
                })
        
        return _response(200, {'query': query_text, 'results': results})
    except Exception as e:
        logger.error(f"Search failed: {str(e)}")
        return _response(500, {'error': str(e)})
    finally:
        db_pool.putconn(conn)