JetBrains AI
Supercharge your tools with AI-powered features inside many JetBrains products
Kotlin으로 AI 에이전트 구축하기 – 2부: 도구에 대한 심층 탐구
이전 글에서는 목록화, 읽기, 쓰기, 편집 기능을 갖춘 기본적인 코딩 에이전트를 어떻게 구축하는지 살펴보았습니다. 오늘은 Koog 프레임워크에서 추가 도구를 만들어 에이전트의 기능을 확장하는 방법에 대해 자세히 살펴봅니다. 일례로, ExecuteShellCommandTool을 구축하여, 에이전트가 코드를 실행하고 실제 엔지니어링에 사용되는 피드백 루프(코드 실행, 실패 관찰, 실제 출력을 기반으로 코드 개선)를 종료하도록 훈련합니다.
LLM은 문법 오류를 피하는 데에는 비교적 강하지만, 통합 단계에서는 어려움을 겪는 경우가 많습니다. 예를 들어, 존재하지 않는 메서드를 호출하거나, import를 누락하거나, 인터페이스를 부분적으로만 구현하는 식입니다. 코드를 즉시 컴파일하고 실행하는 전통적인 접근 방식을 통해 이러한 문제를 빠르게 확인할 수 있습니다. 여기에 약간의 프롬프트를 더하면 LLM에서 이러한 동작을 검증하기 위한 소규모 테스트를 직접 실행하도록 유도할 수 있습니다.
그렇다면 이런 도구는 어떻게 만들 수 있을까요? 기본적인 부분부터 시작해 보겠습니다.
Koog 도구 구조
가장 먼저 추상 클래스인 ai.koog.agents.core.tools.Tool을 상속하여 구현을 시작합니다. 이를 통해 도구 개발에 필요한 핵심 요소를 다음과 같이 정의할 수 있습니다.
- 이름: 저희 팀은 이름 앞뒤에 밑줄을 두 개 붙이는
snake_casing규칙을 따릅니다. 물론 이는 개인적인 취향의 문제입니다. - 설명: 이 필드는 LLM을 위한 주요 문서 역할을 하며, 해당 도구의 기능과 호출 이유를 설명합니다.
Args클래스: 이 클래스는 도구를 호출할 때 LLM이 필요로 하거나 제공할 수 있는 매개변수를 정의합니다.Result클래스: 이 클래스는 LLM에 전달될 메시지로 서식이 지정되는 데이터를 정의합니다. LLM이 읽을 수 있도록 데이터를 문자열 서식으로 지정하는textForLLM()함수를 포함할 수 있습니다.Result클래스는 기본적으로 개발자의 편의를 위한 것으로, 이를 통해 도구 결과를 로그에 기록하거나 UI에서 렌더링하기가 더 쉬워집니다. 에이전트 자체에는 서식이 지정된 문자열만 필요합니다.Execute()메서드: 이 메서드는 Args의 인스턴스를 입력으로 받아Result인스턴스를 반환합니다. LLM이 도구를 호출할 때의 처리 로직을 정의합니다.
이러한 구성 요소는 구축할 대상이 데이터베이스 커넥터든, API 클라이언트든, 오늘 살펴볼 셸 명령 실행기든 상관없이 모든 Koog 도구의 기초가 됩니다. 이제 이러한 원칙이 실제로 어떻게 작동하는지 ExecuteShellCommandTool을 구현해 확인해 보겠습니다.
안전성 확보를 위한 고려 사항
ExecuteShellCommandTool의 구체적인 내용으로 들어가기 전에, 몇 가지 핵심적인 안전 고려 사항을 먼저 짚고 넘어갈 필요가 있습니다.
LLM은 의도적으로 문제를 일으키는 악의적 행위자는 아니지만, 간혹 예기치 못한 실수로 인해 문제가 발생할 수 있습니다. 따라서 LLM에 명령줄 실행 권한을 부여한다면, 그에 따른 실수는 심각한 결과를 초래할 수 있습니다.
하지만 격리된 환경에서 명령어 실행을 샌드박싱하거나 명령어에 부여하는 권한을 제한하는 등, 이 위험을 완화하는 몇 가지 방법이 있습니다. 다만 이러한 방식은 구현이 상당히 복잡할 수 있으니, 여기서는 실제 도구를 만드는 데 중점을 두겠습니다.
이러한 점을 염두에 두고, 가장 간단한 위험 완화 전략을 두 가지 소개해 드립니다.
- 모든 명령어에 대한 명령어 실행 확인. LLM이 실행하려는 모든 명령어를 잠시 시간을 내어 검토하는 것은 가장 안전한 옵션이 될 수 있습니다. 단, 에이전트의 자율성은 심각하게 제한됩니다.
- Brave 모드: 모든 명령어가 자동으로 승인되는 모드입니다. 이 모드는 적절히 샌드박싱되지 않은 경우 작업 중인 시스템에 위험을 초래할 수 있습니다. 여기서는 자체 벤치마크 내부에서만 이를 사용하여 영향 없이 파기할 수 있는 격리된 환경에서 실행합니다.
ExecuteShellCommandTool 구현
ExecuteShellCommandTool은 다음과 같은 구성 요소로 이루어집니다.
- 이름:
__execute_shell_command__. - 설명: “셸 명령어를 실행하고 그 출력을 반환합니다” 등
Args클래스:command,timeoutSeconds,workingDirectoryResult클래스:output,exitCode,command(다소 의외일 수 있지만, 어떤 명령어가 실행되었는지를 로그나 UI에 보고하기에 편리합니다).Execute()메서드: 실행 확인을 요청한 뒤, 명령어를 실행합니다(구현 세부 사항은 아래에서 자세히 살펴봅니다).
Execute() 메서드 구현
이제 execute() 메서드를 어떻게 구현하는지 살펴보겠습니다. 핵심 로직은 다음과 같은 헬퍼 메서드에 위임하여 메서드를 단순하게 유지합니다.
override suspend fun execute(args: Args): Result = when (
val confirmation = confirmationHandler.requestConfirmation(args)
) {
is ShellCommandConfirmation.Approved -> try {
val result = executor.execute(
args.command, args.workingDirectory, args.timeoutSeconds
)
Result(args.command, result.exitCode, result.output)
} catch (e: CancellationException) {
throw e
} catch (e: Exception) {
Result(
args.command, null, "Failed to execute command: ${e.message}"
)
}
is ShellCommandConfirmation.Denied ->
Result(
args.command, null,
"Command execution denied with user response: ${confirmation.userResponse}"
)
}
흐름은 단순합니다. 명령어 실행에 대해 사용자 확인을 요청하고, 명령어 실행기를 통해 승인되면 실행하며, 승인되지 않으면 거부 메시지를 LLM에 반환합니다.
이를 통해 예외를 포착하고 오류 메시지를 LLM에 전달할 수 있으며, LLM이 접근 방식을 조정하거나 대안을 시도하도록 할 수 있습니다.
ConfirmationHandler 구성
ConfirmationHandler는 ExecuteShellCommandTool을 생성할 때 구성할 수 있으며, 다양한 구현을 허용합니다. 현재는 다음의 두 가지를 제공합니다.
PrintShellCommandConfirmationHandler: 명령줄을 통해 사용자에게 확인 메시지를 표시합니다.BraveModeConfirmationHandler: 모든 요청을 자동으로 승인합니다.
두 번째 구현은 아무 조건 없이 승인만 수행하지만, 첫 번째 구현에는 몇 가지 흥미로운 특징이 있습니다.
override suspend fun requestConfirmation(
args: ExecuteShellCommandTool.Args
): ShellCommandConfirmation {
println("Agent wants to execute: ${args.command}")
args.workingDirectory?.let { println("In: $it") }
println("Timeout: ${args.timeoutSeconds}s")
print("Confirm (y / n / reason-for-denying): ")
val userResponse = readln().lowercase()
return when (userResponse) {
"y", "yes" -> ShellCommandConfirmation.Approved
else -> ShellCommandConfirmation.Denied(userResponse)
}
}
사용자는 세 가지 옵션(승인, 거부, 사유를 포함한 거부)으로 인식하지만, 실제 구현상 두 가지 거부 유형은 동일하게 처리됩니다. 두 유형 모두 LLM으로 반환되며, LLM이 각 유형을 적절하게 해석하고 처리합니다.
CommandExecutor 구성
ConfirmationHandler와 마찬가지로 CommandExecutor 역시 구성 가능하지만, 현재는 JVM 구현만 제공합니다. 이론적으로는 Android, iOS, WebAssembly 및 기타 플랫폼용 구현을 생성할 수 있지만, 명확한 수요가 없다면 당분간은 보류합니다.
시간 초과 처리 방법
사용하는 플랫폼과 관계없이 CommandExecutor의 한 가지 측면, 즉 시간 초과 처리는 특별한 주의가 필요합니다. 현재 구현으로는 에이전트가 장시간 실행되는 명령어를 중단할 수 없습니다.
사용자는 때때로 조급해져서 실행을 취소하기 위해 무의식적으로 Ctrl+C를 누르기도 합니다. 하지만 이러한 동작은 에이전트가 멀티스레드 환경에서 사용자의 ‘조급함’을 이해할 것이라고 가정하고 있습니다. 이 상황에서는 더 간단하며 직관적인 더 나은 대안이 있습니다.
LLM에게 최대 실행 시간을 명시하도록 요구하면 이 한도를 초과하는 명령어를 안전하게 중단할 수 있습니다. 이 시간 초과 값은 사용자에게 표시되며, 요청된 기간이 비합리적이거나 과도해 보일 경우 사용자가 실행을 거부할 수 있습니다.
프로세스를 단순히 종료하고 일반적인 시간 초과 메시지를 반환하는 대신, 최대한 많은 출력을 보존하는 것을 목표로 해야 합니다. 불완전한 결과라도 LLM이 유용한 정보를 추출할 수도 있고, 최소한 시간 초과가 발생한 상황을 이해하는 데 도움이 될 수도 있습니다. 신중하게 구현하면 다음을 얻을 수 있습니다.
val stdoutJob = launch {
process.inputStream.bufferedReader().useLines { lines ->
try {
lines.forEach { stdoutBuilder.appendLine(it) }
} catch (_: IOException) {
// Ignore IO exception if the stream is closed and silently stop stream collection
}
}
}
val isCompleted = withTimeoutOrNull(timeoutSeconds * 1000L) {
process.onExit().await()
} != null
if (!isCompleted) {
process.destroyForcibly()
}
stdoutJob.join()
에이전트 수준의 변경 사항
에이전트 수준에서 수정 사항은 매우 미미한 수준이며, 실제 반영된 코드는 단 10여 줄에 불과합니다. diff에서 볼 수 있듯이, 이러한 업데이트의 대부분은 시스템 프롬프트를 확장하는 것과 관련이 있습니다.
단순히 도구를 추가해 변경 범위를 최소화할 수도 있었으나, 에이전트의 성능을 한층 더 끌어올리기 위해 정교한 최적화 작업을 두 가지 더 병행했습니다.
A) BRAVE_MODE 토글
이 토글의 구현은 비교적 단순하며, BRAVE_MODE 환경 변수를 확인하기만 하면 됩니다. 직접 구현하는 과정이 얼마나 간단한지 보여 드리기 위해, Brave 모드 ConfirmationHandler에 람다(lambda) 함수를 적용해 보았습니다. 물론 앞서 소개해 드린 BraveModeConfirmationHandler를 그대로 활용하셔도 무방합니다.
fun createExecuteShellCommandToolFromEnv(): ExecuteShellCommandTool {
return if (System.getenv("BRAVE_MODE")?.lowercase() == "true") {
ExecuteShellCommandTool(JvmShellCommandExecutor()) {
_ -> ShellCommandConfirmation.Approved
}
} else {
ExecuteShellCommandTool(
JvmShellCommandExecutor(),
PrintShellCommandConfirmationHandler(),
)
}
}
B) 프롬프트 내 ‘완료의 정의’ 개선
LLM이 코드 실행이라는 새로운 기능을 효과적으로 활용하도록, 테스트 작성과 실행을 강하게 유도하는 ‘완료의 정의’를 프롬프트에 추가했습니다.
""" ... // Previous prompt from step 01 Production-ready means verified to work—your changes must be proven correct and not introduce regressions. You have shell access to execute commands and run tests. Use this to work with executable feedback instead of assumptions. Establish what correct behavior looks like through tests, then iterate your implementation until tests pass. Validate that existing functionality remains intact. Production-ready means proven through green tests—that's your definition of done. """
이제 다 됐습니다. 최소한의 변경으로 ExecuteShellCommandTool 구성 요소를 완성하고, 이를 에이전트에 통합했습니다. 하지만 여기서 가장 중요한 질문이 남았습니다. 과연 이러한 변화가 실질적인 성능 향상으로 이어졌을까요?
벤치마크 테스트 결과
SWE-bench-verified 데이터 세트 수행 결과, 성능 개선 효과가 확인되었습니다. 성공 사례가 249/500건(50%)에서 279/500건(56%)으로 늘어나며 에이전트의 성능 향상이 확인되었습니다. 현재 리더보드 점수가 70% 내외임을 감안할 때, 이번 결과는 에이전트에게 코드 실행 및 검증 권한을 부여하는 방식이 올바른 발전 방향임을 증명해 줍니다. 다만 에이전트가 어떤 지점에서 고전하고 있는지, 다음 단계에서 무엇을 보완해야 할지 파악하기 위해서는 동작에 대한 가시성 개선이 필요합니다. 바로 그 지점에서 로깅이 중요해집니다.
결론: Koog에서 도구 구축하기
이 글에서는 Koog의 도구 구조가 어떻게 작동하는지 살펴보았습니다. 모든 도구에는 이름, 설명, execute() 메서드가 필요합니다. 또한, ExecuteShellCommandTool을 ConfirmationHandler와 CommandExecutor 구성 요소를 통해 구성 가능하게 설계하였고 복잡한 로직을 교체 가능한 구현으로 위임하는 방법을 보여 드렸습니다.
이러한 동일한 패턴은 데이터베이스 커넥터, API 클라이언트, 파일 프로세서 또는 사용자 정의 통합 등 구축하고자 하는 어떤 도구에도 적용할 수 있습니다. 프레임워크는 LLM 커뮤니케이션을 위한 구조를 제공하고 사용자는 구체적인 기능을 제공하는 방식입니다.
다음 글에서는 에이전트의 동작을 이해하고 개선하는 데 필요한 가시성을 확보하기 위해 에이전트에 로깅과 추적을 추가하는 방법을 살펴봅니다. 반복 처리를 위해서는 에이전트의 의사결정 방식을 파악해야 하며, 적절한 관찰 가능성 도구는 에이전트 자체에 제공하는 도구만큼 중요합니다.
게시물 원문 작성자
Subscribe to JetBrains AI Blog updates
