Skip to content

Migrate Duo Chat Tools: ExplainVulnerability

Overview

Blueprint MR: Adding strategy for migrating Duo Chat tools to... (gitlab-com/content-sites/handbook!7380 - merged)

Proposal

Current behavior

Rails receives from AI Gateway the information about which tool to invoke, generates a prompt and sends it to AI Gateway.

sequenceDiagram
    participant Rails
    participant AI Gateway
    participant LLM

    Rails ->> AI Gateway: POST /v2/chat/agent
    AI Gateway ->> LLM: Creates and sends ReAct Prompt
    LLM -->> AI Gateway: Responds with the right tool to invoke
    AI Gateway -->> Rails: Responds with tool to invoke

    Rails ->> AI Gateway: POST /v1/chat/agent with a prompt
    AI Gateway ->> LLM: Propagates the prompt
    LLM -->> AI Gateway: Response
    AI Gateway -->> Rails: Response

Proposal

Rails receives from AI Gateway the information about which tool to invoke, sends all related data to generate a prompt to AI Gateway. AI Gateway generates a prompt and sends a request to LLM.

sequenceDiagram
    participant Rails
    participant AI Gateway
    participant LLM

    Rails ->> AI Gateway: POST /v2/chat/agent
    AI Gateway ->> LLM: Creates and sends ReAct Prompt
    LLM -->> AI Gateway: Responds with the right tool to invoke
    AI Gateway -->> Rails: Responds with tool to invoke

    Rails ->> AI Gateway: POST /v1/agents/tools/<tool-name> with related data
    AI Gateway ->> LLM: Create a prompt and send it
    LLM -->> AI Gateway: Response
    AI Gateway -->> Rails: Response

When a new version of a prompt is introduced (like ai_gateway/agents/definitions/chat/explain_code/v1), then /v1/agents/tools/<tool-name>/<version> endpoint will be called.

Implementation details

Example: Execute Duo Chat explain code tool via agents (!160252 - merged)

Migrating other tools

Rails side

  • Define prompt_migration_#{unit_primitive} feature flag
  • Include UseAiGatewayAgentPrompt to the executor of the tool
    • If def unit_primitive method is not defined in the executor it will be nil when the feature flag is disabled and if it's enabled
    • If it's defined, it will always contain the defined value
    • It's done to cover this logic
      • If use_ai_gateway_agent_prompt is enabled - agent is called
      • If it's disabled, unit primitive is checked: if unit primitive is present #{BASE_ENDPOINT}/#{unit_primitive} is called; otherwise, ENDPOINT is called
  • Add it_behaves_like 'uses ai gateway agent prompt' to the executor's spec

AI Gateway side

Edited by Igor Drozdov