Skip to content

Commit

Permalink
Revert "Remove cache for now"
Browse files Browse the repository at this point in the history
This reverts commit a804e6d.
  • Loading branch information
dfuchss committed Aug 8, 2024
1 parent 892de10 commit 332d476
Showing 1 changed file with 16 additions and 1 deletion.
17 changes: 16 additions & 1 deletion src/main/kotlin/Translation.kt
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,9 @@ class TranslationService(private val config: TranslationConfig?) {

private val chatModel = config?.let { createChatModel(it) }

private val translationCache = mutableMapOf<String, String>()
private val translationCacheSize = 5

@Synchronized
fun translate(text: String): String {
if (config == null) {
Expand All @@ -22,12 +25,24 @@ class TranslationService(private val config: TranslationConfig?) {
return text
}

var translated = translationCache[text]
if (translated != null) {
return translated
}

try {
return chatModel.generate(config.prompt.replaceFirst("{}", config.model).replaceFirst("{}", text))
translated = chatModel.generate(config.prompt.replaceFirst("{}", config.model).replaceFirst("{}", text))
} catch (e: Exception) {
logger.error("Error while translating text: $text", e)
return text
}

if (translationCache.size >= translationCacheSize) {
translationCache.remove(translationCache.keys.first())
}
translationCache[text] = translated

return translated
}

private fun createChatModel(config: TranslationConfig): ChatLanguageModel {
Expand Down

0 comments on commit 332d476

Please sign in to comment.