RagFlow_Plugs/api/hybrid_search_http.py

#!/usr/bin/env python3
"""
混合检索HTTP服务
使用FastAPI框架实现，提供混合检索的HTTP POST接口
"""

import sys
import os
import requests
from io import BytesIO
from typing import List, Dict, Any
from fastapi import FastAPI, HTTPException, Body
from pydantic import BaseModel
from PIL import Image

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.infinity_util import InfinityVectorDB
from model.multimodal_embedding import Embedding
from conf.config import ModelConfig, VectorDBConfig

# 初始化FastAPI应用
app = FastAPI(
    title="混合检索HTTP服务",
    description="提供混合检索的HTTP POST接口",
    version="1.0.0"
)

# 初始化向量数据库
vector_db = InfinityVectorDB()

# 初始化多模态嵌入模型
embedding_model = Embedding(
    model_name=ModelConfig.get_multimodal_embedding_model_name(),
    api_key=ModelConfig.get_dashscope_api_key()
)

# 定义请求模型
class HybridSearchRequest(BaseModel):
    """混合检索请求模型"""
    text_query: str
    image: str
    topn: int = 2

# 定义响应模型
class HybridSearchResponse(BaseModel):
    """混合检索响应模型"""
    success: bool
    message: str
    output: List[Dict[str, Any]] = []
    total: int = 0

@app.post("/hybrid_search", response_model=HybridSearchResponse)
def hybrid_search(request: HybridSearchRequest = Body(...)):
    """
    混合检索API
    使用文本查询和向量查询进行混合检索
    
    请求参数：
    - text_query: 文本查询
    - image: 图片URL
    - topn: 返回结果数量，默认2
    
    返回结果：
    - success: 是否成功
    - message: 结果消息
    - output: 检索结果列表
    - total: 总命中数
    """
    try:
        # 解析请求参数
        text_query = request.text_query
        image_url = request.image
        topn = request.topn
        
        print(f"开始混合检索，数据库: {VectorDBConfig.get_infinity_database}, 知识库id: {ModelConfig.get_dataset_id()}, 文本查询: {text_query}, 返回数量: {topn}")
        
        # 构建索引名称
        index_name = f"{VectorDBConfig.get_infinity_table_name()}" 
        print(f"开始生成多模态嵌入，文本长度: {len(text_query)}")
        
        # 处理image_url为image: Image.Image
        if isinstance(image_url, str):
                # 下载图片
                response = requests.get(image_url)
                response.raise_for_status()  # 检查HTTP状态码
    
                # 将响应内容转换为字节流
                image_bytes = BytesIO(response.content)
    
                # 创建Image对象
                image = Image.open(image_bytes)
        
        # 生成多模态嵌入向量
        embedding = embedding_model.get_multimodal_embedding(text_query, image)
        
        print(f"多模态嵌入生成完成，向量长度: {len(embedding)}")
        
        # 执行混合检索
        result = vector_db.hybrid_search(
            index_name=index_name,
            match_method="dense",
            vector_field="dense_vector_1024",
            query_vector=embedding,
            element_type="float",
            metric_type="cosine",
            topn=topn,
            text_query=text_query,
            text_field="content"
        )
        
        print(f"混合检索完成，总命中数: {result.get('total', 0)}")
        
        # 返回成功响应
        return HybridSearchResponse(
            success=True,
            message="混合检索成功",
            output=result.get("output", []),
            total=result.get("total", topn)
        )
    except Exception as e:
        print(f"混合检索失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
def health_check():
    """健康检查接口"""
    return {
        "status": "ok",
        "message": "混合检索HTTP服务正常运行"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=18001)