Building a Code Agent

package ai.koog.agents.examples.codeagent.step01

import ai.koog.agents.core.agent.AIAgent
import ai.koog.agents.core.agent.singleRunStrategy
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.ext.tool.file.EditFileTool
import ai.koog.agents.ext.tool.file.ListDirectoryTool
import ai.koog.agents.ext.tool.file.ReadFileTool
import ai.koog.agents.features.eventHandler.feature.handleEvents
import ai.koog.prompt.executor.clients.openai.OpenAIModels
import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
import ai.koog.rag.base.files.JVMFileSystemProvider

val executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))
val agent = AIAgent(
    promptExecutor = executor,
    llmModel = OpenAIModels.Chat.GPT5Codex,
    toolRegistry = ToolRegistry {
        tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly))
        tool(ReadFileTool(JVMFileSystemProvider.ReadOnly))
        tool(EditFileTool(JVMFileSystemProvider.ReadWrite))
    },
    systemPrompt = """
        You are a highly skilled programmer tasked with updating the provided codebase according to the given task.
        Your goal is to deliver production-ready code changes that integrate seamlessly with the existing codebase and solve given task.
    """.trimIndent(),

    strategy = singleRunStrategy(),
    maxIterations = 100
) {
    handleEvents {
        onToolCallStarting { ctx ->
            println(
                "Tool '${ctx.toolName}' called with args:" +
                    " ${ctx.toolArgs.toString().take(100)}"
            )
        }
    }
}

suspend fun main(args: Array<String>) {
    if (args.size < 2) {
        println("Error: Please provide the project absolute path and a task as arguments")
        println("Usage: <absolute_path> <task>")
        return
    }

    val (path, task) = args
    val input = "Project absolute path: $path\n\n## Task\n$task"
    try {
        val result = agent.run(input)
        println(result)
    } finally {
        executor.close()
    }
}

val agent = AIAgent(
    promptExecutor = executor,
    llmModel = OpenAIModels.Chat.GPT5Codex,
    toolRegistry = ToolRegistry {
        tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly))
        tool(ReadFileTool(JVMFileSystemProvider.ReadOnly))
        tool(EditFileTool(JVMFileSystemProvider.ReadWrite))
        tool(createExecuteShellCommandToolFromEnv())  // NEW!
    },
    systemPrompt = """
        You are a highly skilled programmer tasked with updating the provided codebase according to the given task.
        Your goal is to deliver production-ready code changes that integrate seamlessly with the existing codebase and solve given task.
        Ensure minimal possible changes done - that guarantees minimal impact on existing functionality.
        
        You have shell access to execute commands and run tests.
        After investigation, define expected behavior with test scripts, then iterate on your implementation until the tests pass.
        Verify your changes don't break existing functionality through regression testing, but prefer running targeted tests over full test suites.
        Note: the codebase may be fully configured or freshly cloned with no dependencies installed - handle any necessary setup steps.
    """.trimIndent(),
    strategy = singleRunStrategy(),
    maxIterations = 400  // Increased for test iterations
)

fun createExecuteShellCommandToolFromEnv(): ExecuteShellCommandTool {
    return if (System.getenv("BRAVE_MODE")?.lowercase() == "true") {
        // Auto-approve all commands (use with caution!)
        ExecuteShellCommandTool(JvmShellCommandExecutor()) { _ -> ShellCommandConfirmation.Approved }
    } else {
        // Prompt user for confirmation
        ExecuteShellCommandTool(JvmShellCommandExecutor(), PrintShellCommandConfirmationHandler())
    }
}

val agent = AIAgent(
    promptExecutor = executor,
    llmModel = OpenAIModels.Chat.GPT5Codex,
    toolRegistry = ToolRegistry {
        tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly))
        tool(ReadFileTool(JVMFileSystemProvider.ReadOnly))
        tool(EditFileTool(JVMFileSystemProvider.ReadWrite))
        tool(createExecuteShellCommandToolFromEnv())
    },
    systemPrompt = """...""".trimIndent(),
    strategy = singleRunStrategy(),
    maxIterations = 400
) {
    install(OpenTelemetry) {
        setVerbose(true) // Send full strings instead of HIDDEN placeholders
        addLangfuseExporter(
            traceAttributes = listOf(
                CustomAttribute("langfuse.session.id", System.getenv("LANGFUSE_SESSION_ID") ?: "")
            )
        )
    }
    handleEvents {
        onToolCallStarting { ctx ->
            println("Tool '${ctx.toolName}' called with args: ${ctx.toolArgs.toString().take(100)}")
        }
    }
}

