要在 Flask 中同时返回 搜索结果 和 大模型流式生成的文本,可以采用以下方法:

1. 基本思路

你可以将 搜索结果 和 生成的文本 都通过流式响应的方式发送给客户端,确保两者同时返回。在实现时,你的后端需要在流式生成的过程中,既能持续返回生成的文本,又能在生成的每个阶段返回相应的搜索结果。

具体实现的步骤如下:

  1. 先执行搜索,获取相关的文档或信息。
  2. 基于搜索结果启动生成,并在生成的过程中持续返回模型生成的部分内容。
  3. 在整个过程保持 HTTP 连接不关闭,利用 流式响应(yield)向客户端逐步返回搜索结果和生成的内容。

2. 具体实现

2.1 设置 Flask API
from flask import Flask, Response, request
import time

app = Flask(__name__)

# 假设这是搜索模块
def search(query):
    # 模拟搜索过程,返回相关的文档
    time.sleep(1)  # 模拟搜索延迟
    return [f"Document 1 related to {query}", f"Document 2 related to {query}"]

# 假设这是生成模块
def generate_from_search_results(search_results):
    # 模拟流式生成的过程
    for i in range(len(search_results)):
        # 每个搜索结果逐步生成对应的文本
        yield f"Generating content based on: {search_results[i]}\n"
        time.sleep(0.5)  # 模拟生成耗时
        yield f"Generated content part {i+1}\n"
        time.sleep(0.5)  # 继续生成

@app.route('/search_and_generate', methods=['GET'])
def search_and_generate():
    query = request.args.get('query')
    
    # 执行搜索
    search_results = search(query)
    
    # 流式返回搜索结果和生成内容
    def stream_response():
        # 先返回搜索结果
        yield "Search Results:\n"
        for result in search_results:
            yield f"{result}\n"
            time.sleep(0.5)  # 模拟逐个返回搜索结果

        # 接下来返回生成内容
        yield "\nNow, generating content based on search results:\n"
        for part in generate_from_search_results(search_results):
            yield part

    return Response(stream_response(), content_type='text/plain;charset=utf-8')

if __name__ == '__main__':
    app.run(debug=True, threaded=True)
2.2 解释
  • 搜索过程:search(query) 模拟从数据库或搜索引擎中检索与查询相关的文档或数据。这里用 time.sleep 模拟延时。

  • 生成过程:generate_from_search_results(search_results) 模拟生成模型(例如 GPT 或其他大模型)的流式输出。每生成一部分文本,就用 yield 返回。

  • 流式响应:stream_response() 函数使用了 yield 关键字,可以逐步向客户端发送数据。首先返回搜索结果,然后再返回生成的内容。

  • 返回响应:Flask 的 Response 对象通过 stream_response() 持续推送数据到客户端。这使得客户端可以在等待生成的过程中已经接收到搜索结果和部分生成内容。

2.3 异步流式生成

为了提高效率,你可以考虑将搜索和生成过程并行处理。例如,可以利用 asyncio 或线程池来并行执行搜索和生成操作,以避免生成过程中的阻塞。

例如:

import asyncio

async def search(query):
    await asyncio.sleep(1)  # 模拟搜索延迟
    return [f"Document 1 related to {query}", f"Document 2 related to {query}"]

async def generate_from_search_results(search_results):
    for i in range(len(search_results)):
        await asyncio.sleep(0.5)  # 模拟生成延迟
        yield f"Generating content based on: {search_results[i]}\n"
        await asyncio.sleep(0.5)
        yield f"Generated content part {i+1}\n"
        await asyncio.sleep(0.5)

@app.route('/search_and_generate', methods=['GET'])
async def search_and_generate():
    query = request.args.get('query')
    
    # 异步执行搜索和生成
    search_results = await search(query)
    
    async def stream_response():
        # 先返回搜索结果
        yield "Search Results:\n"
        for result in search_results:
            yield f"{result}\n"
            await asyncio.sleep(0.5)

        # 再返回生成内容
        yield "\nNow, generating content based on search results:\n"
        async for part in generate_from_search_results(search_results):
            yield part

    return Response(stream_response(), content_type='text/plain;charset=utf-8')

在这种实现方式下,search 和 generate_from_search_results 变成了异步函数,允许在搜索结果返回后立即并行生成文本。

3. 客户端如何处理流式响应

客户端可以通过流式读取数据并实时显示结果:

const query = "example query";
const response = await fetch(`/search_and_generate?query=${query}`);

const reader = response.body.getReader();
const decoder = new TextDecoder();
let done = false;
let result = '';

while (!done) {
    const { value, done: readerDone } = await reader.read();
    done = readerDone;
    result += decoder.decode(value, { stream: true });
    console.log(result);  // 实时输出生成的内容
}

客户端会逐步接收并显示搜索结果和生成的文本,增强用户体验。

4. 优化与扩展

  • 并发请求优化:如果生成过程非常耗时,可以考虑进一步优化,使用线程池或分布式任务队列(如 Celery)来处理生成任务。

  • 缓存机制:可以对常见查询和搜索结果进行缓存,以减少重复计算。

  • 模型加载优化:如果大模型加载时间较长,可以考虑使用模型加载优化技术(如 模型量化 或 模型蒸馏)以减少推理延迟。

  • 前端界面优化:在前端,可以使用 WebSocket 进行双向通信,实时更新生成的结果。

总结

通过流式响应的方式,Flask 可以在同一个连接中逐步返回搜索结果和生成内容。使用 yield 实现流式输出,同时在生成内容的过程中能够实时返回搜索结果,提升用户体验。这种方法不仅适合流式生成的场景,也能应对高延迟的任务,比如大模型的推理过程。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