diff --git a/server/channels/app/post.go b/server/channels/app/post.go index dd424ee259e..e1a509090ff 100644 --- a/server/channels/app/post.go +++ b/server/channels/app/post.go @@ -3447,6 +3447,20 @@ func (a *App) SendTestMessage(rctx request.CTX, userID string) (*model.Post, *mo return post, nil } +// rewriteResponseJSONSchema is the structured output schema for the LLM rewrite response. +// Defined at package level to avoid re-allocating on every call. +var rewriteResponseJSONSchema = map[string]any{ + "type": "object", + "properties": map[string]any{ + "rewritten_text": map[string]any{ + "type": "string", + "description": "The rewritten version of the message", + }, + }, + "required": []any{"rewritten_text"}, + "additionalProperties": false, +} + // RewriteMessage rewrites a message using AI based on the specified action func (a *App) RewriteMessage( rctx request.CTX, @@ -3495,6 +3509,7 @@ func (a *App) RewriteMessage( {Role: "system", Message: systemPrompt}, {Role: "user", Message: userPrompt}, }, + JSONOutputFormat: rewriteResponseJSONSchema, OperationSubType: normalizeRewriteAction(action), UserID: sessionUserID, } diff --git a/server/channels/app/post_rewrite_test.go b/server/channels/app/post_rewrite_test.go index 571277e7efe..e23e32ddfd8 100644 --- a/server/channels/app/post_rewrite_test.go +++ b/server/channels/app/post_rewrite_test.go @@ -7,6 +7,7 @@ import ( "strings" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/mattermost/mattermost/server/public/model" @@ -26,3 +27,25 @@ func TestBuildRewriteSystemPrompt(t *testing.T) { require.Equal(t, basePrompt, prompt) }) } + +func TestRewriteMessage(t *testing.T) { + t.Run("sets_structured_output_schema_on_bridge_request", func(t *testing.T) { + bridge := &testAgentsBridge{ + completeFn: func(sessionUserID, agentID string, req BridgeCompletionRequest) (string, error) { + return `{"rewritten_text":"Rewritten message"}`, nil + }, + } + + th := Setup(t, WithAgentsBridge(bridge)).InitBasic(t) + ctx := th.Context.WithSession(&model.Session{UserId: th.BasicUser.Id}) + + response, appErr := th.App.RewriteMessage(ctx, model.NewId(), "original message", model.RewriteActionImproveWriting, "", "") + require.Nil(t, appErr) + require.NotNil(t, response) + assert.Equal(t, "Rewritten message", response.RewrittenText) + require.Len(t, bridge.completeCalls, 1) + assert.Equal(t, BridgeOperationRewrite, bridge.completeCalls[0].request.Operation) + assert.Equal(t, string(model.RewriteActionImproveWriting), bridge.completeCalls[0].request.OperationSubType) + assert.Equal(t, rewriteResponseJSONSchema, bridge.completeCalls[0].request.JSONOutputFormat) + }) +}