val findAgent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    llmModel = OpenAIModels.Chat.GPT4_1Mini,
    toolRegistry = ToolRegistry {
        tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly))
        tool(ReadFileTool(JVMFileSystemProvider.ReadOnly))
        tool(RegexSearchTool(JVMFileSystemProvider.ReadOnly))  // NEW!
    },
    systemPrompt = """
        You are an AI assistant specializing in code search.
        Your task is to analyze the user's query and provide clear and specific result.
        
        Break down the query, identify what exactly needs to be found, and note any ambiguities or alternative interpretations.
        If the query is ambiguous or could be improved, provide at least one result for each possible interpretation.
        
        Prioritize accuracy and relevance in your search results.
        * For each result, provide a clear and concise explanation of why it was selected.
        * The explanation should state the specific criteria that led to its selection.
        * If the match is partial or inferred, clearly state the limitations and potential inaccuracies.
        * Ensure to include only relevant snippets in the results.
        
        Ensure to utilize maximum amount of parallelization during the tool calling.
    """.trimIndent(),
    strategy = singleRunStrategy(),
    maxIterations = 100
) {
    setupObservability(agentName = "findAgent")
}

fun createFindAgentTool(): Tool<*, *> {
    return AIAgentService
        .fromAgent(findAgent as GraphAIAgent<String, String>)
        .createAgentTool<String, String>(
            agentName = "__find_in_codebase_agent__",
            agentDescription = """
                This tool is powered by an intelligent micro agent that analyzes and understands code context to find specific elements in your codebase.
                Unlike simple text search (ctrl+F), it intelligently interprets your query to locate classes, functions, variables, or files that best match your intent.
                It requires a detailed query describing what to search for, why you need this information, and an absolute path defining the search scope.
            """.trimIndent(),
            inputDescription = """
                The input contains two components: the absolute_path and the query.
                
                ## Query
                The query is a detailed search query for the intelligent agent to analyze.
                
                Examples of effective queries:
                - Find all implementations of the `UserRepository` interface to understand how data persistence is handled across the application
                - Locate files named `*Service.kt` containing `fun processOrder` because I need to modify the order processing logic
                
                ## absolute_path
                The absolute file system path to the directory where the search should begin.
                
                ## Formatting
                Provide the absolute_path and the query in this format: 'Absolute path for search scope: <absolute_path>\n\n## Query\n<query>'."
            """.trimIndent()
        )
}

import ai.koog.agents.ext.agent.HistoryCompressionConfig
import ai.koog.agents.ext.agent.singleRunStrategyWithHistoryCompression

val agent = AIAgent(
    promptExecutor = multiExecutor,
    llmModel = AnthropicModels.Opus_4_6,
    toolRegistry = ToolRegistry {
        tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly))
        tool(ReadFileTool(JVMFileSystemProvider.ReadOnly))
        tool(EditFileTool(JVMFileSystemProvider.ReadWrite))
        tool(createExecuteShellCommandToolFromEnv())
        tool(createFindAgentTool())
    },
    systemPrompt = """...""".trimIndent(),
    strategy = singleRunStrategyWithHistoryCompression(
        config = HistoryCompressionConfig(
            isHistoryTooBig = CODE_AGENT_HISTORY_TOO_BIG,
            compressionStrategy = CODE_AGENT_COMPRESSION_STRATEGY,
            retrievalModel = OpenAIModels.Chat.GPT4_1Mini
        )
    ),
    maxIterations = 400
)

val CODE_AGENT_HISTORY_TOO_BIG: (HistoryCompressionState) -> Boolean = { state ->
    // Compress when:
    // - More than 100 messages OR
    // - Estimated tokens > 50,000
    state.historySize > 100 || state.estimatedTotalTokens > 50_000
}

val CODE_AGENT_COMPRESSION_STRATEGY: (HistoryCompressionState) -> List<HistoryMessage> = { state ->
    val messages = state.history
    
    // Keep recent messages (last 20)
    val recentMessages = messages.takeLast(20)
    
    // Summarize older messages using LLM
    val olderMessages = messages.dropLast(20)
    val summary = summarizeMessages(olderMessages)
    
    listOf(summary) + recentMessages
}