Migrate Duo Chat Tools: WriteTests
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 benil
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
- If
- If
- Add
it_behaves_like 'uses ai gateway agent prompt'
to the executor's spec
AI Gateway side
- Define prompts in
ai_gateway/prompts/definitions/chat/{tool name}
folder - Example: feat(agents): add prompts for explain code tool (gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!1132 - merged)
Edited by Igor Drozdov