手搓一个 Ollama 本地 SSE 全栈聊天助手

在写的一个小项目中，考虑到调用大模型 API 的经费不足，于是选择使用 Ollama 进行部署

技术栈

后端

Node.js + Express
TS
Ollama
Prisma

前端

React + TS
Tailwindcss

什么是 SSE ?

SSE,全称是 Server-Sent Events,即服务器推送事件，SSE是基于 HTTP 协议的单向通信技术，服务器主动向客户端推送数据,客户端能够通过 EventSource API 或者 Fetch API 进行接收。

代码实现

后端实现

ollama

ts 复制代码

// ollamaService.ts

export class OllamaService {
  private baseURL: string
  private defaultModel: string

  constructor() {
    this.baseURL = process.env.OLLAMA_BASE_URL 
    this.defaultModel = process.env.OLLAMA_MODEL
  }


// 流式输出
  async *chatStream(   // async * 标记为 异步生成器函数
    message: string,
  ): AsyncGenerator<string> {
    try {
      const prompt = this.buildPrompt(message)

      const response = await axios.post(
        `${this.baseURL}/api/generate`,
        {
          model: this.defaultModel,
          prompt: prompt,
          stream: true,   // 启用流式
          options: {
            temperature: 0.7,
            top_p: 0.9,
            num_predict: 512,
          }
        } as OllamaGenerateRequest,
        {
          responseType: 'stream',  // 接收流式响应
          timeout: 60000,
        }
      )

      // 处理流式响应
      for await (const chunk of response.data) {
        const lines = chunk.toString().split('\n').filter((line: string) => line.trim())
        
        for (const line of lines) {
          try {
            const data = JSON.parse(line) as OllamaGenerateResponse
            if (data.response) {
              yield data.response    // 生成器 
            }
          } catch (e) {
            // 忽略解析错误
          }
        }
      }
    } catch (error) {
      //
    }
  }
  
    private buildPrompt(message: string): string {  
        return \`你是 AI 助手。\n\n用户: \${message}\n助手: \`  
    }
    
// 导出单例
export const ollamaService = new OllamaService()

SSE 路由实现

ts 复制代码

// ollama.ts

router.post("/chat/stream", async (req, res) => {
    const { message } = req.body;
    
    // 设置 SSE 响应头
    res.setHeader("Content-Type","text/event-stream");
    res.setHeader("Cache-Control","no-cache");
    res.setHeader("Connection","keep-alive");
    
    try{
    // 流式输出
    // 异步迭代生成器
    for await (const chunk of ollamaService.chatStream(message)){
        res.write(`data:${JSON.stringify({chunk})}\n\n`);
    }
    
    // 发送结束标记
    res.write(`data:${JSON.stringify({done:true})}\n\n`);
    res.end();
    } catch(e){
    res.write(`data:${JSON.stringify({e:"生成失败"})}\n\n`);
    res.end()
        }
})

前端实现

API 封装

ts 复制代码

// ollama.ts
export const chatWithOllamaStream = async (  
    request: { message: string },  
    onChunk: (chunk: string) => void,  
    onError?: (error: Error) => void  
): Promise<void> => {  
    const response = await fetch('http://localhost:5000/api/ollama/chat/stream', {  
    method: 'POST',  
    headers: { 'Content-Type': 'application/json' },  
    body: JSON.stringify(request),  
})  
  
    // 获取 ReadableStream  
    const reader = response.body?.getReader()  
    const decoder = new TextDecoder()  
  
    while (true) {  
    const { done, value } = await reader.read()  
    if (done) break  
  
    // 解析 SSE 格式  
    const text = decoder.decode(value)  
    const lines = text.split('\n').filter(line => line.startsWith('data:'))  
  
    for (const line of lines) {  
    const data = JSON.parse(line.replace('data:', '').trim())  
    if (data.chunk) {  
        onChunk(data.chunk) // 回调更新 UI  
    }  
    if (data.done) return  
        }  
    }  
}

Chat组件

ts 复制代码

// Chat.tsx
const ChatDetailPage = () => {  
const [messages, setMessages] = useState<Message[]>([])  
const [isLoading, setIsLoading] = useState(false)  
  
const handleSendMessage = async (content: string) => {  
// 添加用户消息  
const userMessage = { id: Date.now(), role: 'user', content }  
setMessages(prev => [...prev, userMessage])  
  
// 创建空的 AI 消息  
const aiMessageId = Date.now() + 1  
const aiMessage = { id: aiMessageId, role: 'assistant', content: '' }  
setMessages(prev => [...prev, aiMessage])  
  
// 流式接收 AI 回复  
await chatWithOllamaStream(  
{ message: content },  
(chunk) => {  
// 每收到一个字,更新消息  
setMessages(prev =>  
    prev.map(msg =>   msg.id === aiMessageId  
        ? { ...msg, content: msg.content + chunk }  : msg  )  
)  
    }  
        )  
}  
  
return (  
    <div>  
        {messages.map(msg => (  
            <div key={msg.id}>  
            <strong>{msg.role}:</strong> 
            {msg.content}  
            </div>   ))}  
        <input onSubmit={handleSendMessage} />  
    </div>  
    )  
}