병렬 하위 작업 에이전트 워크플로우

병렬 하위 작업 에이전트 워크플로우 (Parallel Subtask Agent Workflow)

이 노트북에서는 복잡한 작업을 자동으로 여러 개의 하위 작업(subtask)으로 쪼개서 처리하는 범용 에이전트 워크플로우를 만들어 볼 거예요. 쪼개진 하위 작업들은 MistralAI의 LLM을 병렬로 호출해서 처리하고, 필요할 때마다 Tavily API의 실시간 정보로 보강해요. 마지막에는 각 결과를 종합해서 하나의 완성된 답변으로 합칩니다.

출처: 문서

본문

소개 (Introduction)

이 노트북은 복잡한 작업을 자동으로 여러 개의 하위 작업으로 나누는 범용 에이전트 워크플로우를 만드는 방법을 보여줘요.

이 하위 작업들은 MistralAI LLM을 병렬로 호출해서 처리하고, Tavily API의 실시간 정보로 보강합니다.

처리된 결과는 마지막에 하나의 종합적인 응답으로 합쳐져요.

워크플로우 개요 (Workflow Overview)

오케스트레이터(orchestrator) LLM이 메인 작업을 분석해서 서로 독립적인 병렬 하위 작업들로 나눠요.
각 하위 작업은 전문적인 지시사항을 가진 워커(worker) LLM에게 배정됩니다.
워커들은 병렬로 실행되며, 필요할 때 Tavily API를 이용해 최신 정보를 가져와요.
결과는 하나의 통합된 응답으로 종합됩니다.

참고: 하위 작업 처리와 응답 합성에는 MistralAI의 LLM을 쓰고, 최신 실시간 정보를 가져올 때는 Tavily API를 사용할 거예요.

솔루션 아키텍처 (Solution Architecture)

solution architecture

설치 (Installation)

Python Output

!pip install -U mistralai

임포트 (Imports)

import os
import json
import asyncio
import requests
from typing import Any, Optional, Dict, List, Union
from pydantic import Field, BaseModel, ValidationError
from mistralai.client import Mistral
from IPython.display import display, Markdown

import nest_asyncio
nest_asyncio.apply()

API 키 설정 (Set your API keys)

여기서 MistralAI와 Tavily의 API 키를 설정해요. 키는 아래 링크에서 발급받을 수 있어요.

MistralAI: https://console.mistral.ai/api-keys
Tavily: https://app.tavily.com/home

MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY", "<YOUR MISTRAL API KEY>") # Get it from https://console.mistral.ai/api-keys
TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY", "<YOUR TAVILY API KEY>") # Get it from https://app.tavily.com/home

Mistral 클라이언트 초기화 (Initialize Mistral client)

mistral_client = Mistral(api_key=MISTRAL_API_KEY)
MISTRAL_MODEL = "mistral-small-latest"  # Can be configured based on needs

Tavily API 설정 (Tavily API configuration)

TAVILY_API_URL = "https://api.tavily.com/search"
TAVILY_HEADERS = {
    "Authorization": f"Bearer {TAVILY_API_KEY}",
    "Content-Type": "application/json"
}

구조화된 데이터를 위한 Pydantic 모델 (Pydantic Models for Structured Data)

Pydantic 모델은 데이터 검증과 직렬화를 제공해서, LLM에서 받은 데이터가 우리가 기대하는 구조와 일치하도록 보장해요. 덕분에 오케스트레이터와 워커 사이에서 일관성을 유지할 수 있어요.

SubTask: 개별 하위 작업 정의예요. 타입, 설명, 선택적인 검색 쿼리를 가진 하나의 독립적인 작업 단위를 정의해요.

TaskList: 오케스트레이터의 출력 구조예요. 분석 내용과 병렬로 실행될 하위 작업들의 목록을 담고 있어요.

class SubTask(BaseModel):
    """Individual subtask definition"""
    task_id: str
    type: str
    description: str
    search_query: Optional[str]  # Query for Tavily search for the subtask

class TaskList(BaseModel):
    """Structure for orchestrator output"""
    analysis: str
    subtasks: List[SubTask]

API 유틸리티 함수 (API Utility Functions)

API 유틸리티 함수는 외부 API와의 통신을 처리하고 응답을 가공해서, 워크플로우의 나머지 부분에 깔끔한 인터페이스를 제공해요.

fetch_information: 쿼리를 기반으로 Tavily API에서 관련 정보를 가져와 구조화된 결과를 반환해요.

run_mistral_llm: 주어진 프롬프트로 Mistral AI를 표준 호출해서 생성된 콘텐츠를 반환해요.

