feat(routing): inject native ML complexity Router scoring into telegram wrapper header

This commit is contained in:
Tim Z. 2026-04-07 13:56:19 +02:00
parent f521aefa58
commit 3881f69aa0
3 changed files with 30 additions and 14 deletions

View file

@ -1624,11 +1624,9 @@ func (al *AgentLoop) runAgentLoop(
} }
if result.finalContent != "" { if result.finalContent != "" {
var totalTokens, promptTokens, compTokens int var totalTokens int
if usage := ts.GetLastUsage(); usage != nil { if usage := ts.GetLastUsage(); usage != nil {
totalTokens = usage.TotalTokens totalTokens = usage.TotalTokens
promptTokens = usage.PromptTokens
compTokens = usage.CompletionTokens
} }
version := al.GetConfig().BuildInfo.Version version := al.GetConfig().BuildInfo.Version
@ -1644,7 +1642,7 @@ func (al *AgentLoop) runAgentLoop(
Agent: agent.ID, Agent: agent.ID,
Route: providerName, Route: providerName,
Model: modelName, Model: modelName,
Complexity: calculateComplexity(promptTokens, compTokens), Complexity: int(ts.getComplexityScore() * 100),
Tokens: totalTokens, Tokens: totalTokens,
ToolCalls: ts.currentIteration(), ToolCalls: ts.currentIteration(),
Processing: time.Since(ts.startedAt), Processing: time.Since(ts.startedAt),
@ -1850,7 +1848,8 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
ts.ingestMessage(turnCtx, al, rootMsg) ts.ingestMessage(turnCtx, al, rootMsg)
} }
activeCandidates, activeModel, usedLight := al.selectCandidates(ts.agent, ts.userMessage, messages) activeCandidates, activeModel, usedLight, score := al.selectCandidates(ts.agent, ts.userMessage, messages)
ts.setComplexityScore(score)
activeProvider := ts.agent.Provider activeProvider := ts.agent.Provider
if usedLight && ts.agent.LightProvider != nil { if usedLight && ts.agent.LightProvider != nil {
activeProvider = ts.agent.LightProvider activeProvider = ts.agent.LightProvider
@ -2977,12 +2976,18 @@ func (al *AgentLoop) selectCandidates(
agent *AgentInstance, agent *AgentInstance,
userMsg string, userMsg string,
history []providers.Message, history []providers.Message,
) (candidates []providers.FallbackCandidate, model string, usedLight bool) { ) (candidates []providers.FallbackCandidate, model string, usedLight bool, score float64) {
if agent.Router == nil || len(agent.LightCandidates) == 0 { if agent.Router == nil || len(agent.LightCandidates) == 0 {
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, 0
} }
_, usedLight, score := agent.Router.SelectModel(userMsg, history, agent.Model) _, usedLight, score = agent.Router.SelectModel(userMsg, history, agent.Model)
// Cache the native rule-based complexity score to the primary turnState
if ts := turnStateFromContext(context.Background()); ts != nil {
// Note: Since we don't have direct access to ts here, we actually pass it up natively in runTurn
}
if !usedLight { if !usedLight {
logger.DebugCF("agent", "Model routing: primary model selected", logger.DebugCF("agent", "Model routing: primary model selected",
map[string]any{ map[string]any{
@ -2990,7 +2995,7 @@ func (al *AgentLoop) selectCandidates(
"score": score, "score": score,
"threshold": agent.Router.Threshold(), "threshold": agent.Router.Threshold(),
}) })
return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, score
} }
logger.InfoCF("agent", "Model routing: light model selected", logger.InfoCF("agent", "Model routing: light model selected",
@ -3000,7 +3005,7 @@ func (al *AgentLoop) selectCandidates(
"score": score, "score": score,
"threshold": agent.Router.Threshold(), "threshold": agent.Router.Threshold(),
}) })
return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true, score
} }
// resolveContextManager selects the ContextManager implementation based on config. // resolveContextManager selects the ContextManager implementation based on config.

View file

@ -23,7 +23,3 @@ func WrapResponse(rawOutput string, m Metrics) string {
) )
return header + rawOutput return header + rawOutput
} }
func calculateComplexity(promptTokens, completionTokens int) int {
return int(float64(promptTokens)*0.1 + float64(completionTokens)*0.5)
}

View file

@ -103,6 +103,9 @@ type turnState struct {
tokenBudget *atomic.Int64 // Shared token budget counter tokenBudget *atomic.Int64 // Shared token budget counter
lastFinishReason string // Last LLM finish_reason lastFinishReason string // Last LLM finish_reason
lastUsage *providers.UsageInfo // Last LLM usage info lastUsage *providers.UsageInfo // Last LLM usage info
// Telemetry tracking
complexityScore float64 // Cached heuristic difficulty score generated by the router
// Back-reference to the owning AgentLoop (set for SubTurns only, used for hard abort cascade) // Back-reference to the owning AgentLoop (set for SubTurns only, used for hard abort cascade)
al *AgentLoop al *AgentLoop
@ -214,6 +217,18 @@ func (ts *turnState) setIteration(iteration int) {
ts.iteration = iteration ts.iteration = iteration
} }
func (ts *turnState) setComplexityScore(score float64) {
ts.mu.Lock()
defer ts.mu.Unlock()
ts.complexityScore = score
}
func (ts *turnState) getComplexityScore() float64 {
ts.mu.RLock()
defer ts.mu.RUnlock()
return ts.complexityScore
}
func (ts *turnState) currentIteration() int { func (ts *turnState) currentIteration() int {
ts.mu.RLock() ts.mu.RLock()
defer ts.mu.RUnlock() defer ts.mu.RUnlock()