Skip to content

Load prompt and model config from YAML

Background

At the current moment all Agent prompts are being stored in single Python module https://gitlab.com/gitlab-org/duo-workflow/duo-workflow-service/-/blob/64be012bc354241e71d0410a37f5242c0b587d8a/duo_workflow_service/agents/prompts.py#L20. This is not sustainable set up as single Python file will be ever growing entity slowly hindering progress as new prompts are being added

Goal

Enable more efficient iteration process over Agent prompts for Duo Workflow by separating them out of Python code. Driven by iterative spirit Duo Service should start with simple approach and evolve it as the requirements demands.

Implementation

  1. Create new directory to house prompts. This directory structure should mirror https://handbook.gitlab.com/handbook/engineering/architecture/design-documents/prompts_migration/#prompt-definition and follow compatible conventions. With difference being that Duo Workflow does not focus on specific features as a whole, but rather on Agent domain of expertise. Therefore in agents/definitions could for example contain subdirectories like planner or software_engineer
  2. Copy pydantic models used to represent agent definitions from https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/blob/1d6ea7d095f6b3121d31f74ceeef5bcad4b7f9fc/ai_gateway/prompts/config/base.py#L29
    1. Example agent definition: https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/blob/44d8b0e470b91fa1b79a2e5b71dd6a1878aac4d1/ai_gateway/prompts/definitions/chat/write_tests/base.yml#L20
  3. Copy model_factory from gitlab-org/duo-workflow/duo-workflow-service!59 (merged)
  4. Create agent registry class that accepts trailing subdirectory path after agents/definitions eg: planner/v1 or planner/v1/vertex , reads appointed definition and return agent config

Example usage

from prebuild.components.private_scratchpad import State, Agent1To1Reflexion, Toolbox
from langgraph.graph import END, StateGraph

class EngineeringManagerWorkflow(WorkflowProtocol):
    def run(goal: str):
        graph = StateGraph(State)
        prompt = LocalPromptRegistry.get('eng_manager/v1')


        agent = Agent1To1Reflexion(
            prompt_template = propmpt.temlpate,
            model = model_facotry(*prompt.model_params)
            tools = Toolbox.fetch(*prompt.tools)
        )
        agent.attach(graph)
Edited by Mikołaj Wawrzyniak