vllm+go||vllm+go-zero 流式输出

Golang 后端流式输出实现

以下是如何在 Golang 后端实现流式输出的步骤。

1. 后端 (Golang)

首先，你需要创建一个 HTTP 处理器来处理前端的请求，并设置响应头以支持 Server-Sent Events (SSE)。

import (
    "fmt"
    "net/http"
)

func streamHandler(w http.ResponseWriter, r *http.Request) {
    // 设置响应头
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")

    // 创建一个 flusher
    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
        return
    }

    // 模拟从模型获取数据并发送
    for i := 0; i < 10; i++ {
        fmt.Fprintf(w, "data: Message %d\n\n", i)
        flusher.Flush() // 立即将数据发送到客户端
        time.Sleep(time.Second) // 模拟处理时间
    }
}

func main() {
    http.HandleFunc("/stream", streamHandler)
    http.ListenAndServe(":8080", nil)
}

2. 前端 (JavaScript)

在前端，你可以使用 EventSource API 来接收服务器发送的事件。

<!DOCTYPE html>
<html>
<head>
    <title>Streaming Demo</title>
</head>
<body>
    <div id="output"></div>

    <script>
        const output = document.getElementById('output');
        const eventSource = new EventSource('/stream');

        eventSource.onmessage = function(event) {
            const p = document.createElement('p');
            p.textContent = event.data;
            output.appendChild(p);
        };

        eventSource.onerror = function(error) {
            console.error('EventSource failed:', error);
            eventSource.close();
        };
    </script>
</body>
</html>

3. 实际与模型交互

在实际应用中，你可能需要与特定的 AI 模型 API 进行交互。这里以 OpenAI 的 API 为例：

import (
    "context"
    "fmt"
    "io"
    "net/http"

    "github.com/sashabaranov/go-openai"
)

func streamHandler(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")

    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
        return
    }

    client := openai.NewClient("your-api-key")
    ctx := context.Background()

    req := openai.CompletionRequest{
        Model:     openai.GPT3TextDavinci003,
        MaxTokens: 100,
        Prompt:    "Tell me a story",
        Stream:    true,
    }

    stream, err := client.CreateCompletionStream(ctx, req)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    defer stream.Close()

    for {
        response, err := stream.Recv()
        if err == io.EOF {
            break
        }
        if err != nil {
            fmt.Fprintf(w, "data: Error: %v\n\n", err)
            flusher.Flush()
            break
        }

        fmt.Fprintf(w, "data: %s\n\n", response.Choices[0].Text)
        flusher.Flush()
    }
}

这个例子使用了 github.com/sashabaranov/go-openai 库来与 OpenAI API 交互。你需要替换 "your-api-key" 为你的实际 API 密钥。

go-zero 框架下的流式输出实现

以下是如何在 go-zero 框架下实现流式输出的步骤。

1. 定义新的路由

首先，在你的 api 文件中定义新的路由：

service stream-api {
    @handler StreamHandler
    get /stream
}

2. 实现 StreamHandler

然后，在你的 handler 文件中实现 StreamHandler：

package handler

import (
    "net/http"
    "github.com/zeromicro/go-zero/rest/httpx"
    "github.com/sashabaranov/go-openai"
    "context"
    "fmt"
    "io"
)

func StreamHandler(w http.ResponseWriter, r *http.Request) {
    // 设置响应头
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")

    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
        return
    }

    client := openai.NewClient("your-api-key")
    ctx := context.Background()

    req := openai.CompletionRequest{
        Model:     openai.GPT3TextDavinci003,
        MaxTokens: 100,
        Prompt:    "Tell me a story",
        Stream:    true,
    }

    stream, err := client.CreateCompletionStream(ctx, req)
    if err != nil {
        httpx.Error(w, err)
        return
    }
    defer stream.Close()

    for {
        response, err := stream.Recv()
        if err == io.EOF {
            break
        }
        if err != nil {
            fmt.Fprintf(w, "data: Error: %v\n\n", err)
            flusher.Flush()
            break
        }

        fmt.Fprintf(w, "data: %s\n\n", response.Choices[0].Text)
        flusher.Flush()
    }
}

3. 注册新的 handler

在你的 routes.go 文件中注册这个新的 handler：

package handler

import (
    "net/http"

    "github.com/zeromicro/go-zero/rest"
)

func RegisterHandlers(server *rest.Server) {
    server.AddRoutes(
        []rest.Route{
            {
                Method:  http.MethodGet,
                Path:    "/stream",
                Handler: StreamHandler,
            },
        },
    )
}

4. 前端代码保持不变

<!DOCTYPE html>
<html>
<head>
    <title>Streaming Demo</title>
</head>
<body>
    <div id="output"></div>

    <script>
        const output = document.getElementById('output');
        const eventSource = new EventSource('/stream');

        eventSource.onmessage = function(event) {
            const p = document.createElement('p');
            p.textContent = event.data;
            output.appendChild(p);
        };

        eventSource.onerror = function(error) {
            console.error('EventSource failed:', error);
            eventSource.close();
        };
    </script>
</body>
</html>

这样，你就可以在 go-zero 框架中实现流式输出了。记住要替换 "your-api-key" 为你的实际 OpenAI API 密钥。

posted @   ParallelForEach  阅读(393)  评论(0)  编辑  收藏  举报
相关博文：
· 基于 go-zero 框架的项目中集成 WebSocket
· 谷歌统计接入
· 干货：DeepSeek+SpringAI实现流式对话！
· 撸一个前后端分离项目
· 10分钟速成golang
阅读排行：
· 分享4款.NET开源、免费、实用的商城系统
· 全程不用写代码，我用AI程序员写了一个飞机大战
· MongoDB 8.0这个新功能碉堡了，比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· 白话解读 Dapr 1.15：你的「微服务管家」又秀新绝活了
点击右上角即可分享
微信分享提示