diff --git a/server/hooks.go b/server/hooks.go index 999573c..05bed04 100644 --- a/server/hooks.go +++ b/server/hooks.go @@ -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. diff --git a/server/ollama.go b/server/ollama.go index 303858a..ea0d79c 100644 --- a/server/ollama.go +++ b/server/ollama.go @@ -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)