Skip to content

Commit fbe08f0

Browse files
authored
Ollama support for file attachments (#1221)
<!-- Thank you for opening a pull request! Please add a brief description of the proposed change here. Also, please tick the appropriate points in the checklist below. --> ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> ## Breaking Changes <!-- Will users need to update their code or configurations? --> --- #### 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](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md) 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
1 parent 8eadf65 commit fbe08f0

2 files changed

Lines changed: 43 additions & 2 deletions

File tree

  • integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/executor
  • prompt/prompt-executor/prompt-executor-clients/prompt-executor-ollama-client/src/commonMain/kotlin/ai/koog/prompt/executor/ollama/client/dto

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/executor/OllamaExecutorIntegrationTest.kt

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package ai.koog.integration.tests.executor
33
import ai.koog.integration.tests.InjectOllamaTestFixture
44
import ai.koog.integration.tests.OllamaTestFixture
55
import ai.koog.integration.tests.OllamaTestFixtureExtension
6+
import ai.koog.integration.tests.utils.MediaTestScenarios
67
import ai.koog.integration.tests.utils.MediaTestScenarios.ImageTestScenario
78
import ai.koog.integration.tests.utils.MediaTestUtils
89
import ai.koog.integration.tests.utils.MediaTestUtils.checkExecutorMediaResponse
@@ -320,4 +321,27 @@ class OllamaExecutorIntegrationTest : ExecutorIntegrationTestBase() {
320321
}
321322
}
322323
}
324+
325+
@Test
326+
fun `ollama_test txt file processing`() = runTest(timeout = 600.seconds) {
327+
val textFile = MediaTestUtils.createTextFileForScenario(
328+
MediaTestScenarios.TextTestScenario.BASIC_TEXT,
329+
testResourcesDir
330+
)
331+
332+
val prompt = prompt("text-file-test") {
333+
system("You are a helpful assistant that can analyze text files.")
334+
335+
user {
336+
markdown {
337+
+"I'm sending you a text file. Please read its content and summarize it."
338+
}
339+
340+
textFile(KtPath(textFile.pathString), "text/plain")
341+
}
342+
}
343+
344+
val response = executor.execute(prompt, model).single()
345+
response.content.shouldNotBeBlank()
346+
}
323347
}

prompt/prompt-executor/prompt-executor-clients/prompt-executor-ollama-client/src/commonMain/kotlin/ai/koog/prompt/executor/ollama/client/dto/OllamaConverters.kt

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,10 +59,13 @@ internal fun Prompt.toOllamaChatMessages(model: LLModel): List<OllamaChatMessage
5959
}
6060

6161
private fun Message.User.toOllamaChatMessage(model: LLModel): OllamaChatMessageDTO {
62+
val text = StringBuilder()
6263
val images = buildList {
6364
parts.forEach { part ->
6465
when (part) {
65-
is ContentPart.Text -> {}
66+
is ContentPart.Text -> {
67+
text.append(part.text)
68+
}
6669
is ContentPart.Image -> {
6770
require(LLMCapability.Vision.Image in model.capabilities) {
6871
"Model ${model.id} doesn't support images"
@@ -76,14 +79,28 @@ private fun Message.User.toOllamaChatMessage(model: LLModel): OllamaChatMessageD
7679
add(image)
7780
}
7881

82+
is ContentPart.File -> {
83+
val fileContent = when (val actualContent = part.content) {
84+
is AttachmentContent.PlainText -> {
85+
actualContent.text
86+
}
87+
88+
is AttachmentContent.Binary -> actualContent.asBase64()
89+
90+
else -> throw IllegalArgumentException("Unsupported file attachment content: ${content::class}")
91+
}
92+
93+
text.append("\n\n$fileContent")
94+
}
95+
7996
else -> throw IllegalArgumentException("Unsupported attachment type: $part")
8097
}
8198
}
8299
}
83100

84101
return OllamaChatMessageDTO(
85102
role = "user",
86-
content = content,
103+
content = text.toString(),
87104
images = images.takeIf { it.isNotEmpty() }
88105
)
89106
}

0 commit comments

Comments
 (0)