parse_structured_output: Mistral의 구조화된 출력 기능을 이용해 Pydantic 모델에 맞게 응답을 생성하고 파싱해요.

def fetch_information(query: str, max_results: int = 3):
    """Retrieve information from Tavily API"""
    payload = {
        "query": query,
        "search_depth": "advanced",
        "include_answer": True,
        "max_results": max_results
    }

    try:
        response = requests.post(TAVILY_API_URL, json=payload, headers=TAVILY_HEADERS)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching data from Tavily: {e}")
        return {"error": str(e), "results": []}

def run_mistral_llm(prompt: str, system_prompt: Optional[str] = None):
    """Run Mistral LLM with given prompts"""
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    messages.append({"role": "user", "content": prompt})

    response = mistral_client.chat.complete(
        model=MISTRAL_MODEL,
        messages=messages,
        temperature=0.7,
        max_tokens=4000
    )

    return response.choices[0].message.content

def parse_structured_output(prompt: str, response_format: BaseModel, system_prompt: Optional[str] = None):
    """Get structured output from Mistral LLM based on a Pydantic model"""
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    messages.append({"role": "user", "content": prompt})

    response = mistral_client.chat.parse(
        model=MISTRAL_MODEL,
        messages=messages,
        response_format=response_format,
        temperature=0.2
    )

    return json.loads(response.choices[0].message.content)

비동기 워커 함수 (Async Worker Functions)

이 함수들은 하위 작업을 병렬로 실행할 수 있게 해줘요. 여러 컴포넌트를 동시에 처리해서 효율을 크게 높여주죠.

run_task_async: 단일 하위 작업을 비동기적으로 실행하고, 필요할 때 Tavily의 관련 정보로 보강해요.

execute_tasks_in_parallel: 모든 하위 작업의 병렬 실행을 관리해서, 동시에 실행되고 결과가 제대로 수집되도록 보장해요.

async def run_task_async(task: SubTask, original_task: str):
    """Execute a single subtask asynchronously with Tavily enhancement"""

    # Prepare context with Tavily information if a search query is provided
    context = ""
    if task.search_query:
        print(f"Fetching information for: {task.search_query}")
        search_results = fetch_information(task.search_query)

        # Format search results into context
        if "results" in search_results and search_results["results"]:
            context = "### Relevant Information:\n"
            for result in search_results["results"]:
                context += f"- {result.get('content', '')}\n"

        if "answer" in search_results and search_results["answer"]:
            context += f"\n### Summary: {search_results['answer']}\n"

    # Worker prompt with task information and context
    worker_prompt = f"""
    Complete the following subtask based on the given information:

    Original Task: {original_task}

    Subtask Type: {task.type}
    Subtask Description: {task.description}

    {context}

    Please provide a detailed response for this specific subtask only.
    """

    # Use asyncio to run in a thread pool to prevent blocking
    return await asyncio.to_thread(
        run_mistral_llm,
        prompt=worker_prompt,
        system_prompt="You are a specialized agent focused on solving a specific aspect of a larger task."
    )

async def execute_tasks_in_parallel(subtasks: List[SubTask], original_task: str):
    """Execute all subtasks in parallel"""
    tasks = []
    for subtask in subtasks:
        tasks.append(run_task_async(subtask, original_task))

    return await asyncio.gather(*tasks)

메인 워크플로우 함수 (Main Workflow Function)

초기 요청부터 최종 종합 응답까지 전체 병렬 하위 작업 과정을 조율하는 핵심 오케스트레이션 함수예요.

parallel_subtask_workflow: 작업 분해, 하위 작업 병렬 실행, 결과 종합까지 전체 워크플로우를 조율해서 하나의 종합적인 응답으로 만들어요.

단계(Steps):

작업 분석 (Task Analysis): 오케스트레이터가 사용자의 쿼리를 분석해서 서로 다른 하위 작업으로 나눠요.
하위 작업 정의 (Subtask Definition): 각 하위 작업은 고유한 ID, 타입, 설명, 검색 쿼리로 정의돼요.
병렬 실행 (Parallel Execution): 워커 에이전트들이 하위 작업을 동시에 실행해요.
정보 보강 (Information Enhancement): 워커들이 필요할 때 Tavily에서 관련 정보를 가져와요.
결과 수집 (Result Collection): 모든 워커의 출력이 모아져요.
종합 (Synthesis): 개별 결과들이 하나의 종합적인 최종 응답으로 결합돼요.
최종 응답 (Final Response): 개별 분석과 종합된 답변을 포함한 전체 워크플로우 결과가 반환돼요.

