Expose alternatives count on user messages

Kaveh Nejadrequested to merge
feature/expose-message-alternatives-count into master

What does this MR do and why?

Allows the front end to request the new field alternativeCount which returns the count of alternative attempts at each user messages (The total number of user messages that share the same parent_thread_ts - 1 (the current)).

Limited to just user messages as currently you can only retry a user message. A benefit of limiting to user messages, is that often the same checkpoint (therefore thread_ts) can create multiple messages, for example:

Say 1                                                                 thread_ts: 1f18b420-2d24-6f99-8000-bd98b2b46daf  parent_ts: 1f18b420-2d1f-6ba0-bfff-e5880762ca33
    1                                                                 thread_ts: 1f18b420-43e9-63d1-8001-ddf0ad2a7186  parent_ts: 1f18b420-2d24-6f99-8000-bd98b2b46daf
        Which page am I on?                                           thread_ts: 1f18b420-8623-610c-8002-2f358ff2639d  parent_ts: 1f18b420-43e9-63d1-8001-ddf0ad2a7186
            You're on the **Work items** page for th…                 thread_ts: 1f18b420-8623-610c-8002-2f358ff2639d  parent_ts: 1f18b420-43e9-63d1-8001-ddf0ad2a7186

Where 1f18b420-8623-610c-8002-2f358ff2639d is the thread_ts of the user message and LLM response.

By removing the value from all but user messages, we simplify this issue.

The front end will then be able to use this information to request the alternative branches in the next MR !249143 (closed)

Query & query plan

Note: This is my second MR that includes a query plan, and the first was several months ago, if I have missed anything, please let me know.

SELECT DISTINCT
    "p_duo_workflows_checkpoint_blobs"."thread_ts"
FROM
    "p_duo_workflows_checkpoint_blobs"
WHERE
    "p_duo_workflows_checkpoint_blobs"."workflow_id" = 51
    AND "p_duo_workflows_checkpoint_blobs"."workflow_created_at" = '2026-08-11 14:03:50.314807'
    AND "p_duo_workflows_checkpoint_blobs"."project_id" = 1000000
    AND "p_duo_workflows_checkpoint_blobs"."channel" = 'ui_chat_log'

I wasn't sure how to get the exact milisecond for the query, so in console.postgres.ai I altered it to use a range of 1 second.

https://console.postgres.ai/gitlab/projects/gitlab-production-main/sessions/54595/commands/157448

References

Full PoC: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!6124 (closed)

Epic: gitlab-org#21289

Python implementation for retries: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!6356 (merged)

Depends on: !249098 (merged)

Video demo: gitlab-org#21289 (comment 3666671057)

How to set up and validate locally

Query
{
  duoWorkflowWorkflows(workflowId: "gid://gitlab/Ai::DuoWorkflows::Workflow/11") {
    nodes { latestCheckpoint { duoMessages { content messageType threadTs parentTs alternativeCount } } }
  }
}
  1. Enable flags for incremental checkpoints:
    1. Feature.enable(:duo_workflow_incremental_checkpoints)
    2. Feature.enable(:duo_workflow_read_incremental_checkpoints)
    3. Feature.enable(:dw_read_blobs_graphql)
  2. Checkout this branch
    1. From gdk/gitlab run git checkout feature/expose-message-alternatives-count
  3. Restart
    1. gdk restart
  4. Start a conversation "Say 1" => "1" => "Say 2" => "2"
  5. From http://gdk.test:8080/-/graphql-explorer (replace port) run the query above with the workflow ID (You can get the workflow ID from dev tools => Application => Session storage => http://gdk.test:8080 => key)
  6. You should get something similar to
GraphQL response 
{
  "data": {
    "duoWorkflowWorkflows": {
      "nodes": [
        {
          "latestCheckpoint": {
            "duoMessages": [
              {
                "content": "Say 1",
                "messageType": "user",
                "threadTs": "1f194ad0-c695-6bd7-8000-428384bd2199",
                "parentTs": "1f194ad0-c690-6352-bfff-e6d4d04bcd09",
                "alternativeCount": 0
              },
              {
                "content": "1",
                "messageType": "agent",
                "threadTs": "1f194ad0-d74e-65ee-8001-22b4e0499188",
                "parentTs": "1f194ad0-c695-6bd7-8000-428384bd2199",
                "alternativeCount": null
              },
              {
                "content": "Say 2",
                "messageType": "user",
                "threadTs": "1f194ad1-7963-678e-8002-1a67a9ae1e0d",
                "parentTs": "1f194ad0-d74e-65ee-8001-22b4e0499188",
                "alternativeCount": 0
              },
              {
                "content": "2",
                "messageType": "agent",
                "threadTs": "1f194ad1-7963-678e-8002-1a67a9ae1e0d",
                "parentTs": "1f194ad0-d74e-65ee-8001-22b4e0499188",
                "alternativeCount": null
              }
            ]
          }
        }
      ]
    }
  },
  "correlationId": "01KZNP909PEEVDCTV4X2ARYF2G"
}
  1. Copy the parentTs of "Say 2" into gitlab-ai-gateway/duo_workflow_service/server.py replacing resume_checkpoint_ts=start_workflow_request.startRequest.resume_checkpoint_ts with resume_checkpoint_ts="1f194ad0-d74e-65ee-8001-22b4e0499188"
  2. Run gdk restart duo-workflow-service
  3. Send new messages "Say 3" and "Say 4"
    1. 3 will replace 2, and then 4 will replace 3
  4. Run the query again, and each time the alternativeCount of the second user message will increase, and the message will be replaced with the newest message.
GraphQL response 
{
  "data": {
    "duoWorkflowWorkflows": {
      "nodes": [
        {
          "latestCheckpoint": {
            "duoMessages": [
              {
                "content": "Say 1",
                "messageType": "user",
                "threadTs": "1f194ad0-c695-6bd7-8000-428384bd2199",
                "parentTs": "1f194ad0-c690-6352-bfff-e6d4d04bcd09",
                "alternativeCount": 0
              },
              {
                "content": "1",
                "messageType": "agent",
                "threadTs": "1f194ad0-d74e-65ee-8001-22b4e0499188",
                "parentTs": "1f194ad0-c695-6bd7-8000-428384bd2199",
                "alternativeCount": null
              },
              {
                "content": "Say 4",
                "messageType": "user",
                "threadTs": "1f194ad6-5199-68da-8002-73f7b0955386",
                "parentTs": "1f194ad0-d74e-65ee-8001-22b4e0499188",
                "alternativeCount": 2
              },
              {
                "content": "4",
                "messageType": "agent",
                "threadTs": "1f194ad6-5e65-6649-8003-f61a0c372488",
                "parentTs": "1f194ad6-5199-68da-8002-73f7b0955386",
                "alternativeCount": null
              }
            ]
          }
        }
      ]
    }
  },
  "correlationId": "01KZNPCYZJDTV0VH2RW417Y1NH"
}

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Kaveh Nejad

Merge request reports