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 密钥。
标签:Set,http,stream,err,Header,zero,go,vllm From: https://www.cnblogs.com/lwhzj/p/18347541