Code Execution

code_exec:
  type: CodeExecutionStep
  config:
    model: "gemini-2.5-flash"
  input_mapping:
    query: __inputs__#query

nodes:
  # Extract numbers from the query
  extract_data:
    type: LanguageModelStep
    config:
      prompt_template: |
        Extract any numerical data from this query and format as JSON:
        Query: {query}
      model_id: "vertex_ai/gemini-2.0-flash"
    input_mapping:
      query: __inputs__#query

  # Perform calculations
  calculate:
    type: CodeExecutionStep
    config:
      model: "gemini-2.5-flash"
    input_mapping:
      query: calculation_prompt#output

  # Generate final response
  generate:
    type: ResponseGenerationStep
    input_mapping:
      query: __inputs__#query
      calculation_result: calculate#code_result

nodes:
  # Get structured data
  fetch_data:
    type: WebhookStep
    config:
      webhook_url: "https://api.example.com/data"
      method: GET
    input_mapping: {}

  # Transform the data with code
  transform:
    type: CodeExecutionStep
    config:
      model: "gemini-2.5-flash"
    input_mapping:
      query: transform_request#output

  __outputs__:
    type: output
    input_mapping:
      transformed_data: transform#code_result

research:
  type: AgenticResearchStep
  config:
    tools_config:
      - name: search_docs
        description: Search documentation for information.
        step_config:
          type: SearchUnstructuredDataStep
          config:
            top_k: 50

      - name: run_calculation
        description: |
          Execute Python code to perform calculations.
          Use for math, data analysis, or generating charts.
        step_config:
          type: CodeExecutionStep
          config:
            model: "gemini-2.5-flash"

    agent_config:
      agent_loop:
        num_turns: 10
        model_name_or_path: "vertex_ai/claude-sonnet-4-5@20250929"
        research_guidelines_prompt: |
          You have access to:
          - `search_docs` — Find information in documents
          - `run_calculation` — Execute Python code for calculations
          
          Use run_calculation when you need to:
          - Perform mathematical operations
          - Analyze or transform data
          - Generate visualizations

  input_mapping:
    message_history: create_message_history#message_history