Add streaming Ollama response with live post updates
This commit is contained in:
+58
-3
@@ -3,6 +3,7 @@ package main
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mattermost/mattermost/server/public/model"
|
||||
"github.com/mattermost/mattermost/server/public/plugin"
|
||||
@@ -142,8 +143,25 @@ func (p *Plugin) MessageHasBeenPosted(c *plugin.Context, post *model.Post) {
|
||||
return
|
||||
}
|
||||
|
||||
// Send to Ollama.
|
||||
resp, err := p.ollamaClient.ChatCompletion(&ChatRequest{
|
||||
// Create the initial thread post with a placeholder.
|
||||
replyPost := &model.Post{
|
||||
UserId: p.botUserID,
|
||||
ChannelId: post.ChannelId,
|
||||
Message: "…",
|
||||
RootId: post.RootId,
|
||||
}
|
||||
if replyPost.RootId == "" {
|
||||
replyPost.RootId = post.Id
|
||||
}
|
||||
|
||||
createdPost, appErr := p.API.CreatePost(replyPost)
|
||||
if appErr != nil {
|
||||
p.API.LogError("Failed to create reply post", "error", appErr.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Stream the response from Ollama.
|
||||
stream, err := p.ollamaClient.ChatCompletionStream(&ChatRequest{
|
||||
Model: cfg.DefaultModel,
|
||||
Messages: messages,
|
||||
Options: map[string]any{},
|
||||
@@ -153,7 +171,44 @@ func (p *Plugin) MessageHasBeenPosted(c *plugin.Context, post *model.Post) {
|
||||
return
|
||||
}
|
||||
|
||||
p.postReply(post, resp.Message.Content)
|
||||
// Read the stream and update the post periodically.
|
||||
var accumulated strings.Builder
|
||||
updateTicker := time.NewTicker(200 * time.Millisecond)
|
||||
defer updateTicker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case event, ok := <-stream:
|
||||
if !ok {
|
||||
// Stream closed unexpectedly.
|
||||
p.finalisePost(createdPost, accumulated.String())
|
||||
return
|
||||
}
|
||||
if event.Error != nil {
|
||||
p.API.LogError("Stream error", "error", event.Error.Error())
|
||||
p.finalisePost(createdPost, accumulated.String()+"\n\n*Error: response truncated*")
|
||||
return
|
||||
}
|
||||
accumulated.WriteString(event.Token)
|
||||
if event.Done {
|
||||
p.finalisePost(createdPost, accumulated.String())
|
||||
return
|
||||
}
|
||||
case <-updateTicker.C:
|
||||
if accumulated.Len() > 0 {
|
||||
createdPost.Message = accumulated.String()
|
||||
p.API.UpdatePost(createdPost)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// finalisePost updates the post with the final content and persists the session.
|
||||
func (p *Plugin) finalisePost(post *model.Post, content string) {
|
||||
post.Message = content
|
||||
if _, appErr := p.API.UpdatePost(post); appErr != nil {
|
||||
p.API.LogError("Failed to finalise post", "error", appErr.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// postReply creates a reply post in the same thread as the given post.
|
||||
|
||||
@@ -54,6 +54,13 @@ type ListModelsResponse struct {
|
||||
} `json:"models"`
|
||||
}
|
||||
|
||||
// StreamEvent represents a single event from a streaming chat response.
|
||||
type StreamEvent struct {
|
||||
Token string // delta content from the model
|
||||
Done bool // true for the final event
|
||||
Error error // non-nil if the stream failed
|
||||
}
|
||||
|
||||
// ChatCompletion sends a chat request and returns the full response.
|
||||
func (c *OllamaClient) ChatCompletion(req *ChatRequest) (*ChatResponse, error) {
|
||||
req.Stream = false
|
||||
@@ -87,6 +94,72 @@ func (c *OllamaClient) ChatCompletion(req *ChatRequest) (*ChatResponse, error) {
|
||||
return &chatResp, nil
|
||||
}
|
||||
|
||||
// streamingChunk represents a single line from Ollama's ndjson stream.
|
||||
type streamingChunk struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
Done bool `json:"done"`
|
||||
EvalCount int `json:"eval_count"`
|
||||
PromptEvalCount int `json:"prompt_eval_count"`
|
||||
}
|
||||
|
||||
// ChatCompletionStream sends a streaming chat request and returns a channel
|
||||
// of StreamEvents. The caller must drain the channel until Done or Error.
|
||||
func (c *OllamaClient) ChatCompletionStream(req *ChatRequest) (<-chan StreamEvent, error) {
|
||||
req.Stream = true
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequest("POST", c.baseURL+"/api/chat", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http request: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(respBody)))
|
||||
}
|
||||
|
||||
ch := make(chan StreamEvent)
|
||||
go func() {
|
||||
defer resp.Body.Close()
|
||||
defer close(ch)
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
for {
|
||||
var chunk streamingChunk
|
||||
if err := decoder.Decode(&chunk); err != nil {
|
||||
if err == io.EOF {
|
||||
return
|
||||
}
|
||||
ch <- StreamEvent{Error: fmt.Errorf("decode chunk: %w", err)}
|
||||
return
|
||||
}
|
||||
|
||||
ch <- StreamEvent{
|
||||
Token: chunk.Message.Content,
|
||||
Done: chunk.Done,
|
||||
}
|
||||
|
||||
if chunk.Done {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// ListModels fetches the list of available models from Ollama.
|
||||
func (c *OllamaClient) ListModels() ([]string, error) {
|
||||
httpReq, err := http.NewRequest("GET", c.baseURL+"/api/tags", nil)
|
||||
|
||||
Reference in New Issue
Block a user