Saving Conversation History

[samDobsonDev]

I'm working with the Koog framework to build an AI agent service and need help understanding how to properly save and retrieve conversation histories.

After my agent completes its run (including any tool calls, tool call results, assistant responses, etc.), I need to extract and save the complete conversation history to a database. This includes the input user message, tool calls, tool call results and the final assistant message.

I believe the solution involves using LLM sessions through custom nodes, but I'm unclear on the implementation details.

Specific Questions:

Should I save conversation histories through strategy graphs or installed agent features (event handlers/tracing)? I've seen both approaches mentioned, but I'm not sure which is the recommended pattern.

Custom Node Implementation: If using strategy graphs, how do I properly create a custom node that accesses the LLM session to extract message history?

Message Access: What's the correct way to read all message types from the LLM session? I assume I need to use llm.readSession {} but I'm not sure about the exact API.

Database Integration: Where in the agent workflow should I place the history extraction/saving logic? Should it be part of the main strategy flow or a separate post-processing step?

Any code examples or links to the relevant documentation/API reference would be greatly appreciated.

Thanks :)

[Rizzen]

Hey @samDobsonDev! 👋

Looks like the new Agent Persistence feature in the freshly released Koog 0.3.0 might be exactly what you need! You can check out the official docs here.

There are two main approaches available: automatic persistence-restore and manual control. Based on your use case, it sounds like you'd prefer manual control over when and how state is saved and restored.

To implement manual checkpointing, you'll want to install the Persistency feature as described in the docs. Here's a simple example of how to create and store a checkpoint within a custom node:

node<String, String>(name) {
    val input = it
    withPersistency(this) { ctx ->
        createCheckpoint(ctx, name, input, typeOf<String>(), checkpointId)
    }
    return@node "$input\nCheckpoint Created"
}

Alternatively, you can invoke the checkpoint creation from any node using:

withPersistency(this) { ctx ->
    createCheckpoint(ctx, name, input, typeOf<String>(), checkpointId)
}

Later, when you want to restore the checkpoint, you can retrieve it and use it in your LLM session like so:

withPersistency(this) { ctx ->
    val checkpoint = getLatestCheckpoint()

    llm.writeSession {
        val id = prompt.id

        // Option 1: Replace the prompt entirely
        rewritePrompt {
            Prompt.build(id = id) {
                user {
                    checkpoint!!.messageHistory.map {
                        text(it.content)
                    }
                }
            }
        }

        // Option 2: Append to the existing prompt
        updatePrompt {
            user {
                checkpoint!!.messageHistory.map {
                    text(it.content)
                }
            }
        }
    }
}

Hope that helps! Let me know if you run into any problems with that.

[samDobsonDev]

Thank you! 😁