Skip to content

Adds AI Agent detail query to GraphQL

Darby Frey requested to merge graphql-find-ai-agent into master

What does this MR do and why?

This MR adds the ability to look up an AI Agent by ID via GraphQL. For example:

query getAiAgent($fullPath: ID!, $id: AiAgentID!) {
  project(fullPath: $fullPath) {
    id
    aiAgent(id: $id) {
      id
      name
      versions {
        id
        prompt
      }
      latestVersion {
        id
        prompt
      }
    }
  }
}

It also introduces two new scopes on Ai::AgentVersion. See database query and plans:

agent.versions.order_by_agent_id_id_desc produces this query and plan:

SELECT "ai_agent_versions".* 
FROM "ai_agent_versions" 
WHERE "ai_agent_versions"."agent_id" = 1 
ORDER BY agent_id, id DESC
Sort  (cost=0.01..0.02 rows=1 width=104)  
  Sort Key: id DESC  
  ->  Seq Scan on ai_agent_versions  (cost=0.00..0.00 rows=1 width=104)  
        Filter: (agent_id = 1) 

agent.versions.latest_by_agent produces this query and plan:

SELECT DISTINCT ON (agent_id) * 
FROM "ai_agent_versions" 
WHERE "ai_agent_versions"."agent_id" = 1 
ORDER BY agent_id, id DESC
Unique  (cost=0.01..0.02 rows=1 width=104)  
  ->  Sort  (cost=0.01..0.02 rows=1 width=104)  
        Sort Key: id DESC  
        ->  Seq Scan on ai_agent_versions  (cost=0.00..0.00 rows=1 width=104)  
              Filter: (agent_id = 1)  

Contributes to #441959

MR acceptance checklist

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

How to set up and validate locally

  1. Ensure you have an EE license activated

  2. Enable the feature flag

    Feature.enable(:agent_registry)
  3. In the rails console, create an Ai::Agent and multiple Ai::AgentVersions

    agent = Ai::Agent.create!(project: Project.first, name: 'my-agent')
    agent.id # to use in GraphQL query later
    Ai::AgentVersion.create!(agent: agent, prompt: 'my prompt', project: agent.project, model: 'default')
    Ai::AgentVersion.create!(agent: agent, prompt: 'my new prompt', project: agent.project, model: 'default')
  4. Load /-/graphql-explorer in the GDK and issue the following query:

    Query Variables

    {
      "id": "gid://gitlab/Ai::Agent/YOUR_AGENT_ID",
    "fullPath": "YOUR_PROJECT_PATH"
    }

    Query

    query getAiAgent($fullPath: ID!, $id: AiAgentID!) {
      project(fullPath: $fullPath) {
        id
        aiAgent(id: $id) {
          id
          name
          versions {
            id
            prompt
          }
          latestVersion {
            id
            prompt
          }
        }
      }
    }
  5. The response should look something like:

    {
      "data": {
        "project": {
          "id": "gid://gitlab/Project/19",
          "aiAgent": {
            "id": "gid://gitlab/Ai::Agent/31",
            "name": "my-agent",
            "versions": [
              {
                "id": "gid://gitlab/Ai::AgentVersion/16",
                "prompt": "my prompt"
              },
              {
                "id": "gid://gitlab/Ai::AgentVersion/17",
                "prompt": "my new prompt"
              }
            ],
            "latestVersion": {
              "id": "gid://gitlab/Ai::AgentVersion/17",
              "prompt": "my new prompt",
              "model": "default"
            }
          }
        }
      }
    }
Edited by Darby Frey

Merge request reports