fun testPromptTokenizer()

in prompt/prompt-tokenizer/src/commonTest/kotlin/ai/koog/prompt/tokenizer/PromptTokenizerTest.kt [57:112]
44 lines of code
5 McCabe index (conditional complexity)

    fun testPromptTokenizer() {
        // Create a mock tokenizer to track token usage
        val mockTokenizer = MockTokenizer()

        // Create a prompt tokenizer with our mock tokenizer
        val promptTokenizer = OnDemandTokenizer(mockTokenizer)

        // Create a prompt with some messages
        val testPrompt = prompt("test-prompt") {
            system("You are a helpful assistant.")
            user("What is the capital of France?")
            assistant("Paris is the capital of France.")
        }

        // Count tokens in the prompt
        val totalTokens = promptTokenizer.tokenCountFor(testPrompt)

        // Verify that tokens were counted
        assertTrue(totalTokens > 0, "Total tokens should be greater than 0")

        // Verify that the tokenizer was used and counted tokens
        assertTrue(mockTokenizer.totalTokens > 0, "Tokenizer should have counted tokens")

        // Store the tokens counted for the prompt (before counting individual messages)
        val promptTokenCount = mockTokenizer.totalTokens

        // Print the total tokens spent for the prompt
        println("[DEBUG_LOG] Total tokens spent for prompt: $promptTokenCount")

        // Reset tokenizer to count individual messages separately
        mockTokenizer.reset()

        val requestMetainfo = RequestMetaInfo.create(Clock.System)
        val responseMetainfo = ResponseMetaInfo.create(Clock.System)
        // Count tokens for individual messages
        val systemTokens = promptTokenizer.tokenCountFor(
            Message.System("You are a helpful assistant.", requestMetainfo)
        )
        val userTokens = promptTokenizer.tokenCountFor(Message.User("What is the capital of France?", requestMetainfo))
        val assistantTokens = promptTokenizer.tokenCountFor(
            Message.Assistant("Paris is the capital of France.", responseMetainfo)
        )

        // Print token counts for each message
        println("[DEBUG_LOG] System message tokens: $systemTokens")
        println("[DEBUG_LOG] User message tokens: $userTokens")
        println("[DEBUG_LOG] Assistant message tokens: $assistantTokens")

        // Verify that individual messages have positive token counts
        assertTrue(systemTokens > 0, "System message should have tokens")
        assertTrue(userTokens > 0, "User message should have tokens")
        assertTrue(assistantTokens > 0, "Assistant message should have tokens")

        // Verify that the total tokens from prompt counting equals what we stored
        assertEquals(totalTokens, promptTokenCount, "Total tokens should match the stored prompt token count")
    }