Trip Planning Agent

fun createPlannerAgent(
    promptExecutor: PromptExecutor,
    openMeteoClient: OpenMeteoClient,
    googleMapsMcpRegistry: ToolRegistry,
    onToolCallEvent: (String) -> Unit,
    showMessage: suspend (String) -> String,
): AIAgent<UserInput, TripPlan> {
    val googleMapTools = googleMapsMcpRegistry.tools
    val weatherTools = WeatherTools(openMeteoClient)
    val userTools = UserTools(showMessage)

    val toolRegistry = ToolRegistry {
        tool(::addDate)
        tools(weatherTools)
        tools(userTools)
    } + googleMapsMcpRegistry

    val plannerStrategy = plannerStrategy(
        googleMapsTools = googleMapTools,
        addDateTool = ::addDate.asTool(),
        weatherTools = weatherTools,
        userTools = userTools,
    )

    val agentConfig = AIAgentConfig(
        prompt = prompt(
            "planner-agent-prompt",
            params = LLMParams(temperature = 0.2)
        ) {
            system(
                """
                You are a trip planning agent that helps the user to plan their trip.
                Use the information provided by the user to suggest the best possible trip plan.
                """.trimIndent()
            )
        },
        model = OpenAIModels.Chat.GPT4o,
        maxAgentIterations = 200
    )

    return AIAgent<UserInput, TripPlan>(
        promptExecutor = promptExecutor,
        strategy = plannerStrategy,
        agentConfig = agentConfig,
        toolRegistry = toolRegistry,
    ) {
        handleEvents {
            onToolCall { ctx ->
                onToolCallEvent(
                    "Tool ${ctx.tool.name}, args ${
                        ctx.toolArgs.toString().replace('\n', ' ').take(100)
                    }..."
                )
            }
        }
    }
}

private fun plannerStrategy(
    googleMapsTools: List<Tool<*, *>>,
    addDateTool: Tool<*, *>,
    weatherTools: WeatherTools,
    userTools: UserTools
) = strategy<UserInput, TripPlan>("planner-strategy") {
    val userPlanKey = createStorageKey<TripPlan>("user_plan")
    val prevSuggestedPlanKey = createStorageKey<TripPlan>("prev_suggested_plan")

    // Nodes
    val setup by node<UserInput, String> { userInput ->
        llm.writeSession {
            updatePrompt {
                system {
                    +"Today's date is ${userInput.currentDate}."
                }
            }
        }
        userInput.message
    }

    val clarifyUserPlan by subgraphWithTask<String, TripPlan>(
        tools = userTools.asTools() + addDateTool
    ) { initialMessage ->
        xml {
            tag("instructions") {
                +"""
                Clarify a user plan until the locations, dates and additional information, such as user preferences, are provided.    
                """.trimIndent()
            }
            tag("initial_user_message") {
                +initialMessage
            }
        }
    }

    val suggestPlan by subgraphWithTask<SuggestPlanRequest, TripPlan>(
        tools = googleMapsTools + weatherTools.asTools()
    ) { input ->
        xml {
            tag("instructions") {
                markdown {
                    h2("Requirements")
                    bulleted {
                        item("Suggest the plan for ALL days and ALL locations in the user plan, preserving the order.")
                        item("Follow the user plan and provide a detailed step-by-step plan suggestion with multiple options for each date.")
                        item("Consider weather conditions when suggesting places for each date and time to assess how suitable the activity is for the weather.")
                        item("Check detailed information about each place, such as opening hours and reviews, before adding it to the final plan suggestion.")
                    }

                    h2("Tool usage guidelines")
                    +"""
                    ALWAYS use "maps_search_places" tool to search for places, AVOID making your own suggestions.
                    While searching for places, keep search query short and specific:
                    Example DO: "museum", "historical museum", "italian restaurant", "coffee shop", "art gallery"
                    Example DON'T: "interesting cultural sites", "local cuisine restaurants", "restaurant in the city center"
                    """.trimIndent()
                }
            }

            when (input) {
                is SuggestPlanRequest.InitialRequest -> {
                    tag("user_plan") {
                        +input.userPlan.toMarkdownString()
                    }
                }
                is SuggestPlanRequest.CorrectionRequest -> {
                    tag("additional_instructions") {
                        +"User asked for corrections to the previously suggested plan. Provide updated plan according to these corrections."
                    }
                    tag("user_plan") {
                        +input.userPlan.toMarkdownString()
                    }
                    tag("previously_suggested_plan") {
                        +input.prevSuggestedPlan.toMarkdownString()
                    }
                    tag("user_feedback") {
                        +input.userFeedback
                    }
                }
            }
        }
    }

    val saveUserPlan by node<TripPlan, Unit> { plan ->
        storage.set(userPlanKey, plan)
        llm.writeSession {
            replaceHistoryWithTLDR(strategy = HistoryCompressionStrategy.WholeHistory)
        }
    }

    val savePrevSuggestedPlan by node<TripPlan, TripPlan> { plan ->
        storage.set(prevSuggestedPlanKey, plan)
        llm.writeSession {
            replaceHistoryWithTLDR(strategy = HistoryCompressionStrategy.WholeHistory)
        }
        plan
    }

    val showPlanSuggestion by node<String, String> { message ->
        userTools.showMessage(message)
    }

    val processUserFeedback by nodeLLMRequestStructured<PlanSuggestionFeedback>()

    // Edges - define the flow
    nodeStart then setup then clarifyUserPlan then saveUserPlan

    edge(
        savePrevSuggestedPlan forwardTo showPlanSuggestion
            transformed { it.toMarkdownString() }
    )

    edge(showPlanSuggestion forwardTo processUserFeedback)

    // Feedback loop: if not accepted, iterate
    edge(
        processUserFeedback forwardTo createPlanCorrectionRequest
            transformed { it.getOrThrow().structure }
            onCondition { !it.isAccepted }
            transformed { it.message }
    )

    // If accepted, finish
    edge(
        processUserFeedback forwardTo nodeFinish
            transformed { it.getOrThrow().structure }
            onCondition { it.isAccepted }
            transformed { storage.getValue(prevSuggestedPlanKey) }
    )

    edge(createPlanCorrectionRequest forwardTo suggestPlan)
}

class WeatherTools(private val openMeteoClient: OpenMeteoClient) : ToolSet {
    @Tool
    @LLMDescription("Get weather forecast for a location")
    suspend fun getWeather(
        @LLMDescription("Latitude") latitude: Double,
        @LLMDescription("Longitude") longitude: Double,
        @LLMDescription("Start date (YYYY-MM-DD)") startDate: String,
        @LLMDescription("End date (YYYY-MM-DD)") endDate: String
    ): WeatherForecast {
        return openMeteoClient.getWeather(
            latitude, longitude, startDate, endDate
        )
    }
}