feat(router): capture latency + tokens on ROUTER ChatTurnEvent
ChatTurnEvent gains nullable latencyMs and tokensUsed fields (defaults preserve backward-compat; legacy JSON without them deserializes cleanly). emitChatTurn accepts optional metrics; the ROUTER emit site passes inferenceResponse values through; the USER emit site leaves both null.
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
package com.correx.core.events.events
|
package com.correx.core.events.events
|
||||||
|
|
||||||
import com.correx.core.events.types.SessionId
|
import com.correx.core.events.types.SessionId
|
||||||
|
import com.correx.core.inference.TokenUsage
|
||||||
import kotlinx.serialization.SerialName
|
import kotlinx.serialization.SerialName
|
||||||
import kotlinx.serialization.Serializable
|
import kotlinx.serialization.Serializable
|
||||||
|
|
||||||
@@ -12,6 +13,8 @@ data class ChatTurnEvent(
|
|||||||
val role: ChatTurnRole,
|
val role: ChatTurnRole,
|
||||||
val content: String,
|
val content: String,
|
||||||
val timestampMs: Long,
|
val timestampMs: Long,
|
||||||
|
val latencyMs: Long? = null,
|
||||||
|
val tokensUsed: TokenUsage? = null,
|
||||||
) : EventPayload
|
) : EventPayload
|
||||||
|
|
||||||
@Serializable
|
@Serializable
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ import com.correx.core.inference.InferenceRequest
|
|||||||
import com.correx.core.inference.InferenceRouter
|
import com.correx.core.inference.InferenceRouter
|
||||||
import com.correx.core.inference.ModelCapability
|
import com.correx.core.inference.ModelCapability
|
||||||
import com.correx.core.inference.ResponseFormat
|
import com.correx.core.inference.ResponseFormat
|
||||||
|
import com.correx.core.inference.TokenUsage
|
||||||
import com.correx.core.router.l3.L3MemoryEntry
|
import com.correx.core.router.l3.L3MemoryEntry
|
||||||
import com.correx.core.router.l3.L3MemoryStore
|
import com.correx.core.router.l3.L3MemoryStore
|
||||||
import com.correx.core.router.l3.L3Query
|
import com.correx.core.router.l3.L3Query
|
||||||
@@ -148,7 +149,11 @@ class DefaultRouterFacade(
|
|||||||
val content = inferenceResponse.text
|
val content = inferenceResponse.text
|
||||||
|
|
||||||
// Emit ROUTER turn event
|
// Emit ROUTER turn event
|
||||||
emitChatTurn(sessionId, content, ChatTurnRole.ROUTER)
|
emitChatTurn(
|
||||||
|
sessionId, content, ChatTurnRole.ROUTER,
|
||||||
|
latencyMs = inferenceResponse.latencyMs,
|
||||||
|
tokensUsed = inferenceResponse.tokensUsed,
|
||||||
|
)
|
||||||
|
|
||||||
if (mode == ChatMode.STEERING) {
|
if (mode == ChatMode.STEERING) {
|
||||||
val validationError = validateSteering?.invoke(content)
|
val validationError = validateSteering?.invoke(content)
|
||||||
@@ -160,7 +165,13 @@ class DefaultRouterFacade(
|
|||||||
return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING))
|
return RouterResponse(content = content, steeringEmitted = (mode == ChatMode.STEERING))
|
||||||
}
|
}
|
||||||
|
|
||||||
private suspend fun emitChatTurn(sessionId: SessionId, content: String, role: ChatTurnRole): String {
|
private suspend fun emitChatTurn(
|
||||||
|
sessionId: SessionId,
|
||||||
|
content: String,
|
||||||
|
role: ChatTurnRole,
|
||||||
|
latencyMs: Long? = null,
|
||||||
|
tokensUsed: TokenUsage? = null,
|
||||||
|
): String {
|
||||||
val turnId = UUID.randomUUID().toString()
|
val turnId = UUID.randomUUID().toString()
|
||||||
val now = Clock.System.now()
|
val now = Clock.System.now()
|
||||||
eventStore.append(
|
eventStore.append(
|
||||||
@@ -179,6 +190,8 @@ class DefaultRouterFacade(
|
|||||||
role = role,
|
role = role,
|
||||||
content = content,
|
content = content,
|
||||||
timestampMs = now.toEpochMilliseconds(),
|
timestampMs = now.toEpochMilliseconds(),
|
||||||
|
latencyMs = latencyMs,
|
||||||
|
tokensUsed = tokensUsed,
|
||||||
),
|
),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -52,9 +52,11 @@ import com.correx.testing.fixtures.inference.MockTokenizer
|
|||||||
import kotlinx.coroutines.flow.Flow
|
import kotlinx.coroutines.flow.Flow
|
||||||
import kotlinx.coroutines.runBlocking
|
import kotlinx.coroutines.runBlocking
|
||||||
import kotlinx.datetime.Instant
|
import kotlinx.datetime.Instant
|
||||||
|
import kotlinx.serialization.decodeFromString
|
||||||
import org.junit.jupiter.api.Assertions.assertEquals
|
import org.junit.jupiter.api.Assertions.assertEquals
|
||||||
import org.junit.jupiter.api.Assertions.assertFalse
|
import org.junit.jupiter.api.Assertions.assertFalse
|
||||||
import org.junit.jupiter.api.Assertions.assertNotNull
|
import org.junit.jupiter.api.Assertions.assertNotNull
|
||||||
|
import org.junit.jupiter.api.Assertions.assertNull
|
||||||
import org.junit.jupiter.api.Assertions.assertTrue
|
import org.junit.jupiter.api.Assertions.assertTrue
|
||||||
import org.junit.jupiter.api.Test
|
import org.junit.jupiter.api.Test
|
||||||
|
|
||||||
@@ -544,6 +546,75 @@ class RouterFacadeTest {
|
|||||||
assertEquals(2, hits.size)
|
assertEquals(2, hits.size)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --------------------------------------------------------------------------
|
||||||
|
// Task 2.1 — latency + token metrics on ROUTER ChatTurnEvent
|
||||||
|
// --------------------------------------------------------------------------
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `ROUTER ChatTurnEvent carries latencyMs and tokensUsed from inference and USER turn has both null`(): Unit = runBlocking {
|
||||||
|
val capturedRequests = mutableListOf<InferenceRequest>()
|
||||||
|
val knownLatencyMs = 42L
|
||||||
|
val knownTokenUsage = TokenUsage(promptTokens = 10, completionTokens = 20)
|
||||||
|
val metricsProvider = object : InferenceProvider {
|
||||||
|
override val id = ProviderId("mock")
|
||||||
|
override val name = "Mock"
|
||||||
|
override val tokenizer = MockTokenizer()
|
||||||
|
override suspend fun infer(request: InferenceRequest): InferenceResponse {
|
||||||
|
capturedRequests.add(request)
|
||||||
|
return InferenceResponse(
|
||||||
|
requestId = request.requestId,
|
||||||
|
text = "metrics response",
|
||||||
|
finishReason = FinishReason.Stop,
|
||||||
|
tokensUsed = knownTokenUsage,
|
||||||
|
latencyMs = knownLatencyMs,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
|
||||||
|
override fun capabilities(): Set<CapabilityScore> = setOf(CapabilityScore(ModelCapability.General, 1.0))
|
||||||
|
}
|
||||||
|
val eventStore = mockEventStore()
|
||||||
|
val facade = DefaultRouterFacade(
|
||||||
|
routerRepository = object : RouterRepository {
|
||||||
|
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
|
||||||
|
},
|
||||||
|
routerContextBuilder = object : RouterContextBuilder {
|
||||||
|
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
|
||||||
|
},
|
||||||
|
inferenceRouter = object : InferenceRouter {
|
||||||
|
override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider =
|
||||||
|
metricsProvider
|
||||||
|
},
|
||||||
|
eventStore = eventStore,
|
||||||
|
config = RouterConfig(tokenBudget = TokenBudget(limit = 5000)),
|
||||||
|
embedder = NoopEmbedder(dimension = 8),
|
||||||
|
l3MemoryStore = InMemoryL3MemoryStore(),
|
||||||
|
)
|
||||||
|
|
||||||
|
facade.onUserInput(sessionId = SessionId("metrics-session"), input = "hello metrics")
|
||||||
|
|
||||||
|
val chatEvents = eventStore.appendedEvents.map { it.payload }.filterIsInstance<ChatTurnEvent>()
|
||||||
|
assertEquals(2, chatEvents.size)
|
||||||
|
val userEvent = chatEvents.first { it.role == ChatTurnRole.USER }
|
||||||
|
val routerEvent = chatEvents.first { it.role == ChatTurnRole.ROUTER }
|
||||||
|
|
||||||
|
assertNull(userEvent.latencyMs, "USER ChatTurnEvent must have null latencyMs")
|
||||||
|
assertNull(userEvent.tokensUsed, "USER ChatTurnEvent must have null tokensUsed")
|
||||||
|
|
||||||
|
assertEquals(knownLatencyMs, routerEvent.latencyMs)
|
||||||
|
assertNotNull(routerEvent.tokensUsed)
|
||||||
|
assertEquals(30, routerEvent.tokensUsed!!.totalTokens)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun `ChatTurnEvent backward-compat - legacy JSON without latencyMs and tokensUsed deserializes with both null`() {
|
||||||
|
val legacyJson = """{"sessionId":"s1","turnId":"t1","role":"USER","content":"hello","timestampMs":1000}"""
|
||||||
|
val event = com.correx.core.events.serialization.eventJson.decodeFromString<ChatTurnEvent>(legacyJson)
|
||||||
|
assertEquals(SessionId("s1"), event.sessionId)
|
||||||
|
assertEquals("hello", event.content)
|
||||||
|
assertNull(event.latencyMs)
|
||||||
|
assertNull(event.tokensUsed)
|
||||||
|
}
|
||||||
|
|
||||||
// --------------------------------------------------------------------------
|
// --------------------------------------------------------------------------
|
||||||
// Helpers
|
// Helpers
|
||||||
// --------------------------------------------------------------------------
|
// --------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user