async def workflow(user_task: str):
    """Main workflow function to process a task through the parallel subtask agent workflow"""

    print("=== USER TASK ===\n")
    print(user_task)

    # Step 1: Orchestrator breaks down the task into subtasks
    orchestrator_prompt = f"""
    Analyze this task and break it down into 3-5 distinct, specialized subtasks that could be executed in parallel:

    Task: {user_task}

    For each subtask:
    1. Assign a unique task_id
    2. Define a specific type that describes the subtask's focus
    3. Write a clear description explaining what needs to be done
    4. Provide a search query if the subtask requires additional information

    First, provide a brief analysis of your understanding of the task.
    Then, define the subtasks that would collectively solve this problem effectively.

    Remember to make the subtasks complementary, not redundant, focusing on different aspects of the problem.
    """

    orchestrator_system_prompt = """
    You are a task orchestrator that specializes in breaking down complex problems into smaller,
    well-defined subtasks that can be solved independently and in parallel. Think carefully about
    the most logical way to decompose the given task.
    """

    print("\nOrchestrating task decomposition...")
    # Get structured output from orchestrator
    task_breakdown = parse_structured_output(
        prompt=orchestrator_prompt,
        response_format=TaskList,
        system_prompt=orchestrator_system_prompt
    )

    # Display orchestrator output
    print("\n=== ORCHESTRATOR OUTPUT ===")
    print(f"\nANALYSIS:\n{task_breakdown['analysis']}")
    print("\nSUBTASKS:")
    for task in task_breakdown["subtasks"]:
        print(f"- {task['task_id']}: {task['type']} - {task['description'][:100]}...")

    # Step 2: Execute subtasks in parallel
    print("\nExecuting subtasks in parallel...")
    subtask_results = await execute_tasks_in_parallel(
        [SubTask(**task) for task in task_breakdown["subtasks"]],
        user_task
    )

    # Display worker results
    for i, (task, result) in enumerate(zip(task_breakdown["subtasks"], subtask_results)):
        print(f"\n=== WORKER RESULT ({task['type']}) ===")
        print(f"{result[:200]}...\n")

    # Step 3: Synthesize final response
    print("\nSynthesizing final response...")

    # Format worker responses for synthesizer
    worker_responses = ""
    for task, response in zip(task_breakdown["subtasks"], subtask_results):
        worker_responses += f"\n=== SUBTASK: {task['type']} ===\n{response}\n"

    synthesizer_prompt = f"""
    Given the following task: {user_task}

    And these responses from different specialized agents focusing on different aspects of the task:

    {worker_responses}

    Please synthesize a comprehensive, coherent response that addresses the original task.
    Integrate insights from all specialized agents while avoiding redundancy.
    Ensure your response is balanced, considering all the different perspectives provided.
    """

    final_response = run_mistral_llm(
        prompt=synthesizer_prompt,
        system_prompt="You are a synthesis agent that combines specialized analyses into comprehensive responses."
    )

    return {
        "orchestrator_analysis": task_breakdown["analysis"],
        "subtasks": task_breakdown["subtasks"],
        "subtask_results": subtask_results,
        "final_response": final_response
    }

예제 작업으로 워크플로우 실행하기 (Run workflow with an example task)

휴대폰 추천을 비교하는 샘플 예제 작업으로 워크플로우를 실행해 볼게요.

task = "Compare the iPhone 16 Pro, iPhone 15 Pro, and Google Pixel 9 Pro, and recommend which one I should purchase."

Python

result = asyncio.run(workflow(task))

최종 응답 (Final Response)

Python Output

print("\n=== FINAL SYNTHESIZED RESPONSE ===")

display(Markdown(result["final_response"]))

오케스트레이터 분석과 하위 작업 정보, 응답 살펴보기 (Examining Orchestrator Analysis, Subtask information and responses)

오케스트레이터의 분석 내용, 생성된 하위 작업들, 해당 검색 쿼리, 그리고 개별 응답들을 살펴볼 수 있어요.

Python

print("\n=== Orchestrator Analysis ===\n")

display(Markdown(result['orchestrator_analysis']))

Python

print("\n=== SUBTASKS CREATED ===\n")

for subtask in result['subtasks']:
    display(Markdown(f"- {subtask['task_id']}: \n - Task type: {subtask['type']} \n - Task Description: - {subtask['description'][:100]} \n - search_query - {subtask['search_query']}"))

Python

print("\n=== SUBTASKS RESULTS ===\n")

for i, subtask_result in enumerate(result['subtask_results']):

    display(Markdown(f"# Task_{i+1} Result: \n {subtask_result} \n"))

Python

print("\n=== FINAL SYNTHESIZED RESPONSE ===")

display(Markdown(result["final_response"]))

더 알아보기 (Learn more)

원본 문서: Parallel Subtask Agent Workflow