반성(Reflection)으로 오류를 고치고 재시도하기
반성(Reflection)으로 오류를 고치고 재시도하기
LLM이 한 번에 완벽한 결과를 내는 경우는 드물어요. 특히 코드 같은 건 오류가 있기 마련이죠. **반성(Reflection)**은 이 문제를 해결하는 설계 패턴이에요. LLM 생성 결과가 나오면 그 결과를 조건으로 또 다른 LLM 생성(반성)을 수행하는 방식입니다. 예를 들어 코드 작성 작업이 주어졌을 때, 첫 번째 LLM은 코드 스니펫을 만들고 두 번째 LLM은 그 코드에 대한 비평(critique)을 만드는 식이죠.
AutoGen과 에이전트의 맥락에서 반성은 두 에이전트의 쌍으로 구현할 수 있어요. 첫 번째 에이전트가 메시지를 생성하고, 두 번째 에이전트가 그 메시지에 대한 응답을 생성하는 거예요. 두 에이전트는 **정지 조건(예: 최대 반복 횟수나 두 번째 에이전트의 승인)**에 도달할 때까지 계속 상호작용합니다. 이번에는 간단한 반성 패턴을 AutoGen 에이전트로 직접 구현해볼게요.
출처: 공식문서
두 에이전트는 **코더 에이전트(coder)**와 **리뷰어 에이전트(reviewer)**예요. 코더는 코드 스니펫을 생성하고, 리뷰어는 그 코드에 대한 비평을 생성하죠.
메시지 프로토콜
에이전트를 정의하기 전에 먼저 에이전트 간 메시지 프로토콜을 정의해야 해요.
from dataclasses import dataclass
@dataclass
class CodeWritingTask:
task: str
@dataclass
class CodeWritingResult:
task: str
code: str
review: str
@dataclass
class CodeReviewTask:
session_id: str
code_writing_task: str
code_writing_scratchpad: str
code: str
@dataclass
class CodeReviewResult:
review: str
session_id: str
approved: bool
위 메시지 집합은 우리 반성 패턴의 프로토콜을 정의해요.
- 애플리케이션은
CodeWritingTask메시지를 코더 에이전트에게 보냅니다. - 코더 에이전트는
CodeReviewTask메시지를 만들어 리뷰어 에이전트에게 보내요. - 리뷰어 에이전트는
CodeReviewResult메시지를 만들어 코더 에이전트에게 돌려보냅니다. CodeReviewResult에 따라, 코드가 승인되면 코더 에이전트는CodeWritingResult메시지를 애플리케이션에 보내요. 그렇지 않으면 코더 에이전트는 또 다른CodeReviewTask메시지를 리뷰어 에이전트에게 보내고 이 과정이 반복됩니다.
에이전트
이제 반성 패턴의 에이전트를 정의해볼게요.
import json
import re
import uuid
from typing import Dict, List, Union
from autogen_core import MessageContext, RoutedAgent, TopicId, default_subscription, message_handler
from autogen_core.models import (
AssistantMessage,
ChatCompletionClient,
LLMMessage,
SystemMessage,
UserMessage,
)
패턴을 구현하는 데 Broadcast API를 사용해요. 에이전트는 pub/sub 모델을 구현합니다. 코더 에이전트는 CodeWritingTask와 CodeReviewResult 메시지를 구독하고, CodeReviewTask와 CodeWritingResult 메시지를 게시해요.
@default_subscription
class CoderAgent(RoutedAgent):
"""An agent that performs code writing tasks."""
def __init__(self, model_client: ChatCompletionClient) -> None:
super().__init__("A code writing agent.")
self._system_messages: List[LLMMessage] = [
SystemMessage(
content="""You are a proficient coder. You write code to solve problems.
Work with the reviewer to improve your code.
Always put all finished code in a single Markdown code block.
For example:
```python
def hello_world():
print("Hello, World!")
Respond using the following format:
Thoughts:
@message_handler
async def handle_code_writing_task(self, message: CodeWritingTask, ctx: MessageContext) -> None:
# Store the messages in a temporary memory for this request only.
session_id = str(uuid.uuid4())
self._session_memory.setdefault(session_id, []).append(message)
# Generate a response using the chat completion API.
response = await self._model_client.create(
self._system_messages + [UserMessage(content=message.task, source=self.metadata["type"])],
cancellation_token=ctx.cancellation_token,
)
assert isinstance(response.content, str)
# Extract the code block from the response.
code_block = self._extract_code_block(response.content)
if code_block is None:
raise ValueError("Code block not found.")
# Create a code review task.
code_review_task = CodeReviewTask(
session_id=session_id,
code_writing_task=message.task,
code_writing_scratchpad=response.content,
code=code_block,
)
# Store the code review task in the session memory.
self._session_memory[session_id].append(code_review_task)
# Publish a code review task.
await self.publish_message(code_review_task, topic_id=TopicId("default", self.id.key))
@message_handler
async def handle_code_review_result(self, message: CodeReviewResult, ctx: MessageContext) -> None:
# Store the review result in the session memory.
self._session_memory[message.session_id].append(message)
# Obtain the request from previous messages.
review_request = next(
m for m in reversed(self._session_memory[message.session_id]) if isinstance(m, CodeReviewTask)
)
assert review_request is not None
# Check if the code is approved.
if message.approved:
# Publish the code writing result.
await self.publish_message(
CodeWritingResult(
code=review_request.code,
task=review_request.code_writing_task,
review=message.review,
),
topic_id=TopicId("default", self.id.key),
)
print("Code Writing Result:")
print("-" * 80)
print(f"Task:\n{review_request.code_writing_task}")
print("-" * 80)
print(f"Code:\n{review_request.code}")
print("-" * 80)
print(f"Review:\n{message.review}")
print("-" * 80)
else:
# Create a list of LLM messages to send to the model.
messages: List[LLMMessage] = [*self._system_messages]
for m in self._session_memory[message.session_id]:
if isinstance(m, CodeReviewResult):
messages.append(UserMessage(content=m.review, source="Reviewer"))
elif isinstance(m, CodeReviewTask):
messages.append(AssistantMessage(content=m.code_writing_scratchpad, source="Coder"))
elif isinstance(m, CodeWritingTask):
messages.append(UserMessage(content=m.task, source="User"))
else:
raise ValueError(f"Unexpected message type: {m}")
# Generate a revision using the chat completion API.
response = await self._model_client.create(messages, cancellation_token=ctx.cancellation_token)
assert isinstance(response.content, str)
# Extract the code block from the response.
code_block = self._extract_code_block(response.content)
if code_block is None:
raise ValueError("Code block not found.")
# Create a new code review task.
code_review_task = CodeReviewTask(
session_id=message.session_id,
code_writing_task=review_request.code_writing_task,
code_writing_scratchpad=response.content,
code=code_block,
)
# Store the code review task in the session memory.
self._session_memory[message.session_id].append(code_review_task)
# Publish a new code review task.
await self.publish_message(code_review_task, topic_id=TopicId("default", self.id.key))
def _extract_code_block(self, markdown_text: str) -> Union[str, None]:
pattern = r"```(\w+)\n(.*?)\n```"
# Search for the pattern in the markdown text
match = re.search(pattern, markdown_text, re.DOTALL)
# Extract the language and code block if a match is found
if match:
return match.group(2)
return None
`CoderAgent`에 대해 짚어볼 점이 몇 가지 있어요.
- 시스템 메시지에서 chain-of-thought(사고 사슬) 프롬프팅을 사용해요.
- 서로 다른 `CodeWritingTask`에 대한 메시지 히스토리를 딕셔너리에 저장해서, 각 작업이 자기만의 히스토리를 갖게 해요.
- 모델 클라이언트로 LLM 추론 요청을 할 때 메시지 히스토리를 `LLMMessage` 객체 목록으로 변환해 넘겨줘요.
리뷰어 에이전트는 `CodeReviewTask` 메시지를 구독하고 `CodeReviewResult` 메시지를 게시해요.
```python
@default_subscription
class ReviewerAgent(RoutedAgent):
"""An agent that performs code review tasks."""
def __init__(self, model_client: ChatCompletionClient) -> None:
super().__init__("A code reviewer agent.")
self._system_messages: List[LLMMessage] = [
SystemMessage(
content="""You are a code reviewer. You focus on correctness, efficiency and safety of the code.
Respond using the following JSON format:
{
"correctness": "<Your comments>",
"efficiency": "<Your comments>",
"safety": "<Your comments>",
"approval": "<APPROVE or REVISE>",
"suggested_changes": "<Your comments>"
}
""",
)
]
self._session_memory: Dict[str, List[CodeReviewTask | CodeReviewResult]] = {}
self._model_client = model_client
@message_handler
async def handle_code_review_task(self, message: CodeReviewTask, ctx: MessageContext) -> None:
# Format the prompt for the code review.
# Gather the previous feedback if available.
previous_feedback = ""
if message.session_id in self._session_memory:
previous_review = next(
(m for m in reversed(self._session_memory[message.session_id]) if isinstance(m, CodeReviewResult)),
None,
)
if previous_review is not None:
previous_feedback = previous_review.review
# Store the messages in a temporary memory for this request only.
self._session_memory.setdefault(message.session_id, []).append(message)
prompt = f"""The problem statement is: {message.code_writing_task}
The code is:
{message.code}
Previous feedback:
{previous_feedback}
Please review the code. If previous feedback was provided, see if it was addressed.
"""
# Generate a response using the chat completion API.
response = await self._model_client.create(
self._system_messages + [UserMessage(content=prompt, source=self.metadata["type"])],
cancellation_token=ctx.cancellation_token,
json_output=True,
)
assert isinstance(response.content, str)
# TODO: use structured generation library e.g. guidance to ensure the response is in the expected format.
# Parse the response JSON.
review = json.loads(response.content)
# Construct the review text.
review_text = "Code review:\n" + "\n".join([f"{k}: {v}" for k, v in review.items()])
approved = review["approval"].lower().strip() == "approve"
result = CodeReviewResult(
review=review_text,
session_id=message.session_id,
approved=approved,
)
# Store the review result in the session memory.
self._session_memory[message.session_id].append(result)
# Publish the review result.
await self.publish_message(result, topic_id=TopicId("default", self.id.key))
ReviewerAgent는 LLM 추론 요청을 할 때 JSON 모드를 사용하고, 시스템 메시지에서 chain-of-thought 프롬프팅도 사용해요.
로깅
에이전트 간에 주고받는 메시지를 보려면 로깅을 켜세요.
import logging
logging.basicConfig(level=logging.WARNING)
logging.getLogger("autogen_core").setLevel(logging.DEBUG)
패턴 실행하기
코딩 작업으로 이 패턴을 테스트해볼게요. 모든 에이전트가 default_subscription 클래스 데코레이터로 장식돼 있어서, 에이전트가 생성되면 자동으로 기본 토픽을 구독해요. 반성 과정을 시작하려면 기본 토픽에 CodeWritingTask 메시지를 게시하면 됩니다.
from autogen_core import DefaultTopicId, SingleThreadedAgentRuntime
from autogen_ext.models.openai import OpenAIChatCompletionClient
runtime = SingleThreadedAgentRuntime()
model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")
await ReviewerAgent.register(runtime, "ReviewerAgent", lambda: ReviewerAgent(model_client=model_client))
await CoderAgent.register(runtime, "CoderAgent", lambda: CoderAgent(model_client=model_client))
runtime.start()
await runtime.publish_message(
message=CodeWritingTask(task="Write a function to find the sum of all even numbers in a list."),
topic_id=DefaultTopicId(),
)
# Keep processing messages until idle.
await runtime.stop_when_idle()
# Close the model client.
await model_client.close()
로그 메시지에는 코더와 리뷰어 에이전트 사이의 상호작용이 표시돼요. 최종 출력에는 코더 에이전트가 생성한 코드 스니펫과 리뷰어 에이전트가 생성한 비평이 나옵니다.