코드 실행 에이전트 만들기(코딩 에이전트)

코드 실행 에이전트 만들기(코딩 에이전트)

AutoGen을 쓰다 보면 "모델이 코드를 만들고, 그 코드를 직접 실행해서 결과를 확인하는" 워크플로우가 정말 많이 필요해요. 이번에는 두 개의 커스텀 에이전트를 직접 구현해볼게요. 하나는 코드를 작성하는 어시스턴트(Assistant), 다른 하나는 그 코드를 실행하는 실행자(Executor)예요. 프레임워크가 기본 제공하는 AssistantAgent나 CodeExecutorAgent 대신, 직접 만들 수 있다는 걸 보여주기 위해서예요.

출처: 공식문서

준비: 도커 필요

이 예제에서 생성된 코드는 [Docker](https://www.docker.com/) 컨테이너 안에서 실행돼요. 예제를 실행하기 전에 Docker가 [설치](https://docs.docker.com/get-started/get-docker/)되어 있고 실행 중인지 확인하세요. 로컬 코드 실행([LocalCommandLineCodeExecutor](autogen_ext.code_executors.local.LocalCommandLineCodeExecutor))도 가능하지만, LLM이 생성한 코드를 로컬 환경에서 실행하는 위험 때문에 권장하지 않아요.

이 간단한 예제는 테슬라와 엔비디아의 주식 수익률 플롯을 만드는 두 에이전트를 구현해요. 먼저 에이전트 클래스와 각각의 메시지 처리 절차를 정의합니다.

import re
from dataclasses import dataclass
from typing import List

from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler
from autogen_core.code_executor import CodeBlock, CodeExecutor
from autogen_core.models import (
    AssistantMessage,
    ChatCompletionClient,
    LLMMessage,
    SystemMessage,
    UserMessage,
)


@dataclass
class Message:
    content: str


@default_subscription
class Assistant(RoutedAgent):
    def __init__(self, model_client: ChatCompletionClient) -> None:
        super().__init__("An assistant agent.")
        self._model_client = model_client
        self._chat_history: List[LLMMessage] = [
            SystemMessage(
                content="""Write Python script in markdown block, and it will be executed.
Always save figures to file in the current directory. Do not use plt.show(). All code required to complete this task must be contained within a single response.""",
            )
        ]

    @message_handler
    async def handle_message(self, message: Message, ctx: MessageContext) -> None:
        self._chat_history.append(UserMessage(content=message.content, source="user"))
        result = await self._model_client.create(self._chat_history)
        print(f"\n{'-'*80}\nAssistant:\n{result.content}")
        self._chat_history.append(AssistantMessage(content=result.content, source="assistant"))  # type: ignore
        await self.publish_message(Message(content=result.content), DefaultTopicId())  # type: ignore


def extract_markdown_code_blocks(markdown_text: str) -> List[CodeBlock]:
    pattern = re.compile(r"```(?:\s*([\w\+\-]+))?\n([\s\S]*?)```")
    matches = pattern.findall(markdown_text)
    code_blocks: List[CodeBlock] = []
    for match in matches:
        language = match[0].strip() if match[0] else ""
        code_content = match[1]
        code_blocks.append(CodeBlock(code=code_content, language=language))
    return code_blocks


@default_subscription
class Executor(RoutedAgent):
    def __init__(self, code_executor: CodeExecutor) -> None:
        super().__init__("An executor agent.")
        self._code_executor = code_executor

    @message_handler
    async def handle_message(self, message: Message, ctx: MessageContext) -> None:
        code_blocks = extract_markdown_code_blocks(message.content)
        if code_blocks:
            result = await self._code_executor.execute_code_blocks(
                code_blocks, cancellation_token=ctx.cancellation_token
            )
            print(f"\n{'-'*80}\nExecutor:\n{result.output}")
            await self.publish_message(Message(content=result.output), DefaultTopicId())

여기서 눈여겨볼 점이 있어요. 에이전트의 로직(모델을 쓰든 코드 실행기를 쓰든)은 메시지가 어떻게 전달되는지와 완전히 분리돼 있죠. 이것이 이 프레임워크의 핵심 아이디어예요. 프레임워크는 통신 인프라를 제공하고, 에이전트는 자기 자신의 로직만 책임집니다. 이 통신 인프라를 **에이전트 런타임(Agent Runtime)**이라고 불러요.

에이전트 런타임은 이 프레임워크의 핵심 개념이에요. 메시지를 전달하는 것 외에도 에이전트의 생명주기도 관리해요. 그래서 에이전트의 생성도 런타임이 처리합니다.

다음 코드는 로컬에 내장된 에이전트 런타임 구현인 SingleThreadedAgentRuntime을 사용해 에이전트를 등록하고 실행하는 방법을 보여줘요.

import tempfile

from autogen_core import SingleThreadedAgentRuntime
from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor
from autogen_ext.models.openai import OpenAIChatCompletionClient

work_dir = tempfile.mkdtemp()

# Create an local embedded runtime.
runtime = SingleThreadedAgentRuntime()

async with DockerCommandLineCodeExecutor(work_dir=work_dir) as executor:  # type: ignore[syntax]
    # Register the assistant and executor agents by providing
    # their agent types, the factory functions for creating instance and subscriptions.
    model_client = OpenAIChatCompletionClient(
        model="gpt-4o",
        # api_key="YOUR_API_KEY"
    )
    await Assistant.register(
        runtime,
        "assistant",
        lambda: Assistant(model_client=model_client),
    )
    await Executor.register(runtime, "executor", lambda: Executor(executor))

    # Start the runtime and publish a message to the assistant.
    runtime.start()
    await runtime.publish_message(
        Message("Create a plot of NVIDA vs TSLA stock returns YTD from 2024-01-01."), DefaultTopicId()
    )

    # Wait for the runtime to stop when idle.
    await runtime.stop_when_idle()
    # Close the connection to the model client.
    await model_client.close()

에이전트의 출력에서 테슬라와 엔비디아 주식 수익률 플롯이 생성된 걸 확인할 수 있어요.

from IPython.display import Image

Image(filename=f"{work_dir}/nvidia_vs_tesla_ytd_returns.png")  # type: ignore

AutoGen은 분산 에이전트 런타임도 지원해요. 서로 다른 프로세스나 머신에서, 서로 다른 정체성·언어·의존성을 가진 에이전트들을 호스팅할 수 있죠.

에이전트 런타임, 통신, 메시지 처리, 구독에 대해 더 배우고 싶다면 이 퀵스타트 다음 섹션들을 계속 읽어보세요.

더 알아보기 (Learn more)