feat(memory): add token budgeting, fact summaries, and similarity ranking

[fazlerahmanejazi]

Motivation and Context

Support configurable token budgets, summaries, and similarity ranking so memory injections stay concise and relevant.

• add opt-in token budgets, summarisation, and similarity ranking for AgentMemory via the SmartAgentMemoryProvider decorator and shared FactReplayProcessor (uses Embedder, Tokenizer, RankedDocumentStorage; defaults keep legacy behaviour when the decorator isn’t applied).
• add tests (AgentMemoryEnhancementsTest, AgentMemoryEnrichmentTest)

Breaking Changes

No breaking changes—new controls are optional and preserve existing defaults.

---

Type of the changes

• New feature (non-breaking change which adds functionality)
• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update
• Tests improvement
• Refactoring

Checklist

• The pull request has a description of the proposed change
• I read the Contributing Guidelines before opening the pull request
• The pull request uses develop as the base branch
• Tests for the changes have been added
• All new and existing tests passed

Additional steps for pull requests adding a new feature

• An issue describing the proposed change exists
• The pull request includes a link to the issue
• The change was discussed and approved in the issue
• Docs have been added / updated

[aozherelyeva]

Hi @fazlerahmanejazi! Could you please resolve the conflicts? :)

[fazlerahmanejazi]

Hi @fazlerahmanejazi! Could you please resolve the conflicts? :)

Done @aozherelyeva, ready for review!

[Ololoshechkin]

Thanks for your pull request and contribution!

I see the idea as very valuable and practical, though there are some design suggestions that I propose to consider.

AgentMemoryProvider is already an interface that defines how to retrieve facts, and it's being used by AgentMemory feature. So I propose moving all additional filters, embedders and summarizers to some abstract class or implementation of that AgentMemoryProvider interface. Basically, you'd need to override suspend fun load(concept: Concept, subject: MemorySubject, scope: MemoryScope): List<Fact> and suspend fun save(fact: Fact, subject: MemorySubject, scope: MemoryScope) methods for that.

load can do the ranking, summarizing and filtering.

Additionally, please consider using ai.koog.rag.base.RankedDocumentStorage (https://docs.koog.ai/ranked-document-storage/) for ranking -- it's a generic interface that can be implemented via LLM embeddings + local storages as well as with vector databases (ai.koog.rag.vector.EmbeddingBasedDocumentStorage)

[Ololoshechkin]

Could you please elaborate why ai.koog.embeddings.base.Embedder (and classes that implement it) doesn't fit here?

[Ololoshechkin]

Please consider checking ai.koog.prompt.tokenizer.Tokenizer interface (there's already a simple regex-based tokenizer implementation, as well as TiktokenEncoder)

[fazlerahmanejazi]

@Ololoshechkin thanks again for the thoughtful review. I’ve pushed a set of changes that addresses each point:

• Provider-based enrichment/ranking: A SmartAgentMemoryProvider decorator now overrides save/load to run summarisation on save and ranking/budgeting on load. AgentMemory simply forwards the MemoryRequestOptions; plain providers continue to behave exactly as before.

• Replay engine centralised: SmartAgentMemoryProvider delegates the heavy lifting to a reusable FactReplayProcessor, so the enrichment, ranking, and token-budget logic lives in one place. Providers opt in via the lightweight MemoryPostProcessor contract; non-implementers return the raw facts.

• Koog primitives throughout:

  • Replaced the ad-hoc EmbeddingProvider with the existing ai.koog.embeddings.base.Embedder interface.
  • Token estimation now uses the ai.koog.prompt.tokenizer.Tokenizer interface (defaulting to the regex implementation).
  • Ranking now happens via ai.koog.rag.base.RankedDocumentStorage. The processor accepts a storage factory so future integrations can supply persistent/vector-backed implementations; the default is an in-memory adapter that uses the embedder under the hood.

Happy to dig into any of the details!